zenml-nightly 0.70.0.dev20241125__py3-none-any.whl → 0.71.0.dev20241220__py3-none-any.whl

zenml/models/v2/core/step_run.py

@@ -14,7 +14,16 @@
 """Models representing steps runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, ClassVar, Dict, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -41,6 +50,9 @@ if TYPE_CHECKING:
         LogsRequest,
         LogsResponse,
     )
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 
 class StepRunInputResponse(ArtifactVersionResponse):
@@ -142,7 +154,7 @@ class StepRunRequest(WorkspaceScopedRequest):
 class StepRunUpdate(BaseModel):
     """Update model for step runs."""
 
-    outputs: Dict[str, UUID] = Field(
+    outputs: Dict[str, List[UUID]] = Field(
         title="The IDs of the output artifact versions of the step run.",
         default={},
     )
@@ -553,16 +565,6 @@ class StepRunFilter(WorkspaceScopedFilter):
         description="Original id for this step run",
         union_mode="left_to_right",
     )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="User that produced this step run",
-        union_mode="left_to_right",
-    )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None,
-        description="Workspace of this step run",
-        union_mode="left_to_right",
-    )
     model_version_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Model version associated with the step run.",
@@ -576,24 +578,27 @@ class StepRunFilter(WorkspaceScopedFilter):
         default=None,
         description="The run_metadata to filter the step runs by.",
     )
-
     model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
-        self,
+        self, table: Type["AnySchema"]
     ) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
         from sqlmodel import and_
 
         from zenml.zen_stores.schemas import (
             ModelSchema,
             ModelVersionSchema,
+            RunMetadataResourceSchema,
             RunMetadataSchema,
             StepRunSchema,
         )
@@ -612,10 +617,11 @@ class StepRunFilter(WorkspaceScopedFilter):
 
             for key, value in self.run_metadata.items():
                 additional_filter = and_(
-                    RunMetadataSchema.resource_id == StepRunSchema.id,
-                    RunMetadataSchema.resource_type
+                    RunMetadataResourceSchema.resource_id == StepRunSchema.id,
+                    RunMetadataResourceSchema.resource_type
                     == MetadataResourceTypes.STEP_RUN,
-                    RunMetadataSchema.key == key,
+                    RunMetadataResourceSchema.run_metadata_id
+                    == RunMetadataSchema.id,
                     self.generate_custom_query_conditions_for_column(
                         value=value,
                         table=RunMetadataSchema,

zenml/models/v2/core/trigger.py

@@ -13,7 +13,17 @@
 #  permissions and limitations under the License.
 """Collection of all models concerning triggers."""
 
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
 from pydantic import Field, model_validator
@@ -39,6 +49,9 @@ if TYPE_CHECKING:
         ActionResponse,
     )
     from zenml.models.v2.core.event_source import EventSourceResponse
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
 
 
 # ------------------ Request Model ------------------
@@ -358,10 +371,13 @@ class TriggerFilter(WorkspaceScopedFilter):
     )
 
     def get_custom_filters(
-        self,
+        self, table: Type["AnySchema"]
     ) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
+        Args:
+            table: The query table.
+
         Returns:
             A list of custom filters.
         """
@@ -373,7 +389,7 @@ class TriggerFilter(WorkspaceScopedFilter):
             TriggerSchema,
         )
 
-        custom_filters = super().get_custom_filters()
+        custom_filters = super().get_custom_filters(table)
 
         if self.event_source_flavor:
             event_source_flavor_filter = and_(

zenml/models/v2/misc/run_metadata.py

@@ -0,0 +1,38 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Utility classes for modeling run metadata."""
+
+from datetime import datetime
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+
+from zenml.enums import MetadataResourceTypes
+from zenml.metadata.metadata_types import MetadataType
+
+
+class RunMetadataResource(BaseModel):
+    """Utility class to help identify resources to tag metadata to."""
+
+    id: UUID = Field(title="The ID of the resource.")
+    type: MetadataResourceTypes = Field(title="The type of the resource.")
+
+
+class RunMetadataEntry(BaseModel):
+    """Utility class to sort/list run metadata entries."""
+
+    value: MetadataType = Field(title="The value for the run metadata entry")
+    created: datetime = Field(
+        title="The timestamp when this resource was created."
+    )

zenml/orchestrators/base_orchestrator.py

@@ -61,7 +61,7 @@ class BaseOrchestratorConfig(StackComponentConfig):
                     "The 'custom_docker_base_image_name' field has been "
                     "deprecated. To use a custom base container image with your "
                     "orchestrators, please use the DockerSettings in your "
-                    "pipeline (see https://docs.zenml.io/how-to/customize-docker-builds)."
+                    "pipeline (see https://docs.zenml.io/how-to/infrastructure-deployment/customize-docker-builds)."
                 )
 
         return data
@@ -84,6 +84,15 @@ class BaseOrchestratorConfig(StackComponentConfig):
         """
         return False
 
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        return True
+
 
 class BaseOrchestrator(StackComponent, ABC):
     """Base class for all orchestrators.
@@ -205,6 +214,7 @@ class BaseOrchestrator(StackComponent, ABC):
 
         if (
             placeholder_run
+            and self.config.supports_client_side_caching
             and not deployment.schedule
             and not prevent_client_side_caching
         ):
@@ -232,6 +242,8 @@ class BaseOrchestrator(StackComponent, ABC):
                 self._cleanup_run()
                 logger.info("All steps of the pipeline run were cached.")
                 return
+        else:
+            logger.debug("Skipping client-side caching.")
 
         try:
             if metadata_iterator := self.prepare_or_run_pipeline(

zenml/orchestrators/input_utils.py

@@ -20,7 +20,7 @@ from zenml.client import Client
 from zenml.config.step_configurations import Step
 from zenml.enums import ArtifactSaveType, StepRunInputArtifactType
 from zenml.exceptions import InputResolutionError
-from zenml.utils import pagination_utils
+from zenml.utils import pagination_utils, string_utils
 
 if TYPE_CHECKING:
     from zenml.models import PipelineRunResponse
@@ -53,7 +53,8 @@ def resolve_step_inputs(
     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
+            Client().list_run_steps,
+            pipeline_run_id=pipeline_run.id,
         )
     }
 
@@ -66,11 +67,23 @@ def resolve_step_inputs(
                 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
+        )
+
         try:
-            outputs = step_run.outputs[input_.output_name]
+            outputs = step_run.outputs[output_name]
         except KeyError:
             raise InputResolutionError(
-                f"No step output `{input_.output_name}` found for step "
+                f"No step output `{output_name}` found for step "
                 f"`{input_.step_name}`."
             )
 
@@ -83,12 +96,12 @@ def resolve_step_inputs(
             # This should never happen, there can only be a single regular step
             # output for a name
             raise InputResolutionError(
-                f"Too many step outputs for output `{input_.output_name}` of "
+                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 `{input_.output_name}` found for step "
+                f"No step output `{output_name}` found for step "
                 f"`{input_.step_name}`."
             )
 

zenml/orchestrators/output_utils.py

@@ -19,6 +19,7 @@ from uuid import uuid4
 
 from zenml.client import Client
 from zenml.logger import get_logger
+from zenml.utils import string_utils
 
 if TYPE_CHECKING:
     from zenml.artifact_stores import BaseArtifactStore
@@ -75,10 +76,13 @@ def prepare_output_artifact_uris(
     artifact_store = stack.artifact_store
     output_artifact_uris: Dict[str, str] = {}
     for output_name in step.config.outputs.keys():
+        substituted_output_name = string_utils.format_name_template(
+            output_name, substitutions=step_run.config.substitutions
+        )
         artifact_uri = generate_artifact_uri(
             artifact_store=stack.artifact_store,
             step_run=step_run,
-            output_name=output_name,
+            output_name=substituted_output_name,
         )
         if artifact_store.exists(artifact_uri):
             raise RuntimeError("Artifact already exists")

zenml/orchestrators/publish_utils.py

@@ -21,6 +21,7 @@ from zenml.enums import ExecutionStatus, MetadataResourceTypes
 from zenml.models import (
     PipelineRunResponse,
     PipelineRunUpdate,
+    RunMetadataResource,
     StepRunResponse,
     StepRunUpdate,
 )
@@ -32,7 +33,7 @@ if TYPE_CHECKING:
 
 
 def publish_successful_step_run(
-    step_run_id: "UUID", output_artifact_ids: Dict[str, "UUID"]
+    step_run_id: "UUID", output_artifact_ids: Dict[str, List["UUID"]]
 ) -> "StepRunResponse":
     """Publishes a successful step run.
 
@@ -129,8 +130,11 @@ def publish_pipeline_run_metadata(
     for stack_component_id, metadata in pipeline_run_metadata.items():
         client.create_run_metadata(
             metadata=metadata,
-            resource_id=pipeline_run_id,
-            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+            resources=[
+                RunMetadataResource(
+                    id=pipeline_run_id, type=MetadataResourceTypes.PIPELINE_RUN
+                )
+            ],
             stack_component_id=stack_component_id,
         )
 
@@ -150,7 +154,10 @@ def publish_step_run_metadata(
     for stack_component_id, metadata in step_run_metadata.items():
         client.create_run_metadata(
             metadata=metadata,
-            resource_id=step_run_id,
-            resource_type=MetadataResourceTypes.STEP_RUN,
+            resources=[
+                RunMetadataResource(
+                    id=step_run_id, type=MetadataResourceTypes.STEP_RUN
+                )
+            ],
             stack_component_id=stack_component_id,
         )

zenml/orchestrators/step_launcher.py

@@ -179,12 +179,10 @@ class StepLauncher:
                         pipeline_run_id=pipeline_run.id,
                         pipeline_run_metadata=pipeline_run_metadata,
                     )
-
-                pipeline_model_version, pipeline_run = (
-                    step_run_utils.prepare_pipeline_run_model_version(
-                        pipeline_run
-                    )
-                )
+                    if model_version := pipeline_run.model_version:
+                        step_run_utils.log_model_version_dashboard_url(
+                            model_version=model_version
+                        )
 
                 request_factory = step_run_utils.StepRunRequestFactory(
                     deployment=self._deployment,
@@ -209,12 +207,10 @@ class StepLauncher:
                     step_run = Client().zen_store.create_run_step(
                         step_run_request
                     )
-
-                    step_model_version, step_run = (
-                        step_run_utils.prepare_step_run_model_version(
-                            step_run=step_run, pipeline_run=pipeline_run
+                    if model_version := step_run.model_version:
+                        step_run_utils.log_model_version_dashboard_url(
+                            model_version=model_version
                         )
-                    )
 
                 if not step_run.status.is_finished:
                     logger.info(f"Step `{self._step_name}` has started.")
@@ -289,8 +285,8 @@ class StepLauncher:
                         f"Using cached version of step `{self._step_name}`."
                     )
                     if (
-                        model_version := step_model_version
-                        or pipeline_model_version
+                        model_version := step_run.model_version
+                        or pipeline_run.model_version
                     ):
                         step_run_utils.link_output_artifacts_to_model_version(
                             artifacts=step_run.outputs,
@@ -309,8 +305,12 @@ class StepLauncher:
             The created or existing pipeline run,
             and a boolean indicating whether the run was created or reused.
         """
-        run_name = orchestrator_utils.get_run_name(
-            run_name_template=self._deployment.run_name_template
+        start_time = datetime.utcnow()
+        run_name = string_utils.format_name_template(
+            name_template=self._deployment.run_name_template,
+            substitutions=self._deployment.pipeline_configuration._get_full_substitutions(
+                start_time
+            ),
         )
 
         logger.debug("Creating pipeline run %s", run_name)
@@ -329,7 +329,7 @@ class StepLauncher:
             ),
             status=ExecutionStatus.RUNNING,
             orchestrator_environment=get_run_environment_dict(),
-            start_time=datetime.utcnow(),
+            start_time=start_time,
             tags=self._deployment.pipeline_configuration.tags,
         )
         return client.zen_store.get_or_create_run(pipeline_run)