Dynamic pipelines v0

[schustmi]

Describe changes

This PR implements dynamic pipelines for the local/kubernetes orchestrators.

Example

from zenml import step, pipeline

@step
def generate_int() -> int:
  return 3

@step
def do_something(index: int) -> None:
  ...

@pipeline(dynamic=True)
def dynamic_pipeline() -> None:
  count = generate_int()
  # `count` is an artifact, we now load the data
  count_data = count.load()

  for idx in range(count_data):
    # This will run sequentially, like regular python code would.
    # See the features below for an example on how to run steps
    # in parallel.
    do_something(idx)

if __name__ == "__main__":
  dynamic_pipeline()

Features

• Dynamic configuration for steps:

@pipeline(dynamic=True)
def dynamic_pipeline():
  some_step.with_options(enable_cache=False)()

• Specify a step runtime: The runtime configures whether the step will be run in the orchestration environment (runtime=inline) or if the orchestrator will spin up a separate step execution environment (runtime=isolated).

@step(runtime="isolated")
def some_step() -> None:
  ...

• Execute multiple steps in parallel by using step.submit(...). This will either execute the step or launch a new container in a new thread.

@step
def some_step(arg: int) -> int:
  ...

@pipeline(dynamic=True)
def dynamic_pipeline():
  future = some_step.submit(arg=1)
  artifact = future.result()  # wait and get artifact response(s)
  data = future.load()  # wait and load artifact data
  downstream_step(future)  # pass the output to another step

  # Run multiple steps in parallel
  for idx in range(3):
    some_step.submit(arg=idx)

• Specify config templates for steps using depends_on:

# config.yaml
steps:
  some_step:
    parameters:
      arg: 3

# run.py
@step
def some_step(arg: int) -> None:
  ...

@pipeline(dynamic=True, depends_on=[some_step])
def dynamic_pipeline():
  some_step()

if __name__ == "__main__":
  dynamic_pipeline.with_options(config_path="config.yaml")()

Limitations/Known issues

• Our logging storage isn't threadsafe yet, which means logs of parallel steps are mixed up.
• When running multiple steps concurrently, failure in one step does not stop the other steps. Instead, they continue executing until finished.

Pre-requisites

Please ensure you have done the following:

• I have read the CONTRIBUTING.md document.
• I have added tests to cover my changes.
• I have based my new branch on develop and the open PR is targeting develop. If your branch wasn't based on develop read Contribution guide on rebasing branch to develop.
• IMPORTANT: I made sure that my changes are reflected properly in the following resources:
  • ZenML Docs
  • Dashboard: Needs to be communicated to the frontend team.
  • Templates: Might need adjustments (that are not reflected in the template tests) in case of non-breaking changes and deprecations.
  • Projects: Depending on the version dependencies, different projects might get affected.

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Other (add details above)

[Json-Andriopoulos]

Perhaps we should keep a cache dictionary for materialized artifacts, so if (for some weird reason) users do multiple load calls (instead of assigning and reusing) we can return the cached results.

[schustmi]

I like that idea. I'll also add a boolean flag to the load method to prevent this, so users can at least manually disable it for huge artifacts.

[Json-Andriopoulos]

Maybe we should provide a custom implementation for setitem as a better UX/guidance to users than getting an error like this: object does not support item assignment.

[Json-Andriopoulos]

I would recommend adopting one kind of wrapper behavior and stick to it. If we enable users to access futures both dict-style and list/tuple style I think the more we extend the class magic functions the more the behavior may be more confusing regarding what gets executed and what the users expects to be executed.

Plus in general accessing by key is a safer operation and would be a good pattern to enforce. The results may change length, order and accessing by position is a bit un-safe.

[Json-Andriopoulos]

For instance, I can get dict-style (by specifying key) but the __iter__ function I may expect it to return the keys not the values. For better consistency I would say we have 2 options:

1. No magic functions - We provide public methods with clean documentation. Uses can manipulate futures based on the available functions and their signatures. No confusion, users stick to following docstrings.
2. We implement magic functions but we assume implementation wraps one kind of hidden data structure. Users can manipulate it as dict (preferable imo) or as a list/tuple.

[schustmi]

Makes sense, I agree we should only support one way. My vote goes for tuple like behaviour though, as the main purpose of my implementations was to actually support the following: When calling any python function with multiple outputs, there are two cases:

def f():
  return 1, "str"

# Return value is a tuple that can be accessed with `int` keys
tuple_result =  f()
int_result = tuple_result[0]

# Automatic unpacking
int_result, str_result = f()

I think the latter is most common, and is the use-case I think we should support to make it feel as pythonic as possible (and also mirrors how you would call sync steps, in which case the return value will not be a future but instead of tuple of artifacts).

Tuples are also immutable which in our case holds true as well as you can't add outputs to the future result of a step run.

We can then add some helper methods like get_output(key: str) like you suggested to get allow fetching specific outputs by key.

[Json-Andriopoulos]

I see, yes, tuple should do it as well. As long as we are consistent I think we are ok.

[Json-Andriopoulos]

Would __contains__ make sense to also implement here? 🤔

[schustmi]

Oh yes I thought I already did, that definitely makes sense!
Actually if we do implement it as a tuple-like data structure the contains will be with values and might not make much sense, what do you think?

[Json-Andriopoulos]

yy it would make more sense in a dict-like scenario!

[Json-Andriopoulos]

Not in the context of this PR but we can do a better job here, thread-safety wise. In general os.environ should be treated as an immutable value, changing its values may also affect un-intentionally other execution paths (for instance code running in other threads).

I think masking the environment under a custom object/class and making that accessible with context vars resolves the issue (context vars are thread-local, wrapper object loads and freezes state and exposes all the operations we want with relative safety).

[schustmi]

Yes you're right! This is not a problem in this case, but when running multiple steps in parallel we also set different environment variables which is problematic. At least in this case I can switch to using a context var instead

[Json-Andriopoulos]

We can work something here, but different story of course. ContextVars should work perfectly, maybe in combo with a centralized BaseSettings object. Will create the story and we can discuss implementation.

[Json-Andriopoulos]

It is hard to track those environment variables references. Maybe a BaseSettings object to organize those in a single place would be a good idea. Also we wouldn't need this function, pydantic would validate the boolean value for us in one go.

[schustmi]

Yep agreed in general that would be a nice thing. We even have some classes that do this for a subset of env variables (ServerConfiguration and GlobalConfiguration), but not for all of them. This function also treats values like "yes" as True, not sure how pydantic would handle that natively, but we could implement it for sure.

[Json-Andriopoulos]

I think Pydantic captures multiple values as well, I can cross compare. Will open a new story for this just wanted to get your opinion :)