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

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

zenml/utils/function_utils.py

@@ -0,0 +1,260 @@
+#  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.
+"""Utility functions for python functions."""
+
+import inspect
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Callable, Iterator, List, Tuple, TypeVar, Union
+
+import click
+
+from zenml.logger import get_logger
+from zenml.utils.string_utils import random_str
+
+F = TypeVar("F", bound=Callable[..., None])
+
+logger = get_logger(__name__)
+
+_CLI_WRAPPED_SCRIPT_TEMPLATE_HEADER = """
+from zenml.utils.function_utils import _cli_wrapped_function
+
+import sys
+sys.path.append(r"{func_path}")
+
+from {func_module} import {func_name} as func_to_wrap
+
+if entrypoint:=getattr(func_to_wrap, "entrypoint", None):
+    func = _cli_wrapped_function(entrypoint)
+else:
+    func = _cli_wrapped_function(func_to_wrap)
+"""
+_CLI_WRAPPED_MAINS = {
+    "accelerate": """
+if __name__=="__main__":
+    from accelerate import Accelerator
+    import cloudpickle as pickle
+    accelerator = Accelerator()
+    ret = func(standalone_mode=False)
+    if accelerator.is_main_process:
+        pickle.dump(ret, open(r"{output_file}", "wb"))
+"""
+}
+_ALLOWED_TYPES = (str, int, float, bool, Path)
+_ALLOWED_COLLECTIONS = (tuple,)
+_CLICK_TYPES_MAPPER = {
+    str: click.STRING,
+    int: click.INT,
+    float: click.FLOAT,
+    bool: click.BOOL,
+    Path: click.STRING,
+    None: click.STRING,
+}
+
+
+def _cli_arg_name(arg_name: str) -> str:
+    return arg_name.replace("_", "-")
+
+
+def _is_valid_collection_arg(arg_type: Any) -> bool:
+    """Check if the given argument type is a valid collection type.
+
+    Args:
+        arg_type: The type to check.
+
+    Returns:
+        True if the argument type is a valid collection type, False otherwise.
+    """
+    if getattr(arg_type, "__origin__", None) in _ALLOWED_COLLECTIONS:
+        if arg_type.__args__[0] not in _ALLOWED_TYPES:
+            return False
+        return True
+    return False
+
+
+def _is_valid_optional_arg(arg_type: Any) -> bool:
+    """Check if the given argument type is a valid Optional type.
+
+    A valid Optional type is defined as a Union with two arguments, where:
+    - The first argument is either an allowed type or a valid collection type.
+    - The second argument is the NoneType.
+
+    Args:
+        arg_type: The type to check.
+
+    Returns:
+        True if the argument type is a valid Optional type, False otherwise.
+    """
+    if (
+        getattr(arg_type, "_name", None) == "Optional"
+        and getattr(arg_type, "__origin__", None) == Union
+    ):
+        if args := getattr(arg_type, "__args__", None):
+            if len(args) != 2:
+                return False
+            if (
+                args[0] not in _ALLOWED_TYPES
+                and not _is_valid_collection_arg(args[0])
+            ) or args[1] != type(None):
+                return False
+        return True
+    return False
+
+
+def _cli_wrapped_function(func: F) -> F:
+    """Create a decorator to generate the CLI-wrapped function.
+
+    Args:
+        func: The function to decorate.
+
+    Returns:
+        The inner decorator.
+
+    Raises:
+        ValueError: If the function arguments are not valid.
+    """
+    options: List[Any] = []
+    fullargspec = inspect.getfullargspec(func)
+    if fullargspec.defaults is not None:
+        defaults = [None] * (
+            len(fullargspec.args) - len(fullargspec.defaults)
+        ) + list(fullargspec.defaults)
+    else:
+        defaults = [None] * len(fullargspec.args)
+    input_args_dict = (
+        (
+            arg_name,
+            fullargspec.annotations.get(arg_name, None),
+            defaults[i],
+        )
+        for i, arg_name in enumerate(fullargspec.args)
+    )
+    invalid_types = {}
+    for arg_name, arg_type, arg_default in input_args_dict:
+        if _is_valid_optional_arg(arg_type):
+            arg_type = arg_type.__args__[0]
+        arg_name = _cli_arg_name(arg_name)
+        if arg_type == bool:
+            options.append(
+                click.option(
+                    f"--{arg_name}",
+                    type=click.BOOL,
+                    is_flag=True,
+                    default=False,
+                    required=False,
+                )
+            )
+        elif _is_valid_collection_arg(arg_type):
+            member_type = arg_type.__args__[0]
+            options.append(
+                click.option(
+                    f"--{arg_name}",
+                    type=member_type,
+                    default=arg_default,
+                    required=False,
+                    multiple=True,
+                )
+            )
+        elif arg_type in _ALLOWED_TYPES:
+            options.append(
+                click.option(
+                    f"--{arg_name}",
+                    type=_CLICK_TYPES_MAPPER[arg_type],
+                    default=arg_default,
+                    required=False if arg_default is not None else True,
+                )
+            )
+        else:
+            invalid_types[arg_name] = arg_type
+    if invalid_types:
+        raise ValueError(
+            f"Invalid argument types: {invalid_types}. CLI functions only "
+            f"supports: {_ALLOWED_TYPES} types (including Optional) and "
+            f"{_ALLOWED_COLLECTIONS} collections."
+        )
+    options.append(
+        click.command(
+            help="Technical wrapper to pass into the `accelerate launch` command."
+        )
+    )
+
+    def wrapper(function: F) -> F:
+        for option in reversed(options):
+            function = option(function)
+        return function
+
+    func.__doc__ = (
+        f"{func.__doc__}\n\nThis is ZenML-generated " "CLI wrapper function."
+    )
+
+    return wrapper(func)
+
+
+@contextmanager
+def create_cli_wrapped_script(
+    func: F, flavour: str = "accelerate"
+) -> Iterator[Tuple[Path, Path]]:
+    """Create a script with the CLI-wrapped function.
+
+    Args:
+        func: The function to use.
+        flavour: The flavour to use.
+
+    Yields:
+        The paths of the script and the output.
+
+    Raises:
+        ValueError: If the function is not defined in a module.
+    """
+    try:
+        random_name = random_str(20)
+        script_path = Path(random_name + ".py")
+        output_path = Path(random_name + ".out")
+
+        module = inspect.getmodule(func)
+        if module is None:
+            raise ValueError(
+                f"Function `{func.__name__}` must be defined in a "
+                "module to be used with Accelerate."
+            )
+
+        with open(script_path, "w") as f:
+            if path := module.__file__:
+                func_path = str(Path(path).parent)
+                relative_path = path.replace(func_path, "").lstrip(os.sep)
+                relative_path = os.path.splitext(relative_path)[0]
+                clean_module_name = ".".join(relative_path.split(os.sep))
+                script = _CLI_WRAPPED_SCRIPT_TEMPLATE_HEADER.format(
+                    func_path=func_path,
+                    func_module=clean_module_name,
+                    func_name=func.__name__,
+                )
+                script += _CLI_WRAPPED_MAINS[flavour].format(
+                    output_file=str(output_path.absolute())
+                )
+                f.write(script)
+            else:
+                raise ValueError(
+                    f"Cannot find module file path for function `{func.__name__}`."
+                )
+
+        logger.debug(f"Created script:\n\n{script}")
+
+        yield script_path, output_path
+    finally:
+        if script_path.exists():
+            script_path.unlink()
+        if output_path.exists():
+            output_path.unlink()

zenml/utils/json_utils.py

@@ -0,0 +1,131 @@
+#  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.
+"""Carried over version of some functions from the pydantic v1 json module.
+
+Check out the latest version here:
+https://github.com/pydantic/pydantic/blob/v1.10.15/pydantic/json.py
+"""
+
+import datetime
+from collections import deque
+from decimal import Decimal
+from enum import Enum
+from ipaddress import (
+    IPv4Address,
+    IPv4Interface,
+    IPv4Network,
+    IPv6Address,
+    IPv6Interface,
+    IPv6Network,
+)
+from pathlib import Path
+from re import Pattern
+from types import GeneratorType
+from typing import Any, Callable, Dict, Type, Union
+from uuid import UUID
+
+from pydantic import NameEmail, SecretBytes, SecretStr
+from pydantic.color import Color
+
+__all__ = "pydantic_encoder"
+
+
+def isoformat(obj: Union[datetime.date, datetime.time]) -> str:
+    """Function to convert a datetime into iso format.
+
+    Args:
+        obj: input datetime
+
+    Returns:
+        the corresponding time in iso format.
+    """
+    return obj.isoformat()
+
+
+def decimal_encoder(dec_value: Decimal) -> Union[int, float]:
+    """Encodes a Decimal as int of there's no exponent, otherwise float.
+
+    This is useful when we use ConstrainedDecimal to represent Numeric(x,0)
+    where an integer (but not int typed) is used. Encoding this as a float
+    results in failed round-tripping between encode and parse.
+    Our ID type is a prime example of this.
+
+    >>> decimal_encoder(Decimal("1.0"))
+    1.0
+
+    >>> decimal_encoder(Decimal("1"))
+    1
+
+    Args:
+        dec_value: The input Decimal value
+
+    Returns:
+        the encoded result
+    """
+    if dec_value.as_tuple().exponent >= 0:  # type: ignore[operator]
+        return int(dec_value)
+    else:
+        return float(dec_value)
+
+
+ENCODERS_BY_TYPE: Dict[Type[Any], Callable[[Any], Any]] = {
+    bytes: lambda obj: obj.decode(),
+    Color: str,
+    datetime.date: isoformat,
+    datetime.datetime: isoformat,
+    datetime.time: isoformat,
+    datetime.timedelta: lambda td: td.total_seconds(),
+    Decimal: decimal_encoder,
+    Enum: lambda obj: obj.value,
+    frozenset: list,
+    deque: list,
+    GeneratorType: list,
+    IPv4Address: str,
+    IPv4Interface: str,
+    IPv4Network: str,
+    IPv6Address: str,
+    IPv6Interface: str,
+    IPv6Network: str,
+    NameEmail: str,
+    Path: str,
+    Pattern: lambda obj: obj.pattern,
+    SecretBytes: str,
+    SecretStr: str,
+    set: list,
+    UUID: str,
+}
+
+
+def pydantic_encoder(obj: Any) -> Any:
+    from dataclasses import asdict, is_dataclass
+
+    from pydantic import BaseModel
+
+    if isinstance(obj, BaseModel):
+        return obj.model_dump()
+    elif is_dataclass(obj):
+        return asdict(obj)
+
+    # Check the class type and its superclasses for a matching encoder
+    for base in obj.__class__.__mro__[:-1]:
+        try:
+            encoder = ENCODERS_BY_TYPE[base]
+        except KeyError:
+            continue
+        return encoder(obj)
+    else:  # We have exited the for loop without finding a suitable encoder
+        raise TypeError(
+            f"Object of type '{obj.__class__.__name__}' is not JSON "
+            f"serializable."
+        )

zenml/utils/mlstacks_utils.py

@@ -422,9 +422,9 @@ def convert_mlstacks_primitives_to_dicts(
     verify_mlstacks_prerequisites_installation()
 
     # convert to json first to strip out Enums objects
-    stack_dict = json.loads(stack.json())
+    stack_dict = json.loads(stack.model_dump_json())
     components_dicts = [
-        json.loads(component.json()) for component in components
+        json.loads(component.model_dump_json()) for component in components
     ]
 
     return stack_dict, components_dicts

zenml/utils/package_utils.py

@@ -33,7 +33,7 @@ def is_latest_zenml_version() -> bool:
 
     # Get the latest version from PyPI
     try:
-        response = requests.get("https://pypi.org/pypi/zenml/json")
+        response = requests.get("https://pypi.org/pypi/zenml/json", timeout=60)
         response.raise_for_status()
         latest_published_version = response.json()["info"]["version"]
     except Exception as e:

zenml/utils/pipeline_docker_image_builder.py

@@ -119,6 +119,7 @@ class PipelineDockerImageBuilder:
             ValueError: If no Dockerfile and/or custom parent image is
                 specified and the Docker configuration doesn't require an
                 image build.
+            ValueError: If the specified Dockerfile does not exist.
         """
         requirements: Optional[str] = None
         dockerfile: Optional[str] = None
@@ -133,6 +134,14 @@ class PipelineDockerImageBuilder:
             # pipeline?
             return docker_settings.parent_image, dockerfile, requirements
 
+        if docker_settings.dockerfile and not os.path.isfile(
+            docker_settings.dockerfile
+        ):
+            raise ValueError(
+                "Dockerfile at path "
+                f"{os.path.abspath(docker_settings.dockerfile)} not found."
+            )
+
         stack.validate()
         image_builder = stack.image_builder
         if not image_builder: