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

zenml/models/v2/core/artifact_version.py

@@ -20,40 +20,50 @@ from typing import (
     Dict,
     List,
     Optional,
+    Type,
+    TypeVar,
     Union,
 )
 from uuid import UUID
 
-from pydantic import BaseModel, Field
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    field_validator,
+    model_validator,
+)
 
-from zenml.config.source import Source, convert_source_validator
+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.models.v2.base.filter import StrFilter
+from zenml.metadata.metadata_types import MetadataType
+from zenml.models.v2.base.filter import FilterGenerator, StrFilter
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
+    WorkspaceScopedResponseResources,
     WorkspaceScopedTaggableFilter,
 )
 from zenml.models.v2.core.artifact import ArtifactResponse
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+    from sqlalchemy.sql.elements import ColumnElement
 
-    from zenml.model.model import Model
     from zenml.models.v2.core.artifact_visualization import (
         ArtifactVisualizationRequest,
         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
+    from zenml.zen_stores.schemas.base_schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
 
 logger = get_logger(__name__)
 
@@ -63,12 +73,16 @@ logger = get_logger(__name__)
 class ArtifactVersionRequest(WorkspaceScopedRequest):
     """Request model for artifact versions."""
 
-    artifact_id: UUID = Field(
+    artifact_id: Optional[UUID] = Field(
+        default=None,
         title="ID of the artifact to which this version belongs.",
     )
-    version: Union[str, int] = Field(
-        title="Version of the artifact.",
-        max_length=STR_FIELD_MAX_LENGTH,
+    artifact_name: Optional[str] = Field(
+        default=None,
+        title="Name of the artifact to which this version belongs.",
+    )
+    version: Optional[Union[int, str]] = Field(
+        default=None, title="Version of the artifact."
     )
     has_custom_name: bool = Field(
         title="Whether the name is custom (True) or auto-generated (False).",
@@ -82,10 +96,10 @@ class ArtifactVersionRequest(WorkspaceScopedRequest):
     uri: str = Field(
         title="URI of the artifact.", max_length=TEXT_FIELD_MAX_LENGTH
     )
-    materializer: Source = Field(
+    materializer: SourceWithValidator = Field(
         title="Materializer class to use for this artifact.",
     )
-    data_type: Source = Field(
+    data_type: SourceWithValidator = Field(
         title="Data type of the artifact.",
     )
     tags: Optional[List[str]] = Field(
@@ -96,8 +110,55 @@ 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."
+    )
+
+    @field_validator("version")
+    @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
 
-    _convert_source = convert_source_validator("materializer", "data_type")
+        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
+
+    @model_validator(mode="after")
+    def _validate_request(self) -> "ArtifactVersionRequest":
+        """Validate the request values.
+
+        Raises:
+            ValueError: If the request is invalid.
+
+        Returns:
+            The validated request.
+        """
+        if self.artifact_id and self.artifact_name:
+            raise ValueError(
+                "Only one of artifact_name and artifact_id can be set."
+            )
+
+        if not (self.artifact_id or self.artifact_name):
+            raise ValueError(
+                "Either artifact_name or artifact_id must be set."
+            )
+
+        return self
 
 
 # ------------------ Update Model ------------------
@@ -120,34 +181,57 @@ class ArtifactVersionResponseBody(WorkspaceScopedResponseBody):
     artifact: ArtifactResponse = Field(
         title="Artifact to which this version belongs."
     )
-    version: Union[str, int] = Field(
-        title="Version of the artifact.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
+    version: str = Field(title="Version of the artifact.")
     uri: str = Field(
         title="URI of the artifact.", max_length=TEXT_FIELD_MAX_LENGTH
     )
     type: ArtifactType = Field(title="Type of the artifact.")
-    materializer: Source = Field(
+    materializer: SourceWithValidator = Field(
         title="Materializer class to use for this artifact.",
     )
-    data_type: Source = Field(
+    data_type: SourceWithValidator = Field(
         title="Data type of the artifact.",
     )
     tags: List[TagResponse] = Field(
         title="Tags associated with the model",
     )
+    producer_pipeline_run_id: Optional[UUID] = Field(
+        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,
+    )
 
-    _convert_source = convert_source_validator("materializer", "data_type")
+    @field_validator("version")
+    @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
 
 
 class ArtifactVersionResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for artifact versions."""
 
-    artifact_store_id: Optional[UUID] = Field(
-        title="ID of the artifact store in which this artifact is stored.",
-        default=None,
-    )
     producer_step_run_id: Optional[UUID] = Field(
         title="ID of the step run that produced this artifact.",
         default=None,
@@ -155,14 +239,20 @@ 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."
     )
 
 
+class ArtifactVersionResponseResources(WorkspaceScopedResponseResources):
+    """Class for all resource models associated with the artifact version entity."""
+
+
 class ArtifactVersionResponse(
     WorkspaceScopedResponse[
-        ArtifactVersionResponseBody, ArtifactVersionResponseMetadata
+        ArtifactVersionResponseBody,
+        ArtifactVersionResponseMetadata,
+        ArtifactVersionResponseResources,
     ]
 ):
     """Response model for artifact versions."""
@@ -223,6 +313,24 @@ class ArtifactVersionResponse(
         """
         return self.get_body().tags
 
+    @property
+    def producer_pipeline_run_id(self) -> Optional[UUID]:
+        """The `producer_pipeline_run_id` property.
+
+        Returns:
+            the value of the property.
+        """
+        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.
@@ -230,7 +338,7 @@ class ArtifactVersionResponse(
         Returns:
             the value of the property.
         """
-        return self.get_metadata().artifact_store_id
+        return self.get_body().artifact_store_id
 
     @property
     def producer_step_run_id(self) -> Optional[UUID]:
@@ -253,7 +361,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:
@@ -321,17 +429,31 @@ class ArtifactVersionResponse(
 
         return load_artifact_from_response(self)
 
-    def read(self) -> Any:
-        """(Deprecated) Materializes (loads) the data stored in this artifact.
+    def download_files(self, path: str, overwrite: bool = False) -> None:
+        """Downloads data for an artifact with no materializing.
 
-        Returns:
-            The materialized data.
+        Any artifacts will be saved as a zip file to the given path.
+
+        Args:
+            path: The path to save the binary data to.
+            overwrite: Whether to overwrite the file if it already exists.
+
+        Raises:
+            ValueError: If the path does not end with '.zip'.
         """
-        logger.warning(
-            "`artifact.read()` is deprecated and will be removed in a future "
-            "release. Please use `artifact.load()` instead."
+        if not path.endswith(".zip"):
+            raise ValueError(
+                "The path should end with '.zip' to save the binary data."
+            )
+        from zenml.artifacts.utils import (
+            download_artifact_files_from_response,
+        )
+
+        download_artifact_files_from_response(
+            self,
+            path=path,
+            overwrite=overwrite,
         )
-        return self.load()
 
     def visualize(self, title: Optional[str] = None) -> None:
         """Visualize the artifact in notebook environments.
@@ -350,18 +472,20 @@ class ArtifactVersionResponse(
 class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of artifact versions."""
 
-    # `name` and `only_unused` refer to properties related to other entities
-    #  rather than a field in the db, hence they needs to be handled
-    #  explicitly
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
         "name",
         "only_unused",
         "has_custom_name",
+        "model",
+        "pipeline_run",
+        "model_version_id",
+        "run_metadata",
     ]
     artifact_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="ID of the artifact to which this version belongs.",
+        union_mode="left_to_right",
     )
     name: Optional[str] = Field(
         default=None,
@@ -374,6 +498,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
     version_number: Optional[Union[int, str]] = Field(
         default=None,
         description="Version of the artifact if it is an integer",
+        union_mode="left_to_right",
     )
     uri: Optional[str] = Field(
         default=None,
@@ -392,13 +517,15 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         description="Datatype of the artifact",
     )
     artifact_store_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Artifact store for this artifact"
-    )
-    workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace for this artifact"
+        default=None,
+        description="Artifact store for this artifact",
+        union_mode="left_to_right",
     )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User that produced this artifact"
+    model_version_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="ID of the model version that is associated with this "
+        "artifact version.",
+        union_mode="left_to_right",
     )
     only_unused: Optional[bool] = Field(
         default=False, description="Filter only for unused artifacts"
@@ -407,27 +534,54 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
         default=None,
         description="Filter only artifacts with/without custom names.",
     )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the artifact version.",
+    )
+    model: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the model that is associated with this "
+        "artifact version.",
+    )
+    pipeline_run: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of a pipeline run that is associated with this "
+        "artifact version.",
+    )
+    run_metadata: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="The run_metadata to filter the artifact versions by.",
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
-        self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+        self, table: Type["AnySchema"]
+    ) -> List[Union["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 sqlalchemy import and_
-        from sqlmodel import select
+        from sqlmodel import and_, or_, select
 
-        from zenml.zen_stores.schemas.artifact_schemas import (
+        from zenml.zen_stores.schemas import (
             ArtifactSchema,
             ArtifactVersionSchema,
-        )
-        from zenml.zen_stores.schemas.step_run_schemas import (
+            ModelSchema,
+            ModelVersionArtifactSchema,
+            ModelVersionSchema,
+            PipelineRunSchema,
+            RunMetadataResourceSchema,
+            RunMetadataSchema,
             StepRunInputArtifactSchema,
             StepRunOutputArtifactSchema,
+            StepRunSchema,
         )
 
         if self.name:
@@ -437,7 +591,7 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
                 column="name",
                 value=value,
             )
-            artifact_name_filter = and_(  # type: ignore[type-var]
+            artifact_name_filter = and_(
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
                 filter_.generate_query_conditions(ArtifactSchema),
             )
@@ -454,13 +608,81 @@ class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(unused_filter)
 
+        if self.model_version_id:
+            value, operator = self._resolve_operator(self.model_version_id)
+
+            model_version_filter = and_(
+                ArtifactVersionSchema.id
+                == ModelVersionArtifactSchema.artifact_version_id,
+                ModelVersionArtifactSchema.model_version_id
+                == ModelVersionSchema.id,
+                FilterGenerator(ModelVersionSchema)
+                .define_filter(column="id", value=value, operator=operator)
+                .generate_query_conditions(ModelVersionSchema),
+            )
+            custom_filters.append(model_version_filter)
+
         if self.has_custom_name is not None:
-            custom_name_filter = and_(  # type: ignore[type-var]
+            custom_name_filter = and_(
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
                 ArtifactSchema.has_custom_name == self.has_custom_name,
             )
             custom_filters.append(custom_name_filter)
 
+        if self.model:
+            model_filter = and_(
+                ArtifactVersionSchema.id
+                == ModelVersionArtifactSchema.artifact_version_id,
+                ModelVersionArtifactSchema.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.pipeline_run:
+            pipeline_run_filter = and_(
+                or_(
+                    and_(
+                        ArtifactVersionSchema.id
+                        == StepRunOutputArtifactSchema.artifact_id,
+                        StepRunOutputArtifactSchema.step_id
+                        == StepRunSchema.id,
+                    ),
+                    and_(
+                        ArtifactVersionSchema.id
+                        == StepRunInputArtifactSchema.artifact_id,
+                        StepRunInputArtifactSchema.step_id == StepRunSchema.id,
+                    ),
+                ),
+                StepRunSchema.pipeline_run_id == PipelineRunSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.pipeline_run, table=PipelineRunSchema
+                ),
+            )
+            custom_filters.append(pipeline_run_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
+                    == ArtifactVersionSchema.id,
+                    RunMetadataResourceSchema.resource_type
+                    == MetadataResourceTypes.ARTIFACT_VERSION,
+                    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
 
 
@@ -475,9 +697,10 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
     """
 
     id: Optional[UUID] = None  # type: ignore[assignment]
-    _lazy_load_name: Optional[str] = None
-    _lazy_load_version: Optional[str] = None
-    _lazy_load_model: "Model"
+    lazy_load_name: Optional[str] = None
+    lazy_load_version: 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.
@@ -498,7 +721,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:
@@ -507,7 +730,8 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
         from zenml.metadata.lazy_load import RunMetadataLazyGetter
 
         return RunMetadataLazyGetter(  # type: ignore[return-value]
-            self._lazy_load_model,
-            self._lazy_load_name,
-            self._lazy_load_version,
+            self.lazy_load_model_name,
+            self.lazy_load_model_version,
+            self.lazy_load_name,
+            self.lazy_load_version,
         )

zenml/models/v2/core/artifact_visualization.py

@@ -17,10 +17,11 @@ from uuid import UUID
 
 from zenml.enums import VisualizationType
 from zenml.models.v2.base.base import (
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
-    BaseResponseBody,
     BaseResponseMetadata,
+    BaseResponseResources,
 )
 
 # ------------------ Request Model ------------------
@@ -40,7 +41,7 @@ class ArtifactVisualizationRequest(BaseRequest):
 # ------------------ Response Model ------------------
 
 
-class ArtifactVisualizationResponseBody(BaseResponseBody):
+class ArtifactVisualizationResponseBody(BaseDatedResponseBody):
     """Response body for artifact visualizations."""
 
     type: VisualizationType
@@ -53,10 +54,15 @@ class ArtifactVisualizationResponseMetadata(BaseResponseMetadata):
     artifact_version_id: UUID
 
 
+class ArtifactVisualizationResponseResources(BaseResponseResources):
+    """Class for all resource models associated with the artifact visualization."""
+
+
 class ArtifactVisualizationResponse(
-    BaseResponse[
+    BaseIdentifiedResponse[
         ArtifactVisualizationResponseBody,
         ArtifactVisualizationResponseMetadata,
+        ArtifactVisualizationResponseResources,
     ]
 ):
     """Response model for artifact visualizations."""

zenml/models/v2/core/code_reference.py

@@ -19,10 +19,11 @@ from uuid import UUID
 from pydantic import Field
 
 from zenml.models.v2.base.base import (
+    BaseDatedResponseBody,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
-    BaseResponseBody,
     BaseResponseMetadata,
+    BaseResponseResources,
 )
 
 if TYPE_CHECKING:
@@ -52,7 +53,7 @@ class CodeReferenceRequest(BaseRequest):
 # ------------------ Response Model ------------------
 
 
-class CodeReferenceResponseBody(BaseResponseBody):
+class CodeReferenceResponseBody(BaseDatedResponseBody):
     """Response body for code references."""
 
     commit: str = Field(description="The commit of the code reference.")
@@ -68,8 +69,16 @@ class CodeReferenceResponseMetadata(BaseResponseMetadata):
     """Response metadata for code references."""
 
 
+class CodeReferenceResponseResources(BaseResponseResources):
+    """Class for all resource models associated with the code reference entity."""
+
+
 class CodeReferenceResponse(
-    BaseResponse[CodeReferenceResponseBody, CodeReferenceResponseMetadata]
+    BaseIdentifiedResponse[
+        CodeReferenceResponseBody,
+        CodeReferenceResponseMetadata,
+        CodeReferenceResponseResources,
+    ]
 ):
     """Response model for code references."""