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

zenml/steps/base_step.py

@@ -34,7 +34,7 @@ from typing import (
     cast,
 )
 
-from pydantic import BaseModel, Extra, ValidationError
+from pydantic import BaseModel, ConfigDict, ValidationError
 
 from zenml.client_lazy_loader import ClientLazyLoader
 from zenml.config.retry_config import StepRetryConfig
@@ -59,6 +59,7 @@ from zenml.utils import (
     settings_utils,
     source_code_utils,
     source_utils,
+    typing_utils,
 )
 
 if TYPE_CHECKING:
@@ -513,17 +514,17 @@ class BaseStep(metaclass=BaseStepMeta):
                     )
             elif isinstance(value, LazyArtifactVersionResponse):
                 model_artifacts_or_metadata[key] = ModelVersionDataLazyLoader(
-                    model=value._lazy_load_model,
-                    artifact_name=value._lazy_load_name,
-                    artifact_version=value._lazy_load_version,
+                    model=value.lazy_load_model,
+                    artifact_name=value.lazy_load_name,
+                    artifact_version=value.lazy_load_version,
                     metadata_name=None,
                 )
             elif isinstance(value, LazyRunMetadataResponse):
                 model_artifacts_or_metadata[key] = ModelVersionDataLazyLoader(
-                    model=value._lazy_load_model,
-                    artifact_name=value._lazy_load_artifact_name,
-                    artifact_version=value._lazy_load_artifact_version,
-                    metadata_name=value._lazy_load_metadata_name,
+                    model=value.lazy_load_model,
+                    artifact_name=value.lazy_load_artifact_name,
+                    artifact_version=value.lazy_load_artifact_version,
+                    metadata_name=value.lazy_load_metadata_name,
                 )
             elif isinstance(value, ClientLazyLoader):
                 client_lazy_loaders[key] = value
@@ -583,7 +584,7 @@ class BaseStep(metaclass=BaseStepMeta):
         from zenml.new.pipelines.pipeline import Pipeline
 
         if not Pipeline.ACTIVE_PIPELINE:
-            # The step is being called outside of the context of a pipeline,
+            # The step is being called outside the context of a pipeline,
             # we simply call the entrypoint
             return self.call_entrypoint(*args, **kwargs)
 
@@ -645,12 +646,15 @@ class BaseStep(metaclass=BaseStepMeta):
         try:
             validated_args = pydantic_utils.validate_function_args(
                 self.entrypoint,
-                {"arbitrary_types_allowed": True, "smart_union": True},
+                ConfigDict(arbitrary_types_allowed=True),
                 *args,
                 **kwargs,
             )
         except ValidationError as e:
-            raise StepInterfaceError("Invalid entrypoint arguments.") from e
+            raise StepInterfaceError(
+                "Invalid step function entrypoint arguments. Check out the "
+                "pydantic error above for more details."
+            ) from e
 
         return self.entrypoint(**validated_args)
 
@@ -796,7 +800,7 @@ class BaseStep(metaclass=BaseStepMeta):
             success_hook_source = resolve_and_validate_hook(on_success)
 
         if isinstance(parameters, BaseParameters):
-            parameters = parameters.dict()
+            parameters = parameters.model_dump()
 
         values = dict_utils.remove_none_values(
             {
@@ -1111,12 +1115,6 @@ To avoid this consider setting step parameters only in one place (config or code
                 output_name, PartialArtifactConfiguration()
             )
 
-            from pydantic.typing import (
-                get_origin,
-                is_none_type,
-                is_union,
-            )
-
             from zenml.steps.utils import get_args
 
             if not output.materializer_source:
@@ -1129,13 +1127,15 @@ To avoid this consider setting step parameters only in one place (config or code
                     )
                     continue
 
-                if is_union(
-                    get_origin(output_annotation.resolved_annotation)
+                if typing_utils.is_union(
+                    typing_utils.get_origin(
+                        output_annotation.resolved_annotation
+                    )
                     or output_annotation.resolved_annotation
                 ):
                     output_types = tuple(
                         type(None)
-                        if is_none_type(output_type)
+                        if typing_utils.is_none_type(output_type)
                         else output_type
                         for output_type in get_args(
                             output_annotation.resolved_annotation
@@ -1169,7 +1169,7 @@ To avoid this consider setting step parameters only in one place (config or code
         config = StepConfigurationUpdate(**values)
         self._apply_configuration(config)
 
-        self._configuration = self._configuration.copy(
+        self._configuration = self._configuration.model_copy(
             update={
                 "caching_parameters": self.caching_parameters,
                 "external_input_artifacts": external_artifacts,
@@ -1178,7 +1178,9 @@ To avoid this consider setting step parameters only in one place (config or code
             }
         )
 
-        return StepConfiguration.parse_obj(self._configuration)
+        return StepConfiguration.model_validate(
+            self._configuration.model_dump()
+        )
 
     def _finalize_parameters(self) -> Dict[str, Any]:
         """Finalizes the config parameters for running this step.
@@ -1199,7 +1201,7 @@ To avoid this consider setting step parameters only in one place (config or code
                 # Make sure we have all necessary values to instantiate the
                 # pydantic model later
                 model = annotation(**value)
-                params[key] = model.dict()
+                params[key] = model.model_dump()
             else:
                 params[key] = value
 
@@ -1254,7 +1256,7 @@ To avoid this consider setting step parameters only in one place (config or code
         for (
             name,
             field,
-        ) in self.entrypoint_definition.legacy_params.annotation.__fields__.items():
+        ) in self.entrypoint_definition.legacy_params.annotation.model_fields.items():
             if name in self.configuration.parameters:
                 # a value for this parameter has been set already
                 values[name] = self.configuration.parameters[name]
@@ -1262,7 +1264,7 @@ To avoid this consider setting step parameters only in one place (config or code
                 # a value for this parameter has been set in the "new" way
                 # already
                 values[name] = params_defined_in_new_way[name]
-            elif field.required:
+            elif field.is_required():
                 # this field has no default value set and therefore needs
                 # to be passed via an initialized config object
                 missing_keys.append(name)
@@ -1278,8 +1280,12 @@ To avoid this consider setting step parameters only in one place (config or code
             )
 
         if (
-            self.entrypoint_definition.legacy_params.annotation.__config__.extra
-            == Extra.allow
+            getattr(
+                self.entrypoint_definition.legacy_params.annotation.model_config,
+                "extra",
+                None,
+            )
+            == "allow"
         ):
             # Add all parameters for the config class for backwards
             # compatibility if the config class allows extra attributes

zenml/steps/entrypoint_function_utils.py

@@ -27,7 +27,7 @@ from typing import (
     Union,
 )
 
-from pydantic import BaseConfig, ValidationError, create_model
+from pydantic import ConfigDict, ValidationError, create_model
 
 from zenml.constants import ENFORCE_TYPE_ANNOTATIONS
 from zenml.exceptions import StepInterfaceError
@@ -235,16 +235,14 @@ class EntrypointFunctionDefinition(NamedTuple):
             parameter: The function parameter for which the value was provided.
             value: The input value.
         """
-
-        class ModelConfig(BaseConfig):
-            arbitrary_types_allowed = False
+        config_dict = ConfigDict(arbitrary_types_allowed=False)
 
         # Create a pydantic model with just a single required field with the
         # type annotation of the parameter to verify the input type including
         # pydantics type coercion
         validation_model_class = create_model(
             "input_validation_model",
-            __config__=ModelConfig,
+            __config__=config_dict,
             value=(parameter.annotation, ...),
         )
         validation_model_class(value=value)

zenml/steps/utils.py

@@ -21,7 +21,6 @@ import textwrap
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 from uuid import UUID
 
-import pydantic.typing as pydantic_typing
 from pydantic import BaseModel
 from typing_extensions import Annotated
 
@@ -32,7 +31,7 @@ from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.new.steps.step_context import get_step_context
 from zenml.steps.step_output import Output
-from zenml.utils import source_code_utils
+from zenml.utils import source_code_utils, typing_utils
 
 logger = get_logger(__name__)
 
@@ -42,8 +41,8 @@ SINGLE_RETURN_OUT_NAME = "output"
 class OutputSignature(BaseModel):
     """The signature of an output artifact."""
 
-    resolved_annotation: Any
-    artifact_config: Optional[ArtifactConfig]
+    resolved_annotation: Any = None
+    artifact_config: Optional[ArtifactConfig] = None
     has_custom_name: bool = False
 
 
@@ -60,8 +59,7 @@ def get_args(obj: Any) -> Tuple[Any, ...]:
         The args of the annotation.
     """
     return tuple(
-        pydantic_typing.get_origin(v) or v
-        for v in pydantic_typing.get_args(obj)
+        typing_utils.get_origin(v) or v for v in typing_utils.get_args(obj)
     )
 
 
@@ -123,11 +121,11 @@ def parse_return_type_annotations(
             for output_name, output_type in return_annotation.items()
         }
 
-    elif pydantic_typing.get_origin(return_annotation) is tuple:
+    elif typing_utils.get_origin(return_annotation) is tuple:
         requires_multiple_artifacts = has_tuple_return(func)
         if requires_multiple_artifacts:
             output_signature: Dict[str, Any] = {}
-            args = pydantic_typing.get_args(return_annotation)
+            args = typing_utils.get_args(return_annotation)
             if args[-1] is Ellipsis:
                 raise RuntimeError(
                     "Variable length output annotations are not allowed."
@@ -179,12 +177,12 @@ def resolve_type_annotation(obj: Any) -> Any:
     Returns:
         The non-generic class for generic aliases of the typing module.
     """
-    origin = pydantic_typing.get_origin(obj) or obj
+    origin = typing_utils.get_origin(obj) or obj
 
     if origin is Annotated:
-        annotation, *_ = pydantic_typing.get_args(obj)
+        annotation, *_ = typing_utils.get_args(obj)
         return resolve_type_annotation(annotation)
-    elif pydantic_typing.is_union(origin):
+    elif typing_utils.is_union(origin):
         return obj
 
     return origin
@@ -212,10 +210,10 @@ def get_artifact_config_from_annotation_metadata(
     Returns:
         The artifact config.
     """
-    if (pydantic_typing.get_origin(annotation) or annotation) is not Annotated:
+    if (typing_utils.get_origin(annotation) or annotation) is not Annotated:
         return None
 
-    annotation, *metadata = pydantic_typing.get_args(annotation)
+    annotation, *metadata = typing_utils.get_args(annotation)
 
     error_message = (
         "Artifact annotation should only contain two elements: the artifact "
@@ -251,7 +249,7 @@ def get_artifact_config_from_annotation_metadata(
         if not artifact_config:
             artifact_config = ArtifactConfig(name=output_name)
         elif not artifact_config.name:
-            artifact_config = artifact_config.copy()
+            artifact_config = artifact_config.model_copy()
             artifact_config.name = output_name
 
     if artifact_config and artifact_config.name == "":

zenml/utils/cuda_utils.py

@@ -0,0 +1,50 @@
+# Apache Software License 2.0
+#
+# Copyright (c) ZenML GmbH 2024. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Utilities for managing GPU memory."""
+
+import gc
+
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def cleanup_gpu_memory(force: bool = False) -> None:
+    """Clean up GPU memory.
+
+    Args:
+        force: whether to force the cleanup of GPU memory (must be passed explicitly)
+    """
+    if not force:
+        logger.warning(
+            "This will clean up all GPU memory on current physical machine. "
+            "This action is considered to be dangerous by default, since "
+            "it might affect other processes running in the same environment. "
+            "If this is intended, please explicitly pass `force=True`."
+        )
+    else:
+        try:
+            import torch
+        except ModuleNotFoundError:
+            logger.warning(
+                "No PyTorch installed. Skipping GPU memory cleanup."
+            )
+            return
+
+        logger.info("Cleaning up GPU memory...")
+        while gc.collect():
+            torch.cuda.empty_cache()

zenml/utils/deprecation_utils.py

@@ -14,14 +14,12 @@
 """Deprecation utilities."""
 
 import warnings
-from typing import TYPE_CHECKING, Any, Dict, Set, Tuple, Type, Union
+from typing import Any, Dict, Set, Tuple, Type, Union
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, model_validator
 
 from zenml.logger import get_logger
-
-if TYPE_CHECKING:
-    AnyClassMethod = classmethod[Any]  # type: ignore[type-arg]
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -30,7 +28,7 @@ PREVIOUS_DEPRECATION_WARNINGS_ATTRIBUTE = "__previous_deprecation_warnings"
 
 def deprecate_pydantic_attributes(
     *attributes: Union[str, Tuple[str, str]],
-) -> "AnyClassMethod":
+) -> Any:
     """Utility function for deprecating and migrating pydantic attributes.
 
     **Usage**:
@@ -55,22 +53,24 @@ def deprecate_pydantic_attributes(
     Args:
         *attributes: List of attributes to deprecate. This is either the name
             of the attribute to deprecate, or a tuple containing the name of
-            the deprecated attribute and it's replacement.
+            the deprecated attribute, and it's replacement.
 
     Returns:
         Pydantic validator class method to be used on BaseModel subclasses
         to deprecate or migrate attributes.
     """
 
-    @root_validator(pre=True, allow_reuse=True)
+    @model_validator(mode="before")  # type: ignore[misc]
+    @classmethod
+    @before_validator_handler
     def _deprecation_validator(
-        cls: Type[BaseModel], values: Dict[str, Any]
+        cls: Type[BaseModel], data: Dict[str, Any]
     ) -> Dict[str, Any]:
         """Pydantic validator function for deprecating pydantic attributes.
 
         Args:
             cls: The class on which the attributes are defined.
-            values: All values passed at model initialization.
+            data: All values passed at model initialization.
 
         Raises:
             AssertionError: If either the deprecated or replacement attribute
@@ -110,14 +110,14 @@ def deprecate_pydantic_attributes(
                 deprecated_attribute, replacement_attribute = attribute
 
                 assert (
-                    replacement_attribute in cls.__fields__
+                    replacement_attribute in cls.model_fields
                 ), f"Unable to find attribute {replacement_attribute}."
 
             assert (
-                deprecated_attribute in cls.__fields__
+                deprecated_attribute in cls.model_fields
             ), f"Unable to find attribute {deprecated_attribute}."
 
-            if cls.__fields__[deprecated_attribute].required:
+            if cls.model_fields[deprecated_attribute].is_required():
                 raise TypeError(
                     f"Unable to deprecate attribute '{deprecated_attribute}' "
                     f"of class {cls.__name__}. In order to deprecate an "
@@ -126,7 +126,7 @@ def deprecate_pydantic_attributes(
                     "annotation."
                 )
 
-            if values.get(deprecated_attribute, None) is None:
+            if data.get(deprecated_attribute, None) is None:
                 continue
 
             if replacement_attribute is None:
@@ -144,17 +144,15 @@ def deprecate_pydantic_attributes(
                 attribute=deprecated_attribute,
             )
 
-            if values.get(replacement_attribute, None) is None:
+            if data.get(replacement_attribute, None) is None:
                 logger.debug(
                     "Migrating value of deprecated attribute %s to "
                     "replacement attribute %s.",
                     deprecated_attribute,
                     replacement_attribute,
                 )
-                values[replacement_attribute] = values.pop(
-                    deprecated_attribute
-                )
-            elif values[deprecated_attribute] != values[replacement_attribute]:
+                data[replacement_attribute] = data.pop(deprecated_attribute)
+            elif data[deprecated_attribute] != data[replacement_attribute]:
                 raise ValueError(
                     "Got different values for deprecated attribute "
                     f"{deprecated_attribute} and replacement "
@@ -170,6 +168,6 @@ def deprecate_pydantic_attributes(
             previous_deprecation_warnings,
         )
 
-        return values
+        return data
 
     return _deprecation_validator

zenml/utils/dict_utils.py

@@ -17,7 +17,7 @@ import base64
 import json
 from typing import Any, Dict
 
-from pydantic.json import pydantic_encoder
+from zenml.utils.json_utils import pydantic_encoder
 
 
 def recursive_update(

zenml/utils/filesync_model.py

@@ -13,11 +13,16 @@
 #  permissions and limitations under the License.
 """Filesync utils for ZenML."""
 
-import json
 import os
 from typing import Any, Optional
 
-from pydantic import BaseModel
+from pydantic import (
+    BaseModel,
+    ValidationError,
+    ValidationInfo,
+    ValidatorFunctionWrapHandler,
+    model_validator,
+)
 
 from zenml.io import fileio
 from zenml.logger import get_logger
@@ -40,29 +45,69 @@ class FileSyncModel(BaseModel):
     _config_file: str
     _config_file_timestamp: Optional[float] = None
 
-    def __init__(self, config_file: str, **kwargs: Any) -> None:
-        """Create a FileSyncModel instance synchronized with a configuration file on disk.
+    @model_validator(mode="wrap")
+    @classmethod
+    def config_validator(
+        cls,
+        data: Any,
+        handler: ValidatorFunctionWrapHandler,
+        info: ValidationInfo,
+    ) -> "FileSyncModel":
+        """Wrap model validator to infer the config_file during initialization.
 
         Args:
-            config_file: configuration file path. If the file exists, the model
-                will be initialized with the values from the file.
-            **kwargs: additional keyword arguments to pass to the Pydantic model
-                constructor. If supplied, these values will override those
-                loaded from the configuration file.
+            data: The raw data that is provided before the validation.
+            handler: The actual validation function pydantic would use for the
+                built-in validation function.
+            info: The context information during the execution of this
+                validation function.
+
+        Returns:
+            the actual instance after the validation
+
+        Raises:
+            ValidationError: if you try to validate through a JSON string. You
+                need to provide a config_file path when you create a
+                FileSyncModel.
+            AssertionError: if the raw input does not include a config_file
+                path for the configuration file.
         """
-        config_dict = {}
-        if fileio.exists(config_file):
-            config_dict = yaml_utils.read_yaml(config_file)
+        # Disable json validation
+        if info.mode == "json":
+            raise ValidationError(
+                "You can not instantiate filesync models using the JSON mode."
+            )
 
-        self._config_file = config_file
-        self._config_file_timestamp = None
+        if isinstance(data, dict):
+            # Assert that the config file is defined
+            assert (
+                "config_file" in data
+            ), "You have to provide a path for the configuration file."
 
-        config_dict.update(kwargs)
-        super(FileSyncModel, self).__init__(**config_dict)
+            config_file = data.pop("config_file")
 
-        # write the configuration file to disk, to reflect new attributes
-        # and schema changes
-        self.write_config()
+            # Load the current values and update with new values
+            config_dict = {}
+            if fileio.exists(config_file):
+                config_dict = yaml_utils.read_yaml(config_file)
+            config_dict.update(data)
+
+            # Execute the regular validation
+            model = handler(config_dict)
+
+            assert isinstance(model, cls)
+
+            # Assign the private attribute and save the config
+            model._config_file = config_file
+            model.write_config()
+
+        else:
+            # If the raw value is not a dict, apply proper validation.
+            model = handler(data)
+
+            assert isinstance(model, cls)
+
+        return model
 
     def __setattr__(self, key: str, value: Any) -> None:
         """Sets an attribute on the model and persists it in the configuration file.
@@ -91,8 +136,7 @@ class FileSyncModel(BaseModel):
 
     def write_config(self) -> None:
         """Writes the model to the configuration file."""
-        config_dict = json.loads(self.json())
-        yaml_utils.write_yaml(self._config_file, config_dict)
+        yaml_utils.write_yaml(self._config_file, self.model_dump(mode="json"))
         self._config_file_timestamp = os.path.getmtime(self._config_file)
 
     def load_config(self) -> None:
@@ -115,10 +159,3 @@ class FileSyncModel(BaseModel):
             super(FileSyncModel, self).__setattr__(key, value)
 
         self._config_file_timestamp = file_timestamp
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True