zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/models/v2/core/step_run.py

@@ -14,31 +14,64 @@
 """Models representing steps runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+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,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
+    WorkspaceScopedResponseResources,
 )
+from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
+from zenml.models.v2.core.model_version import ModelVersionResponse
 
 if TYPE_CHECKING:
-    from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
+    from sqlalchemy.sql.elements import ColumnElement
+
     from zenml.models.v2.core.logs import (
         LogsRequest,
         LogsResponse,
     )
-    from zenml.models.v2.core.run_metadata import (
-        RunMetadataResponse,
-    )
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
+
+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 ------------------
@@ -93,11 +126,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.",
@@ -106,6 +139,13 @@ class StepRunRequest(WorkspaceScopedRequest):
     deployment: UUID = Field(
         title="The deployment associated with the step run."
     )
+    model_version_id: Optional[UUID] = Field(
+        title="The ID of the model version that was "
+        "configured by this step run explicitly.",
+        default=None,
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
 
 
 # ------------------ Update Model ------------------
@@ -114,14 +154,10 @@ 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={},
     )
-    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={},
@@ -134,6 +170,12 @@ class StepRunUpdate(BaseModel):
         title="The end time of the step run.",
         default=None,
     )
+    model_version_id: Optional[UUID] = Field(
+        title="The ID of the model version that was "
+        "configured by this step run explicitly.",
+        default=None,
+    )
+    model_config = ConfigDict(protected_namespaces=())
 
 
 # ------------------ Response Model ------------------
@@ -141,14 +183,28 @@ class StepRunResponseBody(WorkspaceScopedResponseBody):
     """Response body for step runs."""
 
     status: ExecutionStatus = Field(title="The status of the step.")
-    inputs: Dict[str, "ArtifactVersionResponse"] = Field(
+    start_time: Optional[datetime] = Field(
+        title="The start time of the step run.",
+        default=None,
+    )
+    end_time: Optional[datetime] = Field(
+        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={},
+        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 "
+        "configured by this step run explicitly.",
+        default=None,
     )
+    model_config = ConfigDict(protected_namespaces=())
 
 
 class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
@@ -180,16 +236,6 @@ class StepRunResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=TEXT_FIELD_MAX_LENGTH,
     )
 
-    # Timestamps
-    start_time: Optional[datetime] = Field(
-        title="The start time of the step run.",
-        default=None,
-    )
-    end_time: Optional[datetime] = Field(
-        title="The end time of the step run.",
-        default=None,
-    )
-
     # References
     logs: Optional["LogsResponse"] = Field(
         title="Logs associated with this step run.",
@@ -209,14 +255,30 @@ 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={},
     )
 
 
+class StepRunResponseResources(WorkspaceScopedResponseResources):
+    """Class for all resource models associated with the step run entity."""
+
+    model_version: Optional[ModelVersionResponse] = None
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
+
+
 class StepRunResponse(
-    WorkspaceScopedResponse[StepRunResponseBody, StepRunResponseMetadata]
+    WorkspaceScopedResponse[
+        StepRunResponseBody, StepRunResponseMetadata, StepRunResponseResources
+    ]
 ):
     """Response model for step runs."""
 
@@ -237,7 +299,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:
@@ -256,7 +318,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:
@@ -267,12 +329,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
@@ -285,7 +350,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:
@@ -294,7 +359,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:
@@ -302,6 +367,15 @@ class StepRunResponse(
         """
         return self.get_body().outputs
 
+    @property
+    def model_version_id(self) -> Optional[UUID]:
+        """The `model_version_id` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().model_version_id
+
     @property
     def config(self) -> "StepConfiguration":
         """The `config` property.
@@ -363,7 +437,7 @@ class StepRunResponse(
         Returns:
             the value of the property.
         """
-        return self.get_metadata().start_time
+        return self.get_body().start_time
 
     @property
     def end_time(self) -> Optional[datetime]:
@@ -372,7 +446,7 @@ class StepRunResponse(
         Returns:
             the value of the property.
         """
-        return self.get_metadata().end_time
+        return self.get_body().end_time
 
     @property
     def logs(self) -> Optional["LogsResponse"]:
@@ -420,7 +494,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:
@@ -428,6 +502,15 @@ class StepRunResponse(
         """
         return self.get_metadata().run_metadata
 
+    @property
+    def model_version(self) -> Optional[ModelVersionResponse]:
+        """The `model_version` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().model_version
+
 
 # ------------------ Filter Model ------------------
 
@@ -435,6 +518,12 @@ class StepRunResponse(
 class StepRunFilter(WorkspaceScopedFilter):
     """Model to enable advanced filtering of step runs."""
 
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        "model",
+        "run_metadata",
+    ]
+
     name: Optional[str] = Field(
         default=None,
         description="Name of the step run",
@@ -452,20 +541,93 @@ class StepRunFilter(WorkspaceScopedFilter):
         description="Status of the Step Run",
     )
     start_time: Optional[Union[datetime, str]] = Field(
-        default=None, description="Start time for this run"
+        default=None,
+        description="Start time for this run",
+        union_mode="left_to_right",
     )
     end_time: Optional[Union[datetime, str]] = Field(
-        default=None, description="End time for this run"
+        default=None,
+        description="End time for this run",
+        union_mode="left_to_right",
     )
     pipeline_run_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Pipeline run of this step run"
+        default=None,
+        description="Pipeline run of this step run",
+        union_mode="left_to_right",
+    )
+    deployment_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Deployment of this step run",
+        union_mode="left_to_right",
     )
     original_step_run_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Original id for this step run"
+        default=None,
+        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"
+    model_version_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Model version associated with the step run.",
+        union_mode="left_to_right",
     )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of this step run"
+    model: Optional[Union[UUID, str]] = Field(
+        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=())
+
+    def get_custom_filters(
+        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(table)
+
+        from sqlmodel import and_
+
+        from zenml.zen_stores.schemas import (
+            ModelSchema,
+            ModelVersionSchema,
+            RunMetadataResourceSchema,
+            RunMetadataSchema,
+            StepRunSchema,
+        )
+
+        if self.model:
+            model_filter = and_(
+                StepRunSchema.model_version_id == ModelVersionSchema.id,
+                ModelVersionSchema.model_id == ModelSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.model, table=ModelSchema
+                ),
+            )
+            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_(
+                    RunMetadataResourceSchema.resource_id == StepRunSchema.id,
+                    RunMetadataResourceSchema.resource_type
+                    == MetadataResourceTypes.STEP_RUN,
+                    RunMetadataResourceSchema.run_metadata_id
+                    == RunMetadataSchema.id,
+                    self.generate_custom_query_conditions_for_column(
+                        value=value,
+                        table=RunMetadataSchema,
+                        column="value",
+                    ),
+                )
+                custom_filters.append(additional_filter)
+
+        return custom_filters

zenml/models/v2/core/tag.py

@@ -21,10 +21,11 @@ from pydantic import BaseModel, Field
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.enums import ColorVariants
 from zenml.models.v2.base.base import (
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
-    BaseResponseBody,
     BaseResponseMetadata,
+    BaseResponseResources,
 )
 from zenml.models.v2.base.filter import BaseFilter
 
@@ -50,14 +51,14 @@ class TagRequest(BaseRequest):
 class TagUpdate(BaseModel):
     """Update model for tags."""
 
-    name: Optional[str]
-    color: Optional[ColorVariants]
+    name: Optional[str] = None
+    color: Optional[ColorVariants] = None
 
 
 # ------------------ Response Model ------------------
 
 
-class TagResponseBody(BaseResponseBody):
+class TagResponseBody(BaseDatedResponseBody):
     """Response body for tags."""
 
     color: ColorVariants = Field(
@@ -69,7 +70,15 @@ class TagResponseBody(BaseResponseBody):
     )
 
 
-class TagResponse(BaseResponse[TagResponseBody, BaseResponseMetadata]):
+class TagResponseResources(BaseResponseResources):
+    """Class for all resource models associated with the tag entity."""
+
+
+class TagResponse(
+    BaseIdentifiedResponse[
+        TagResponseBody, BaseResponseMetadata, TagResponseResources
+    ]
+):
     """Response model for tags."""
 
     name: str = Field(
@@ -112,5 +121,9 @@ class TagResponse(BaseResponse[TagResponseBody, BaseResponseMetadata]):
 class TagFilter(BaseFilter):
     """Model to enable advanced filtering of all tags."""
 
-    name: Optional[str]
-    color: Optional[ColorVariants]
+    name: Optional[str] = Field(
+        description="The unique title of the tag.", default=None
+    )
+    color: Optional[ColorVariants] = Field(
+        description="The color variant assigned to the tag.", default=None
+    )

zenml/models/v2/core/tag_resource.py

@@ -17,10 +17,11 @@ from uuid import UUID
 
 from zenml.enums import TaggableResourceTypes
 from zenml.models.v2.base.base import (
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
-    BaseResponseBody,
     BaseResponseMetadata,
+    BaseResponseResources,
 )
 
 # ------------------ Request Model ------------------
@@ -41,7 +42,7 @@ class TagResourceRequest(BaseRequest):
 # ------------------ Response Model ------------------
 
 
-class TagResourceResponseBody(BaseResponseBody):
+class TagResourceResponseBody(BaseDatedResponseBody):
     """Response body for the links between tags and resources."""
 
     tag_id: UUID
@@ -49,8 +50,16 @@ class TagResourceResponseBody(BaseResponseBody):
     resource_type: TaggableResourceTypes
 
 
+class TagResourceResponseResources(BaseResponseResources):
+    """Class for all resource models associated with the tag resource entity."""
+
+
 class TagResourceResponse(
-    BaseResponse[TagResourceResponseBody, BaseResponseMetadata]
+    BaseIdentifiedResponse[
+        TagResourceResponseBody,
+        BaseResponseMetadata,
+        TagResourceResponseResources,
+    ]
 ):
     """Response model for the links between tags and resources."""