zenml-nightly 0.68.1.dev20241106__py3-none-any.whl → 0.70.0.dev20241116__py3-none-any.whl

zenml/models/v2/core/step_run.py

@@ -21,7 +21,8 @@ from pydantic import BaseModel, 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
+from zenml.enums import ExecutionStatus, StepRunInputArtifactType
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -30,19 +31,35 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
+from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
 from zenml.models.v2.core.model_version import ModelVersionResponse
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
 
-    from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.logs import (
         LogsRequest,
         LogsResponse,
     )
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
+
+
+class StepRunInputResponse(ArtifactVersionResponse):
+    """Response model for step run inputs."""
+
+    input_type: StepRunInputArtifactType
+
+    def get_hydrated_version(self) -> "StepRunInputResponse":
+        """Get the hydrated version of this step run input.
+
+        Returns:
+            an instance of the same entity with the metadata field attached.
+        """
+        from zenml.client import Client
+
+        return StepRunInputResponse(
+            input_type=self.input_type,
+            **Client().zen_store.get_artifact_version(self.id).model_dump(),
+        )
 
 
 # ------------------ Request Model ------------------
@@ -97,11 +114,11 @@ class StepRunRequest(WorkspaceScopedRequest):
     )
     inputs: Dict[str, UUID] = Field(
         title="The IDs of the input artifact versions of the step run.",
-        default={},
+        default_factory=dict,
     )
-    outputs: Dict[str, UUID] = Field(
+    outputs: Dict[str, List[UUID]] = Field(
         title="The IDs of the output artifact versions of the step run.",
-        default={},
+        default_factory=dict,
     )
     logs: Optional["LogsRequest"] = Field(
         title="Logs associated with this step run.",
@@ -129,10 +146,6 @@ class StepRunUpdate(BaseModel):
         title="The IDs of the output artifact versions of the step run.",
         default={},
     )
-    saved_artifact_versions: Dict[str, UUID] = Field(
-        title="The IDs of artifact versions that were saved by this step run.",
-        default={},
-    )
     loaded_artifact_versions: Dict[str, UUID] = Field(
         title="The IDs of artifact versions that were loaded by this step run.",
         default={},
@@ -166,13 +179,13 @@ class StepRunResponseBody(WorkspaceScopedResponseBody):
         title="The end time of the step run.",
         default=None,
     )
-    inputs: Dict[str, "ArtifactVersionResponse"] = Field(
+    inputs: Dict[str, StepRunInputResponse] = Field(
         title="The input artifact versions of the step run.",
-        default={},
+        default_factory=dict,
     )
-    outputs: Dict[str, "ArtifactVersionResponse"] = Field(
+    outputs: Dict[str, List[ArtifactVersionResponse]] = Field(
         title="The output artifact versions of the step run.",
-        default={},
+        default_factory=dict,
     )
     model_version_id: Optional[UUID] = Field(
         title="The ID of the model version that was "
@@ -230,7 +243,7 @@ class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
         title="The IDs of the parent steps of this step run.",
         default_factory=list,
     )
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         title="Metadata associated with this step run.",
         default={},
     )
@@ -274,7 +287,7 @@ class StepRunResponse(
 
     # Helper properties
     @property
-    def input(self) -> "ArtifactVersionResponse":
+    def input(self) -> ArtifactVersionResponse:
         """Returns the input artifact that was used to run this step.
 
         Returns:
@@ -293,7 +306,7 @@ class StepRunResponse(
         return next(iter(self.inputs.values()))
 
     @property
-    def output(self) -> "ArtifactVersionResponse":
+    def output(self) -> ArtifactVersionResponse:
         """Returns the output artifact that was written by this step.
 
         Returns:
@@ -304,12 +317,15 @@ class StepRunResponse(
         """
         if not self.outputs:
             raise ValueError(f"Step {self.name} has no outputs.")
-        if len(self.outputs) > 1:
+        if len(self.outputs) > 1 or (
+            len(self.outputs) == 1
+            and len(next(iter(self.outputs.values()))) > 1
+        ):
             raise ValueError(
                 f"Step {self.name} has multiple outputs, so `Step.output` is "
                 "ambiguous. Please use `Step.outputs` instead."
             )
-        return next(iter(self.outputs.values()))
+        return next(iter(self.outputs.values()))[0]
 
     # Body and metadata properties
     @property
@@ -322,7 +338,7 @@ class StepRunResponse(
         return self.get_body().status
 
     @property
-    def inputs(self) -> Dict[str, "ArtifactVersionResponse"]:
+    def inputs(self) -> Dict[str, StepRunInputResponse]:
         """The `inputs` property.
 
         Returns:
@@ -331,7 +347,7 @@ class StepRunResponse(
         return self.get_body().inputs
 
     @property
-    def outputs(self) -> Dict[str, "ArtifactVersionResponse"]:
+    def outputs(self) -> Dict[str, List[ArtifactVersionResponse]]:
         """The `outputs` property.
 
         Returns:
@@ -466,7 +482,7 @@ class StepRunResponse(
         return self.get_metadata().parent_step_ids
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `run_metadata` property.
 
         Returns:
@@ -493,6 +509,7 @@ class StepRunFilter(WorkspaceScopedFilter):
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
         "model",
+        "run_metadata",
     ]
 
     name: Optional[str] = Field(
@@ -555,6 +572,10 @@ class StepRunFilter(WorkspaceScopedFilter):
         default=None,
         description="Name/ID of the model associated with the step run.",
     )
+    run_metadata: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="The run_metadata to filter the step runs by.",
+    )
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -573,6 +594,7 @@ class StepRunFilter(WorkspaceScopedFilter):
         from zenml.zen_stores.schemas import (
             ModelSchema,
             ModelVersionSchema,
+            RunMetadataSchema,
             StepRunSchema,
         )
 
@@ -585,5 +607,21 @@ class StepRunFilter(WorkspaceScopedFilter):
                 ),
             )
             custom_filters.append(model_filter)
+        if self.run_metadata is not None:
+            from zenml.enums import MetadataResourceTypes
+
+            for key, value in self.run_metadata.items():
+                additional_filter = and_(
+                    RunMetadataSchema.resource_id == StepRunSchema.id,
+                    RunMetadataSchema.resource_type
+                    == MetadataResourceTypes.STEP_RUN,
+                    RunMetadataSchema.key == key,
+                    self.generate_custom_query_conditions_for_column(
+                        value=value,
+                        table=RunMetadataSchema,
+                        column="value",
+                    ),
+                )
+                custom_filters.append(additional_filter)
 
         return custom_filters

zenml/orchestrators/base_orchestrator.py

@@ -15,6 +15,7 @@
 
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
+from uuid import UUID
 
 from pydantic import model_validator
 
@@ -186,7 +187,17 @@ class BaseOrchestrator(StackComponent, ABC):
         """
         self._prepare_run(deployment=deployment)
 
-        environment = get_config_environment_vars(deployment=deployment)
+        pipeline_run_id: Optional[UUID] = None
+        schedule_id: Optional[UUID] = None
+        if deployment.schedule:
+            schedule_id = deployment.schedule.id
+        if placeholder_run:
+            pipeline_run_id = placeholder_run.id
+
+        environment = get_config_environment_vars(
+            schedule_id=schedule_id,
+            pipeline_run_id=pipeline_run_id,
+        )
 
         prevent_client_side_caching = handle_bool_env_var(
             ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING, default=False

zenml/orchestrators/input_utils.py

@@ -18,17 +18,19 @@ from uuid import UUID
 
 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
 
 if TYPE_CHECKING:
-    from zenml.models import ArtifactVersionResponse, PipelineRunResponse
+    from zenml.models import PipelineRunResponse
+    from zenml.models.v2.core.step_run import StepRunInputResponse
 
 
 def resolve_step_inputs(
     step: "Step",
     pipeline_run: "PipelineRunResponse",
-) -> Tuple[Dict[str, "ArtifactVersionResponse"], List[UUID]]:
+) -> Tuple[Dict[str, "StepRunInputResponse"], List[UUID]]:
     """Resolves inputs for the current step.
 
     Args:
@@ -45,7 +47,8 @@ def resolve_step_inputs(
         The IDs of the input artifact versions and the IDs of parent steps of
             the current step.
     """
-    from zenml.models import ArtifactVersionResponse, RunMetadataResponse
+    from zenml.models import ArtifactVersionResponse
+    from zenml.models.v2.core.step_run import StepRunInputResponse
 
     current_run_steps = {
         run_step.name: run_step
@@ -54,7 +57,7 @@ def resolve_step_inputs(
         )
     }
 
-    input_artifacts: Dict[str, "ArtifactVersionResponse"] = {}
+    input_artifacts: Dict[str, StepRunInputResponse] = {}
     for name, input_ in step.spec.inputs.items():
         try:
             step_run = current_run_steps[input_.step_name]
@@ -64,22 +67,44 @@ def resolve_step_inputs(
             )
 
         try:
-            artifact = step_run.outputs[input_.output_name]
+            outputs = step_run.outputs[input_.output_name]
         except KeyError:
             raise InputResolutionError(
-                f"No output `{input_.output_name}` found for step "
+                f"No step output `{input_.output_name}` found for step "
                 f"`{input_.step_name}`."
             )
 
-        input_artifacts[name] = artifact
+        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
+            raise InputResolutionError(
+                f"Too many step outputs for output `{input_.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"`{input_.step_name}`."
+            )
+
+        input_artifacts[name] = StepRunInputResponse(
+            input_type=StepRunInputArtifactType.STEP_OUTPUT,
+            **step_outputs[0].model_dump(),
+        )
 
     for (
         name,
         external_artifact,
     ) in step.config.external_input_artifacts.items():
         artifact_version_id = external_artifact.get_artifact_version_id()
-        input_artifacts[name] = Client().get_artifact_version(
-            artifact_version_id
+        input_artifacts[name] = StepRunInputResponse(
+            input_type=StepRunInputArtifactType.EXTERNAL,
+            **Client().get_artifact_version(artifact_version_id).model_dump(),
         )
 
     for name, config_ in step.config.model_artifacts_or_metadata.items():
@@ -98,9 +123,7 @@ def resolve_step_inputs(
             ):
                 # metadata values should go directly in parameters, as primitive types
                 step.config.parameters[name] = (
-                    context_model_version.run_metadata[
-                        config_.metadata_name
-                    ].value
+                    context_model_version.run_metadata[config_.metadata_name]
                 )
             elif config_.artifact_name is None:
                 err_msg = (
@@ -112,14 +135,15 @@ def resolve_step_inputs(
                     config_.artifact_name, config_.artifact_version
                 ):
                     if config_.metadata_name is None:
-                        input_artifacts[name] = artifact_
+                        input_artifacts[name] = StepRunInputResponse(
+                            input_type=StepRunInputArtifactType.LAZY_LOADED,
+                            **artifact_.model_dump(),
+                        )
                     elif config_.metadata_name:
                         # metadata values should go directly in parameters, as primitive types
                         try:
                             step.config.parameters[name] = (
-                                artifact_.run_metadata[
-                                    config_.metadata_name
-                                ].value
+                                artifact_.run_metadata[config_.metadata_name]
                             )
                         except KeyError:
                             err_msg = (
@@ -141,9 +165,10 @@ def resolve_step_inputs(
     for name, cll_ in step.config.client_lazy_loaders.items():
         value_ = cll_.evaluate()
         if isinstance(value_, ArtifactVersionResponse):
-            input_artifacts[name] = value_
-        elif isinstance(value_, RunMetadataResponse):
-            step.config.parameters[name] = value_.value
+            input_artifacts[name] = StepRunInputResponse(
+                input_type=StepRunInputArtifactType.LAZY_LOADED,
+                **value_.model_dump(),
+            )
         else:
             step.config.parameters[name] = value_
 

zenml/orchestrators/step_launcher.py

@@ -33,13 +33,13 @@ from zenml.environment import get_run_environment_dict
 from zenml.logger import get_logger
 from zenml.logging import step_logging
 from zenml.models import (
-    ArtifactVersionResponse,
     LogsRequest,
     PipelineDeploymentResponse,
     PipelineRunRequest,
     PipelineRunResponse,
     StepRunResponse,
 )
+from zenml.models.v2.core.step_run import StepRunInputResponse
 from zenml.orchestrators import output_utils, publish_utils, step_run_utils
 from zenml.orchestrators import utils as orchestrator_utils
 from zenml.orchestrators.step_runner import StepRunner
@@ -422,7 +422,8 @@ class StepLauncher:
             )
         )
         environment = orchestrator_utils.get_config_environment_vars(
-            deployment=self._deployment
+            pipeline_run_id=step_run_info.run_id,
+            step_run_id=step_run_info.step_run_id,
         )
         if last_retry:
             environment[ENV_ZENML_IGNORE_FAILURE_HOOK] = str(False)
@@ -442,7 +443,7 @@ class StepLauncher:
         pipeline_run: PipelineRunResponse,
         step_run: StepRunResponse,
         step_run_info: StepRunInfo,
-        input_artifacts: Dict[str, ArtifactVersionResponse],
+        input_artifacts: Dict[str, StepRunInputResponse],
         output_artifact_uris: Dict[str, str],
         last_retry: bool,
     ) -> None:

zenml/orchestrators/step_run_utils.py

@@ -14,12 +14,12 @@
 """Utilities for creating step runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, Mapping, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Set, Tuple
 
 from zenml.client import Client
 from zenml.config.step_configurations import ArtifactConfiguration, Step
 from zenml.constants import CODE_HASH_PARAMETER_NAME, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ExecutionStatus
+from zenml.enums import ArtifactSaveType, ExecutionStatus
 from zenml.logger import get_logger
 from zenml.model.utils import link_artifact_version_to_model_version
 from zenml.models import (
@@ -104,6 +104,7 @@ class StepRunRequestFactory:
             input_name: artifact.id
             for input_name, artifact in input_artifacts.items()
         }
+
         request.inputs = input_artifact_ids
         request.parent_step_ids = parent_step_ids
 
@@ -142,8 +143,8 @@ class StepRunRequestFactory:
 
                 request.original_step_run_id = cached_step_run.id
                 request.outputs = {
-                    output_name: artifact.id
-                    for output_name, artifact in cached_step_run.outputs.items()
+                    output_name: [artifact.id for artifact in artifacts]
+                    for output_name, artifacts in cached_step_run.outputs.items()
                 }
 
                 request.status = ExecutionStatus.CACHED
@@ -551,7 +552,7 @@ def link_pipeline_run_to_model_version(
 
 
 def link_output_artifacts_to_model_version(
-    artifacts: Dict[str, ArtifactVersionResponse],
+    artifacts: Dict[str, List[ArtifactVersionResponse]],
     output_configurations: Mapping[str, ArtifactConfiguration],
     model_version: ModelVersionResponse,
 ) -> None:
@@ -562,13 +563,16 @@ def link_output_artifacts_to_model_version(
         output_configurations: The output configurations for the step.
         model_version: The model version to link.
     """
-    for output_name, output_artifact in artifacts.items():
-        artifact_config = None
-        if output_config := output_configurations.get(output_name, None):
-            artifact_config = output_config.artifact_config
-
-        link_artifact_version_to_model_version(
-            artifact_version=output_artifact,
-            model_version=model_version,
-            artifact_config=artifact_config,
-        )
+    for output_name, output_artifacts in artifacts.items():
+        for output_artifact in output_artifacts:
+            artifact_config = None
+            if output_artifact.save_type == ArtifactSaveType.STEP_OUTPUT and (
+                output_config := output_configurations.get(output_name, None)
+            ):
+                artifact_config = output_config.artifact_config
+
+            link_artifact_version_to_model_version(
+                artifact_version=output_artifact,
+                model_version=model_version,
+                artifact_config=artifact_config,
+            )

zenml/orchestrators/step_runner.py

@@ -28,7 +28,8 @@ from typing import (
 )
 
 from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
-from zenml.artifacts.utils import save_artifact
+from zenml.artifacts.utils import _store_artifact_data_and_prepare_request
+from zenml.client import Client
 from zenml.config.step_configurations import StepConfiguration
 from zenml.config.step_run_info import StepRunInfo
 from zenml.constants import (
@@ -36,10 +37,12 @@ from zenml.constants import (
     ENV_ZENML_IGNORE_FAILURE_HOOK,
     handle_bool_env_var,
 )
+from zenml.enums import ArtifactSaveType
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.logging.step_logging import StepLogsStorageContext, redirected
 from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.models.v2.core.step_run import StepRunInputResponse
 from zenml.orchestrators.publish_utils import (
     publish_step_run_metadata,
     publish_successful_step_run,
@@ -98,7 +101,7 @@ class StepRunner:
         self,
         pipeline_run: "PipelineRunResponse",
         step_run: "StepRunResponse",
-        input_artifacts: Dict[str, "ArtifactVersionResponse"],
+        input_artifacts: Dict[str, StepRunInputResponse],
         output_artifact_uris: Dict[str, str],
         step_run_info: StepRunInfo,
     ) -> None:
@@ -241,7 +244,9 @@ class StepRunner:
                             from zenml.orchestrators import step_run_utils
 
                             step_run_utils.link_output_artifacts_to_model_version(
-                                artifacts=output_artifacts,
+                                artifacts={
+                                    k: [v] for k, v in output_artifacts.items()
+                                },
                                 output_configurations=step_run.config.outputs,
                                 model_version=model_version,
                             )
@@ -302,7 +307,7 @@ class StepRunner:
         self,
         args: List[str],
         annotations: Dict[str, Any],
-        input_artifacts: Dict[str, "ArtifactVersionResponse"],
+        input_artifacts: Dict[str, StepRunInputResponse],
     ) -> Dict[str, Any]:
         """Parses the inputs for a step entrypoint function.
 
@@ -534,7 +539,7 @@ class StepRunner:
             The IDs of the published output artifacts.
         """
         step_context = get_step_context()
-        output_artifacts: Dict[str, "ArtifactVersionResponse"] = {}
+        artifact_requests = []
 
         for output_name, return_value in output_data.items():
             data_type = type(return_value)
@@ -595,22 +600,25 @@ class StepRunner:
             # Get full set of tags
             tags = step_context.get_output_tags(output_name)
 
-            artifact = save_artifact(
+            artifact_request = _store_artifact_data_and_prepare_request(
                 name=artifact_name,
                 data=return_value,
-                materializer=materializer_class,
+                materializer_class=materializer_class,
                 uri=uri,
-                extract_metadata=artifact_metadata_enabled,
-                include_visualizations=artifact_visualization_enabled,
+                store_metadata=artifact_metadata_enabled,
+                store_visualizations=artifact_visualization_enabled,
                 has_custom_name=has_custom_name,
                 version=version,
                 tags=tags,
-                user_metadata=user_metadata,
-                manual_save=False,
+                save_type=ArtifactSaveType.STEP_OUTPUT,
+                metadata=user_metadata,
             )
-            output_artifacts[output_name] = artifact
+            artifact_requests.append(artifact_request)
 
-        return output_artifacts
+        responses = Client().zen_store.batch_create_artifact_versions(
+            artifact_requests
+        )
+        return dict(zip(output_data.keys(), responses))
 
     def load_and_run_hook(
         self,