zenml-nightly 0.68.1.dev20241107__py3-none-any.whl → 0.68.1.dev20241108__py3-none-any.whl

zenml/models/v2/core/artifact_version.py

@@ -34,7 +34,7 @@ from pydantic import (
 
 from zenml.config.source import Source, SourceWithValidator
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
-from zenml.enums import ArtifactType, GenericFilterOps
+from zenml.enums import ArtifactSaveType, ArtifactType, GenericFilterOps
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.filter import StrFilter
@@ -57,9 +57,6 @@ if TYPE_CHECKING:
         ArtifactVisualizationResponse,
     )
     from zenml.models.v2.core.pipeline_run import PipelineRunResponse
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
     from zenml.models.v2.core.step_run import StepRunResponse
 
 logger = get_logger(__name__)
@@ -107,6 +104,9 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
     visualizations: Optional[List["ArtifactVisualizationRequest"]] = Field(
         default=None, title="Visualizations of the artifact."
     )
+    save_type: ArtifactSaveType = Field(
+        title="The save type of the artifact version.",
+    )
     metadata: Optional[Dict[str, MetadataType]] = Field(
         default=None, title="Metadata of the artifact version."
     )
@@ -193,6 +193,9 @@ class ArtifactVersionResponseBody(WorkspaceScopedResponseBody):
         title="The ID of the pipeline run that generated this artifact version.",
         default=None,
     )
+    save_type: ArtifactSaveType = Field(
+        title="The save type of the artifact version.",
+    )
     artifact_store_id: Optional[UUID] = Field(
         title="ID of the artifact store in which this artifact is stored.",
         default=None,
@@ -230,7 +233,7 @@ class ArtifactVersionResponseMetadata(WorkspaceScopedResponseMetadata):
     visualizations: Optional[List["ArtifactVisualizationResponse"]] = Field(
         default=None, title="Visualizations of the artifact."
     )
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         default={}, title="Metadata of the artifact."
     )
 
@@ -313,6 +316,15 @@ class ArtifactVersionResponse(
         """
         return self.get_body().producer_pipeline_run_id
 
+    @property
+    def save_type(self) -> ArtifactSaveType:
+        """The `save_type` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().save_type
+
     @property
     def artifact_store_id(self) -> Optional[UUID]:
         """The `artifact_store_id` property.
@@ -343,7 +355,7 @@ class ArtifactVersionResponse(
         return self.get_metadata().visualizations
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `metadata` property.
 
         Returns:
@@ -671,7 +683,7 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
         )
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `metadata` property in lazy loading mode.
 
         Returns:

zenml/models/v2/core/model_version.py

@@ -29,6 +29,7 @@ from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ModelStages
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.filter import AnyQuery
 from zenml.models.v2.base.page import Page
 from zenml.models.v2.base.scoped import (
@@ -49,9 +50,6 @@ if TYPE_CHECKING:
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.model import ModelResponse
     from zenml.models.v2.core.pipeline_run import PipelineRunResponse
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
     from zenml.zen_stores.schemas import BaseSchema
 
     AnySchema = TypeVar("AnySchema", bound=BaseSchema)
@@ -193,7 +191,7 @@ class ModelVersionResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=TEXT_FIELD_MAX_LENGTH,
         default=None,
     )
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         description="Metadata linked to the model version",
         default={},
     )
@@ -304,7 +302,7 @@ class ModelVersionResponse(
         return self.get_metadata().description
 
     @property
-    def run_metadata(self) -> Optional[Dict[str, "RunMetadataResponse"]]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `run_metadata` property.
 
         Returns:

zenml/models/v2/core/pipeline_run.py

@@ -30,6 +30,7 @@ from pydantic import BaseModel, ConfigDict, Field
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import ExecutionStatus
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
@@ -51,9 +52,6 @@ if TYPE_CHECKING:
     from zenml.models.v2.core.pipeline_build import (
         PipelineBuildResponse,
     )
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
     from zenml.models.v2.core.schedule import ScheduleResponse
     from zenml.models.v2.core.stack import StackResponse
     from zenml.models.v2.core.step_run import StepRunResponse
@@ -190,7 +188,7 @@ class PipelineRunResponseBody(WorkspaceScopedResponseBody):
 class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for pipeline runs."""
 
-    run_metadata: Dict[str, "RunMetadataResponse"] = Field(
+    run_metadata: Dict[str, MetadataType] = Field(
         default={},
         title="Metadata associated with this pipeline run.",
     )
@@ -450,7 +448,7 @@ class PipelineRunResponse(
         return self.get_body().model_version_id
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, MetadataType]:
         """The `run_metadata` property.
 
         Returns:

zenml/models/v2/core/run_metadata.py

@@ -13,21 +13,15 @@
 #  permissions and limitations under the License.
 """Models representing run metadata."""
 
-from typing import Any, Dict, Optional, Union
+from typing import Dict, Optional
 from uuid import UUID
 
-from pydantic import Field, field_validator
+from pydantic import Field
 
-from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import MetadataResourceTypes
 from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum
 from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
     WorkspaceScopedRequest,
-    WorkspaceScopedResponse,
-    WorkspaceScopedResponseBody,
-    WorkspaceScopedResponseMetadata,
-    WorkspaceScopedResponseResources,
 )
 
 # ------------------ Request Model ------------------
@@ -51,212 +45,3 @@ class RunMetadataRequest(WorkspaceScopedRequest):
     types: Dict[str, "MetadataTypeEnum"] = Field(
         title="The types of the metadata to be created.",
     )
-
-
-# ------------------ Update Model ------------------
-
-# There is no update model for run metadata.
-
-# ------------------ Response Model ------------------
-
-
-class RunMetadataResponseBody(WorkspaceScopedResponseBody):
-    """Response body for run metadata."""
-
-    key: str = Field(title="The key of the metadata.")
-    value: MetadataType = Field(
-        title="The value of the metadata.", union_mode="smart"
-    )
-    type: MetadataTypeEnum = Field(title="The type of the metadata.")
-
-    @field_validator("key", "type")
-    @classmethod
-    def str_field_max_length_check(cls, value: Any) -> Any:
-        """Checks if the length of the value exceeds the maximum str length.
-
-        Args:
-            value: the value set in the field
-
-        Returns:
-            the value itself.
-
-        Raises:
-            AssertionError: if the length of the field is longer than the
-                maximum threshold.
-        """
-        assert len(str(value)) < STR_FIELD_MAX_LENGTH, (
-            "The length of the value for this field can not "
-            f"exceed {STR_FIELD_MAX_LENGTH}"
-        )
-        return value
-
-    @field_validator("value")
-    @classmethod
-    def text_field_max_length_check(cls, value: Any) -> Any:
-        """Checks if the length of the value exceeds the maximum text length.
-
-        Args:
-            value: the value set in the field
-
-        Returns:
-            the value itself.
-
-        Raises:
-            AssertionError: if the length of the field is longer than the
-                maximum threshold.
-        """
-        assert len(str(value)) < TEXT_FIELD_MAX_LENGTH, (
-            "The length of the value for this field can not "
-            f"exceed {TEXT_FIELD_MAX_LENGTH}"
-        )
-        return value
-
-
-class RunMetadataResponseMetadata(WorkspaceScopedResponseMetadata):
-    """Response metadata for run metadata."""
-
-    resource_id: UUID = Field(
-        title="The ID of the resource that this metadata belongs to.",
-    )
-    resource_type: MetadataResourceTypes = Field(
-        title="The type of the resource that this metadata belongs to.",
-    )
-    stack_component_id: Optional[UUID] = Field(
-        title="The ID of the stack component that this metadata belongs to."
-    )
-
-
-class RunMetadataResponseResources(WorkspaceScopedResponseResources):
-    """Class for all resource models associated with the run metadata entity."""
-
-
-class RunMetadataResponse(
-    WorkspaceScopedResponse[
-        RunMetadataResponseBody,
-        RunMetadataResponseMetadata,
-        RunMetadataResponseResources,
-    ]
-):
-    """Response model for run metadata."""
-
-    def get_hydrated_version(self) -> "RunMetadataResponse":
-        """Get the hydrated version of this run metadata.
-
-        Returns:
-            an instance of the same entity with the metadata field attached.
-        """
-        from zenml.client import Client
-
-        return Client().zen_store.get_run_metadata(self.id)
-
-    # Body and metadata properties
-    @property
-    def key(self) -> str:
-        """The `key` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().key
-
-    @property
-    def value(self) -> MetadataType:
-        """The `value` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().value
-
-    @property
-    def type(self) -> MetadataTypeEnum:
-        """The `type` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_body().type
-
-    @property
-    def resource_id(self) -> UUID:
-        """The `resource_id` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().resource_id
-
-    @property
-    def resource_type(self) -> MetadataResourceTypes:
-        """The `resource_type` property.
-
-        Returns:
-            the value of the property.
-        """
-        return MetadataResourceTypes(self.get_metadata().resource_type)
-
-    @property
-    def stack_component_id(self) -> Optional[UUID]:
-        """The `stack_component_id` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().stack_component_id
-
-
-# ------------------ Filter Model ------------------
-
-
-class RunMetadataFilter(WorkspaceScopedFilter):
-    """Model to enable advanced filtering of run metadata."""
-
-    resource_id: Optional[Union[str, UUID]] = Field(
-        default=None, union_mode="left_to_right"
-    )
-    resource_type: Optional[MetadataResourceTypes] = None
-    stack_component_id: Optional[Union[str, UUID]] = Field(
-        default=None, union_mode="left_to_right"
-    )
-    key: Optional[str] = None
-    type: Optional[Union[str, MetadataTypeEnum]] = Field(
-        default=None, union_mode="left_to_right"
-    )
-
-
-# -------------------- Lazy Loader --------------------
-
-
-class LazyRunMetadataResponse(RunMetadataResponse):
-    """Lazy run metadata response.
-
-    Used if the run metadata is accessed from the model in
-    a pipeline context available only during pipeline compilation.
-    """
-
-    id: Optional[UUID] = None  # type: ignore[assignment]
-    lazy_load_artifact_name: Optional[str] = None
-    lazy_load_artifact_version: Optional[str] = None
-    lazy_load_metadata_name: Optional[str] = None
-    lazy_load_model_name: str
-    lazy_load_model_version: Optional[str] = None
-
-    def get_body(self) -> None:  # type: ignore[override]
-        """Protects from misuse of the lazy loader.
-
-        Raises:
-            RuntimeError: always
-        """
-        raise RuntimeError(
-            "Cannot access run metadata body before pipeline runs."
-        )
-
-    def get_metadata(self) -> None:  # type: ignore[override]
-        """Protects from misuse of the lazy loader.
-
-        Raises:
-            RuntimeError: always
-        """
-        raise RuntimeError(
-            "Cannot access run metadata metadata before pipeline runs."
-        )

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:

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_