zenml-nightly 0.58.2.dev20240614__py3-none-any.whl → 0.58.2.dev20240622__py3-none-any.whl

zenml/models/__init__.py

@@ -20,11 +20,12 @@ from zenml.models.v2.base.base import (
     BaseZenModel,
     BaseRequest,
     BaseResponse,
+    BaseUpdate,
     BaseIdentifiedResponse,
     BaseResponseBody,
     BaseResponseMetadata,
     BaseResponseResources,
-    BaseDatedResponseBody
+    BaseDatedResponseBody,
 )
 from zenml.models.v2.base.scoped import (
     UserScopedRequest,
@@ -37,7 +38,8 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
-    WorkspaceScopedTaggableFilter
+    WorkspaceScopedResponseResources,
+    WorkspaceScopedTaggableFilter,
 )
 from zenml.models.v2.base.filter import (
     BaseFilter,
@@ -82,6 +84,7 @@ from zenml.models.v2.core.artifact_version import (
     ArtifactVersionResponseBody,
     ArtifactVersionResponseMetadata,
     ArtifactVersionUpdate,
+    LazyArtifactVersionResponse
 )
 from zenml.models.v2.core.artifact_visualization import (
     ArtifactVisualizationRequest,
@@ -187,10 +190,12 @@ from zenml.models.v2.core.pipeline import (
     PipelineResponse,
     PipelineResponseBody,
     PipelineResponseMetadata,
-    PipelineNamespaceFilter,
-    PipelineNamespaceResponse,
+)
+from zenml.models.v2.core.pipeline_namespace import (
     PipelineNamespaceResponseBody,
     PipelineNamespaceResponseMetadata,
+    PipelineNamespaceResponse,
+    PipelineNamespaceFilter,
 )
 from zenml.models.v2.core.pipeline_build import (
     PipelineBuildBase,
@@ -219,6 +224,7 @@ from zenml.models.v2.core.pipeline_run import (
 )
 from zenml.models.v2.base.base_plugin_flavor import BasePluginFlavorResponse
 from zenml.models.v2.core.run_metadata import (
+    LazyRunMetadataResponse,
     RunMetadataRequest,
     RunMetadataFilter,
     RunMetadataResponse,
@@ -319,7 +325,7 @@ from zenml.models.v2.core.trigger import (
     TriggerResponse,
     TriggerResponseBody,
     TriggerResponseMetadata,
-    TriggerResponseResources
+    TriggerResponseResources,
 )
 from zenml.models.v2.core.trigger_execution import (
     TriggerExecutionRequest,
@@ -327,7 +333,7 @@ from zenml.models.v2.core.trigger_execution import (
     TriggerExecutionResponse,
     TriggerExecutionResponseBody,
     TriggerExecutionResponseMetadata,
-    TriggerExecutionResponseResources
+    TriggerExecutionResponseResources,
 )
 from zenml.models.v2.core.event_source import (
     EventSourceRequest,
@@ -336,7 +342,7 @@ from zenml.models.v2.core.event_source import (
     EventSourceResponse,
     EventSourceResponseBody,
     EventSourceResponseMetadata,
-    EventSourceResponseResources
+    EventSourceResponseResources,
 )
 from zenml.models.v2.misc.user_auth import UserAuthModel
 from zenml.models.v2.misc.build_item import BuildItem
@@ -375,203 +381,69 @@ from zenml.models.v2.core.server_settings import (
 # ----------------------------- Forward References -----------------------------
 
 # V2
-APIKeyResponseBody.update_forward_refs(
-    ServiceAccountResponse=ServiceAccountResponse,
-)
-ArtifactVersionRequest.update_forward_refs(
-    ArtifactVisualizationRequest=ArtifactVisualizationRequest,
-)
-ArtifactVersionResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-ArtifactVersionResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-    ArtifactVisualizationResponse=ArtifactVisualizationResponse,
-    RunMetadataResponse=RunMetadataResponse,
-)
-CodeReferenceResponseBody.update_forward_refs(
-    CodeRepositoryResponse=CodeRepositoryResponse,
-)
-CodeRepositoryResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-CodeRepositoryResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-ComponentResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-ComponentResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-    ServiceConnectorResponse=ServiceConnectorResponse,
-)
-EventSourceResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-EventSourceResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-EventSourceResponseResources.update_forward_refs(
-    TriggerResponse=TriggerResponse,
-)
-FlavorResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-FlavorResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-ServiceResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-ServiceResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-ServiceResponseResources.update_forward_refs(
-    ModelVersionResponse=ModelVersionResponse,
-)
-ModelResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-    TagResponse=TagResponse,
-)
-ModelResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-ModelVersionResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-    ModelResponse=ModelResponse,
-    RunMetadataResponse=RunMetadataResponse,
-)
-ModelVersionResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-    RunMetadataResponse=RunMetadataResponse,
-)
-ModelVersionResponseResources.update_forward_refs(
-    ServiceResponse=ServiceResponse,
-)
-ModelVersionArtifactResponseBody.update_forward_refs(
-    ArtifactVersionResponse=ArtifactVersionResponse,
-)
-ModelVersionPipelineRunResponseBody.update_forward_refs(
-    PipelineRunResponse=PipelineRunResponse
-)
-OAuthDeviceResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-PipelineResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-PipelineResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-PipelineBuildBase.update_forward_refs(
-    BuildItem=BuildItem,
-)
-PipelineBuildResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-PipelineBuildResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-    PipelineResponse=PipelineResponse,
-    StackResponse=StackResponse,
-    BuildItem=BuildItem,
-)
-PipelineDeploymentRequest.update_forward_refs(
-    CodeReferenceRequest=CodeReferenceRequest,
-)
-PipelineDeploymentResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-PipelineDeploymentResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-    PipelineResponse=PipelineResponse,
-    StackResponse=StackResponse,
-    PipelineBuildResponse=PipelineBuildResponse,
-    ScheduleResponse=ScheduleResponse,
-    CodeReferenceResponse=CodeReferenceResponse,
-)
-PipelineDeploymentResponseResources.update_forward_refs(
-    TriggerResponse=TriggerResponse,
-)
-PipelineRunResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-    PipelineResponse=PipelineResponse,
-    StackResponse=StackResponse,
-    PipelineBuildResponse=PipelineBuildResponse,
-    ScheduleResponse=ScheduleResponse,
-    CodeReferenceResponse=CodeReferenceResponse,
-    TriggerExecutionResponse=TriggerExecutionResponse,
-)
-PipelineRunResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-    RunMetadataResponse=RunMetadataResponse,
-    StepRunResponse=StepRunResponse,
-)
-RunMetadataResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-RunMetadataResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-ScheduleResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-ScheduleResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-SecretResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-SecretResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-ServiceConnectorResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-ServiceConnectorResponseMetadata.update_forward_refs(
-    ServiceConnectorTypeModel=ServiceConnectorTypeModel,
-    WorkspaceResponse=WorkspaceResponse,
-    ComponentResponse=ComponentResponse,
-)
-StackResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-StackResponseMetadata.update_forward_refs(
-    ComponentResponse=ComponentResponse,
-    WorkspaceResponse=WorkspaceResponse,
-)
-StepRunRequest.update_forward_refs(
-    LogsRequest=LogsRequest,
-)
-StepRunResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-    ArtifactVersionResponse=ArtifactVersionResponse,
-)
-StepRunResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-    LogsResponse=LogsResponse,
-    RunMetadataResponse=RunMetadataResponse,
-)
-TriggerExecutionResponseResources.update_forward_refs(
-    TriggerResponse=TriggerResponse
-)
-TriggerResponseBody.update_forward_refs(
-    UserResponse=UserResponse,
-)
-TriggerResponseMetadata.update_forward_refs(
-    WorkspaceResponse=WorkspaceResponse,
-)
-TriggerResponseResources.update_forward_refs(
-    EventSourceResponse=EventSourceResponse,
-)
-TriggerResponseResources.update_forward_refs(
-    EventSourceResponse=EventSourceResponse,
-    TriggerExecutionResponse=TriggerExecutionResponse,
-)
+APIKeyResponseBody.model_rebuild()
+ArtifactVersionRequest.model_rebuild()
+ArtifactVersionResponseBody.model_rebuild()
+ArtifactVersionResponseMetadata.model_rebuild()
+CodeReferenceResponseBody.model_rebuild()
+CodeRepositoryResponseBody.model_rebuild()
+CodeRepositoryResponseMetadata.model_rebuild()
+ComponentResponseBody.model_rebuild()
+ComponentResponseMetadata.model_rebuild()
+EventSourceResponseBody.model_rebuild()
+EventSourceResponseMetadata.model_rebuild()
+EventSourceResponseResources.model_rebuild()
+FlavorResponseBody.model_rebuild()
+FlavorResponseMetadata.model_rebuild()
+LazyArtifactVersionResponse.model_rebuild()
+LazyRunMetadataResponse.model_rebuild()
+ModelResponseBody.model_rebuild()
+ModelResponseMetadata.model_rebuild()
+ModelVersionResponseBody.model_rebuild()
+ModelVersionResponseMetadata.model_rebuild()
+ModelVersionResponseResources.model_rebuild()
+ModelVersionArtifactResponseBody.model_rebuild()
+ModelVersionPipelineRunResponseBody.model_rebuild()
+OAuthDeviceResponseBody.model_rebuild()
+PipelineResponseBody.model_rebuild()
+PipelineResponseMetadata.model_rebuild()
+PipelineBuildBase.model_rebuild()
+PipelineBuildResponseBody.model_rebuild()
+PipelineBuildResponseMetadata.model_rebuild()
+PipelineDeploymentRequest.model_rebuild()
+PipelineDeploymentResponseBody.model_rebuild()
+PipelineDeploymentResponseMetadata.model_rebuild()
+PipelineDeploymentResponseResources.model_rebuild()
+PipelineRunResponseBody.model_rebuild()
+PipelineRunResponseMetadata.model_rebuild()
+RunMetadataResponseBody.model_rebuild()
+RunMetadataResponseMetadata.model_rebuild()
+ScheduleResponseBody.model_rebuild()
+ScheduleResponseMetadata.model_rebuild()
+SecretResponseBody.model_rebuild()
+SecretResponseMetadata.model_rebuild()
+ServiceResponseBody.model_rebuild()
+ServiceResponseMetadata.model_rebuild()
+ServiceResponseResources.model_rebuild()
+ServiceConnectorResponseBody.model_rebuild()
+ServiceConnectorResponseMetadata.model_rebuild()
+StackResponseBody.model_rebuild()
+StackResponseMetadata.model_rebuild()
+StepRunRequest.model_rebuild()
+StepRunResponseBody.model_rebuild()
+StepRunResponseMetadata.model_rebuild()
+TriggerExecutionResponseResources.model_rebuild()
+TriggerResponseBody.model_rebuild()
+TriggerResponseMetadata.model_rebuild()
+TriggerResponseResources.model_rebuild()
+TriggerResponseResources.model_rebuild()
+
 
 __all__ = [
     # V2 Base
     "BaseRequest",
     "BaseResponse",
+    "BaseUpdate",
     "BaseIdentifiedResponse",
     "BaseResponseBody",
     "BaseResponseMetadata",
@@ -589,6 +461,7 @@ __all__ = [
     "WorkspaceScopedResponse",
     "WorkspaceScopedResponseBody",
     "WorkspaceScopedResponseMetadata",
+    "WorkspaceScopedResponseResources",
     "WorkspaceScopedTaggableFilter",
     "BaseFilter",
     "StrFilter",
@@ -732,6 +605,13 @@ __all__ = [
     "SecretResponseBody",
     "SecretResponseMetadata",
     "SecretUpdate",
+    "ServiceResponse",
+    "ServiceResponseBody",
+    "ServiceResponseMetadata",
+    "ServiceUpdate",
+    "ServiceFilter",
+    "ServiceRequest",
+    "ServiceResponseResources",
     "ServerActivationRequest",
     "ServerSettingsResponse",
     "ServerSettingsResponseResources",
@@ -802,13 +682,6 @@ __all__ = [
     "WorkspaceResponse",
     "WorkspaceResponseBody",
     "WorkspaceResponseMetadata",
-    "ServiceResponse",
-    "ServiceResponseBody",
-    "ServiceResponseMetadata",
-    "ServiceUpdate",
-    "ServiceFilter",
-    "ServiceRequest",
-    "ServiceResponseResources",
     # V2 Misc
     "AuthenticationMethodModel",
     "ServiceConnectorResourcesModel",

zenml/models/v2/base/base.py

@@ -17,8 +17,7 @@ from datetime import datetime
 from typing import Any, Dict, Generic, Optional, TypeVar
 from uuid import UUID
 
-from pydantic import Extra, Field, SecretStr
-from pydantic.generics import GenericModel
+from pydantic import ConfigDict, Field
 
 from zenml.analytics.models import AnalyticsTrackedModelMixin
 from zenml.enums import ResponseUpdateStrategy
@@ -35,26 +34,16 @@ class BaseZenModel(YAMLSerializationMixin, AnalyticsTrackedModelMixin):
     """Base model class for all ZenML models.
 
     This class is used as a base class for all ZenML models. It provides
-    functionality for tracking analytics events and proper encoding of
-    SecretStr values.
+    functionality for tracking analytics events.
     """
 
-    class Config:
-        """Pydantic configuration class."""
-
-        # This is needed to allow the REST client and server to unpack SecretStr
-        # values correctly.
-        json_encoders = {
-            SecretStr: lambda v: v.get_secret_value()
-            if v is not None
-            else None
-        }
-
+    model_config = ConfigDict(
         # Ignore extras on all models to support forwards and backwards
         # compatibility (e.g. new fields in newer versions of ZenML servers
         # are allowed to be passed to older versions of ZenML clients and
         # vice versa but will be ignored).
-        extra = "ignore"
+        extra="ignore",
+    )
 
 
 # -------------------- Request Model --------------------
@@ -67,6 +56,21 @@ class BaseRequest(BaseZenModel):
     """
 
 
+# -------------------- Update Model --------------------
+
+
+class BaseUpdate(BaseZenModel):
+    """Base update model.
+
+    Used as a base class for all update models.
+    """
+
+    model_config = ConfigDict(
+        # Ignore extras on all update models.
+        extra="ignore",
+    )
+
+
 # -------------------- Response Model --------------------
 
 
@@ -87,10 +91,7 @@ class BaseResponseResources(BaseZenModel):
     Used as a base class for all resource models associated with responses.
     """
 
-    class Config:
-        """Allows additional resources to be added."""
-
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
 
 AnyBody = TypeVar("AnyBody", bound=BaseResponseBody)
@@ -98,9 +99,7 @@ AnyMetadata = TypeVar("AnyMetadata", bound=BaseResponseMetadata)
 AnyResources = TypeVar("AnyResources", bound=BaseResponseResources)
 
 
-class BaseResponse(
-    GenericModel, Generic[AnyBody, AnyMetadata, AnyResources], BaseZenModel
-):
+class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
     """Base domain model for all responses."""
 
     # Body and metadata pair
@@ -140,7 +139,7 @@ class BaseResponse(
             )
 
         # Check if the name has changed
-        if "name" in self.__fields__:
+        if "name" in self.model_fields:
             original_name = getattr(self, "name")
             hydrated_name = getattr(hydrated_model, "name")
 
@@ -178,7 +177,7 @@ class BaseResponse(
                     )
 
         # Check all the fields in the body
-        for field in self.get_body().__fields__:
+        for field in self.get_body().model_fields:
             original_value = getattr(self.get_body(), field)
             hydrated_value = getattr(hydrated_model.get_body(), field)
 
@@ -253,9 +252,21 @@ class BaseResponse(
         """
         if self.metadata is None:
             # If the metadata is not there, check the class first.
-            metadata_type = self.__fields__["metadata"].type_
+            metadata_annotation = self.model_fields["metadata"].annotation
+            assert metadata_annotation is not None, (
+                "For each response model, an annotated metadata"
+                "field should exist."
+            )
+
+            # metadata is defined as:
+            #   metadata: Optional[....ResponseMetadata] = Field(default=None)
+            # We need to find the actual class inside the Optional annotation.
+            from zenml.utils.typing_utils import get_args
+
+            metadata_type = get_args(metadata_annotation)[0]
+            assert issubclass(metadata_type, BaseResponseMetadata)
 
-            if len(metadata_type.__fields__):
+            if len(metadata_type.model_fields):
                 # If the metadata class defines any fields, fetch the metadata
                 # through the hydrated version.
                 hydrated_version = self.get_hydrated_version()
@@ -281,9 +292,21 @@ class BaseResponse(
         """
         if self.resources is None:
             # If the resources are not there, check the class first.
-            resources_type = self.__fields__["resources"].type_
+            resources_annotation = self.model_fields["resources"].annotation
+            assert resources_annotation is not None, (
+                "For each response model, an annotated resources"
+                "field should exist."
+            )
+
+            # metadata is defined as:
+            #   metadata: Optional[....ResponseMetadata] = Field(default=None)
+            # We need to find the actual class inside the Optional annotation.
+            from zenml.utils.typing_utils import get_args
+
+            resources_type = get_args(resources_annotation)[0]
+            assert issubclass(resources_type, BaseResponseResources)
 
-            if len(resources_type.__fields__):
+            if len(resources_type.model_fields):
                 # If the resources class defines any fields, fetch the resources
                 # through the hydrated version.
                 hydrated_version = self.get_hydrated_version()
@@ -292,7 +315,7 @@ class BaseResponse(
             else:
                 # Otherwise, use the resources class to create an empty
                 # resources object.
-                self.metadata = resources_type()
+                self.resources = resources_type()
 
         if self.resources is None:
             raise RuntimeError(
@@ -327,17 +350,7 @@ class BaseIdentifiedResponse(
     """Base domain model for resources with DB representation."""
 
     id: UUID = Field(title="The unique resource id.")
-    body: Optional["AnyDatedBody"] = Field(
-        title="The body of the resource, "
-        "containing at the minimum "
-        "creation and updated fields."
-    )
-    metadata: Optional["AnyMetadata"] = Field(
-        title="The metadata related to this resource."
-    )
-    resources: Optional["AnyResources"] = Field(
-        title="The resources related to this resource."
-    )
+
     permission_denied: bool = False
 
     # Helper functions

zenml/models/v2/base/base_plugin_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import Generic, TypeVar
 
-from pydantic import Extra, Field
+from pydantic import ConfigDict, Field
 
 from zenml.enums import PluginSubType, PluginType
 from zenml.models.v2.base.base import (
@@ -56,11 +56,7 @@ class BasePluginFlavorResponse(
     name: str = Field(title="Name of the flavor.")
     type: PluginType = Field(title="Type of the plugin.")
     subtype: PluginSubType = Field(title="Subtype of the plugin.")
-
-    class Config:
-        """Configuration for base plugin flavor response."""
-
-        extra = Extra.ignore
+    model_config = ConfigDict(extra="ignore")
 
     def get_hydrated_version(
         self,