Dynamic config

Author: schustmi

src/zenml/pipelines/dynamic/pipeline_definition.py

@@ -50,17 +50,19 @@
 class DynamicPipeline(Pipeline):
     """ZenML pipeline class."""
 
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self._depends_on = kwargs.pop("depends_on", None) or []
+        if self._depends_on:
+            # TODO: This doesn't really work, as `step.with_options`
+            sources = [step.resolve().import_path for step in self._depends_on]
+            if len(sources) != len(set(sources)):
+                raise ValueError("Duplicate steps in depends_on.")
+
+        super().__init__(*args, **kwargs)
+
     @property
     def depends_on(self) -> List["BaseStep"]:
-        # TODO: Even with this, it will not be possible to define all potential
-        # docker builds:
-        # If a step will be called multiple times, once without step operator
-        # and once with, it might require multiple docker images. Even if this
-        # list had two copies of the same step, how would we map at runtime
-        # which one to use?
-        # TODO: maybe this needs to be a dict, and steps can select which
-        # "template" config (including the docker image) to use?
-        return getattr(self, "_depends_on", [])
+        return self._depends_on
 
     @property
     def is_dynamic(self) -> bool:
@@ -148,7 +150,7 @@ def pipeline_(param_name: str):
         self._parameters = validated_args
         self._invocations = {}
         with self:
-            for step in self.depends_on:
+            for step in self._depends_on:
                 self.add_step_invocation(
                     step,
                     input_artifacts={},
@@ -181,7 +183,7 @@ def add_dynamic_invocation(
             external_artifacts=external_artifacts,
             model_artifacts_or_metadata={},
             client_lazy_loaders={},
-            parameters={},
+            parameters=step.configuration.parameters,
             default_parameters={},
             upstream_steps=upstream_steps or set(),
             pipeline=self,

src/zenml/pipelines/dynamic/runner.py

@@ -18,7 +18,7 @@
 from zenml import ExternalArtifact
 from zenml.client import Client
 from zenml.config.compiler import Compiler
-from zenml.config.step_configurations import Step, StepConfiguration
+from zenml.config.step_configurations import Step
 from zenml.exceptions import RunStoppedException
 from zenml.logger import get_logger
 from zenml.logging.step_logging import setup_pipeline_logging
@@ -281,14 +281,15 @@ def _compile_step(
                 annotation=Any,
                 pipeline=pipeline,
             )
-        elif isinstance(value, ArtifactVersionResponse):
+        elif isinstance(value, (ArtifactVersionResponse, ExternalArtifact)):
             external_artifacts[name] = value
         else:
             external_artifacts[name] = ExternalArtifact(value=value)
 
-    if config := get_dynamic_step_configuration(snapshot, step):
-        step._configuration = config
+    if template := get_static_step_template(snapshot, step):
+        step._configuration = template.config
 
+    step._apply_dynamic_configuration()
     invocation_id = pipeline.add_dynamic_invocation(
         step=step,
         custom_id=id,
@@ -413,14 +414,14 @@ def _runs_in_process(step: "Step") -> bool:
     return True
 
 
-def get_dynamic_step_configuration(
+def get_static_step_template(
     snapshot: "PipelineSnapshotResponse",
     step: "BaseStep",
-) -> Optional["StepConfiguration"]:
+) -> Optional["Step"]:
     step_source = step.resolve().import_path
 
     for compiled_step in snapshot.step_configurations.values():
         if compiled_step.spec.source.import_path == step_source:
-            return compiled_step.config
+            return compiled_step
 
     return None

src/zenml/pipelines/pipeline_decorator.py

@@ -35,6 +35,7 @@
     from zenml.config.retry_config import StepRetryConfig
     from zenml.model.model import Model
     from zenml.pipelines.pipeline_definition import Pipeline
+    from zenml.steps.base_step import BaseStep
     from zenml.types import HookSpecification, InitHookSpecification
     from zenml.utils.tag_utils import Tag
 
@@ -51,6 +52,7 @@ def pipeline(_func: "F") -> "Pipeline": ...
 def pipeline(
     *,
     name: Optional[str] = None,
+    depends_on: Optional[List["BaseStep"]] = None,
     enable_cache: Optional[bool] = None,
     enable_artifact_metadata: Optional[bool] = None,
     enable_step_logs: Optional[bool] = None,
@@ -77,6 +79,7 @@ def pipeline(
     _func: Optional["F"] = None,
     *,
     name: Optional[str] = None,
+    depends_on: Optional[List["BaseStep"]] = None,
     enable_cache: Optional[bool] = None,
     enable_artifact_metadata: Optional[bool] = None,
     enable_step_logs: Optional[bool] = None,
@@ -142,6 +145,7 @@ def inner_decorator(func: "F") -> "Pipeline":
 
         p = DynamicPipeline(
             name=name or func.__name__,
+            depends_on=depends_on,
             entrypoint=func,
             enable_cache=enable_cache,
             enable_artifact_metadata=enable_artifact_metadata,

src/zenml/steps/base_step.py

@@ -201,6 +201,7 @@ def __init__(
             )
 
         self._configuration = PartialStepConfiguration(name=name)
+        self._dynamic_configuration: Optional["StepConfigurationUpdate"] = None
         self.configure(
             enable_cache=enable_cache,
             enable_artifact_metadata=enable_artifact_metadata,
@@ -897,15 +898,34 @@ def _apply_configuration(
                 or not. See the `BaseStep.configure(...)` method for a detailed
                 explanation.
         """
+        from zenml.pipelines.dynamic.context import DynamicPipelineRunContext
+
         self._validate_configuration(config, runtime_parameters)
 
+        if DynamicPipelineRunContext.get():
+            if self._dynamic_configuration is None:
+                self._dynamic_configuration = config
+            else:
+                self._dynamic_configuration = pydantic_utils.update_model(
+                    self._dynamic_configuration, update=config, recursive=merge
+                )
+            return
+
         self._configuration = pydantic_utils.update_model(
             self._configuration, update=config, recursive=merge
         )
 
         logger.debug("Updated step configuration:")
         logger.debug(self._configuration)
 
+    def _apply_dynamic_configuration(self) -> None:
+        if self._dynamic_configuration:
+            self._configuration = pydantic_utils.update_model(
+                self._configuration,
+                update=self._dynamic_configuration,
+                recursive=True,
+            )
+
     def _validate_configuration(
         self,
         config: "StepConfigurationUpdate",