zenml-nightly 0.58.2.dev20240615__py3-none-any.whl → 0.58.2.dev20240623__py3-none-any.whl

zenml/models/v2/core/logs.py

@@ -13,10 +13,10 @@
 #  permissions and limitations under the License.
 """Models representing logs."""
 
-from typing import Optional, Union
+from typing import Any, Optional, Union
 from uuid import UUID
 
-from pydantic import Field
+from pydantic import Field, field_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.models.v2.base.base import (
@@ -33,15 +33,55 @@ from zenml.models.v2.base.base import (
 class LogsRequest(BaseRequest):
     """Request model for logs."""
 
-    uri: str = Field(
-        title="The uri of the logs file",
-        max_length=TEXT_FIELD_MAX_LENGTH,
-    )
+    uri: str = Field(title="The uri of the logs file")
+
     artifact_store_id: Union[str, UUID] = Field(
         title="The artifact store ID to associate the logs with.",
-        max_length=STR_FIELD_MAX_LENGTH,
+        union_mode="left_to_right",
     )
 
+    @field_validator("uri")
+    @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
+
+    @field_validator("artifact_store_id")
+    @classmethod
+    def str_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)) < STR_FIELD_MAX_LENGTH, (
+            "The length of the value for this field can not "
+            f"exceed {STR_FIELD_MAX_LENGTH}"
+        )
+        return value
+
 
 # ------------------ Update Model ------------------
 
@@ -66,17 +106,37 @@ class LogsResponseMetadata(BaseResponseMetadata):
         title="Step ID to associate the logs with.",
         default=None,
         description="When this is set, pipeline_run_id should be set to None.",
+        union_mode="left_to_right",
     )
     pipeline_run_id: Optional[Union[str, UUID]] = Field(
         title="Pipeline run ID to associate the logs with.",
         default=None,
         description="When this is set, step_run_id should be set to None.",
+        union_mode="left_to_right",
     )
     artifact_store_id: Union[str, UUID] = Field(
         title="The artifact store ID to associate the logs with.",
-        max_length=STR_FIELD_MAX_LENGTH,
+        union_mode="left_to_right",
     )
 
+    @field_validator("artifact_store_id")
+    @classmethod
+    def str_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)) < STR_FIELD_MAX_LENGTH
+        return value
+
 
 class LogsResponseResources(BaseResponseResources):
     """Class for all resource models associated with the Logs entity."""

zenml/models/v2/core/model.py

@@ -117,8 +117,8 @@ class ModelResponseBody(WorkspaceScopedResponseBody):
     tags: List["TagResponse"] = Field(
         title="Tags associated with the model",
     )
-    latest_version_name: Optional[str]
-    latest_version_id: Optional[UUID]
+    latest_version_name: Optional[str] = None
+    latest_version_id: Optional[UUID] = None
 
 
 class ModelResponseMetadata(WorkspaceScopedResponseMetadata):
@@ -322,10 +322,14 @@ class ModelFilter(WorkspaceScopedTaggableFilter):
         description="Name of the Model",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of the Model"
+        default=None,
+        description="Workspace of the Model",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User of the Model"
+        default=None,
+        description="User of the Model",
+        union_mode="left_to_right",
     )
 
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [

zenml/models/v2/core/model_version.py

@@ -16,7 +16,7 @@
 from typing import TYPE_CHECKING, Dict, List, Optional, Type, TypeVar, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field, PrivateAttr, validator
+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
@@ -70,12 +70,14 @@ class ModelVersionRequest(WorkspaceScopedRequest):
 
     number: Optional[int] = Field(
         description="The number of the model version",
+        default=None,
     )
     model: UUID = Field(
         description="The ID of the model containing version",
     )
     tags: Optional[List[str]] = Field(
         title="Tags associated with the model version",
+        default=None,
     )
 
 
@@ -89,7 +91,9 @@ class ModelVersionUpdate(BaseModel):
         description="The ID of the model containing version",
     )
     stage: Optional[Union[str, ModelStages]] = Field(
-        description="Target model version stage to be set", default=None
+        description="Target model version stage to be set",
+        default=None,
+        union_mode="left_to_right",
     )
     force: bool = Field(
         description="Whether existing model version in target stage should be "
@@ -113,7 +117,8 @@ class ModelVersionUpdate(BaseModel):
         default=None,
     )
 
-    @validator("stage")
+    @field_validator("stage")
+    @classmethod
     def _validate_stage(cls, stage: str) -> str:
         stage = getattr(stage, "value", stage)
         if stage is not None and stage not in [
@@ -160,6 +165,14 @@ class ModelVersionResponseBody(WorkspaceScopedResponseBody):
         title="Tags associated with the model version", default=[]
     )
 
+    # 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 ModelVersionResponseMetadata(WorkspaceScopedResponseMetadata):
     """Response metadata for model versions."""
@@ -578,13 +591,19 @@ class ModelVersionFilter(WorkspaceScopedTaggableFilter):
         description="The number of the Model Version",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The workspace of the Model Version"
+        default=None,
+        description="The workspace of the Model Version",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The user of the Model Version"
+        default=None,
+        description="The user of the Model Version",
+        union_mode="left_to_right",
     )
     stage: Optional[Union[str, ModelStages]] = Field(
-        description="The model version stage", default=None
+        description="The model version stage",
+        default=None,
+        union_mode="left_to_right",
     )
 
     _model_id: UUID = PrivateAttr(None)

zenml/models/v2/core/model_version_artifact.py

@@ -13,11 +13,10 @@
 #  permissions and limitations under the License.
 """Models representing the link between model versions and artifacts."""
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 from uuid import UUID
 
-from pydantic import Field, validator
-from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+from pydantic import ConfigDict, Field, model_validator
 
 from zenml.enums import GenericFilterOps
 from zenml.models.v2.base.base import (
@@ -33,6 +32,8 @@ from zenml.models.v2.base.scoped import (
 )
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
 
 
@@ -48,17 +49,22 @@ class ModelVersionArtifactRequest(WorkspaceScopedRequest):
     is_model_artifact: bool = False
     is_deployment_artifact: bool = False
 
-    @validator("is_deployment_artifact")
-    def _validate_is_endpoint_artifact(
-        cls, is_deployment_artifact: bool, values: Dict[str, Any]
-    ) -> bool:
-        is_model_artifact = values.get("is_model_artifact", False)
-        if is_model_artifact and is_deployment_artifact:
+    # 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=())
+
+    @model_validator(mode="after")
+    def _validate_is_endpoint_artifact(self) -> "ModelVersionArtifactRequest":
+        if self.is_model_artifact and self.is_deployment_artifact:
             raise ValueError(
                 "Artifact cannot be a model artifact and deployment artifact "
                 "at the same time."
             )
-        return is_deployment_artifact
+        return self
 
 
 # ------------------ Update Model ------------------
@@ -77,6 +83,14 @@ class ModelVersionArtifactResponseBody(BaseDatedResponseBody):
     is_model_artifact: bool = False
     is_deployment_artifact: bool = False
 
+    # 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 ModelVersionArtifactResponseResources(BaseResponseResources):
     """Class for all resource models associated with the model version artifact entity."""
@@ -168,19 +182,29 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
     ]
 
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The workspace of the Model Version"
+        default=None,
+        description="The workspace of the Model Version",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The user of the Model Version"
+        default=None,
+        description="The user of the Model Version",
+        union_mode="left_to_right",
     )
     model_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by model ID"
+        default=None,
+        description="Filter by model ID",
+        union_mode="left_to_right",
     )
     model_version_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by model version ID"
+        default=None,
+        description="Filter by model version ID",
+        union_mode="left_to_right",
     )
     artifact_version_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by artifact ID"
+        default=None,
+        description="Filter by artifact ID",
+        union_mode="left_to_right",
     )
     artifact_name: Optional[str] = Field(
         default=None,
@@ -191,9 +215,15 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
     only_deployment_artifacts: Optional[bool] = False
     has_custom_name: Optional[bool] = None
 
-    def get_custom_filters(
-        self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+    # 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=())
+
+    def get_custom_filters(self) -> List[Union["ColumnElement[bool]"]]:
         """Get custom filters.
 
         Returns:
@@ -201,7 +231,7 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
         """
         custom_filters = super().get_custom_filters()
 
-        from sqlalchemy import and_
+        from sqlmodel import and_
 
         from zenml.zen_stores.schemas.artifact_schemas import (
             ArtifactSchema,
@@ -218,7 +248,7 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
                 column="name",
                 value=value,
             )
-            artifact_name_filter = and_(  # type: ignore[type-var]
+            artifact_name_filter = and_(
                 ModelVersionArtifactSchema.artifact_version_id
                 == ArtifactVersionSchema.id,
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
@@ -246,7 +276,7 @@ class ModelVersionArtifactFilter(WorkspaceScopedFilter):
             custom_filters.append(deployment_artifact_filter)
 
         if self.has_custom_name is not None:
-            custom_name_filter = and_(  # type: ignore[type-var]
+            custom_name_filter = and_(
                 ModelVersionArtifactSchema.artifact_version_id
                 == ArtifactVersionSchema.id,
                 ArtifactVersionSchema.artifact_id == ArtifactSchema.id,

zenml/models/v2/core/model_version_pipeline_run.py

@@ -13,11 +13,11 @@
 #  permissions and limitations under the License.
 """Models representing the link between model versions and pipeline runs."""
 
-from typing import Any, List, Optional, Union
+from typing import List, Optional, Union
 from uuid import UUID
 
-from pydantic import Field
-from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+from pydantic import ConfigDict, Field
+from sqlalchemy.sql.elements import ColumnElement
 
 from zenml.enums import GenericFilterOps
 from zenml.models.v2.base.base import (
@@ -43,6 +43,14 @@ class ModelVersionPipelineRunRequest(WorkspaceScopedRequest):
     model_version: UUID
     pipeline_run: UUID
 
+    # 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=())
+
 
 # ------------------ Update Model ------------------
 
@@ -58,6 +66,14 @@ class ModelVersionPipelineRunResponseBody(BaseDatedResponseBody):
     model_version: UUID
     pipeline_run: PipelineRunResponse
 
+    # 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 ModelVersionPipelineRunResponseResources(BaseResponseResources):
     """Class for all resource models associated with the model version pipeline run entity."""
@@ -123,28 +139,44 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
     ]
 
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The workspace of the Model Version"
+        default=None,
+        description="The workspace of the Model Version",
+        union_mode="left_to_right",
     )
     user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="The user of the Model Version"
+        default=None,
+        description="The user of the Model Version",
+        union_mode="left_to_right",
     )
     model_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by model ID"
+        default=None,
+        description="Filter by model ID",
+        union_mode="left_to_right",
     )
     model_version_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by model version ID"
+        default=None,
+        description="Filter by model version ID",
+        union_mode="left_to_right",
     )
     pipeline_run_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Filter by pipeline run ID"
+        default=None,
+        description="Filter by pipeline run ID",
+        union_mode="left_to_right",
     )
     pipeline_run_name: Optional[str] = Field(
         default=None,
         description="Name of the pipeline run",
     )
 
-    def get_custom_filters(
-        self,
-    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+    # 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=())
+
+    def get_custom_filters(self) -> List["ColumnElement[bool]"]:
         """Get custom filters.
 
         Returns:
@@ -152,7 +184,7 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
         """
         custom_filters = super().get_custom_filters()
 
-        from sqlalchemy import and_
+        from sqlmodel import and_
 
         from zenml.zen_stores.schemas.model_schemas import (
             ModelVersionPipelineRunSchema,
@@ -170,7 +202,7 @@ class ModelVersionPipelineRunFilter(WorkspaceScopedFilter):
                 column="name",
                 value=value,
             )
-            pipeline_run_name_filter = and_(  # type: ignore[type-var]
+            pipeline_run_name_filter = and_(
                 ModelVersionPipelineRunSchema.pipeline_run_id
                 == PipelineRunSchema.id,
                 filter_.generate_query_conditions(PipelineRunSchema),

zenml/models/v2/core/pipeline.py

@@ -21,27 +21,18 @@ from pydantic import Field
 from zenml.config.pipeline_spec import PipelineSpec
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ExecutionStatus
-from zenml.models import (
-    BaseFilter,
-    BaseResponseMetadata,
+from zenml.models.v2.base.base import BaseUpdate
+from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
+    WorkspaceScopedResponseResources,
 )
-from zenml.models.v2.base.base import (
-    BaseResponse,
-    BaseResponseBody,
-    BaseResponseResources,
-)
-from zenml.models.v2.base.scoped import WorkspaceScopedResponseResources
-from zenml.models.v2.base.update import update_model
 
 if TYPE_CHECKING:
-    from zenml.models.v2.core.pipeline_run import (
-        PipelineRunResponse,
-    )
+    from zenml.models.v2.core.pipeline_run import PipelineRunResponse
 
 
 # ------------------ Request Model ------------------
@@ -65,6 +56,7 @@ class PipelineRequest(WorkspaceScopedRequest):
     docstring: Optional[str] = Field(
         title="The docstring of the pipeline.",
         max_length=TEXT_FIELD_MAX_LENGTH,
+        default=None,
     )
     spec: PipelineSpec = Field(title="The spec of the pipeline.")
 
@@ -72,10 +64,34 @@ class PipelineRequest(WorkspaceScopedRequest):
 # ------------------ Update Model ------------------
 
 
-@update_model
-class PipelineUpdate(PipelineRequest):
+class PipelineUpdate(BaseUpdate):
     """Update model for pipelines."""
 
+    name: Optional[str] = Field(
+        title="The name of the pipeline.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    version: Optional[str] = Field(
+        title="The version of the pipeline.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    version_hash: Optional[str] = Field(
+        title="The version hash of the pipeline.",
+        max_length=STR_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    docstring: Optional[str] = Field(
+        title="The docstring of the pipeline.",
+        max_length=TEXT_FIELD_MAX_LENGTH,
+        default=None,
+    )
+    spec: Optional[PipelineSpec] = Field(
+        title="The spec of the pipeline.",
+        default=None,
+    )
+
 
 # ------------------ Response Model ------------------
 
@@ -103,6 +119,7 @@ class PipelineResponseMetadata(WorkspaceScopedResponseMetadata):
     docstring: Optional[str] = Field(
         title="The docstring of the pipeline.",
         max_length=TEXT_FIELD_MAX_LENGTH,
+        default=None,
     )
 
 
@@ -277,64 +294,12 @@ class PipelineFilter(WorkspaceScopedFilter):
         description="Docstring of the Pipeline",
     )
     workspace_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="Workspace of the Pipeline"
-    )
-    user_id: Optional[Union[UUID, str]] = Field(
-        default=None, description="User of the Pipeline"
-    )
-
-
-# ------------------ Pipeline Namespaces ------------------
-
-
-class PipelineNamespaceResponseBody(BaseResponseBody):
-    """Response body for pipeline namespaces."""
-
-    latest_run_id: Optional[UUID] = Field(
-        default=None,
-        title="The ID of the latest run of the pipeline namespace.",
-    )
-    latest_run_status: Optional[ExecutionStatus] = Field(
         default=None,
-        title="The status of the latest run of the pipeline namespace.",
+        description="Workspace of the Pipeline",
+        union_mode="left_to_right",
     )
-
-
-class PipelineNamespaceResponseMetadata(BaseResponseMetadata):
-    """Response metadata for pipeline namespaces."""
-
-
-class PipelineNamespaceResponseResources(BaseResponseResources):
-    """Class for all resource models associated with the pipeline namespace entity."""
-
-
-class PipelineNamespaceResponse(
-    BaseResponse[
-        PipelineNamespaceResponseBody,
-        PipelineNamespaceResponseMetadata,
-        PipelineNamespaceResponseResources,
-    ]
-):
-    """Response model for pipeline namespaces."""
-
-    name: str = Field(
-        title="The name of the pipeline namespace.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
-
-    def get_hydrated_version(self) -> "PipelineNamespaceResponse":
-        """Get the hydrated version of this pipeline namespace.
-
-        Returns:
-            an instance of the same entity with the metadata field attached.
-        """
-        return self
-
-
-class PipelineNamespaceFilter(BaseFilter):
-    """Pipeline namespace filter model."""
-
-    name: Optional[str] = Field(
+    user_id: Optional[Union[UUID, str]] = Field(
         default=None,
-        description="Name of the pipeline namespace.",
+        description="User of the Pipeline",
+        union_mode="left_to_right",
     )