zenml-nightly 0.58.2.dev20240618__py3-none-any.whl → 0.58.2.dev20240619__py3-none-any.whl

zenml/orchestrators/step_runner.py

@@ -28,8 +28,6 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic.typing import get_origin, is_union
-
 from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
 from zenml.artifacts.utils import save_artifact
 from zenml.config.step_configurations import StepConfiguration
@@ -63,6 +61,7 @@ from zenml.steps.utils import (
     resolve_type_annotation,
 )
 from zenml.utils import materializer_utils, source_utils
+from zenml.utils.typing_utils import get_origin, is_union
 
 if TYPE_CHECKING:
     from zenml.config.source import Source
@@ -396,7 +395,7 @@ class StepRunner:
 
             # Parse the parameters
             if issubclass(arg_type, BaseParameters):
-                step_params = arg_type.parse_obj(
+                step_params = arg_type.model_validate(
                     self.configuration.parameters[arg]
                 )
                 function_params[arg] = step_params
@@ -440,7 +439,7 @@ class StepRunner:
         # Skip materialization for `UnmaterializedArtifact`.
         if data_type == UnmaterializedArtifact:
             return UnmaterializedArtifact(
-                **artifact.get_hydrated_version().dict()
+                **artifact.get_hydrated_version().model_dump()
             )
 
         if data_type is Any or is_union(get_origin(data_type)):
@@ -513,8 +512,6 @@ class StepRunner:
                 f"(return values: {return_values})."
             )
 
-        from pydantic.typing import get_origin, is_union
-
         from zenml.steps.utils import get_args
 
         validated_outputs: Dict[str, Any] = {}

zenml/orchestrators/utils.py

@@ -225,7 +225,7 @@ def _link_cached_artifacts_to_model(
         if output_name_ in output_annotations:
             annotation = output_annotations.get(output_name_, None)
             if annotation and annotation.artifact_config is not None:
-                artifact_config_ = annotation.artifact_config.copy()
+                artifact_config_ = annotation.artifact_config.model_copy()
         # no artifact config found or artifact was produced by `save_artifact`
         # inside the step body, so was never in annotations
         if artifact_config_ is None:

zenml/plugins/base_plugin_flavor.py

@@ -16,7 +16,7 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, ClassVar, Type
 
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, ConfigDict
 
 from zenml.config.global_config import GlobalConfiguration
 from zenml.enums import PluginSubType, PluginType
@@ -29,16 +29,12 @@ if TYPE_CHECKING:
 class BasePluginConfig(BaseModel, ABC):
     """Allows configuring of Event Source and Filter configuration."""
 
-    class Config:
-        """Pydantic configuration class."""
-
-        # public attributes are immutable
-        allow_mutation = True
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+    model_config = ConfigDict(
+        # public attributes are mutable
+        frozen=False,
         # ignore extra attributes during model initialization
-        extra = Extra.ignore
+        extra="ignore",
+    )
 
 
 class BasePlugin(ABC):

zenml/plugins/plugin_flavor_registry.py

@@ -16,7 +16,7 @@
 import math
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Type
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.enums import PluginSubType, PluginType
 from zenml.integrations.registry import integration_registry
@@ -35,12 +35,8 @@ class RegistryEntry(BaseModel):
     """Registry Entry Class for the Plugin Registry."""
 
     flavor_class: Type[BasePluginFlavor]
-    plugin_instance: Optional[BasePlugin]
-
-    class Config:
-        """Pydantic configuration class."""
-
-        arbitrary_types_allowed = True
+    plugin_instance: Optional[BasePlugin] = None
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
 class PluginFlavorRegistry:

zenml/secret/base_secret.py

@@ -15,7 +15,7 @@
 
 from typing import Any, Dict, List
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class BaseSecretSchema(BaseModel):
@@ -31,7 +31,7 @@ class BaseSecretSchema(BaseModel):
         Returns:
             A list of all attribute names that are part of the schema.
         """
-        return list(cls.__fields__.keys())
+        return list(cls.model_fields.keys())
 
     def get_values(self) -> Dict[str, Any]:
         """Get all values of the secret schema.
@@ -39,12 +39,11 @@ class BaseSecretSchema(BaseModel):
         Returns:
             A dictionary of all attribute names and their corresponding values.
         """
-        return self.dict(exclude_none=True)
-
-    class Config:
-        """Pydantic configuration class."""
+        return self.model_dump(exclude_none=True)
 
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
+        validate_assignment=True,
         # report extra attributes as validation failures
-        extra = "forbid"
+        extra="forbid",
+    )

zenml/service_connectors/docker_service_connector.py

@@ -23,7 +23,7 @@ from typing import Any, List, Optional
 
 from docker.client import DockerClient
 from docker.errors import DockerException
-from pydantic import Field, SecretStr
+from pydantic import Field
 
 from zenml.constants import DOCKER_REGISTRY_RESOURCE_TYPE
 from zenml.exceptions import AuthorizationException
@@ -39,6 +39,7 @@ from zenml.service_connectors.service_connector import (
 )
 from zenml.utils import docker_utils
 from zenml.utils.enum_utils import StrEnum
+from zenml.utils.secret_utils import PlainSerializedSecretStr
 
 logger = get_logger(__name__)
 
@@ -46,10 +47,10 @@ logger = get_logger(__name__)
 class DockerCredentials(AuthenticationConfig):
     """Docker client authentication credentials."""
 
-    username: SecretStr = Field(
+    username: PlainSerializedSecretStr = Field(
         title="Username",
     )
-    password: SecretStr = Field(
+    password: PlainSerializedSecretStr = Field(
         title="Password",
     )
 

zenml/service_connectors/service_connector.py

@@ -34,7 +34,7 @@ from pydantic import (
     SecretStr,
     ValidationError,
 )
-from pydantic.main import ModelMetaclass
+from pydantic._internal._model_construction import ModelMetaclass
 
 from zenml.client import Client
 from zenml.constants import (
@@ -71,7 +71,7 @@ class AuthenticationConfig(BaseModel):
         """
         return {
             k: v
-            for k, v in self.dict(exclude_none=True).items()
+            for k, v in self.model_dump(exclude_none=True).items()
             if isinstance(v, SecretStr)
         }
 
@@ -84,7 +84,7 @@ class AuthenticationConfig(BaseModel):
         """
         return {
             k: v
-            for k, v in self.dict(exclude_none=True).items()
+            for k, v in self.model_dump(exclude_none=True).items()
             if not isinstance(v, SecretStr)
         }
 
@@ -95,7 +95,7 @@ class AuthenticationConfig(BaseModel):
         Returns:
             A dictionary of all values in the configuration.
         """
-        return self.dict(exclude_none=True)
+        return self.model_dump(exclude_none=True)
 
 
 class ServiceConnectorMeta(ModelMetaclass):
@@ -487,7 +487,7 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
         ):
             return self
 
-        copy = self.copy()
+        copy = self.model_copy()
         copy.resource_type = resource_type
         copy.resource_id = resource_id
         return copy
@@ -1401,10 +1401,3 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
         )
 
         return connector_client
-
-    class Config:
-        """Connector configuration."""
-
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True

zenml/service_connectors/service_connector_registry.py

@@ -22,7 +22,6 @@ if TYPE_CHECKING:
     from zenml.models import (
         ServiceConnectorRequest,
         ServiceConnectorResponse,
-        ServiceConnectorUpdate,
     )
     from zenml.service_connectors.service_connector import ServiceConnector
 logger = get_logger(__name__)
@@ -90,7 +89,7 @@ class ServiceConnectorRegistry:
                 f"Please make sure the corresponding packages and/or ZenML "
                 f"integration are installed and try again."
             )
-        return self.service_connector_types[connector_type].copy()
+        return self.service_connector_types[connector_type].model_copy()
 
     def __getitem__(self, key: str) -> ServiceConnectorTypeModel:
         """Get a service connector type by its connector type identifier.
@@ -155,7 +154,7 @@ class ServiceConnectorRegistry:
                     or auth_method in service_connector_type.auth_method_dict
                 )
             ):
-                matches.append(service_connector_type.copy())
+                matches.append(service_connector_type.model_copy())
 
         return matches
 
@@ -164,7 +163,6 @@ class ServiceConnectorRegistry:
         model: Union[
             "ServiceConnectorRequest",
             "ServiceConnectorResponse",
-            "ServiceConnectorUpdate",
         ],
     ) -> "ServiceConnector":
         """Validate a service connector model and create an instance from it.

zenml/services/container/container_service.py

@@ -289,7 +289,7 @@ class ContainerService(BaseService):
         assert self.status.log_file is not None
 
         with open(self.status.config_file, "w") as f:
-            f.write(self.json(indent=4))
+            f.write(self.model_dump_json(indent=4))
         pathlib.Path(self.status.log_file).touch()
 
         command = [

zenml/services/container/container_service_endpoint.py

@@ -80,7 +80,7 @@ class ContainerServiceEndpoint(BaseServiceEndpoint):
     )
     monitor: Optional[
         Union[HTTPEndpointHealthMonitor, TCPEndpointHealthMonitor]
-    ] = Field(..., discriminator="type")
+    ] = Field(..., discriminator="type", union_mode="left_to_right")
 
     def _lookup_free_port(self) -> int:
         """Search for a free TCP port for the service endpoint.

zenml/services/local/local_service.py

@@ -331,7 +331,7 @@ class LocalDaemonService(BaseService):
         assert self.status.pid_file is not None
 
         with open(self.status.config_file, "w") as f:
-            f.write(self.json(indent=4))
+            f.write(self.model_dump_json(indent=4))
 
         # delete the previous PID file, in case a previous daemon process
         # crashed and left a stale PID file

zenml/services/local/local_service_endpoint.py

@@ -83,7 +83,7 @@ class LocalDaemonServiceEndpoint(BaseServiceEndpoint):
     )
     monitor: Optional[
         Union[HTTPEndpointHealthMonitor, TCPEndpointHealthMonitor]
-    ] = Field(..., discriminator="type")
+    ] = Field(..., discriminator="type", union_mode="left_to_right")
 
     def _lookup_free_port(self) -> int:
         """Search for a free TCP port for the service endpoint.

zenml/services/service.py

@@ -31,6 +31,8 @@ from typing import (
 )
 from uuid import UUID
 
+from pydantic import ConfigDict
+
 from zenml.console import console
 from zenml.logger import get_logger
 from zenml.services.service_endpoint import BaseServiceEndpoint
@@ -117,6 +119,14 @@ class ServiceConfig(BaseTypedModel):
     model_version: str = ""
     service_name: str = ""
 
+    # 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 __init__(self, **data: Any):
         """Initialize the service configuration.
 
@@ -142,7 +152,7 @@ class ServiceConfig(BaseTypedModel):
             a dictionary of service labels.
         """
         labels = {}
-        for k, v in self.dict().items():
+        for k, v in self.model_dump().items():
             label = f"zenml_{k}".upper()
             labels[label] = str(v)
         return labels
@@ -172,7 +182,7 @@ class BaseService(BaseTypedModel):
     config: ServiceConfig
     status: ServiceStatus
     # TODO [ENG-703]: allow multiple endpoints per service
-    endpoint: Optional[BaseServiceEndpoint]
+    endpoint: Optional[BaseServiceEndpoint] = None
 
     def __init__(
         self,
@@ -209,7 +219,7 @@ class BaseService(BaseTypedModel):
             admin_state=model.admin_state,
             config=model.config,
             status=model.status,
-            service_type=model.service_type.dict(),
+            service_type=model.service_type.model_dump(),
             endpoint=model.endpoint,
         )
 
@@ -512,14 +522,10 @@ class BaseService(BaseTypedModel):
         """
         return self.__repr__()
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # validate attribute assignments
-        validate_assignment = True
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+        validate_assignment=True,
+    )
 
 
 class BaseDeploymentService(BaseService):

zenml/services/service_type.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of a ZenML ServiceType class."""
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class ServiceType(BaseModel):
@@ -33,8 +33,7 @@ class ServiceType(BaseModel):
     description: str = ""
     logo_url: str = ""
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # make the service type immutable and hashable
-        frozen = True
+        frozen=True
+    )

zenml/services/terraform/terraform_service.py

@@ -176,7 +176,7 @@ class TerraformService(BaseService):
         assert self.status.config_file is not None
 
         with open(self.status.config_file, "w") as f:
-            f.write(self.json(indent=4))
+            f.write(self.model_dump_json(indent=4))
 
     def _write_vars_to_file(self, vars: Dict[str, Any]) -> None:
         """Write variables to the variables file.

zenml/stack/flavor.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Base ZenML Flavor implementation."""
 
-import json
 from abc import abstractmethod
 from typing import Any, Dict, Optional, Type, cast
 
@@ -102,10 +101,7 @@ class Flavor:
         Returns:
             The config schema.
         """
-        config_schema: Dict[str, Any] = json.loads(
-            self.config_class.schema_json()
-        )
-        return config_schema
+        return self.config_class.model_json_schema()
 
     @property
     def service_connector_requirements(

zenml/stack/flavor_registry.py

@@ -122,8 +122,8 @@ class FlavorRegistry:
                 if len(existing_flavor) == 0:
                     store.create_flavor(flavor_request_model)
                 else:
-                    flavor_update_model = FlavorUpdate.parse_obj(
-                        flavor_request_model
+                    flavor_update_model = FlavorUpdate.model_validate(
+                        dict(flavor_request_model)
                     )
                     store.update_flavor(
                         existing_flavor[0].id, flavor_update_model
@@ -154,8 +154,8 @@ class FlavorRegistry:
                         if len(existing_flavor) == 0:
                             store.create_flavor(flavor_request_model)
                         else:
-                            flavor_update_model = FlavorUpdate.parse_obj(
-                                flavor_request_model
+                            flavor_update_model = FlavorUpdate.model_validate(
+                                dict(flavor_request_model)
                             )
                             store.update_flavor(
                                 existing_flavor[0].id, flavor_update_model

zenml/stack/stack.py

@@ -15,6 +15,7 @@
 
 import functools
 import itertools
+import json
 import os
 from datetime import datetime
 from typing import (
@@ -481,7 +482,9 @@ class Stack:
             A dictionary containing the stack components.
         """
         component_dict = {
-            component_type.value: component.config.json(sort_keys=True)
+            component_type.value: json.dumps(
+                component.config.model_dump(mode="json"), sort_keys=True
+            )
             for component_type, component in self.components.items()
         }
         component_dict.update({"name": self.name})

zenml/stack/stack_component.py

@@ -21,8 +21,7 @@ from inspect import isclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Type, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Extra, root_validator
-from pydantic.typing import get_origin
+from pydantic import BaseModel, ConfigDict, model_validator
 
 from zenml.config.build_configuration import BuildConfiguration
 from zenml.config.step_configurations import Step
@@ -31,7 +30,12 @@ from zenml.enums import StackComponentType
 from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.models import ServiceConnectorRequirements, StepRunResponse
-from zenml.utils import secret_utils, settings_utils
+from zenml.utils import (
+    pydantic_utils,
+    secret_utils,
+    settings_utils,
+    typing_utils,
+)
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import BaseSettings
@@ -65,7 +69,8 @@ class StackComponentConfig(BaseModel, ABC):
         custom pydantic validation are set as secret references.
 
         Args:
-            warn_about_plain_text_secrets: If true, then warns about using plain-text secrets.
+            warn_about_plain_text_secrets: If true, then warns about using
+                plain-text secrets.
             **kwargs: Arguments to initialize this stack component.
 
         Raises:
@@ -75,7 +80,7 @@ class StackComponentConfig(BaseModel, ABC):
         """
         for key, value in kwargs.items():
             try:
-                field = self.__class__.__fields__[key]
+                field = self.__class__.model_fields[key]
             except KeyError:
                 # Value for a private attribute or non-existing field, this
                 # will fail during the upcoming pydantic validation
@@ -101,8 +106,9 @@ class StackComponentConfig(BaseModel, ABC):
                     )
                 continue
 
-            requires_validation = field.pre_validators or field.post_validators
-            if requires_validation:
+            if pydantic_utils.has_validators(
+                pydantic_class=self.__class__, field_name=key
+            ):
                 raise ValueError(
                     f"Passing the stack component attribute `{key}` as a "
                     "secret reference is not allowed as additional validation "
@@ -120,7 +126,7 @@ class StackComponentConfig(BaseModel, ABC):
         """
         return {
             secret_utils.parse_secret_reference(v)
-            for v in self.dict().values()
+            for v in self.model_dump().values()
             if secret_utils.is_secret_reference(v)
         }
 
@@ -246,52 +252,70 @@ class StackComponentConfig(BaseModel, ABC):
         # (see https://github.com/python/mypy/issues/13319).
         __getattribute__ = __custom_getattribute__
 
-    @root_validator(pre=True)
-    def _convert_json_strings(cls, values: Dict[str, Any]) -> Any:
+    @model_validator(mode="before")
+    @classmethod
+    @pydantic_utils.before_validator_handler
+    def _convert_json_strings(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Converts potential JSON strings.
 
         Args:
-            values: The model values.
+            data: The model data.
 
         Returns:
-            The potentially converted values.
+            The potentially converted data.
 
         Raises:
             ValueError: If any of the values is an invalid JSON string.
         """
-        for key, field in cls.__fields__.items():
-            value = values.get(key, None)
+        for key, field in cls.model_fields.items():
+            if not field.annotation:
+                continue
+
+            value = data.get(key, None)
 
             if isinstance(value, str):
-                if get_origin(field.outer_type_) in {
+                if typing_utils.is_optional(field.annotation):
+                    args = list(typing_utils.get_args(field.annotation))
+                    if str in args:
+                        # Don't do any type coercion in case str is in the
+                        # possible types of the field
+                        continue
+
+                    # Remove `NoneType` from the arguments
+                    NoneType = type(None)
+                    if NoneType in args:
+                        args.remove(NoneType)
+
+                    # We just choose the first arg and match against this
+                    annotation = args[0]
+                else:
+                    annotation = field.annotation
+
+                if typing_utils.get_origin(annotation) in {
                     dict,
                     list,
                     Mapping,
                     Sequence,
                 }:
                     try:
-                        values[key] = json.loads(value)
+                        data[key] = json.loads(value)
                     except json.JSONDecodeError as e:
                         raise ValueError(
                             f"Invalid json string '{value}'"
                         ) from e
-                elif isclass(field.type_) and issubclass(
-                    field.type_, BaseModel
-                ):
-                    values[key] = field.type_.parse_raw(value).dict()
-
-        return values
+                elif isclass(annotation) and issubclass(annotation, BaseModel):
+                    data[key] = annotation.model_validate_json(
+                        value
+                    ).model_dump()
 
-    class Config:
-        """Pydantic configuration class."""
+        return data
 
+    model_config = ConfigDict(
         # public attributes are immutable
-        allow_mutation = False
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+        frozen=True,
         # prevent extra attributes during model initialization
-        extra = Extra.forbid
+        extra="forbid",
+    )
 
 
 class StackComponent:
@@ -480,7 +504,7 @@ class StackComponent:
         )
 
         if key in all_settings:
-            return self.settings_class.parse_obj(all_settings[key])
+            return self.settings_class.model_validate(dict(all_settings[key]))
         else:
             return self.settings_class()
 
@@ -810,7 +834,7 @@ class StackComponent:
             A string representation of the stack component.
         """
         attribute_representation = ", ".join(
-            f"{key}={value}" for key, value in self.config.dict().items()
+            f"{key}={value}" for key, value in self.config.model_dump().items()
         )
         return (
             f"{self.__class__.__qualname__}(type={self.type}, "