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

zenml/config/docker_settings.py

@@ -16,11 +16,13 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import Extra, root_validator
+from pydantic import Field, model_validator
+from pydantic_settings import SettingsConfigDict
 
 from zenml.config.base_settings import BaseSettings
 from zenml.logger import get_logger
 from zenml.utils import deprecation_utils
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -70,7 +72,7 @@ class DockerSettings(BaseSettings):
     --------------
     * No `dockerfile` specified: If any of the options regarding
     requirements, environment variables or copying files require us to build an
-    image, ZenML will build this image. Otherwise the `parent_image` will be
+    image, ZenML will build this image. Otherwise, the `parent_image` will be
     used to run the pipeline.
     * `dockerfile` specified: ZenML will first build an image based on the
     specified Dockerfile. If any of the options regarding
@@ -190,8 +192,10 @@ class DockerSettings(BaseSettings):
     python_package_installer_args: Dict[str, Any] = {}
     replicate_local_python_environment: Optional[
         Union[List[str], PythonEnvironmentExportMethod]
-    ] = None
-    requirements: Union[None, str, List[str]] = None
+    ] = Field(default=None, union_mode="left_to_right")
+    requirements: Union[None, str, List[str]] = Field(
+        default=None, union_mode="left_to_right"
+    )
     required_integrations: List[str] = []
     required_hub_plugins: List[str] = []
     install_stack_requirements: bool = True
@@ -208,41 +212,40 @@ class DockerSettings(BaseSettings):
         "copy_files", "copy_global_config"
     )
 
-    @root_validator(pre=True)
-    def _migrate_copy_files(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _migrate_copy_files(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Migrates the value from the old copy_files attribute.
 
         Args:
-            values: The settings values.
+            data: The settings values.
 
         Returns:
             The migrated settings values.
         """
-        copy_files = values.get("copy_files", None)
+        copy_files = data.get("copy_files", None)
 
         if copy_files is None:
-            return values
+            return data
 
-        if values.get("source_files", None):
+        if data.get("source_files", None):
             # Ignore the copy files value in favor of the new source files
             logger.warning(
                 "Both `copy_files` and `source_files` specified for the "
                 "DockerSettings, ignoring the `copy_files` value."
             )
         elif copy_files is True:
-            values["source_files"] = SourceFileMode.INCLUDE
+            data["source_files"] = SourceFileMode.INCLUDE
         elif copy_files is False:
-            values["source_files"] = SourceFileMode.IGNORE
+            data["source_files"] = SourceFileMode.IGNORE
 
-        return values
+        return data
 
-    @root_validator
-    def _validate_skip_build(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def _validate_skip_build(self) -> "DockerSettings":
         """Ensures that a parent image is passed when trying to skip the build.
 
-        Args:
-            values: The settings values.
-
         Returns:
             The validated settings values.
 
@@ -250,10 +253,7 @@ class DockerSettings(BaseSettings):
             ValueError: If the build should be skipped but no parent image
                 was specified.
         """
-        skip_build = values.get("skip_build", False)
-        parent_image = values.get("parent_image")
-
-        if skip_build and not parent_image:
+        if self.skip_build and not self.parent_image:
             raise ValueError(
                 "Docker settings that specify `skip_build=True` must always "
                 "contain a `parent_image`. This parent image will be used "
@@ -261,12 +261,11 @@ class DockerSettings(BaseSettings):
                 "Docker builds on top of it."
             )
 
-        return values
-
-    class Config:
-        """Pydantic configuration class."""
+        return self
 
+    model_config = SettingsConfigDict(
         # public attributes are immutable
-        allow_mutation = False
+        frozen=True,
         # prevent extra attributes during model initialization
-        extra = Extra.forbid
+        extra="forbid",
+    )

zenml/config/global_config.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Functionality to support ZenML GlobalConfiguration."""
 
-import json
 import os
 import uuid
 from pathlib import Path
@@ -21,8 +20,15 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, cast
 from uuid import UUID
 
 from packaging import version
-from pydantic import BaseModel, Field, SecretStr, ValidationError, validator
-from pydantic.main import ModelMetaclass
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    SerializeAsAny,
+    ValidationError,
+    field_validator,
+)
+from pydantic._internal._model_construction import ModelMetaclass
 
 from zenml import __version__
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
@@ -108,17 +114,16 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
         store: Store configuration.
         active_stack_id: The ID of the active stack.
         active_workspace_name: The name of the active workspace.
-        jwt_secret_key: The secret key used to sign and verify JWT tokens.
     """
 
     user_id: uuid.UUID = Field(default_factory=uuid.uuid4)
     user_email: Optional[str] = None
     user_email_opt_in: Optional[bool] = None
     analytics_opt_in: bool = True
-    version: Optional[str]
-    store: Optional[StoreConfiguration]
-    active_stack_id: Optional[uuid.UUID]
-    active_workspace_name: Optional[str]
+    version: Optional[str] = None
+    store: Optional[SerializeAsAny[StoreConfiguration]] = None
+    active_stack_id: Optional[uuid.UUID] = None
+    active_workspace_name: Optional[str] = None
 
     _zen_store: Optional["BaseZenStore"] = None
     _active_workspace: Optional["WorkspaceResponse"] = None
@@ -169,12 +174,13 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
         if config:
             config._write_config()
 
-    @validator("version")
-    def _validate_version(cls, v: Optional[str]) -> Optional[str]:
+    @field_validator("version")
+    @classmethod
+    def _validate_version(cls, value: Optional[str]) -> Optional[str]:
         """Validate the version attribute.
 
         Args:
-            v: The version attribute value.
+            value: The version attribute value.
 
         Returns:
             The version attribute value.
@@ -182,18 +188,18 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
         Raises:
             RuntimeError: If the version parsing fails.
         """
-        if v is None:
-            return v
+        if value is None:
+            return value
 
-        if not isinstance(version.parse(v), version.Version):
+        if not isinstance(version.parse(value), version.Version):
             # If the version parsing fails, it returns a `LegacyVersion`
             # instead. Check to make sure it's an actual `Version` object
             # which represents a valid version.
             raise RuntimeError(
-                f"Invalid version in global configuration: {v}."
+                f"Invalid version in global configuration: {value}."
             )
 
-        return v
+        return value
 
     def __setattr__(self, key: str, value: Any) -> None:
         """Sets an attribute and persists it in the global configuration.
@@ -222,7 +228,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
             The attribute value.
         """
         value = super().__getattribute__(key)
-        if key.startswith("_") or key not in type(self).__fields__:
+        if key.startswith("_") or key not in type(self).model_fields:
             return value
 
         environment_variable_name = f"{CONFIG_ENV_VAR_PREFIX}{key.upper()}"
@@ -321,7 +327,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
             return
 
         config_file = self._config_file
-        yaml_dict = json.loads(self.json(exclude_none=True))
+        yaml_dict = self.model_dump(mode="json", exclude_none=True)
         logger.debug(f"Writing config to {config_file}")
 
         if not fileio.exists(config_file):
@@ -436,7 +442,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
         """
         environment_vars = {}
 
-        for key in self.__fields__.keys():
+        for key in self.model_fields.keys():
             if key == "store":
                 # The store configuration uses its own environment variable
                 # naming scheme
@@ -448,7 +454,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
                     value
                 )
 
-        store_dict = self.store_configuration.dict(exclude_none=True)
+        store_dict = self.store_configuration.model_dump(exclude_none=True)
 
         # The secrets store and backup secrets store configurations use their
         # own environment variables naming scheme
@@ -556,7 +562,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
                 logger.debug(
                     "Using environment variables to update the default store"
                 )
-                store = store.copy(update=env_store_config, deep=True)
+                store = store.model_copy(update=env_store_config, deep=True)
 
         # Step 3: Replace or update the baseline secrets store configuration
         # with the environment variables. This only applies to SQL stores.
@@ -586,7 +592,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
                         "Using environment variables to update the secrets "
                         "store"
                     )
-                    store.secrets_store = store.secrets_store.copy(
+                    store.secrets_store = store.secrets_store.model_copy(
                         update=env_secrets_store_config, deep=True
                     )
 
@@ -605,7 +611,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
                         "secrets store"
                     )
                     store.backup_secrets_store = (
-                        store.backup_secrets_store.copy(
+                        store.backup_secrets_store.model_copy(
                             update=env_backup_secrets_store_config, deep=True
                         )
                     )
@@ -769,21 +775,11 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
 
         return self.active_stack_id
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # Validate attributes when assigning them. We need to set this in order
         # to have a mix of mutable and immutable attributes
-        validate_assignment = True
+        validate_assignment=True,
         # Allow extra attributes from configs of previous ZenML versions to
         # permit downgrading
-        extra = "allow"
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
-
-        # This is needed to allow correct handling of SecretStr values during
-        # serialization.
-        json_encoders = {
-            SecretStr: lambda v: v.get_secret_value() if v else None
-        }
+        extra="allow",
+    )

zenml/config/pipeline_configurations.py

@@ -15,11 +15,11 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
-from pydantic import validator
+from pydantic import SerializeAsAny, field_validator
 
 from zenml.config.constants import DOCKER_SETTINGS_KEY
 from zenml.config.retry_config import StepRetryConfig
-from zenml.config.source import Source, convert_source_validator
+from zenml.config.source import SourceWithValidator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.model.model import Model
 
@@ -38,25 +38,22 @@ class PipelineConfigurationUpdate(StrictBaseModel):
     enable_artifact_metadata: Optional[bool] = None
     enable_artifact_visualization: Optional[bool] = None
     enable_step_logs: Optional[bool] = None
-    settings: Dict[str, BaseSettings] = {}
+    settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
     extra: Dict[str, Any] = {}
-    failure_hook_source: Optional[Source] = None
-    success_hook_source: Optional[Source] = None
+    failure_hook_source: Optional[SourceWithValidator] = None
+    success_hook_source: Optional[SourceWithValidator] = None
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
     retry: Optional[StepRetryConfig] = None
 
-    _convert_source = convert_source_validator(
-        "failure_hook_source", "success_hook_source"
-    )
-
 
 class PipelineConfiguration(PipelineConfigurationUpdate):
     """Pipeline configuration class."""
 
     name: str
 
-    @validator("name")
+    @field_validator("name")
+    @classmethod
     def ensure_pipeline_name_allowed(cls, name: str) -> str:
         """Ensures the pipeline name is allowed.
 
@@ -88,4 +85,4 @@ class PipelineConfiguration(PipelineConfigurationUpdate):
         model_or_dict: SettingsOrDict = self.settings.get(
             DOCKER_SETTINGS_KEY, {}
         )
-        return DockerSettings.parse_obj(model_or_dict)
+        return DockerSettings.model_validate(model_or_dict)

zenml/config/pipeline_run_configuration.py

@@ -16,6 +16,8 @@
 from typing import Any, Dict, Optional, Union
 from uuid import UUID
 
+from pydantic import Field, SerializeAsAny
+
 from zenml.config.base_settings import BaseSettings
 from zenml.config.retry_config import StepRetryConfig
 from zenml.config.schedule import Schedule
@@ -37,9 +39,11 @@ class PipelineRunConfiguration(
     enable_artifact_visualization: Optional[bool] = None
     enable_step_logs: Optional[bool] = None
     schedule: Optional[Schedule] = None
-    build: Union[PipelineBuildBase, UUID, None] = None
+    build: Union[PipelineBuildBase, UUID, None] = Field(
+        default=None, union_mode="left_to_right"
+    )
     steps: Dict[str, StepConfigurationUpdate] = {}
-    settings: Dict[str, BaseSettings] = {}
+    settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
     extra: Dict[str, Any] = {}
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None

zenml/config/pipeline_spec.py

@@ -16,11 +16,10 @@
 import json
 from typing import Any, Dict, List, Optional
 
-from pydantic.json import pydantic_encoder
-
 from zenml.config.source import Source
 from zenml.config.step_configurations import StepSpec
 from zenml.config.strict_base_model import StrictBaseModel
+from zenml.utils.json_utils import pydantic_encoder
 
 
 class PipelineSpec(StrictBaseModel):
@@ -61,13 +60,13 @@ class PipelineSpec(StrictBaseModel):
         """
         from packaging import version
 
-        dict_ = self.dict()
+        dict_ = self.model_dump()
 
         if self.source:
             dict_["source"] = self.source.import_path
 
         for step_dict in dict_["steps"]:
-            step_dict["source"] = Source.parse_obj(
+            step_dict["source"] = Source.model_validate(
                 step_dict["source"]
             ).import_path
 

zenml/config/resource_settings.py

@@ -16,7 +16,8 @@
 from enum import Enum
 from typing import Optional, Union
 
-from pydantic import Extra, Field, NonNegativeInt, PositiveFloat
+from pydantic import Field, NonNegativeInt, PositiveFloat
+from pydantic_settings import SettingsConfigDict
 
 from zenml.config.base_settings import BaseSettings
 
@@ -70,7 +71,7 @@ class ResourceSettings(BaseSettings):
 
     cpu_count: Optional[PositiveFloat] = None
     gpu_count: Optional[NonNegativeInt] = None
-    memory: Optional[str] = Field(regex=MEMORY_REGEX)
+    memory: Optional[str] = Field(pattern=MEMORY_REGEX, default=None)
 
     @property
     def empty(self) -> bool:
@@ -82,7 +83,7 @@ class ResourceSettings(BaseSettings):
         # To detect whether this config is empty (= no values specified), we
         # check if there are any attributes which are explicitly set to any
         # value other than `None`.
-        return len(self.dict(exclude_unset=True, exclude_none=True)) == 0
+        return len(self.model_dump(exclude_unset=True, exclude_none=True)) == 0
 
     def get_memory(
         self, unit: Union[str, ByteUnit] = ByteUnit.GB
@@ -114,11 +115,9 @@ class ResourceSettings(BaseSettings):
             # Should never happen due to the regex validation
             raise ValueError(f"Unable to parse memory unit from '{memory}'.")
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = SettingsConfigDict(
         # public attributes are immutable
-        allow_mutation = False
-
+        frozen=True,
         # prevent extra attributes during model initialization
-        extra = Extra.forbid
+        extra="forbid",
+    )

zenml/config/schedule.py

@@ -14,9 +14,9 @@
 """Class for defining a pipeline schedule."""
 
 import datetime
-from typing import Any, Dict, Optional
+from typing import Optional
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, model_validator
 
 from zenml.logger import get_logger
 
@@ -55,15 +55,10 @@ class Schedule(BaseModel):
     catchup: bool = False
     run_once_start_time: Optional[datetime.datetime] = None
 
-    @root_validator
-    def _ensure_cron_or_periodic_schedule_configured(
-        cls, values: Dict[str, Any]
-    ) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def _ensure_cron_or_periodic_schedule_configured(self) -> "Schedule":
         """Ensures that the cron expression or start time + interval are set.
 
-        Args:
-            values: All attributes of the schedule.
-
         Returns:
             All schedule attributes.
 
@@ -71,13 +66,9 @@ class Schedule(BaseModel):
             ValueError: If no cron expression or start time + interval were
                 provided.
         """
-        cron_expression = values.get("cron_expression")
-        periodic_schedule = values.get("start_time") and values.get(
-            "interval_second"
-        )
-        run_once_starts_at = values.get("run_once_start_time")
+        periodic_schedule = self.start_time and self.interval_second
 
-        if cron_expression and periodic_schedule:
+        if self.cron_expression and periodic_schedule:
             logger.warning(
                 "This schedule was created with a cron expression as well as "
                 "values for `start_time` and `interval_seconds`. The resulting "
@@ -85,17 +76,22 @@ class Schedule(BaseModel):
                 "but will usually ignore the interval and use the cron "
                 "expression."
             )
-            return values
-        elif cron_expression and run_once_starts_at:
+
+            return self
+        elif self.cron_expression and self.run_once_start_time:
             logger.warning(
                 "This schedule was created with a cron expression as well as "
                 "a value for `run_once_start_time`. The resulting behavior "
                 "depends on the concrete orchestrator implementation but will "
                 "usually ignore the `run_once_start_time`."
             )
-            return values
-        elif cron_expression or periodic_schedule or run_once_starts_at:
-            return values
+            return self
+        elif (
+            self.cron_expression
+            or periodic_schedule
+            or self.run_once_start_time
+        ):
+            return self
         else:
             raise ValueError(
                 "Either a cron expression, a start time and interval seconds "

zenml/config/secret_reference_mixin.py

@@ -19,6 +19,7 @@ from pydantic import BaseModel
 
 from zenml.logger import get_logger
 from zenml.utils import secret_utils
+from zenml.utils.pydantic_utils import has_validators
 
 logger = get_logger(__name__)
 
@@ -45,7 +46,7 @@ class SecretReferenceMixin(BaseModel):
         """
         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
@@ -75,7 +76,9 @@ class SecretReferenceMixin(BaseModel):
                     "not allowed."
                 )
 
-            requires_validation = field.pre_validators or field.post_validators
+            requires_validation = has_validators(
+                pydantic_class=self.__class__, field_name=key
+            )
             if requires_validation:
                 raise ValueError(
                     f"Passing the attribute `{key}` as a secret reference is "
@@ -146,6 +149,6 @@ class SecretReferenceMixin(BaseModel):
         """
         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)
         }

zenml/config/secrets_store_config.py

@@ -13,9 +13,9 @@
 #  permissions and limitations under the License.
 """Functionality to support ZenML secrets store configurations."""
 
-from typing import Any, Dict, Optional
+from typing import Optional
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, ConfigDict, model_validator
 
 from zenml.enums import SecretsStoreType
 from zenml.logger import get_logger
@@ -41,45 +41,38 @@ class SecretsStoreConfiguration(BaseModel):
     type: SecretsStoreType
     class_path: Optional[str] = None
 
-    @root_validator
-    def validate_custom(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def validate_custom(self) -> "SecretsStoreConfiguration":
         """Validate that class_path is set for custom secrets stores.
 
-        Args:
-            values: Dict representing user-specified runtime settings.
-
         Returns:
             Validated settings.
 
         Raises:
-            ValueError: If class_path is not set when using an custom secrets
+            ValueError: If class_path is not set when using a custom secrets
                 store.
         """
-        if not values.get("type"):
-            return values
-        if values["type"] == SecretsStoreType.CUSTOM:
-            if values["class_path"] is None:
+        if not self.type:
+            return self
+        if self.type == SecretsStoreType.CUSTOM:
+            if self.class_path is None:
                 raise ValueError(
                     "A class_path must be set when using a custom secrets "
                     "store implementation."
                 )
-        elif values["class_path"] is not None:
+        elif self.class_path is not None:
             raise ValueError(
                 f"The class_path attribute is not supported for the "
-                f"{values['type']} secrets store type."
+                f"{self.type} secrets store type."
             )
 
-        return values
-
-    class Config:
-        """Pydantic configuration class."""
+        return self
 
+    model_config = ConfigDict(
         # Validate attributes when assigning them. We need to set this in order
         # to have a mix of mutable and immutable attributes
-        validate_assignment = True
+        validate_assignment=True,
         # Allow extra attributes to be set in the base class. The concrete
         # classes are responsible for validating the attributes.
-        extra = "allow"
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+        extra="allow",
+    )