zenml-nightly 0.83.0.dev20250623__py3-none-any.whl → 0.83.1.dev20250625__py3-none-any.whl

zenml/orchestrators/base_orchestrator.py

@@ -14,7 +14,16 @@
 """Base orchestrator class."""
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    Optional,
+    Type,
+    cast,
+)
 from uuid import UUID
 
 from pydantic import model_validator
@@ -24,9 +33,13 @@ from zenml.constants import (
     handle_bool_env_var,
 )
 from zenml.enums import ExecutionStatus, StackComponentType
+from zenml.exceptions import RunMonitoringError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
-from zenml.orchestrators.publish_utils import publish_pipeline_run_metadata
+from zenml.orchestrators.publish_utils import (
+    publish_pipeline_run_metadata,
+    publish_schedule_metadata,
+)
 from zenml.orchestrators.step_launcher import StepLauncher
 from zenml.orchestrators.utils import get_config_environment_vars
 from zenml.stack import Flavor, Stack, StackComponent, StackComponentConfig
@@ -39,6 +52,26 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
+class SubmissionResult:
+    """Result of submitting a pipeline run."""
+
+    def __init__(
+        self,
+        wait_for_completion: Optional[Callable[[], None]] = None,
+        metadata: Optional[Dict[str, MetadataType]] = None,
+    ):
+        """Initialize a submission result.
+
+        Args:
+            wait_for_completion: A function that waits for the pipeline run to
+                complete. If provided, this will be called after the pipeline
+                run was submitted successfully.
+            metadata: Metadata for the pipeline run or schedule.
+        """
+        self.wait_for_completion = wait_for_completion
+        self.metadata = metadata
+
+
 class BaseOrchestratorConfig(StackComponentConfig):
     """Base orchestrator config."""
 
@@ -95,22 +128,7 @@ class BaseOrchestratorConfig(StackComponentConfig):
 
 
 class BaseOrchestrator(StackComponent, ABC):
-    """Base class for all orchestrators.
-
-    In order to implement an orchestrator you will need to subclass from this
-    class.
-
-    How it works:
-    -------------
-    The `run(...)` method is the entrypoint that is executed when the
-    pipeline's run method is called within the user code
-    (`pipeline_instance.run(...)`).
-
-    This method will do some internal preparation and then call the
-    `prepare_or_run_pipeline(...)` method. BaseOrchestrator subclasses must
-    implement this method and either run the pipeline steps directly or deploy
-    the pipeline to some remote infrastructure.
-    """
+    """Base class for all orchestrators."""
 
     _active_deployment: Optional["PipelineDeploymentResponse"] = None
 
@@ -134,7 +152,32 @@ class BaseOrchestrator(StackComponent, ABC):
             The orchestrator run id.
         """
 
-    @abstractmethod
+    def submit_pipeline(
+        self,
+        deployment: "PipelineDeploymentResponse",
+        stack: "Stack",
+        environment: Dict[str, str],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
+
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
+
+        Args:
+            deployment: The pipeline deployment to submit.
+            stack: The stack the pipeline will run on.
+            environment: Environment variables to set in the orchestration
+                environment. These don't need to be set if running locally.
+            placeholder_run: An optional placeholder run for the deployment.
+
+        Returns:
+            Optional submission result.
+        """
+        return None
+
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -142,34 +185,7 @@ class BaseOrchestrator(StackComponent, ABC):
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Optional[Iterator[Dict[str, MetadataType]]]:
-        """The method needs to be implemented by the respective orchestrator.
-
-        Depending on the type of orchestrator you'll have to perform slightly
-        different operations.
-
-        Simple Case:
-        ------------
-        The Steps are run directly from within the same environment in which
-        the orchestrator code is executed. In this case you will need to
-        deal with implementation-specific runtime configurations (like the
-        schedule) and then iterate through the steps and finally call
-        `self.run_step(...)` to execute each step.
-
-        Advanced Case:
-        --------------
-        Most orchestrators will not run the steps directly. Instead, they
-        build some intermediate representation of the pipeline that is then
-        used to create and run the pipeline and its steps on the target
-        environment. For such orchestrators this method will have to build
-        this representation and deploy it.
-
-        Regardless of the implementation details, the orchestrator will need
-        to run each step in the target environment. For this the
-        `self.run_step(...)` method should be used.
-
-        The easiest way to make this work is by using an entrypoint
-        configuration to run single steps (`zenml.entrypoints.step_entrypoint_configuration.StepEntrypointConfiguration`)
-        or entire pipelines (`zenml.entrypoints.pipeline_entrypoint_configuration.PipelineEntrypointConfiguration`).
+        """DEPRECATED: Prepare or run a pipeline.
 
         Args:
             deployment: The pipeline deployment to prepare or run.
@@ -177,9 +193,6 @@ class BaseOrchestrator(StackComponent, ABC):
             environment: Environment variables to set in the orchestration
                 environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
-
-        Yields:
-            Metadata for the pipeline run.
         """
 
     def run(
@@ -195,6 +208,10 @@ class BaseOrchestrator(StackComponent, ABC):
             stack: The stack on which to run the pipeline.
             placeholder_run: An optional placeholder run for the deployment.
                 This will be deleted in case the pipeline deployment failed.
+
+        Raises:
+            RunMonitoringError: If a failure happened while monitoring the
+                pipeline run.
         """
         self._prepare_run(deployment=deployment)
 
@@ -237,24 +254,78 @@ class BaseOrchestrator(StackComponent, ABC):
             logger.debug("Skipping client-side caching.")
 
         try:
-            if metadata_iterator := self.prepare_or_run_pipeline(
-                deployment=deployment,
-                stack=stack,
-                environment=environment,
-                placeholder_run=placeholder_run,
+            if (
+                getattr(self.submit_pipeline, "__func__", None)
+                is BaseOrchestrator.submit_pipeline
             ):
-                for metadata_dict in metadata_iterator:
-                    try:
-                        if placeholder_run:
-                            publish_pipeline_run_metadata(
-                                pipeline_run_id=placeholder_run.id,
-                                pipeline_run_metadata={self.id: metadata_dict},
+                logger.warning(
+                    "The orchestrator '%s' is still using the deprecated "
+                    "`prepare_or_run_pipeline(...)` method which will be "
+                    "removed in the future. Please implement the replacement "
+                    "`submit_pipeline(...)` method for your custom "
+                    "orchestrator.",
+                    self.name,
+                )
+                if metadata_iterator := self.prepare_or_run_pipeline(
+                    deployment=deployment,
+                    stack=stack,
+                    environment=environment,
+                    placeholder_run=placeholder_run,
+                ):
+                    for metadata_dict in metadata_iterator:
+                        try:
+                            if placeholder_run:
+                                publish_pipeline_run_metadata(
+                                    pipeline_run_id=placeholder_run.id,
+                                    pipeline_run_metadata={
+                                        self.id: metadata_dict
+                                    },
+                                )
+                        except Exception as e:
+                            logger.debug(
+                                "Something went went wrong trying to publish the"
+                                f"run metadata: {e}"
                             )
-                    except Exception as e:
-                        logger.debug(
-                            "Something went went wrong trying to publish the"
-                            f"run metadata: {e}"
-                        )
+            else:
+                submission_result = self.submit_pipeline(
+                    deployment=deployment,
+                    stack=stack,
+                    environment=environment,
+                    placeholder_run=placeholder_run,
+                )
+
+                if submission_result:
+                    if submission_result.metadata:
+                        if placeholder_run:
+                            try:
+                                publish_pipeline_run_metadata(
+                                    pipeline_run_id=placeholder_run.id,
+                                    pipeline_run_metadata={
+                                        self.id: submission_result.metadata
+                                    },
+                                )
+                            except Exception as e:
+                                logger.debug(
+                                    "Error publishing run metadata: %s", e
+                                )
+                        elif deployment.schedule:
+                            try:
+                                publish_schedule_metadata(
+                                    schedule_id=deployment.schedule.id,
+                                    schedule_metadata={
+                                        self.id: submission_result.metadata
+                                    },
+                                )
+                            except Exception as e:
+                                logger.debug(
+                                    "Error publishing schedule metadata: %s", e
+                                )
+
+                    if submission_result.wait_for_completion:
+                        try:
+                            submission_result.wait_for_completion()
+                        except BaseException as e:
+                            raise RunMonitoringError(original_exception=e)
         finally:
             self._cleanup_run()
 

zenml/orchestrators/input_utils.py

@@ -14,8 +14,7 @@
 """Utilities for inputs."""
 
 import json
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
-from uuid import UUID
+from typing import TYPE_CHECKING, Dict, Optional
 
 from zenml.client import Client
 from zenml.config.step_configurations import Step
@@ -32,7 +31,7 @@ def resolve_step_inputs(
     step: "Step",
     pipeline_run: "PipelineRunResponse",
     step_runs: Optional[Dict[str, "StepRunResponse"]] = None,
-) -> Tuple[Dict[str, "StepRunInputResponse"], List[UUID]]:
+) -> Dict[str, "StepRunInputResponse"]:
     """Resolves inputs for the current step.
 
     Args:
@@ -49,16 +48,14 @@ def resolve_step_inputs(
             resolved in runtime due to missing object.
 
     Returns:
-        The IDs of the input artifact versions and the IDs of parent steps of
-            the current step.
+        The input artifact versions.
     """
     from zenml.models import ArtifactVersionResponse
     from zenml.models.v2.core.step_run import StepRunInputResponse
 
     step_runs = step_runs or {}
 
-    steps_to_fetch = set(step.spec.upstream_steps)
-    steps_to_fetch.update(
+    steps_to_fetch = set(
         input_.step_name for input_ in step.spec.inputs.values()
     )
     # Remove all the step runs that we've already fetched.
@@ -205,9 +202,4 @@ def resolve_step_inputs(
         else:
             step.config.parameters[name] = value_
 
-    parent_step_ids = [
-        step_runs[upstream_step].id
-        for upstream_step in step.spec.upstream_steps
-    ]
-
-    return input_artifacts, parent_step_ids
+    return input_artifacts

zenml/orchestrators/local/local_orchestrator.py

@@ -14,14 +14,15 @@
 """Implementation of the ZenML local orchestrator."""
 
 import time
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type
+from typing import TYPE_CHECKING, Dict, Optional, Type
 from uuid import uuid4
 
 from zenml.logger import get_logger
-from zenml.orchestrators import BaseOrchestrator
-from zenml.orchestrators.base_orchestrator import (
+from zenml.orchestrators import (
+    BaseOrchestrator,
     BaseOrchestratorConfig,
     BaseOrchestratorFlavor,
+    SubmissionResult,
 )
 from zenml.stack import Stack
 from zenml.utils import string_utils
@@ -41,21 +42,29 @@ class LocalOrchestrator(BaseOrchestrator):
 
     _orchestrator_run_id: Optional[str] = None
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
-        """Iterates through all steps and executes them sequentially.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
+
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Args:
-            deployment: The pipeline deployment to prepare or run.
-            stack: The stack on which the pipeline is deployed.
+            deployment: The pipeline deployment to submit.
+            stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
+
+        Returns:
+            Optional submission result.
         """
         if deployment.schedule:
             logger.warning(
@@ -87,6 +96,7 @@ class LocalOrchestrator(BaseOrchestrator):
             string_utils.get_human_readable_time(run_duration),
         )
         self._orchestrator_run_id = None
+        return None
 
     def get_orchestrator_run_id(self) -> str:
         """Returns the active orchestrator run id.

zenml/orchestrators/local_docker/local_docker_orchestrator.py

@@ -34,6 +34,7 @@ from zenml.orchestrators import (
     BaseOrchestratorConfig,
     BaseOrchestratorFlavor,
     ContainerizedOrchestrator,
+    SubmissionResult,
 )
 from zenml.stack import Stack, StackValidator
 from zenml.utils import docker_utils, string_utils
@@ -91,24 +92,32 @@ class LocalDockerOrchestrator(ContainerizedOrchestrator):
                 f"{ENV_ZENML_DOCKER_ORCHESTRATOR_RUN_ID}."
             )
 
-    def prepare_or_run_pipeline(
+    def submit_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
-        """Sequentially runs all pipeline steps in local Docker containers.
+    ) -> Optional[SubmissionResult]:
+        """Submits a pipeline to the orchestrator.
+
+        This method should only submit the pipeline and not wait for it to
+        complete. If the orchestrator is configured to wait for the pipeline run
+        to complete, a function that waits for the pipeline run to complete can
+        be passed as part of the submission result.
 
         Args:
-            deployment: The pipeline deployment to prepare or run.
+            deployment: The pipeline deployment to submit.
             stack: The stack the pipeline will run on.
             environment: Environment variables to set in the orchestration
-                environment.
+                environment. These don't need to be set if running locally.
             placeholder_run: An optional placeholder run for the deployment.
 
         Raises:
             RuntimeError: If a step fails.
+
+        Returns:
+            Optional submission result.
         """
         if deployment.schedule:
             logger.warning(
@@ -194,6 +203,7 @@ class LocalDockerOrchestrator(ContainerizedOrchestrator):
             "Pipeline run has finished in `%s`.",
             string_utils.get_human_readable_time(run_duration),
         )
+        return None
 
 
 class LocalDockerOrchestratorSettings(BaseSettings):

zenml/orchestrators/publish_utils.py

@@ -161,3 +161,27 @@ def publish_step_run_metadata(
             ],
             stack_component_id=stack_component_id,
         )
+
+
+def publish_schedule_metadata(
+    schedule_id: "UUID",
+    schedule_metadata: Dict["UUID", Dict[str, "MetadataType"]],
+) -> None:
+    """Publishes the given schedule metadata.
+
+    Args:
+        schedule_id: The ID of the schedule.
+        schedule_metadata: A dictionary mapping stack component IDs to the
+            metadata they created.
+    """
+    client = Client()
+    for stack_component_id, metadata in schedule_metadata.items():
+        client.create_run_metadata(
+            metadata=metadata,
+            resources=[
+                RunMetadataResource(
+                    id=schedule_id, type=MetadataResourceTypes.SCHEDULE
+                )
+            ],
+            stack_component_id=stack_component_id,
+        )

zenml/orchestrators/step_run_utils.py

@@ -95,7 +95,7 @@ class StepRunRequestFactory:
         """
         step = self.deployment.step_configurations[request.name]
 
-        input_artifacts, parent_step_ids = input_utils.resolve_step_inputs(
+        input_artifacts = input_utils.resolve_step_inputs(
             step=step,
             pipeline_run=self.pipeline_run,
             step_runs=step_runs,
@@ -108,7 +108,6 @@ class StepRunRequestFactory:
         request.inputs = {
             name: [artifact.id] for name, artifact in input_artifacts.items()
         }
-        request.parent_step_ids = parent_step_ids
 
         cache_key = cache_utils.generate_cache_key(
             step=step,

zenml/pipelines/run_utils.py

@@ -12,6 +12,7 @@ from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.source import Source, SourceType
 from zenml.config.step_configurations import StepConfigurationUpdate
 from zenml.enums import ExecutionStatus
+from zenml.exceptions import RunMonitoringError
 from zenml.logger import get_logger
 from zenml.models import (
     FlavorFilter,
@@ -136,8 +137,9 @@ def deploy_pipeline(
         stack: The stack on which to run the deployment.
         placeholder_run: An optional placeholder run for the deployment.
 
+    # noqa: DAR401
     Raises:
-        Exception: Any exception that happened while deploying or running
+        BaseException: Any exception that happened while deploying or running
             (in case it happens synchronously) the pipeline.
     """
     # Prevent execution of nested pipelines which might lead to
@@ -150,16 +152,19 @@ def deploy_pipeline(
             deployment=deployment,
             placeholder_run=placeholder_run,
         )
-    except Exception as e:
+    except RunMonitoringError as e:
+        # Don't mark the run as failed if the error happened during monitoring
+        # of the run.
+        raise e.original_exception from None
+    except BaseException as e:
         if (
             placeholder_run
-            and Client()
+            and not Client()
             .get_pipeline_run(placeholder_run.id, hydrate=False)
-            .status
-            == ExecutionStatus.INITIALIZING
+            .status.is_finished
         ):
-            # The run failed during the initialization phase -> We change it's
-            # status to `Failed`
+            # We failed during/before the submission of the run, so we mark the
+            # run as failed if it is still in an initializing/running state.
             publish_failed_pipeline_run(placeholder_run.id)
 
         raise e

zenml/step_operators/step_operator_entrypoint_configuration.py

@@ -90,7 +90,7 @@ class StepOperatorEntrypointConfiguration(StepEntrypointConfiguration):
         )
 
         stack = Client().active_stack
-        input_artifacts, _ = input_utils.resolve_step_inputs(
+        input_artifacts = input_utils.resolve_step_inputs(
             step=step, pipeline_run=pipeline_run
         )
         output_artifact_uris = output_utils.prepare_output_artifact_uris(

zenml/zen_server/dashboard/assets/{404-DmJUgorp.js → 404-B5eko6XL.js}

@@ -1 +1 @@
-import{j as e}from"./@radix-Cdvw4jJ8.js";import{B as s,r as t}from"./index-BFqbGSck.js";import{E as r}from"./EmptyState-CjrgDtVk.js";import{S as a}from"./help-Co6aedki.js";import{L as o}from"./@react-router-DeDfXbUF.js";import"./@tanstack-5gTMR7G2.js";import"./@reactflow-8OCk19Fi.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@radix-Cdvw4jJ8.js";import{B as s,r as t}from"./index-BgEfQ3_G.js";import{E as r}from"./EmptyState-CjrgDtVk.js";import{S as a}from"./help-Co6aedki.js";import{L as o}from"./@react-router-DeDfXbUF.js";import"./@tanstack-5gTMR7G2.js";import"./@reactflow-B_iCtR7X.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};