misc

Author: schustmi

src/zenml/config/compiler.py

@@ -35,7 +35,6 @@
 from zenml.config.step_configurations import (
     InputSpec,
     Step,
-    StepConfiguration,
     StepConfigurationUpdate,
     StepSpec,
 )
@@ -48,7 +47,6 @@
 if TYPE_CHECKING:
     from zenml.pipelines.pipeline_definition import Pipeline
     from zenml.stack import Stack, StackComponent
-    from zenml.steps.base_step import BaseStep
     from zenml.steps.step_invocation import StepInvocation
 
 from zenml.logger import get_logger
@@ -128,29 +126,17 @@ def compile(
                 merge=False,
             )
 
-        if pipeline.is_dynamic:
-            step_templates = {
-                step.name: self._compile_config_template(
-                    step=step, stack=stack
-                )
-                for step in pipeline.depends_on
-            }
-            steps = {}
-        else:
-            step_templates = None
-            steps = {
-                invocation_id: self._compile_step_invocation(
-                    invocation=invocation,
-                    stack=stack,
-                    step_config=(run_configuration.steps or {}).get(
-                        invocation_id
-                    ),
-                    pipeline_configuration=pipeline.configuration,
-                )
-                for invocation_id, invocation in self._get_sorted_invocations(
-                    pipeline=pipeline
-                )
-            }
+        steps = {
+            invocation_id: self._compile_step_invocation(
+                invocation=invocation,
+                stack=stack,
+                step_config=(run_configuration.steps or {}).get(invocation_id),
+                pipeline_configuration=pipeline.configuration,
+            )
+            for invocation_id, invocation in self._get_sorted_invocations(
+                pipeline=pipeline
+            )
+        }
 
         self._ensure_required_stack_components_exist(stack=stack, steps=steps)
 
@@ -170,7 +156,6 @@ def compile(
             is_dynamic=pipeline.is_dynamic,
             pipeline_configuration=pipeline.configuration,
             step_configurations=steps,
-            step_configuration_templates=step_templates,
             client_environment=get_run_environment_dict(),
             client_version=client_version,
             server_version=server_version,
@@ -536,48 +521,6 @@ def _compile_step_invocation(
             step_config_overrides=step_configuration_overrides,
         )
 
-    def _compile_config_template(
-        self,
-        step: "BaseStep",
-        stack: "Stack",
-        step_config: Optional["StepConfigurationUpdate"],
-    ) -> StepConfiguration:
-        """Compiles a ZenML step.
-
-        Args:
-            invocation: The step invocation to compile.
-            stack: The stack on which the pipeline will be run.
-            step_config: Run configuration for the step.
-            pipeline_configuration: Configuration for the pipeline.
-
-        Returns:
-            The compiled step.
-        """
-        if step_config:
-            step._apply_configuration(step_config)
-
-        convert_component_shortcut_settings_keys(
-            step.configuration.settings, stack=stack
-        )
-        step_secrets = secret_utils.resolve_and_verify_secrets(
-            step.configuration.secrets
-        )
-        step_settings = self._filter_and_validate_settings(
-            settings=step.configuration.settings,
-            configuration_level=ConfigurationLevel.STEP,
-            stack=stack,
-        )
-        step.configure(
-            secrets=step_secrets,
-            settings=step_settings,
-            merge=False,
-        )
-
-        # TODO: apply pipeline config
-        return StepConfiguration.model_validate(
-            step.configuration.model_dump()
-        )
-
     def _get_sorted_invocations(
         self,
         pipeline: "Pipeline",

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -857,6 +857,7 @@ def run_dynamic_out_of_process_step(
         settings = cast(
             KubernetesOrchestratorSettings, self.get_settings(step_run_info)
         )
+        # TODO: Way to fetch the correct docker image here.
         image_name = step_run_info.get_image(key=ORCHESTRATOR_DOCKER_IMAGE_KEY)
         command = StepOperatorEntrypointConfiguration.get_entrypoint_command()
         args = StepOperatorEntrypointConfiguration.get_entrypoint_arguments(

src/zenml/models/v2/core/pipeline_snapshot.py

@@ -31,7 +31,7 @@
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.pipeline_spec import PipelineSpec
-from zenml.config.step_configurations import Step, StepConfiguration
+from zenml.config.step_configurations import Step
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.models.v2.base.base import BaseUpdate, BaseZenModel
@@ -82,12 +82,6 @@ class PipelineSnapshotBase(BaseZenModel):
     step_configurations: Dict[str, Step] = Field(
         default={}, title="The step configurations for this snapshot."
     )
-    step_configuration_templates: Optional[Dict[str, StepConfiguration]] = (
-        Field(
-            default=None,
-            title="The step configuration templates of the snapshot.",
-        )
-    )
     client_environment: Dict[str, Any] = Field(
         default={}, title="The client environment for this snapshot."
     )

src/zenml/orchestrators/containerized_orchestrator.py

@@ -120,24 +120,6 @@ def get_docker_builds(
                 builds.append(pipeline_build)
                 included_pipeline_build = True
 
-        for name, step_config in snapshot.step_configuration_templates.items():
-            step_settings = step_config.docker_settings
-
-            if step_settings != pipeline_settings:
-                build = BuildConfiguration(
-                    key=ORCHESTRATOR_DOCKER_IMAGE_KEY,
-                    settings=step_settings,
-                    step_name=name,
-                )
-                builds.append(build)
-            elif not included_pipeline_build:
-                pipeline_build = BuildConfiguration(
-                    key=ORCHESTRATOR_DOCKER_IMAGE_KEY,
-                    settings=pipeline_settings,
-                )
-                builds.append(pipeline_build)
-                included_pipeline_build = True
-
         if not included_pipeline_build and self.should_build_pipeline_image(
             snapshot
         ):

src/zenml/pipelines/dynamic/context.py

@@ -69,7 +69,3 @@ def __enter__(self) -> Self:
                 "Calling a pipeline within a dynamic pipeline is not allowed."
             )
         return super().__enter__()
-
-
-def get_dynamic_pipeline_run_context() -> Optional[DynamicPipelineRunContext]:
-    return DynamicPipelineRunContext.get()

src/zenml/pipelines/dynamic/pipeline_definition.py

@@ -30,11 +30,9 @@
 from zenml import ExternalArtifact
 from zenml.client import Client
 from zenml.logger import get_logger
-from zenml.logging.step_logging import setup_pipeline_logging
 from zenml.models import ArtifactVersionResponse, PipelineRunResponse
 from zenml.pipelines.pipeline_definition import Pipeline
 from zenml.pipelines.run_utils import (
-    create_placeholder_run,
     should_prevent_pipeline_execution,
 )
 from zenml.steps.step_invocation import StepInvocation
@@ -54,7 +52,15 @@ class DynamicPipeline(Pipeline):
 
     @property
     def depends_on(self) -> List["BaseStep"]:
-        return []
+        # 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", [])
 
     @property
     def is_dynamic(self) -> bool:
@@ -141,6 +147,18 @@ def pipeline_(param_name: str):
 
         self._parameters = validated_args
         self._invocations = {}
+        with self:
+            for step in self.depends_on:
+                self.add_step_invocation(
+                    step,
+                    input_artifacts={},
+                    external_artifacts={},
+                    model_artifacts_or_metadata={},
+                    client_lazy_loaders={},
+                    parameters={},
+                    default_parameters={},
+                    upstream_steps=set(),
+                )
 
     def add_dynamic_invocation(
         self,
@@ -179,28 +197,13 @@ def __call__(
             return
 
         stack = Client().active_stack
-
         if not stack.orchestrator.supports_dynamic_pipelines:
             raise RuntimeError(
                 f"The {stack.orchestrator.__class__.__name__} does not support dynamic pipelines. "
             )
 
         self.prepare(*args, **kwargs)
-        snapshot = self._create_snapshot(**self._run_args)
-
-        with setup_pipeline_logging(
-            source="client", snapshot=snapshot
-        ) as logs_request:
-            run = create_placeholder_run(
-                snapshot=snapshot,
-                logs=logs_request,
-            )
-            stack.orchestrator.run(
-                snapshot=snapshot,
-                stack=stack,
-                placeholder_run=run,
-            )
-        return Client().get_pipeline_run(run.id)
+        return self._run()
 
     def _call_entrypoint(self, *args: Any, **kwargs: Any) -> None:
         """Calls the pipeline entrypoint function with the given arguments.
@@ -227,6 +230,9 @@ def _call_entrypoint(self, *args: Any, **kwargs: Any) -> None:
                 "Check out the pydantic error above for more details."
             ) from e
 
+        # Clear the invocations as they might still contain invocations from
+        # the compilation phase.
+        self._invocations = {}
         self.entrypoint(**validated_args)
 
     def _compute_output_schema(self) -> Optional[Dict[str, Any]]:

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
+from zenml.config.step_configurations import Step, StepConfiguration
 from zenml.exceptions import RunStoppedException
 from zenml.logger import get_logger
 from zenml.logging.step_logging import setup_pipeline_logging
@@ -174,9 +174,10 @@ def run_step_sync(
             "StepRunResultFuture", Sequence["StepRunResultFuture"], None
         ] = None,
     ) -> StepRunResult:
+        step = step.copy()
         inputs, upstream_steps = _prepare_step_run(step, args, kwargs, after)
         compiled_step, _ = _compile_step(
-            self.pipeline, step, id, upstream_steps, inputs
+            self._snapshot, self.pipeline, step, id, upstream_steps, inputs
         )
         step_run = _run_step_sync(
             snapshot=self._snapshot,
@@ -197,9 +198,10 @@ def run_step_in_thread(
             "StepRunResultFuture", Sequence["StepRunResultFuture"], None
         ] = None,
     ) -> StepRunResultFuture:
+        step = step.copy()
         inputs, upstream_steps = _prepare_step_run(step, args, kwargs, after)
         compiled_step, invocation_id = _compile_step(
-            self.pipeline, step, id, upstream_steps, inputs
+            self._snapshot, self.pipeline, step, id, upstream_steps, inputs
         )
 
         def _run() -> StepRunResult:
@@ -266,6 +268,7 @@ def _await_and_validate_input(input: Any):
 
 
 def _compile_step(
+    snapshot: "PipelineSnapshotResponse",
     pipeline: "DynamicPipeline",
     step: "BaseStep",
     id: Optional[str],
@@ -287,6 +290,9 @@ def _compile_step(
         else:
             external_artifacts[name] = ExternalArtifact(value=value)
 
+    if config := get_dynamic_step_configuration(snapshot, step):
+        step._configuration = config
+
     invocation_id = pipeline.add_dynamic_invocation(
         step=step,
         custom_id=id,
@@ -409,3 +415,16 @@ def _runs_in_process(step: "Step") -> bool:
         return False
 
     return True
+
+
+def get_dynamic_step_configuration(
+    snapshot: "PipelineSnapshotResponse",
+    step: "BaseStep",
+) -> Optional["StepConfiguration"]:
+    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 None

src/zenml/steps/base_step.py

@@ -200,14 +200,12 @@ def __init__(
                 },
             )
 
-        self._configuration = PartialStepConfiguration(
-            name=name,
+        self._configuration = PartialStepConfiguration(name=name)
+        self.configure(
             enable_cache=enable_cache,
             enable_artifact_metadata=enable_artifact_metadata,
             enable_artifact_visualization=enable_artifact_visualization,
             enable_step_logs=enable_step_logs,
-        )
-        self.configure(
             experiment_tracker=experiment_tracker,
             step_operator=step_operator,
             output_materializers=output_materializers,
@@ -1061,6 +1059,7 @@ def _finalize_configuration(
         external_artifacts: Dict[str, "ExternalArtifactConfiguration"],
         model_artifacts_or_metadata: Dict[str, "ModelVersionDataLazyLoader"],
         client_lazy_loaders: Dict[str, "ClientLazyLoader"],
+        skip_input_validation: bool = False,
     ) -> "StepConfiguration":
         """Finalizes the configuration after the step was called.
 
@@ -1075,6 +1074,7 @@ def _finalize_configuration(
             model_artifacts_or_metadata: The model artifacts or metadata of
                 this step.
             client_lazy_loaders: The client lazy loaders of this step.
+            skip_input_validation: If True, will skip the input validation.
 
         Raises:
             StepInterfaceError: If explicit materializers were specified for an
@@ -1164,12 +1164,13 @@ def _finalize_configuration(
 
         parameters = self._finalize_parameters()
         self.configure(parameters=parameters, merge=False)
-        self._validate_inputs(
-            input_artifacts=input_artifacts,
-            external_artifacts=external_artifacts,
-            model_artifacts_or_metadata=model_artifacts_or_metadata,
-            client_lazy_loaders=client_lazy_loaders,
-        )
+        if not skip_input_validation:
+            self._validate_inputs(
+                input_artifacts=input_artifacts,
+                external_artifacts=external_artifacts,
+                model_artifacts_or_metadata=model_artifacts_or_metadata,
+                client_lazy_loaders=client_lazy_loaders,
+            )
 
         values = dict_utils.remove_none_values({"outputs": outputs or None})
         config = StepConfigurationUpdate(**values)

src/zenml/steps/step_invocation.py

@@ -119,4 +119,6 @@ def finalize(self, parameters_to_ignore: Set[str]) -> "StepConfiguration":
             external_artifacts=external_artifacts,
             model_artifacts_or_metadata=self.model_artifacts_or_metadata,
             client_lazy_loaders=self.client_lazy_loaders,
+            # Skip input validation for dynamic pipeline step templates
+            skip_input_validation=self.pipeline.is_dynamic,
         )