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

zenml/utils/pydantic_utils.py

@@ -12,17 +12,31 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Utilities for pydantic models."""
+
 import inspect
 import json
+from json.decoder import JSONDecodeError
 from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union, cast
 
 import yaml
-from pydantic import BaseModel
-from pydantic.decorator import ValidatedFunction
-from pydantic.json import pydantic_encoder
-from pydantic.utils import sequence_like
-
-from zenml.utils import dict_utils, yaml_utils
+from pydantic import (
+    AfterValidator,
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    PlainValidator,
+    ValidationInfo,
+    WrapValidator,
+    validate_call,
+)
+from pydantic._internal import _repr as pydantic_repr
+from pydantic.v1.utils import sequence_like
+
+from zenml.logger import get_logger
+from zenml.utils import dict_utils, typing_utils, yaml_utils
+from zenml.utils.json_utils import pydantic_encoder
+
+logger = get_logger(__name__)
 
 M = TypeVar("M", bound="BaseModel")
 
@@ -31,7 +45,7 @@ def update_model(
     original: M,
     update: Union["BaseModel", Dict[str, Any]],
     recursive: bool = True,
-    exclude_none: bool = True,
+    exclude_none: bool = False,
 ) -> M:
     """Updates a pydantic model.
 
@@ -39,8 +53,7 @@ def update_model(
         original: The model to update.
         update: The update values.
         recursive: If `True`, dictionary values will be updated recursively.
-        exclude_none: If `True`, `None` values in the update dictionary
-            will be removed.
+        exclude_none: If `True`, `None` values in the update will be removed.
 
     Returns:
         The updated model.
@@ -53,15 +66,17 @@ def update_model(
         else:
             update_dict = update
     else:
-        update_dict = update.dict(exclude_unset=True)
+        update_dict = update.model_dump(
+            exclude_unset=True, exclude_none=exclude_none
+        )
 
-    original_dict = original.dict(exclude_unset=True)
+    original_dict = original.model_dump(exclude_unset=True)
     if recursive:
         values = dict_utils.recursive_update(original_dict, update_dict)
     else:
         values = {**original_dict, **update_dict}
 
-    return original.__class__(**values)
+    return original.__class__.model_validate(values)
 
 
 class TemplateGenerator:
@@ -93,7 +108,7 @@ class TemplateGenerator:
                 self.instance_or_class
             )
 
-        # Convert to json in an intermediate step so we can leverage Pydantic's
+        # Convert to json in an intermediate step, so we can leverage Pydantic's
         # encoder to support types like UUID and datetime
         json_string = json.dumps(template, default=pydantic_encoder)
         return cast(Dict[str, Any], json.loads(json_string))
@@ -109,7 +124,7 @@ class TemplateGenerator:
         """
         template = self._generate_template_for_model_class(model.__class__)
 
-        for name in model.__fields_set__:
+        for name in model.model_fields_set:
             value = getattr(model, name)
             template[name] = self._generate_template_for_value(value)
 
@@ -129,19 +144,25 @@ class TemplateGenerator:
         """
         template: Dict[str, Any] = {}
 
-        for name, field in model_class.__fields__.items():
-            if self._is_model_class(field.outer_type_):
-                template[name] = self._generate_template_for_model_class(
-                    field.outer_type_
-                )
-            elif field.outer_type_ is Optional and self._is_model_class(
-                field.type_
-            ):
-                template[name] = self._generate_template_for_model_class(
-                    field.type_
-                )
-            else:
-                template[name] = field._type_display()
+        for name, field in model_class.model_fields.items():
+            annotation = field.annotation
+
+            if annotation is not None:
+                if self._is_model_class(annotation):
+                    template[name] = self._generate_template_for_model_class(
+                        annotation
+                    )
+
+                elif typing_utils.is_optional(
+                    annotation
+                ) and self._is_model_class(
+                    typing_utils.get_args(annotation)[0]
+                ):
+                    template[name] = self._generate_template_for_model_class(
+                        typing_utils.get_args(annotation)[0]
+                    )
+                else:
+                    template[name] = pydantic_repr.display_as_type(annotation)
 
         return template
 
@@ -187,12 +208,16 @@ class YAMLSerializationMixin(BaseModel):
 
         Args:
             sort_keys: Whether to sort the keys in the YAML representation.
-            **kwargs: Kwargs to pass to the pydantic json(...) method.
+            **kwargs: Kwargs to pass to the pydantic model_dump(...) method.
 
         Returns:
             YAML string representation.
         """
-        dict_ = json.loads(self.json(**kwargs, sort_keys=sort_keys))
+        dict_ = json.loads(
+            json.dumps(
+                self.model_dump(mode="json", **kwargs), sort_keys=sort_keys
+            )
+        )
         return yaml.dump(dict_, sort_keys=sort_keys)
 
     @classmethod
@@ -206,12 +231,12 @@ class YAMLSerializationMixin(BaseModel):
             The model instance.
         """
         dict_ = yaml_utils.read_yaml(path)
-        return cls.parse_obj(dict_)
+        return cls.model_validate(dict_)
 
 
 def validate_function_args(
     __func: Callable[..., Any],
-    __config: Dict[str, Any],
+    __config: Optional[ConfigDict],
     *args: Any,
     **kwargs: Any,
 ) -> Dict[str, Any]:
@@ -232,40 +257,225 @@ def validate_function_args(
     Returns:
         The validated arguments.
     """
-    parameter_prefix = "zenml__"
-
     signature = inspect.signature(__func)
-    parameters = [
-        param.replace(name=f"{parameter_prefix}{param.name}")
-        for param in signature.parameters.values()
-    ]
-    signature = signature.replace(parameters=parameters)
-
-    def f() -> None:
-        pass
-
-    # We create a dummy function with the original function signature, but
-    # add a prefix to all arguments to avoid potential clashes with pydantic
-    # BaseModel attributes
+
+    validated_args = ()
+    validated_kwargs = {}
+
+    def f(*args: Any, **kwargs: Dict[Any, Any]) -> None:
+        nonlocal validated_args
+        nonlocal validated_kwargs
+
+        validated_args = args
+        validated_kwargs = kwargs
+
+    # We create a dummy function with the original function signature to run
+    # pydantic validation without actually running the function code
     f.__signature__ = signature  # type: ignore[attr-defined]
-    f.__annotations__ = {
-        f"{parameter_prefix}{key}": annotation
-        for key, annotation in __func.__annotations__.items()
-    }
-
-    validation_func = ValidatedFunction(f, config=__config)
-
-    kwargs = {
-        f"{parameter_prefix}{key}": value for key, value in kwargs.items()
-    }
-    model = validation_func.init_model_instance(*args, **kwargs)
-
-    validated_args = {
-        k[len(parameter_prefix) :]: v
-        for k, v in model._iter()
-        if k in model.__fields_set__
-        or model.__fields__[k].default_factory
-        or model.__fields__[k].default
-    }
-
-    return validated_args
+    f.__annotations__ = __func.__annotations__
+
+    validated_function = validate_call(config=__config, validate_return=False)(
+        f
+    )
+
+    # This raises a pydantic.ValidatonError in case the arguments are not valid
+    validated_function(*args, **kwargs)
+
+    return signature.bind(*validated_args, **validated_kwargs).arguments
+
+
+def model_validator_data_handler(
+    raw_data: Any,
+    base_class: Type[BaseModel],
+    validation_info: ValidationInfo,
+) -> Dict[str, Any]:
+    """Utility function to parse raw input data of varying types to a dict.
+
+    With the change to pydantic v2, validators which operate with "before"
+    (or previously known as the "pre" parameter) are getting "Any" types of raw
+    input instead of a "Dict[str, Any]" as before. Depending on the use-case,
+    this can create conflicts after the migration and this function will be
+    used as a helper function to handle different types of raw input data.
+
+    A code snippet to showcase how the behavior changes. The "before" validator
+    prints the type of the input:
+
+        class Base(BaseModel):
+            a: int = 3
+
+        class MyClass(Base):
+            @model_validator(mode="before")
+            @classmethod
+            def before_validator(cls, data: Any) -> Any:
+                print(type(data))
+                return {}
+
+        one = MyClass() # prints "<class 'dict'>"
+        MyClass.model_validate(one)  # prints NOTHING, it is already validated
+        MyClass.model_validate("asdf")  # prints "<class 'str'>", fails without the modified return.
+        MyClass.model_validate(RandomClass())  # prints "<class 'RandomClass'>", fails without the modified return.
+        MyClass.model_validate(Base())  # prints "<class 'Base'>", fails without the modified return.
+        MyClass.model_validate_json(json.dumps("aria"))  # prints "<class 'str'>", fails without the modified return.
+        MyClass.model_validate_json(json.dumps([1]))  # prints "<class 'list'>", fails without the modified return.
+        MyClass.model_validate_json(one.model_dump_json())  # prints "<class 'dict'>"
+
+    Args:
+        raw_data: The raw data passed to the validator, can be "Any" type.
+        base_class: The class that the validator belongs to
+        validation_info: Extra information about the validation process.
+
+    Raises:
+        TypeError: if the type of the data is not processable.
+        ValueError: in case of an unknown validation mode.
+
+    Returns:
+        A dictionary which will be passed to the eventual validator of pydantic.
+    """
+    if validation_info.mode == "python":
+        # This is mode is only active if people validate objects using pythonic
+        # raw data such as MyClass(...) or MyClass.model_validate()
+
+        if isinstance(raw_data, dict):
+            # In most cases, this is the behavior as the raw input is a dict
+            return raw_data
+
+        elif isinstance(raw_data, base_class):
+            # In some cases, we pass the same object type to the validation
+            # in such cases, it is critical we keep the original structure of
+            # fields that are already set.
+            return dict(raw_data)
+
+        elif issubclass(base_class, raw_data.__class__):
+            # There are a few occurrences where the annotation of the field is
+            # denoted by a subclass, and we use the instance of its super class
+            # as the raw input. In such cases we will use the same approach as
+            # before, while raising a debug message.
+            logger.debug(
+                f"During the validation of a `{base_class}` object, an instance"
+                f"of `{raw_data.__class__}` (super class of `{base_class}`) "
+                f"has been passed as raw input. This might lead to unexpected "
+                f"behavior in case `{base_class}` have features which can not"
+                f"be extracted from an instance of a `{raw_data.__class__}`."
+            )
+            return dict(raw_data)
+
+        elif isinstance(raw_data, str):
+            # If the raw input is a raw string, we can try to use the `json`
+            # module to parse it. The resulting data needs to be a proper
+            # dict for us to pass it to the validation process.
+            try:
+                json_data = json.loads(raw_data)
+
+                if isinstance(json_data, dict):
+                    return json_data
+                else:
+                    raise TypeError("The resulting json data is not a dict!")
+
+            except (TypeError, JSONDecodeError) as e:
+                raise TypeError(
+                    "The raw json input string can not be converted to a "
+                    f"dict: {e}"
+                )
+        else:
+            raise TypeError(
+                "Unsupported type of raw input data for the `python` validation"
+                "mode of the pydantic class. Please consider changing the way "
+                f"you are creating using the `{base_class}` or instead use"
+                f"`{base_class}.model_validate_json()`."
+            )
+
+    elif validation_info.mode == "json":
+        # This is mode is only active if people validate objects using json
+        # input data such as MyClass.model_validate_json()
+        if isinstance(raw_data, dict):
+            return raw_data
+        else:
+            raise TypeError(
+                f"The resulting JSON data {raw_data} is not a dict, therefore"
+                f"can not be used by the validation process."
+            )
+    else:
+        # Unknown validation mode
+        raise ValueError(f"Unknown validation mode. {validation_info.mode}")
+
+
+def before_validator_handler(
+    method: Callable[..., Any],
+) -> Callable[[Any, Any, Any], Any]:
+    """Decorator to handle the raw input data for pydantic model validators.
+
+    Args:
+        method: the class method with the actual validation logic.
+
+    Returns:
+        the validator method
+    """
+
+    def before_validator(
+        cls: Type[BaseModel], data: Any, validation_info: ValidationInfo
+    ) -> Any:
+        """Wrapper method to handle the raw data.
+
+        Args:
+            cls: the class handler
+            data: the raw input data
+            validation_info: the context of the validation.
+
+        Returns:
+            the validated data
+        """
+        data = model_validator_data_handler(
+            raw_data=data, base_class=cls, validation_info=validation_info
+        )
+        return method(cls=cls, data=data)
+
+    return before_validator
+
+
+def has_validators(
+    pydantic_class: Type[BaseModel],
+    field_name: Optional[str] = None,
+) -> bool:
+    """Function to check if a Pydantic model or a pydantic field has validators.
+
+    Args:
+        pydantic_class: The class defining the pydantic model.
+        field_name: Optional, field info. If specified, this function will focus
+            on a singular field within the class. If not specified, it will
+            check model validators.
+
+    Returns:
+        Whether the specified field or class has a validator or not.
+    """
+    # If field is not specified check model validators
+    if field_name is None:
+        if pydantic_class.__pydantic_decorators__.model_validators:
+            return True
+
+    # Else, check field validators
+    else:
+        # 1. Field validators can be defined through @field_validator decorators
+        f_validators = pydantic_class.__pydantic_decorators__.field_validators
+
+        for name, f_v in f_validators.items():
+            if field_name in f_v.info.fields:
+                return True
+
+        # 2. Field validators can be defined through the Annotation[.....]
+        field_info = pydantic_class.model_fields[field_name]
+        if metadata := field_info.metadata:
+            if any(
+                isinstance(
+                    m,
+                    (
+                        AfterValidator,
+                        BeforeValidator,
+                        PlainValidator,
+                        WrapValidator,
+                    ),
+                )
+                for m in metadata
+            ):
+                return True
+
+    return False

zenml/utils/requirements_utils.py

@@ -0,0 +1,71 @@
+#  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:
+#
+#       https://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.
+"""Requirement utils."""
+
+from typing import TYPE_CHECKING, List, Set, Tuple
+
+from zenml.integrations.utils import get_integration_for_module
+
+if TYPE_CHECKING:
+    from zenml.models import ComponentResponse, StackResponse
+
+
+def get_requirements_for_stack(
+    stack: "StackResponse",
+) -> Tuple[List[str], List[str]]:
+    """Get requirements for a stack model.
+
+    Args:
+        stack: The stack for which to get the requirements.
+
+    Returns:
+        Tuple of PyPI and APT requirements of the stack.
+    """
+    pypi_requirements: Set[str] = set()
+    apt_packages: Set[str] = set()
+
+    for component_list in stack.components.values():
+        assert len(component_list) == 1
+        component = component_list[0]
+        (
+            component_pypi_requirements,
+            component_apt_packages,
+        ) = get_requirements_for_component(component=component)
+        pypi_requirements = pypi_requirements.union(
+            component_pypi_requirements
+        )
+        apt_packages = apt_packages.union(component_apt_packages)
+
+    return sorted(pypi_requirements), sorted(apt_packages)
+
+
+def get_requirements_for_component(
+    component: "ComponentResponse",
+) -> Tuple[List[str], List[str]]:
+    """Get requirements for a component model.
+
+    Args:
+        component: The component for which to get the requirements.
+
+    Returns:
+        Tuple of PyPI and APT requirements of the component.
+    """
+    integration = get_integration_for_module(
+        module_name=component.flavor.source
+    )
+
+    if integration:
+        return integration.get_requirements(), integration.APT_PACKAGES
+    else:
+        return [], []

zenml/utils/secret_utils.py

@@ -12,19 +12,33 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Utility functions for secrets and secret references."""
+
 import re
 from typing import TYPE_CHECKING, Any, NamedTuple
 
-from pydantic import Field
+from pydantic import Field, PlainSerializer, SecretStr
+from typing_extensions import Annotated
+
+from zenml.logger import get_logger
 
 if TYPE_CHECKING:
-    from pydantic.fields import ModelField
+    from pydantic.fields import FieldInfo
 
 _secret_reference_expression = re.compile(r"\{\{\s*\S+?\.\S+\s*\}\}")
 
 PYDANTIC_SENSITIVE_FIELD_MARKER = "sensitive"
 PYDANTIC_CLEAR_TEXT_FIELD_MARKER = "prevent_secret_reference"
 
+PlainSerializedSecretStr = Annotated[
+    SecretStr,
+    PlainSerializer(
+        lambda v: v.get_secret_value() if v is not None else None,
+        when_used="json",
+    ),
+]
+
+logger = get_logger(__name__)
+
 
 def is_secret_reference(value: Any) -> bool:
     """Checks whether any value is a secret reference.
@@ -86,8 +100,9 @@ def SecretField(*args: Any, **kwargs: Any) -> Any:
     Returns:
         Pydantic field info.
     """
-    kwargs[PYDANTIC_SENSITIVE_FIELD_MARKER] = True
-    return Field(*args, **kwargs)  # type: ignore[pydantic-field]
+    json_schema_extra = kwargs.get("json_schema_extra", {})
+    json_schema_extra.update({PYDANTIC_SENSITIVE_FIELD_MARKER: True})
+    return Field(json_schema_extra=json_schema_extra, *args, **kwargs)  # type: ignore[pydantic-field]
 
 
 def ClearTextField(*args: Any, **kwargs: Any) -> Any:
@@ -102,11 +117,12 @@ def ClearTextField(*args: Any, **kwargs: Any) -> Any:
     Returns:
         Pydantic field info.
     """
-    kwargs[PYDANTIC_CLEAR_TEXT_FIELD_MARKER] = True
-    return Field(*args, **kwargs)  # type: ignore[pydantic-field]
+    json_schema_extra = kwargs.get("json_schema_extra", {})
+    json_schema_extra.update({PYDANTIC_CLEAR_TEXT_FIELD_MARKER: True})
+    return Field(json_schema_extra=json_schema_extra, *args, **kwargs)  # type: ignore[pydantic-field]
 
 
-def is_secret_field(field: "ModelField") -> bool:
+def is_secret_field(field: "FieldInfo") -> bool:
     """Returns whether a pydantic field contains sensitive information or not.
 
     Args:
@@ -115,10 +131,29 @@ def is_secret_field(field: "ModelField") -> bool:
     Returns:
         `True` if the field contains sensitive information, `False` otherwise.
     """
-    return field.field_info.extra.get(PYDANTIC_SENSITIVE_FIELD_MARKER, False)  # type: ignore[no-any-return]
-
-
-def is_clear_text_field(field: "ModelField") -> bool:
+    if field.json_schema_extra is not None:
+        if isinstance(field.json_schema_extra, dict):
+            if marker := field.json_schema_extra.get(
+                PYDANTIC_SENSITIVE_FIELD_MARKER
+            ):
+                assert isinstance(marker, bool), (
+                    f"The parameter `{PYDANTIC_SENSITIVE_FIELD_MARKER}` in the "
+                    f"field definition can only be a boolean value."
+                )
+                return marker
+
+        else:
+            logger.warning(
+                f"The 'json_schema_extra' of the field '{field.title}' is "
+                "not defined as a dict. This might lead to unexpected "
+                "behavior as we are checking it is a secret text field. "
+                "Returning 'False' as default..."
+            )
+
+    return False
+
+
+def is_clear_text_field(field: "FieldInfo") -> bool:
     """Returns whether a pydantic field prevents secret references or not.
 
     Args:
@@ -127,4 +162,23 @@ def is_clear_text_field(field: "ModelField") -> bool:
     Returns:
         `True` if the field prevents secret references, `False` otherwise.
     """
-    return field.field_info.extra.get(PYDANTIC_CLEAR_TEXT_FIELD_MARKER, False)  # type: ignore[no-any-return]
+    if field.json_schema_extra is not None:
+        if isinstance(field.json_schema_extra, dict):
+            if marker := field.json_schema_extra.get(
+                PYDANTIC_CLEAR_TEXT_FIELD_MARKER
+            ):
+                assert isinstance(marker, bool), (
+                    f"The parameter `{PYDANTIC_CLEAR_TEXT_FIELD_MARKER}` in the "
+                    f"field definition can only be a boolean value."
+                )
+                return marker
+
+        else:
+            logger.warning(
+                f"The 'json_schema_extra' of the field '{field.title}' is "
+                "not defined as a dict. This might lead to unexpected "
+                "behavior as we are checking it is a clear text field. "
+                "Returning 'False' as default..."
+            )
+
+    return False

zenml/utils/settings_utils.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Utility functions for ZenML settings."""
+
 import re
 from typing import TYPE_CHECKING, Dict, Sequence, Type
 
@@ -24,7 +25,7 @@ if TYPE_CHECKING:
     from zenml.stack.flavor import Flavor
 
 STACK_COMPONENT_REGEX = re.compile(
-    "(" + "|".join(StackComponentType.values()) + r")\..*"
+    "(" + "|".join(StackComponentType.values()) + r")(\..*)?"
 )
 
 

zenml/utils/singleton.py

@@ -64,9 +64,21 @@ class SingletonMetaClass(type):
 
         return cls.__singleton_instance
 
-    def _clear(cls) -> None:
-        """Clear the singleton instance."""
-        cls.__singleton_instance = None
+    def _clear(cls, instance: Optional["SingletonMetaClass"] = None) -> None:
+        """Clear or replace the singleton instance.
+
+        Args:
+            instance: The new singleton instance.
+        """
+        cls.__singleton_instance = instance
+
+    def _instance(cls) -> Optional["SingletonMetaClass"]:
+        """Get the singleton instance.
+
+        Returns:
+            The singleton instance.
+        """
+        return cls.__singleton_instance
 
     def _exists(cls) -> bool:
         """Check if the singleton instance exists.

zenml/utils/source_code_utils.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Utilities for getting the source code of objects."""
+
 import hashlib
 import inspect
 import sys