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/config/step_run_info.py

@@ -18,7 +18,11 @@
 
 from zenml.config.frozen_base_model import FrozenBaseModel
 from zenml.config.pipeline_configurations import PipelineConfiguration
-from zenml.config.step_configurations import StepConfiguration
+from zenml.config.step_configurations import StepConfiguration, StepSpec
+from zenml.logger import get_logger
+from zenml.models import PipelineSnapshotResponse
+
+logger = get_logger(__name__)
 
 
 class StepRunInfo(FrozenBaseModel):
@@ -30,7 +34,9 @@ class StepRunInfo(FrozenBaseModel):
     pipeline_step_name: str
 
     config: StepConfiguration
+    spec: StepSpec
     pipeline: PipelineConfiguration
+    snapshot: PipelineSnapshotResponse
 
     force_write_logs: Callable[..., Any]
 
@@ -46,15 +52,32 @@ def get_image(self, key: str) -> str:
         Returns:
             The image name or digest.
         """
-        from zenml.client import Client
-
-        run = Client().get_pipeline_run(self.run_id)
-        if not run.build:
+        if not self.snapshot.build:
             raise RuntimeError(
-                f"Missing build for run {run.id}. This is probably because "
-                "the build was manually deleted."
+                f"Missing build for snapshot {self.snapshot.id}. This is "
+                "probably because the build was manually deleted."
             )
 
-        return run.build.get_image(
-            component_key=key, step=self.pipeline_step_name
-        )
+        if self.snapshot.is_dynamic:
+            # TODO: better way for this
+            for (
+                invocation_id,
+                compiled_step,
+            ) in self.snapshot.step_configurations.values():
+                if (
+                    compiled_step.spec.source.import_path
+                    == self.spec.source.import_path
+                ):
+                    step_key = invocation_id
+                    break
+            else:
+                logger.warning(
+                    "Unable to find config template for step %s. Falling "
+                    "back to the pipeline image.",
+                    self.pipeline_step_name,
+                )
+                step_key = None
+        else:
+            step_key = self.pipeline_step_name
+
+        return self.snapshot.build.get_image(component_key=key, step=step_key)

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

@@ -175,6 +175,10 @@ class PipelineRunUpdate(BaseUpdate):
         max_length=STR_FIELD_MAX_LENGTH,
     )
     end_time: Optional[datetime] = None
+    completed: Optional[bool] = Field(
+        default=None,
+        title="Whether the pipeline run is completed.",
+    )
     orchestrator_run_id: Optional[str] = None
     # TODO: we should maybe have a different update model here, the upper
     #  three attributes should only be for internal use

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/orchestrators/publish_utils.py

@@ -131,6 +131,7 @@ def publish_successful_pipeline_run(
         run_update=PipelineRunUpdate(
             status=ExecutionStatus.COMPLETED,
             end_time=utc_now(),
+            completed=True,
         ),
     )
 
@@ -151,6 +152,7 @@ def publish_failed_pipeline_run(
         run_update=PipelineRunUpdate(
             status=ExecutionStatus.FAILED,
             end_time=utc_now(),
+            completed=True,
         ),
     )
 

src/zenml/orchestrators/step_launcher.py

@@ -431,12 +431,14 @@ def _run_step(
         """
         step_run_info = StepRunInfo(
             config=self._step.config,
+            spec=self._step.spec,
             pipeline=self._snapshot.pipeline_configuration,
             run_name=pipeline_run.name,
             pipeline_step_name=self._invocation_id,
             run_id=pipeline_run.id,
             step_run_id=step_run.id,
             force_write_logs=force_write_logs,
+            snapshot=self._snapshot,
         )
 
         output_artifact_uris = output_utils.prepare_output_artifact_uris(

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]]: