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

zenml/config/server_config.py

@@ -19,7 +19,7 @@ from secrets import token_hex
 from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field, SecretStr, root_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from zenml.constants import (
     DEFAULT_ZENML_JWT_TOKEN_ALGORITHM,
@@ -46,6 +46,7 @@ from zenml.constants import (
 from zenml.enums import AuthScheme
 from zenml.logger import get_logger
 from zenml.models import ServerDeploymentType
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -270,30 +271,41 @@ class ServerConfiguration(BaseModel):
     login_rate_limit_minute: int = DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE
     login_rate_limit_day: int = DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY
 
-    secure_headers_server: Union[bool, str] = True
-    secure_headers_hsts: Union[bool, str] = (
-        DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS
+    secure_headers_server: Union[bool, str] = Field(
+        default=True,
+        union_mode="left_to_right",
     )
-    secure_headers_xfo: Union[bool, str] = (
-        DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO
+    secure_headers_hsts: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS,
+        union_mode="left_to_right",
     )
-    secure_headers_xxp: Union[bool, str] = (
-        DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP
+    secure_headers_xfo: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO,
+        union_mode="left_to_right",
     )
-    secure_headers_content: Union[bool, str] = (
-        DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT
+    secure_headers_xxp: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP,
+        union_mode="left_to_right",
     )
-    secure_headers_csp: Union[bool, str] = (
-        DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP
+    secure_headers_content: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT,
+        union_mode="left_to_right",
     )
-    secure_headers_referrer: Union[bool, str] = (
-        DEFAULT_ZENML_SERVER_SECURE_HEADERS_REFERRER
+    secure_headers_csp: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP,
+        union_mode="left_to_right",
     )
-    secure_headers_cache: Union[bool, str] = (
-        DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE
+    secure_headers_referrer: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_REFERRER,
+        union_mode="left_to_right",
     )
-    secure_headers_permissions: Union[bool, str] = (
-        DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS
+    secure_headers_cache: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
+        union_mode="left_to_right",
+    )
+    secure_headers_permissions: Union[bool, str] = Field(
+        default=DEFAULT_ZENML_SERVER_SECURE_HEADERS_PERMISSIONS,
+        union_mode="left_to_right",
     )
     use_legacy_dashboard: bool = DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD
 
@@ -306,12 +318,14 @@ class ServerConfiguration(BaseModel):
 
     _deployment_id: Optional[UUID] = None
 
-    @root_validator(pre=True)
-    def _validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _validate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Validate the server configuration.
 
         Args:
-            values: The server configuration values.
+            data: The server configuration values.
 
         Returns:
             The validated server configuration values.
@@ -319,10 +333,10 @@ class ServerConfiguration(BaseModel):
         Raises:
             ValueError: If the server configuration is invalid.
         """
-        if values.get("auth_scheme") == AuthScheme.EXTERNAL:
+        if data.get("auth_scheme") == AuthScheme.EXTERNAL:
             # If the authentication scheme is set to `EXTERNAL`, the
             # external authenticator URLs must be specified.
-            if not values.get("external_login_url") or not values.get(
+            if not data.get("external_login_url") or not data.get(
                 "external_user_info_url"
             ):
                 raise ValueError(
@@ -333,22 +347,22 @@ class ServerConfiguration(BaseModel):
 
             # If the authentication scheme is set to `EXTERNAL`, the
             # external cookie name must be specified.
-            if not values.get("external_cookie_name"):
+            if not data.get("external_cookie_name"):
                 raise ValueError(
                     "The external cookie name must be specified when "
                     "using the EXTERNAL authentication scheme."
                 )
 
-        if cors_allow_origins := values.get("cors_allow_origins"):
+        if cors_allow_origins := data.get("cors_allow_origins"):
             origins = cors_allow_origins.split(",")
-            values["cors_allow_origins"] = origins
+            data["cors_allow_origins"] = origins
         else:
-            values["cors_allow_origins"] = ["*"]
+            data["cors_allow_origins"] = ["*"]
 
         # if metadata is a string, convert it to a dictionary
-        if isinstance(values.get("metadata"), str):
+        if isinstance(data.get("metadata"), str):
             try:
-                values["metadata"] = json.loads(values["metadata"])
+                data["metadata"] = json.loads(data["metadata"])
             except json.JSONDecodeError as e:
                 raise ValueError(
                     f"The server metadata is not a valid JSON string: {e}"
@@ -356,15 +370,15 @@ class ServerConfiguration(BaseModel):
 
         # if one of the secure headers options is set to a boolean value, set
         # the corresponding value
-        for k, v in values.copy().items():
+        for k, v in data.copy().items():
             if k.startswith("secure_headers_") and isinstance(v, str):
                 if v.lower() in ["disabled", "no", "none", "false", "off", ""]:
-                    values[k] = False
+                    data[k] = False
                 if v.lower() in ["enabled", "yes", "true", "on"]:
                     # Revert to the default value if the header is enabled
-                    del values[k]
+                    del data[k]
 
-        return values
+        return data
 
     @property
     def deployment_id(self) -> UUID:
@@ -490,18 +504,8 @@ class ServerConfiguration(BaseModel):
 
         return ServerConfiguration(**env_server_config)
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # 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/settings_resolver.py

@@ -122,7 +122,7 @@ class SettingsResolver:
         Returns:
             The converted settings.
         """
-        settings_dict = self._settings.dict()
+        settings_dict = dict(self._settings)
         try:
             return target_class(**settings_dict)
         except ValidationError:

zenml/config/source.py

@@ -14,10 +14,17 @@
 """Source classes."""
 
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional
 from uuid import UUID
 
-from pydantic import BaseModel, Extra, validator
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    SerializeAsAny,
+    field_validator,
+)
+from typing_extensions import Annotated
 
 from zenml.logger import get_logger
 
@@ -130,10 +137,29 @@ class Source(BaseModel):
         """
         return self.attribute is None
 
-    class Config:
-        """Pydantic config class."""
+    model_config = ConfigDict(extra="allow")
 
-        extra = Extra.allow
+    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
+        """Dump the source as a dictionary.
+
+        Args:
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The source as a dictionary.
+        """
+        return super().model_dump(serialize_as_any=True, **kwargs)
+
+    def model_dump_json(self, **kwargs: Any) -> str:
+        """Dump the source as a JSON string.
+
+        Args:
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The source as a JSON string.
+        """
+        return super().model_dump_json(serialize_as_any=True, **kwargs)
 
 
 class DistributionPackageSource(Source):
@@ -148,7 +174,8 @@ class DistributionPackageSource(Source):
     version: Optional[str] = None
     type: SourceType = SourceType.DISTRIBUTION_PACKAGE
 
-    @validator("type")
+    @field_validator("type")
+    @classmethod
     def _validate_type(cls, value: SourceType) -> SourceType:
         """Validate the source type.
 
@@ -182,7 +209,8 @@ class CodeRepositorySource(Source):
     subdirectory: str
     type: SourceType = SourceType.CODE_REPOSITORY
 
-    @validator("type")
+    @field_validator("type")
+    @classmethod
     def _validate_type(cls, value: SourceType) -> SourceType:
         """Validate the source type.
 
@@ -201,39 +229,21 @@ class CodeRepositorySource(Source):
         return value
 
 
-def convert_source_validator(*attributes: str) -> "AnyClassMethod":
-    """Function to convert pydantic fields containing legacy class paths.
-
-    In older versions, sources (sometimes also called class paths) like
-    `zenml.materializers.BuiltInMaterializer` were stored as strings in our
-    configuration classes. These strings got replaced by a separate class, and
-    this function returns a validator to convert those old strings to the new
-    classes.
+def convert_source(source: Any) -> Any:
+    """Converts an old source string to a source object.
 
     Args:
-        *attributes: List of attributes to convert.
+        source: Source string or object.
 
     Returns:
-        Pydantic validator class method to be used on BaseModel subclasses
-        to convert source fields.
+        The converted source.
     """
+    if isinstance(source, str):
+        source = Source.from_import_path(source)
 
-    @validator(*attributes, pre=True, allow_reuse=True)
-    def _convert_source(
-        cls: Type[BaseModel], value: Union[Source, str, None]
-    ) -> Optional[Source]:
-        """Converts an old source string to a source object.
-
-        Args:
-            cls: The class on which the attributes are defined.
-            value: Source string or object.
+    return source
 
-        Returns:
-            The converted source.
-        """
-        if isinstance(value, str):
-            value = Source.from_import_path(value)
-
-        return value
 
-    return _convert_source
+SourceWithValidator = Annotated[
+    SerializeAsAny[Source], BeforeValidator(convert_source)
+]

zenml/config/step_configurations.py

@@ -24,7 +24,12 @@ from typing import (
     Union,
 )
 
-from pydantic import root_validator, validator
+from pydantic import (
+    ConfigDict,
+    SerializeAsAny,
+    field_validator,
+    model_validator,
+)
 
 from zenml.artifacts.external_artifact_config import (
     ExternalArtifactConfiguration,
@@ -33,12 +38,13 @@ from zenml.client_lazy_loader import ClientLazyLoader
 from zenml.config.base_settings import BaseSettings, SettingsOrDict
 from zenml.config.constants import DOCKER_SETTINGS_KEY, RESOURCE_SETTINGS_KEY
 from zenml.config.retry_config import StepRetryConfig
-from zenml.config.source import Source, convert_source_validator
+from zenml.config.source import Source, SourceWithValidator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.logger import get_logger
 from zenml.model.lazy_load import ModelVersionDataLazyLoader
 from zenml.model.model import Model
 from zenml.utils import deprecation_utils
+from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.config import DockerSettings, ResourceSettings
@@ -54,25 +60,30 @@ class PartialArtifactConfiguration(StrictBaseModel):
     # for all steps/outputs
     default_materializer_source: Optional[Source] = None
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
     def _remove_deprecated_attributes(
-        cls, values: Dict[str, Any]
+        cls, data: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Removes deprecated attributes from the values dict.
 
         Args:
-            values: The values dict used to instantiate the model.
+            data: The values dict used to instantiate the model.
 
         Returns:
             The values dict without deprecated attributes.
         """
         deprecated_attributes = ["artifact_source"]
+
         for deprecated_attribute in deprecated_attributes:
-            if deprecated_attribute in values:
-                values.pop(deprecated_attribute)
-        return values
+            if deprecated_attribute in data:
+                data.pop(deprecated_attribute)
+
+        return data
 
-    @validator("materializer_source", pre=True)
+    @field_validator("materializer_source", mode="before")
+    @classmethod
     def _convert_source(
         cls,
         value: Union[None, Source, Dict[str, Any], str, Tuple[Source, ...]],
@@ -88,7 +99,7 @@ class PartialArtifactConfiguration(StrictBaseModel):
         if isinstance(value, str):
             value = (Source.from_import_path(value),)
         elif isinstance(value, dict):
-            value = (Source.parse_obj(value),)
+            value = (Source.model_validate(value),)
         elif isinstance(value, Source):
             value = (value,)
 
@@ -100,7 +111,8 @@ class ArtifactConfiguration(PartialArtifactConfiguration):
 
     materializer_source: Tuple[Source, ...]
 
-    @validator("materializer_source", pre=True)
+    @field_validator("materializer_source", mode="before")
+    @classmethod
     def _convert_source(
         cls, value: Union[Source, Dict[str, Any], str, Tuple[Source, ...]]
     ) -> Tuple[Source, ...]:
@@ -115,7 +127,7 @@ class ArtifactConfiguration(PartialArtifactConfiguration):
         if isinstance(value, str):
             value = (Source.from_import_path(value),)
         elif isinstance(value, dict):
-            value = (Source.parse_obj(value),)
+            value = (Source.model_validate(value),)
         elif isinstance(value, Source):
             value = (value,)
 
@@ -133,18 +145,15 @@ class StepConfigurationUpdate(StrictBaseModel):
     step_operator: Optional[str] = None
     experiment_tracker: Optional[str] = None
     parameters: Dict[str, Any] = {}
-    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
     retry: Optional[StepRetryConfig] = None
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 
-    _convert_source = convert_source_validator(
-        "failure_hook_source", "success_hook_source"
-    )
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
         "name"
     )
@@ -160,27 +169,37 @@ class PartialStepConfiguration(StepConfigurationUpdate):
     client_lazy_loaders: Mapping[str, ClientLazyLoader] = {}
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 
+    # 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=())
+
     # Override the deprecation validator as we do not want to deprecate the
     # `name`` attribute on this class.
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes()
 
-    @root_validator(pre=True)
-    def _remove_deprecated_attributes(
-        cls, values: Dict[str, Any]
-    ) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _remove_deprecated_attributes(cls, data: Any) -> Any:
         """Removes deprecated attributes from the values dict.
 
         Args:
-            values: The values dict used to instantiate the model.
+            data: The values dict used to instantiate the model.
 
         Returns:
             The values dict without deprecated attributes.
         """
         deprecated_attributes = ["docstring", "inputs"]
+
         for deprecated_attribute in deprecated_attributes:
-            if deprecated_attribute in values:
-                values.pop(deprecated_attribute)
-        return values
+            if deprecated_attribute in data:
+                data.pop(deprecated_attribute)
+
+        return data
 
 
 class StepConfiguration(PartialStepConfiguration):
@@ -200,7 +219,10 @@ class StepConfiguration(PartialStepConfiguration):
         model_or_dict: SettingsOrDict = self.settings.get(
             RESOURCE_SETTINGS_KEY, {}
         )
-        return ResourceSettings.parse_obj(model_or_dict)
+
+        if isinstance(model_or_dict, BaseSettings):
+            model_or_dict = model_or_dict.model_dump()
+        return ResourceSettings.model_validate(model_or_dict)
 
     @property
     def docker_settings(self) -> "DockerSettings":
@@ -214,7 +236,9 @@ class StepConfiguration(PartialStepConfiguration):
         model_or_dict: SettingsOrDict = self.settings.get(
             DOCKER_SETTINGS_KEY, {}
         )
-        return DockerSettings.parse_obj(model_or_dict)
+        if isinstance(model_or_dict, BaseSettings):
+            model_or_dict = model_or_dict.model_dump()
+        return DockerSettings.model_validate(model_or_dict)
 
 
 class InputSpec(StrictBaseModel):
@@ -227,14 +251,12 @@ class InputSpec(StrictBaseModel):
 class StepSpec(StrictBaseModel):
     """Specification of a pipeline."""
 
-    source: Source
+    source: SourceWithValidator
     upstream_steps: List[str]
     inputs: Dict[str, InputSpec] = {}
     # The default value is to ensure compatibility with specs of version <0.2
     pipeline_parameter_name: str = ""
 
-    _convert_source = convert_source_validator("source")
-
     def __eq__(self, other: Any) -> bool:
         """Returns whether the other object is referring to the same step.
 

zenml/config/step_run_info.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Step run info."""
 
+from typing import Any, Callable
 from uuid import UUID
 
 from zenml.config.pipeline_configurations import PipelineConfiguration
@@ -31,6 +32,8 @@ class StepRunInfo(StrictBaseModel):
     config: StepConfiguration
     pipeline: PipelineConfiguration
 
+    force_write_logs: Callable[..., Any]
+
     def get_image(self, key: str) -> str:
         """Gets the Docker image for the given key.
 

zenml/config/store_config.py

@@ -15,11 +15,12 @@
 
 from typing import Any, Dict, Optional
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, ConfigDict, SerializeAsAny, model_validator
 
 from zenml.config.secrets_store_config import SecretsStoreConfiguration
 from zenml.enums import StoreType
 from zenml.logger import get_logger
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -43,8 +44,10 @@ class StoreConfiguration(BaseModel):
 
     type: StoreType
     url: str
-    secrets_store: Optional[SecretsStoreConfiguration] = None
-    backup_secrets_store: Optional[SecretsStoreConfiguration] = None
+    secrets_store: Optional[SerializeAsAny[SecretsStoreConfiguration]] = None
+    backup_secrets_store: Optional[
+        SerializeAsAny[SecretsStoreConfiguration]
+    ] = None
 
     @classmethod
     def supports_url_scheme(cls, url: str) -> bool:
@@ -61,38 +64,36 @@ class StoreConfiguration(BaseModel):
         """
         return True
 
-    @root_validator(pre=True)
-    def validate_secrets_store(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def validate_store_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Validate the secrets store configuration.
 
         Args:
-            values: The values of the store configuration.
+            data: The values of the store configuration.
 
         Returns:
             The values of the store configuration.
         """
-        if values.get("secrets_store") is None:
-            return values
+        if data.get("secrets_store") is None:
+            return data
 
         # Remove the legacy REST secrets store configuration since it is no
         # longer supported/needed
-        secrets_store = values["secrets_store"]
+        secrets_store = data["secrets_store"]
         if isinstance(secrets_store, dict):
             secrets_store_type = secrets_store.get("type")
             if secrets_store_type == "rest":
-                del values["secrets_store"]
-
-        return values
+                del data["secrets_store"]
 
-    class Config:
-        """Pydantic configuration class."""
+        return data
 
+    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",
+    )

zenml/config/strict_base_model.py

@@ -13,14 +13,10 @@
 #  permissions and limitations under the License.
 """Strict immutable pydantic model."""
 
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, ConfigDict
 
 
 class StrictBaseModel(BaseModel):
     """Immutable pydantic model which prevents extra attributes."""
 
-    class Config:
-        """Pydantic config class."""
-
-        allow_mutation = False
-        extra = Extra.forbid
+    model_config = ConfigDict(frozen=True, extra="forbid")

zenml/constants.py

@@ -69,6 +69,30 @@ def handle_json_env_var(
         return default
 
 
+def is_true_string_value(value: Any) -> bool:
+    """Checks if the given value is a string representation of 'True'.
+
+    Args:
+        value: the value to check.
+
+    Returns:
+        Whether the input value represents a string version of 'True'.
+    """
+    return value in ["1", "y", "yes", "True", "true"]
+
+
+def is_false_string_value(value: Any) -> bool:
+    """Checks if the given value is a string representation of 'False'.
+
+    Args:
+        value: the value to check.
+
+    Returns:
+        Whether the input value represents a string version of 'False'.
+    """
+    return value in ["0", "n", "no", "False", "false"]
+
+
 def handle_bool_env_var(var: str, default: bool = False) -> bool:
     """Converts normal env var to boolean.
 
@@ -80,9 +104,9 @@ def handle_bool_env_var(var: str, default: bool = False) -> bool:
         The converted value.
     """
     value = os.getenv(var)
-    if value in ["1", "y", "yes", "True", "true"]:
+    if is_true_string_value(value):
         return True
-    elif value in ["0", "n", "no", "False", "false"]:
+    elif is_false_string_value(value):
         return False
     return default