zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/config/source.py

@@ -12,11 +12,19 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """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
 
@@ -34,6 +42,7 @@ class SourceType(Enum):
     INTERNAL = "internal"
     DISTRIBUTION_PACKAGE = "distribution_package"
     CODE_REPOSITORY = "code_repository"
+    NOTEBOOK = "notebook"
     UNKNOWN = "unknown"
 
 
@@ -129,10 +138,29 @@ class Source(BaseModel):
         """
         return self.attribute is None
 
-    class Config:
-        """Pydantic config class."""
+    model_config = ConfigDict(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.
 
-        extra = Extra.allow
+        Returns:
+            The source as a JSON string.
+        """
+        return super().model_dump_json(serialize_as_any=True, **kwargs)
 
 
 class DistributionPackageSource(Source):
@@ -147,7 +175,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.
 
@@ -181,7 +210,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.
 
@@ -200,39 +230,75 @@ class CodeRepositorySource(Source):
         return value
 
 
-def convert_source_validator(*attributes: str) -> "AnyClassMethod":
-    """Function to convert pydantic fields containing legacy class paths.
+class NotebookSource(Source):
+    """Source representing an object defined in a notebook.
 
-    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.
+    Attributes:
+        replacement_module: Name of the module from which this source should
+            be loaded in case the code is not running in a notebook.
+        artifact_store_id: ID of the artifact store in which the replacement
+            module code is stored.
+    """
 
-    Args:
-        *attributes: List of attributes to convert.
+    replacement_module: Optional[str] = None
+    artifact_store_id: Optional[UUID] = None
+    type: SourceType = SourceType.NOTEBOOK
 
-    Returns:
-        Pydantic validator class method to be used on BaseModel subclasses
-        to convert source fields.
-    """
+    @field_validator("type")
+    @classmethod
+    def _validate_type(cls, value: SourceType) -> SourceType:
+        """Validate the source type.
+
+        Args:
+            value: The source type.
+
+        Raises:
+            ValueError: If the source type is not `NOTEBOOK`.
+
+        Returns:
+            The source type.
+        """
+        if value != SourceType.NOTEBOOK:
+            raise ValueError("Invalid source type.")
 
-    @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.
+        return value
+
+    @field_validator("module")
+    @classmethod
+    def _validate_module(cls, value: str) -> str:
+        """Validate the module.
 
         Args:
-            cls: The class on which the attributes are defined.
-            value: Source string or object.
+            value: The module.
+
+        Raises:
+            ValueError: If the module is not `__main__`.
 
         Returns:
-            The converted source.
+            The module.
         """
-        if isinstance(value, str):
-            value = Source.from_import_path(value)
+        if value != "__main__":
+            raise ValueError("Invalid module for notebook source.")
 
         return value
 
-    return _convert_source
+
+def convert_source(source: Any) -> Any:
+    """Converts an old source string to a source object.
+
+    Args:
+        source: Source string or object.
+
+    Returns:
+        The converted source.
+    """
+    if isinstance(source, str):
+        source = Source.from_import_path(source)
+
+    return source
+
+
+SourceWithValidator = Annotated[
+    SerializeAsAny[Source],
+    BeforeValidator(convert_source),
+]

zenml/config/step_configurations.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Pipeline configuration classes."""
 
+from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -24,22 +25,32 @@ from typing import (
     Union,
 )
 
-from pydantic import root_validator, validator
+from pydantic import (
+    ConfigDict,
+    SerializeAsAny,
+    field_validator,
+    model_validator,
+)
 
+from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.artifacts.external_artifact_config import (
     ExternalArtifactConfiguration,
 )
+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.source import Source, convert_source_validator
+from zenml.config.retry_config import StepRetryConfig
+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
+    from zenml.config.pipeline_configurations import PipelineConfiguration
 
 logger = get_logger(__name__)
 
@@ -47,30 +58,36 @@ logger = get_logger(__name__)
 class PartialArtifactConfiguration(StrictBaseModel):
     """Class representing a partial input/output artifact configuration."""
 
-    materializer_source: Optional[Tuple[Source, ...]] = None
+    materializer_source: Optional[Tuple[SourceWithValidator, ...]] = None
     # TODO: This could be moved to the `PipelineDeployment` as it's the same
     # for all steps/outputs
-    default_materializer_source: Optional[Source] = None
+    default_materializer_source: Optional[SourceWithValidator] = None
+    artifact_config: Optional[ArtifactConfig] = 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, ...]],
@@ -86,7 +103,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,)
 
@@ -96,9 +113,10 @@ class PartialArtifactConfiguration(StrictBaseModel):
 class ArtifactConfiguration(PartialArtifactConfiguration):
     """Class representing a complete input/output artifact configuration."""
 
-    materializer_source: Tuple[Source, ...]
+    materializer_source: Tuple[SourceWithValidator, ...]
 
-    @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, ...]:
@@ -113,7 +131,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,)
 
@@ -123,7 +141,6 @@ class ArtifactConfiguration(PartialArtifactConfiguration):
 class StepConfigurationUpdate(StrictBaseModel):
     """Class for step configuration updates."""
 
-    name: Optional[str] = None
     enable_cache: Optional[bool] = None
     enable_artifact_metadata: Optional[bool] = None
     enable_artifact_visualization: Optional[bool] = None
@@ -131,21 +148,16 @@ 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
+    substitutions: Dict[str, str] = {}
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 
-    _convert_source = convert_source_validator(
-        "failure_hook_source", "success_hook_source"
-    )
-    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
-        "name"
-    )
-
 
 class PartialStepConfiguration(StepConfigurationUpdate):
     """Class representing a partial step configuration."""
@@ -154,29 +166,40 @@ class PartialStepConfiguration(StepConfigurationUpdate):
     caching_parameters: Mapping[str, Any] = {}
     external_input_artifacts: Mapping[str, ExternalArtifactConfiguration] = {}
     model_artifacts_or_metadata: Mapping[str, ModelVersionDataLazyLoader] = {}
+    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):
@@ -196,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":
@@ -210,7 +236,27 @@ 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)
+
+    def _get_full_substitutions(
+        self,
+        pipeline_config: "PipelineConfiguration",
+        start_time: Optional[datetime],
+    ) -> Dict[str, str]:
+        """Get the full set of substitutions for this step configuration.
+
+        Args:
+            pipeline_config: The pipeline configuration.
+            start_time: The start time of the pipeline run.
+
+        Returns:
+            The full set of substitutions for this step configuration.
+        """
+        ret = pipeline_config._get_full_substitutions(start_time)
+        ret.update(self.substitutions)
+        return ret
 
 
 class InputSpec(StrictBaseModel):
@@ -223,14 +269,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

@@ -13,14 +13,14 @@
 #  permissions and limitations under the License.
 """Functionality to support ZenML store configurations."""
 
-from pathlib import PurePath
 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__)
 
@@ -44,42 +44,10 @@ class StoreConfiguration(BaseModel):
 
     type: StoreType
     url: str
-    secrets_store: Optional[SecretsStoreConfiguration] = None
-    backup_secrets_store: Optional[SecretsStoreConfiguration] = None
-
-    @classmethod
-    def copy_configuration(
-        cls,
-        config: "StoreConfiguration",
-        config_path: str,
-        load_config_path: Optional[PurePath] = None,
-    ) -> "StoreConfiguration":
-        """Create a copy of the store config using a different configuration path.
-
-        This method is used to create a copy of the store configuration that can
-        be loaded using a different configuration path or in the context of a
-        new environment, such as a container image.
-
-        The configuration files accompanying the store configuration are also
-        copied to the new configuration path (e.g. certificates etc.).
-
-        Args:
-            config: The store configuration to copy.
-            config_path: new path where the configuration copy will be loaded
-                from.
-            load_config_path: absolute path that will be used to load the copied
-                configuration. This can be set to a value different from
-                `config_path` if the configuration copy will be loaded from
-                a different environment, e.g. when the configuration is copied
-                to a container image and loaded using a different absolute path.
-                This will be reflected in the paths and URLs encoded in the
-                copied configuration.
-
-        Returns:
-            A new store configuration object that reflects the new configuration
-            path.
-        """
-        return config.copy()
+    secrets_store: Optional[SerializeAsAny[SecretsStoreConfiguration]] = None
+    backup_secrets_store: Optional[
+        SerializeAsAny[SecretsStoreConfiguration]
+    ] = None
 
     @classmethod
     def supports_url_scheme(cls, url: str) -> bool:
@@ -96,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/console.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """ZenML console implementation."""
 
-
 from rich.console import Console
 from rich.style import Style
 from rich.theme import Theme