zenml-nightly 0.82.1.dev20250527__py3-none-any.whl → 0.82.1.dev20250528__py3-none-any.whl

zenml/models/v2/core/step_run.py

@@ -30,7 +30,11 @@ from pydantic import ConfigDict, Field
 
 from zenml.config.step_configurations import StepConfiguration, StepSpec
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ExecutionStatus, StepRunInputArtifactType
+from zenml.enums import (
+    ArtifactSaveType,
+    ExecutionStatus,
+    StepRunInputArtifactType,
+)
 from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
@@ -126,7 +130,7 @@ class StepRunRequest(ProjectScopedRequest):
         title="The IDs of the parent steps of this step run.",
         default_factory=list,
     )
-    inputs: Dict[str, UUID] = Field(
+    inputs: Dict[str, List[UUID]] = Field(
         title="The IDs of the input artifact versions of the step run.",
         default_factory=dict,
     )
@@ -180,19 +184,15 @@ class StepRunResponseBody(ProjectScopedResponseBody):
         title="The end time of the step run.",
         default=None,
     )
-    inputs: Dict[str, StepRunInputResponse] = Field(
-        title="The input artifact versions of the step run.",
-        default_factory=dict,
-    )
-    outputs: Dict[str, List[ArtifactVersionResponse]] = Field(
-        title="The output artifact versions of the step run.",
-        default_factory=dict,
-    )
     model_version_id: Optional[UUID] = Field(
         title="The ID of the model version that was "
         "configured by this step run explicitly.",
         default=None,
     )
+    substitutions: Dict[str, str] = Field(
+        title="The substitutions of the step run.",
+        default={},
+    )
     model_config = ConfigDict(protected_namespaces=())
 
 
@@ -254,6 +254,14 @@ class StepRunResponseResources(ProjectScopedResponseResources):
     """Class for all resource models associated with the step run entity."""
 
     model_version: Optional[ModelVersionResponse] = None
+    inputs: Dict[str, List[StepRunInputResponse]] = Field(
+        title="The input artifact versions of the step run.",
+        default_factory=dict,
+    )
+    outputs: Dict[str, List[ArtifactVersionResponse]] = Field(
+        title="The output artifact versions of the step run.",
+        default_factory=dict,
+    )
 
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.
@@ -288,7 +296,7 @@ class StepRunResponse(
 
     # Helper properties
     @property
-    def input(self) -> ArtifactVersionResponse:
+    def input(self) -> StepRunInputResponse:
         """Returns the input artifact that was used to run this step.
 
         Returns:
@@ -299,12 +307,14 @@ class StepRunResponse(
         """
         if not self.inputs:
             raise ValueError(f"Step {self.name} has no inputs.")
-        if len(self.inputs) > 1:
+        if len(self.inputs) > 1 or (
+            len(self.inputs) == 1 and len(next(iter(self.inputs.values()))) > 1
+        ):
             raise ValueError(
                 f"Step {self.name} has multiple inputs, so `Step.input` is "
                 "ambiguous. Please use `Step.inputs` instead."
             )
-        return next(iter(self.inputs.values()))
+        return next(iter(self.inputs.values()))[0]
 
     @property
     def output(self) -> ArtifactVersionResponse:
@@ -328,6 +338,71 @@ class StepRunResponse(
             )
         return next(iter(self.outputs.values()))[0]
 
+    @property
+    def regular_inputs(self) -> Dict[str, StepRunInputResponse]:
+        """Returns the regular step inputs of the step run.
+
+        Regular step inputs are the inputs that are defined in the step function
+        signature, and are not manually loaded during the step execution.
+
+        Raises:
+            ValueError: If there were multiple regular input artifacts for the
+                same input name.
+
+        Returns:
+            The regular step inputs.
+        """
+        result = {}
+
+        for input_name, input_artifacts in self.inputs.items():
+            filtered = [
+                input_artifact
+                for input_artifact in input_artifacts
+                if input_artifact.input_type != StepRunInputArtifactType.MANUAL
+            ]
+            if len(filtered) > 1:
+                raise ValueError(
+                    f"Expected 1 regular input artifact for {input_name}, got "
+                    f"{len(filtered)}."
+                )
+            if filtered:
+                result[input_name] = filtered[0]
+
+        return result
+
+    @property
+    def regular_outputs(self) -> Dict[str, ArtifactVersionResponse]:
+        """Returns the regular step outputs of the step run.
+
+        Regular step outputs are the outputs that are defined in the step
+        function signature, and are not manually saved during the step
+        execution.
+
+        Raises:
+            ValueError: If there were multiple regular output artifacts for the
+                same output name.
+
+        Returns:
+            The regular step outputs.
+        """
+        result = {}
+
+        for output_name, output_artifacts in self.outputs.items():
+            filtered = [
+                output_artifact
+                for output_artifact in output_artifacts
+                if output_artifact.save_type == ArtifactSaveType.STEP_OUTPUT
+            ]
+            if len(filtered) > 1:
+                raise ValueError(
+                    f"Expected 1 regular output artifact for {output_name}, "
+                    f"got {len(filtered)}."
+                )
+            if filtered:
+                result[output_name] = filtered[0]
+
+        return result
+
     # Body and metadata properties
     @property
     def status(self) -> ExecutionStatus:
@@ -339,13 +414,13 @@ class StepRunResponse(
         return self.get_body().status
 
     @property
-    def inputs(self) -> Dict[str, StepRunInputResponse]:
+    def inputs(self) -> Dict[str, List[StepRunInputResponse]]:
         """The `inputs` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_body().inputs
+        return self.get_resources().inputs
 
     @property
     def outputs(self) -> Dict[str, List[ArtifactVersionResponse]]:
@@ -354,7 +429,7 @@ class StepRunResponse(
         Returns:
             the value of the property.
         """
-        return self.get_body().outputs
+        return self.get_resources().outputs
 
     @property
     def model_version_id(self) -> Optional[UUID]:
@@ -365,6 +440,15 @@ class StepRunResponse(
         """
         return self.get_body().model_version_id
 
+    @property
+    def substitutions(self) -> Dict[str, str]:
+        """The `substitutions` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().substitutions
+
     @property
     def config(self) -> "StepConfiguration":
         """The `config` property.

zenml/models/v2/core/tag.py

@@ -125,14 +125,15 @@ class TagResponseBody(UserScopedResponseBody):
     exclusive: bool = Field(
         description="The flag signifying whether the tag is an exclusive tag."
     )
-    tagged_count: int = Field(
-        description="The count of resources tagged with this tag."
-    )
 
 
 class TagResponseMetadata(UserScopedResponseMetadata):
     """Response metadata for tags."""
 
+    tagged_count: int = Field(
+        description="The count of resources tagged with this tag."
+    )
+
 
 class TagResponseResources(UserScopedResponseResources):
     """Class for all resource models associated with the tag entity."""
@@ -185,7 +186,7 @@ class TagResponse(
         Returns:
             the value of the property.
         """
-        return self.get_body().tagged_count
+        return self.get_metadata().tagged_count
 
 
 # ------------------ Filter Model ------------------

zenml/models/v2/misc/pipeline_run_dag.py

@@ -0,0 +1,46 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Pipeline run DAG models."""
+
+from typing import Any, Dict, List, Optional
+from uuid import UUID
+
+from pydantic import BaseModel
+
+from zenml.enums import ExecutionStatus
+
+
+class PipelineRunDAG(BaseModel):
+    """Pipeline run DAG."""
+
+    id: UUID
+    status: ExecutionStatus
+    nodes: List["Node"]
+    edges: List["Edge"]
+
+    class Node(BaseModel):
+        """Node in the pipeline run DAG."""
+
+        node_id: str
+        type: str
+        id: Optional[UUID] = None
+        name: str
+        metadata: Dict[str, Any] = {}
+
+    class Edge(BaseModel):
+        """Edge in the pipeline run DAG."""
+
+        source: str
+        target: str
+        metadata: Dict[str, Any] = {}

zenml/orchestrators/base_orchestrator.py

@@ -220,29 +220,18 @@ class BaseOrchestrator(StackComponent, ABC):
             and not deployment.schedule
             and not prevent_client_side_caching
         ):
-            from zenml.orchestrators import step_run_utils
+            from zenml.orchestrators import cache_utils
 
-            cached_invocations = step_run_utils.create_cached_step_runs(
-                deployment=deployment,
-                pipeline_run=placeholder_run,
-                stack=stack,
+            run_required = (
+                cache_utils.create_cached_step_runs_and_prune_deployment(
+                    deployment=deployment,
+                    pipeline_run=placeholder_run,
+                    stack=stack,
+                )
             )
 
-            for invocation_id in cached_invocations:
-                # Remove the cached step invocations from the deployment so
-                # the orchestrator does not try to run them
-                deployment.step_configurations.pop(invocation_id)
-
-            for step in deployment.step_configurations.values():
-                for invocation_id in cached_invocations:
-                    if invocation_id in step.spec.upstream_steps:
-                        step.spec.upstream_steps.remove(invocation_id)
-
-            if len(deployment.step_configurations) == 0:
-                # All steps were cached, we update the pipeline run status and
-                # don't actually use the orchestrator to run the pipeline
+            if not run_required:
                 self._cleanup_run()
-                logger.info("All steps of the pipeline run were cached.")
                 return
         else:
             logger.debug("Skipping client-side caching.")

zenml/orchestrators/cache_utils.py

@@ -19,13 +19,20 @@ from typing import TYPE_CHECKING, Dict, Optional
 from zenml.client import Client
 from zenml.enums import ExecutionStatus, SorterOps
 from zenml.logger import get_logger
+from zenml.orchestrators import step_run_utils
 
 if TYPE_CHECKING:
     from uuid import UUID
 
     from zenml.artifact_stores import BaseArtifactStore
     from zenml.config.step_configurations import Step
-    from zenml.models import StepRunResponse
+    from zenml.models import (
+        PipelineDeploymentResponse,
+        PipelineRunResponse,
+        StepRunResponse,
+    )
+    from zenml.stack import Stack
+
 
 logger = get_logger(__name__)
 
@@ -127,3 +134,43 @@ def get_cached_step_run(cache_key: str) -> Optional["StepRunResponse"]:
     if cache_candidates:
         return cache_candidates[0]
     return None
+
+
+def create_cached_step_runs_and_prune_deployment(
+    deployment: "PipelineDeploymentResponse",
+    pipeline_run: "PipelineRunResponse",
+    stack: "Stack",
+) -> bool:
+    """Create cached step runs and prune the cached steps from the deployment.
+
+    Args:
+        deployment: The deployment of the pipeline run.
+        pipeline_run: The pipeline run for which to create the step runs.
+        stack: The stack on which the pipeline run is happening.
+
+    Returns:
+        Whether an actual pipeline run is still required.
+    """
+    cached_invocations = step_run_utils.create_cached_step_runs(
+        deployment=deployment,
+        pipeline_run=pipeline_run,
+        stack=stack,
+    )
+
+    for invocation_id in cached_invocations:
+        # Remove the cached step invocations from the deployment so
+        # the orchestrator does not try to run them
+        deployment.step_configurations.pop(invocation_id)
+
+    for step in deployment.step_configurations.values():
+        for invocation_id in cached_invocations:
+            if invocation_id in step.spec.upstream_steps:
+                step.spec.upstream_steps.remove(invocation_id)
+
+    if len(deployment.step_configurations) == 0:
+        # All steps were cached, we update the pipeline run status and
+        # don't actually use the orchestrator to run the pipeline
+        logger.info("All steps of the pipeline run were cached.")
+        return False
+
+    return True

zenml/orchestrators/input_utils.py

@@ -13,29 +13,34 @@
 #  permissions and limitations under the License.
 """Utilities for inputs."""
 
-from typing import TYPE_CHECKING, Dict, List, Tuple
+import json
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 from uuid import UUID
 
 from zenml.client import Client
 from zenml.config.step_configurations import Step
-from zenml.enums import ArtifactSaveType, StepRunInputArtifactType
+from zenml.enums import StepRunInputArtifactType
 from zenml.exceptions import InputResolutionError
 from zenml.utils import pagination_utils, string_utils
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineRunResponse
+    from zenml.models import PipelineRunResponse, StepRunResponse
     from zenml.models.v2.core.step_run import StepRunInputResponse
 
 
 def resolve_step_inputs(
     step: "Step",
     pipeline_run: "PipelineRunResponse",
+    step_runs: Optional[Dict[str, "StepRunResponse"]] = None,
 ) -> Tuple[Dict[str, "StepRunInputResponse"], List[UUID]]:
     """Resolves inputs for the current step.
 
     Args:
         step: The step for which to resolve the inputs.
         pipeline_run: The current pipeline run.
+        step_runs: A dictionary of already fetched step runs to use for input
+            resolution. This will be updated in-place with newly fetched step
+            runs.
 
     Raises:
         InputResolutionError: If input resolving failed due to a missing
@@ -50,65 +55,56 @@ def resolve_step_inputs(
     from zenml.models import ArtifactVersionResponse
     from zenml.models.v2.core.step_run import StepRunInputResponse
 
-    current_run_steps = {
-        run_step.name: run_step
-        for run_step in pagination_utils.depaginate(
-            Client().list_run_steps,
-            pipeline_run_id=pipeline_run.id,
-            project=pipeline_run.project.id,
+    step_runs = step_runs or {}
+
+    steps_to_fetch = set(step.spec.upstream_steps)
+    steps_to_fetch.update(
+        input_.step_name for input_ in step.spec.inputs.values()
+    )
+    # Remove all the step runs that we've already fetched.
+    steps_to_fetch.difference_update(step_runs.keys())
+
+    if steps_to_fetch:
+        step_runs.update(
+            {
+                run_step.name: run_step
+                for run_step in pagination_utils.depaginate(
+                    Client().list_run_steps,
+                    pipeline_run_id=pipeline_run.id,
+                    project=pipeline_run.project_id,
+                    name="oneof:" + json.dumps(list(steps_to_fetch)),
+                )
+            }
         )
-    }
 
     input_artifacts: Dict[str, StepRunInputResponse] = {}
     for name, input_ in step.spec.inputs.items():
         try:
-            step_run = current_run_steps[input_.step_name]
+            step_run = step_runs[input_.step_name]
         except KeyError:
             raise InputResolutionError(
                 f"No step `{input_.step_name}` found in current run."
             )
 
-        # Try to get the substitutions from the pipeline run first, as we
-        # already have a hydrated version of that. In the unlikely case
-        # that the pipeline run is outdated, we fetch it from the step
-        # run instead which will costs us one hydration call.
-        substitutions = (
-            pipeline_run.step_substitutions.get(step_run.name)
-            or step_run.config.substitutions
-        )
         output_name = string_utils.format_name_template(
-            input_.output_name, substitutions=substitutions
+            input_.output_name, substitutions=step_run.substitutions
         )
 
         try:
-            outputs = step_run.outputs[output_name]
+            output = step_run.regular_outputs[output_name]
         except KeyError:
             raise InputResolutionError(
                 f"No step output `{output_name}` found for step "
                 f"`{input_.step_name}`."
             )
-
-        step_outputs = [
-            output
-            for output in outputs
-            if output.save_type == ArtifactSaveType.STEP_OUTPUT
-        ]
-        if len(step_outputs) > 2:
-            # This should never happen, there can only be a single regular step
-            # output for a name
+        except ValueError:
             raise InputResolutionError(
-                f"Too many step outputs for output `{output_name}` of "
-                f"step `{input_.step_name}`."
-            )
-        elif len(step_outputs) == 0:
-            raise InputResolutionError(
-                f"No step output `{output_name}` found for step "
-                f"`{input_.step_name}`."
+                f"Expected 1 regular output artifact for {output_name}."
             )
 
         input_artifacts[name] = StepRunInputResponse(
             input_type=StepRunInputArtifactType.STEP_OUTPUT,
-            **step_outputs[0].model_dump(),
+            **output.model_dump(),
         )
 
     for (
@@ -187,7 +183,7 @@ def resolve_step_inputs(
             step.config.parameters[name] = value_
 
     parent_step_ids = [
-        current_run_steps[upstream_step].id
+        step_runs[upstream_step].id
         for upstream_step in step.spec.upstream_steps
     ]
 

zenml/orchestrators/step_launcher.py

@@ -381,7 +381,7 @@ class StepLauncher:
                     pipeline_run=pipeline_run,
                     step_run=step_run,
                     step_run_info=step_run_info,
-                    input_artifacts=step_run.inputs,
+                    input_artifacts=step_run.regular_inputs,
                     output_artifact_uris=output_artifact_uris,
                     last_retry=last_retry,
                 )

zenml/orchestrators/step_run_utils.py

@@ -28,6 +28,7 @@ from zenml.models import (
     PipelineDeploymentResponse,
     PipelineRunResponse,
     StepRunRequest,
+    StepRunResponse,
 )
 from zenml.orchestrators import cache_utils, input_utils, utils
 from zenml.stack import Stack
@@ -79,24 +80,34 @@ class StepRunRequestFactory:
             project=Client().active_project.id,
         )
 
-    def populate_request(self, request: StepRunRequest) -> None:
+    def populate_request(
+        self,
+        request: StepRunRequest,
+        step_runs: Optional[Dict[str, "StepRunResponse"]] = None,
+    ) -> None:
         """Populate a step run request with additional information.
 
         Args:
             request: The request to populate.
+            step_runs: A dictionary of already fetched step runs to use for
+                input resolution. This will be updated in-place with newly
+                fetched step runs.
         """
         step = self.deployment.step_configurations[request.name]
 
         input_artifacts, parent_step_ids = input_utils.resolve_step_inputs(
             step=step,
             pipeline_run=self.pipeline_run,
+            step_runs=step_runs,
         )
         input_artifact_ids = {
             input_name: artifact.id
             for input_name, artifact in input_artifacts.items()
         }
 
-        request.inputs = input_artifact_ids
+        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(
@@ -128,8 +139,8 @@ class StepRunRequestFactory:
                 cache_key=cache_key
             ):
                 request.inputs = {
-                    input_name: artifact.id
-                    for input_name, artifact in cached_step_run.inputs.items()
+                    input_name: [artifact.id for artifact in artifacts]
+                    for input_name, artifacts in cached_step_run.inputs.items()
                 }
 
                 request.original_step_run_id = cached_step_run.id
@@ -283,6 +294,9 @@ def create_cached_step_runs(
     request_factory = StepRunRequestFactory(
         deployment=deployment, pipeline_run=pipeline_run, stack=stack
     )
+    # This is used to cache the step runs that we created to avoid unnecessary
+    # server requests.
+    step_runs: Dict[str, "StepRunResponse"] = {}
 
     while (
         cache_candidates := find_cacheable_invocation_candidates(
@@ -296,16 +310,13 @@ def create_cached_step_runs(
         for invocation_id in cache_candidates:
             visited_invocations.add(invocation_id)
 
-            # Make sure the request factory has the most up to date pipeline
-            # run to avoid hydration calls
-            request_factory.pipeline_run = Client().get_pipeline_run(
-                pipeline_run.id
-            )
             try:
                 step_run_request = request_factory.create_request(
                     invocation_id
                 )
-                request_factory.populate_request(step_run_request)
+                request_factory.populate_request(
+                    step_run_request, step_runs=step_runs
+                )
             except Exception as e:
                 # We failed to create/populate the step run. This might be due
                 # to some input resolution error, or an error importing the step
@@ -325,6 +336,11 @@ def create_cached_step_runs(
 
             step_run = Client().zen_store.create_run_step(step_run_request)
 
+            # Include the newly created step run in the step runs dictionary to
+            # avoid fetching it again later when downstream steps need it for
+            # input resolution.
+            step_runs[invocation_id] = step_run
+
             if (
                 model_version := step_run.model_version
                 or pipeline_run.model_version

zenml/pipelines/pipeline_definition.py

@@ -1065,16 +1065,16 @@ To avoid this consider setting pipeline parameters only in one place (config or
                     if not source.is_internal:
                         custom_materializer = True
 
-        stack_creator = Client().get_stack(stack.id).user
+        stack_creator = Client().get_stack(stack.id).user_id
         active_user = Client().active_user
-        own_stack = stack_creator and stack_creator.id == active_user.id
+        own_stack = stack_creator and stack_creator == active_user.id
 
         stack_metadata = {
             component_type.value: component.flavor
             for component_type, component in stack.components.items()
         }
         return {
-            "project_id": deployment.project.id,
+            "project_id": deployment.project_id,
             "store_type": Client().zen_store.type.value,
             **stack_metadata,
             "total_steps": len(self.invocations),

zenml/pipelines/run_utils.py

@@ -64,10 +64,9 @@ def create_placeholder_run(
     Returns:
         The placeholder run or `None` if no run was created.
     """
-    assert deployment.user
-
     if deployment.schedule:
         return None
+
     start_time = utc_now()
     run_request = PipelineRunRequest(
         name=string_utils.format_name_template(
@@ -84,7 +83,7 @@ def create_placeholder_run(
         # running.
         start_time=start_time,
         orchestrator_run_id=None,
-        project=deployment.project.id,
+        project=deployment.project_id,
         deployment=deployment.id,
         pipeline=deployment.pipeline.id if deployment.pipeline else None,
         status=ExecutionStatus.INITIALIZING,