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

zenml/stack/stack_component.py

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

zenml/step_operators/step_operator_entrypoint_configuration.py

@@ -86,6 +86,7 @@ class StepOperatorEntrypointConfiguration(StepEntrypointConfiguration):
             pipeline_step_name=self.entrypoint_args[STEP_NAME_OPTION],
             run_id=pipeline_run.id,
             step_run_id=step_run_id,
+            force_write_logs=lambda: None,
         )
 
         stack = Client().active_stack

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(