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

zenml/utils/archivable.py

@@ -0,0 +1,178 @@
+#  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.
+"""Archivable mixin."""
+
+import io
+import tarfile
+import zipfile
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import IO, Any, Dict, Optional
+
+from zenml.io import fileio
+from zenml.utils.enum_utils import StrEnum
+
+
+class ArchiveType(StrEnum):
+    """Archive types supported by the ZenML build context."""
+
+    TAR = "tar"
+    TAR_GZ = "tar.gz"
+    ZIP = "zip"
+
+
+class Archivable(ABC):
+    """Archivable mixin class."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize the object.
+
+        Args:
+            *args: Unused args for subclasses.
+            **kwargs: Unused keyword args for subclasses.
+        """
+        self._extra_files: Dict[str, str] = {}
+
+    def add_file(self, source: str, destination: str) -> None:
+        """Adds a file to the archive.
+
+        Args:
+            source: The source of the file to add. This can either be a path
+                or the file content.
+            destination: The path inside the archive where the file
+                should be added.
+        """
+        if fileio.exists(source):
+            with fileio.open(source) as f:
+                self._extra_files[destination] = f.read()
+        else:
+            self._extra_files[destination] = source
+
+    def add_directory(self, source: str, destination: str) -> None:
+        """Adds a directory to the archive.
+
+        Args:
+            source: Path to the directory.
+            destination: The path inside the build context where the directory
+                should be added.
+
+        Raises:
+            ValueError: If `source` does not point to a directory.
+        """
+        if not fileio.isdir(source):
+            raise ValueError(
+                f"Can't add directory {source} to the build context as it "
+                "does not exist or is not a directory."
+            )
+
+        for dir, _, files in fileio.walk(source):
+            dir_path = Path(fileio.convert_to_str(dir))
+            for file_name in files:
+                file_name = fileio.convert_to_str(file_name)
+                file_source = dir_path / file_name
+                file_destination = (
+                    Path(destination)
+                    / dir_path.relative_to(source)
+                    / file_name
+                )
+
+                with file_source.open("r") as f:
+                    self._extra_files[file_destination.as_posix()] = f.read()
+
+    def write_archive(
+        self,
+        output_file: IO[bytes],
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
+    ) -> None:
+        """Writes an archive of the build context to the given file.
+
+        Args:
+            output_file: The file to write the archive to.
+            archive_type: The type of archive to create.
+        """
+        files = self.get_files()
+        extra_files = self.get_extra_files()
+        close_fileobj: Optional[Any] = None
+        fileobj: Any = output_file
+
+        if archive_type == ArchiveType.ZIP:
+            fileobj = zipfile.ZipFile(output_file, "w", zipfile.ZIP_DEFLATED)
+        else:
+            if archive_type == ArchiveType.TAR_GZ:
+                from gzip import GzipFile
+
+                # We don't use the builtin gzip functionality of the `tarfile`
+                # library as that one includes the tar filename and creation
+                # timestamp in the archive which causes the hash of the resulting
+                # file to be different each time. We use this hash to avoid
+                # duplicate uploads, which is why we pass empty values for filename
+                # and mtime here.
+                close_fileobj = fileobj = GzipFile(
+                    filename="", mode="wb", fileobj=output_file, mtime=0.0
+                )
+            fileobj = tarfile.open(mode="w", fileobj=fileobj)
+
+        try:
+            with fileobj as af:
+                for archive_path, file_path in files.items():
+                    if archive_path in extra_files:
+                        continue
+                    if archive_type == ArchiveType.ZIP:
+                        assert isinstance(af, zipfile.ZipFile)
+                        af.write(file_path, arcname=archive_path)
+                    else:
+                        assert isinstance(af, tarfile.TarFile)
+                        if info := af.gettarinfo(
+                            file_path, arcname=archive_path
+                        ):
+                            if info.isfile():
+                                with open(file_path, "rb") as f:
+                                    af.addfile(info, f)
+                            else:
+                                af.addfile(info, None)
+
+                for archive_path, contents in extra_files.items():
+                    contents_encoded = contents.encode("utf-8")
+
+                    if archive_type == ArchiveType.ZIP:
+                        assert isinstance(af, zipfile.ZipFile)
+                        af.writestr(archive_path, contents_encoded)
+                    else:
+                        assert isinstance(af, tarfile.TarFile)
+                        info = tarfile.TarInfo(archive_path)
+                        info.size = len(contents_encoded)
+                        af.addfile(info, io.BytesIO(contents_encoded))
+        finally:
+            if close_fileobj:
+                close_fileobj.close()
+
+        output_file.seek(0)
+
+    @abstractmethod
+    def get_files(self) -> Dict[str, str]:
+        """Gets all regular files that should be included in the archive.
+
+        Returns:
+            A dict {path_in_archive: path_on_filesystem} for all regular files
+            in the archive.
+        """
+
+    def get_extra_files(self) -> Dict[str, str]:
+        """Gets all extra files that should be included in the archive.
+
+        Returns:
+            A dict {path_in_archive: file_content} for all extra files in the
+            archive.
+        """
+        return self._extra_files.copy()

zenml/utils/callback_registry.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.
+"""Callback registry implementation."""
+
+from typing import Any, Callable, Dict, List, Tuple
+
+from typing_extensions import ParamSpec
+
+from zenml.logger import get_logger
+
+P = ParamSpec("P")
+
+logger = get_logger(__name__)
+
+
+class CallbackRegistry:
+    """Callback registry class."""
+
+    def __init__(self) -> None:
+        """Initializes the callback registry."""
+        self._callbacks: List[
+            Tuple[Callable[P, Any], Tuple[Any], Dict[str, Any]]
+        ] = []
+
+    def register_callback(
+        self, callback: Callable[P, Any], *args: P.args, **kwargs: P.kwargs
+    ) -> None:
+        """Register a callback.
+
+        Args:
+            callback: The callback to register.
+            *args: Arguments to call the callback with.
+            **kwargs: Keyword arguments to call the callback with.
+        """
+        self._callbacks.append((callback, args, kwargs))  # type: ignore[arg-type]
+
+    def reset(self) -> None:
+        """Reset the callbacks."""
+        self._callbacks = []
+
+    def execute_callbacks(self, raise_on_exception: bool) -> None:
+        """Execute all registered callbacks.
+
+        Args:
+            raise_on_exception: If True, exceptions raised during the execution
+                of the callbacks will be raised. If False, a warning with the
+                exception will be logged instead.
+
+        Raises:
+            Exception: Exceptions raised in any of the callbacks if
+                `raise_on_exception` is set to True.
+        """
+        for callback, args, kwargs in self._callbacks:
+            try:
+                callback(*args, **kwargs)
+            except Exception as e:
+                if raise_on_exception:
+                    raise e
+                else:
+                    logger.warning("Failed to run callback: %s", str(e))

zenml/utils/code_repository_utils.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Utilities for code repositories."""
+
 import os
 from typing import (
     Dict,

zenml/utils/code_utils.py

@@ -0,0 +1,346 @@
+#  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.
+"""Code utilities."""
+
+import hashlib
+import os
+import shutil
+import sys
+import tempfile
+from pathlib import Path
+from typing import IO, TYPE_CHECKING, Dict, Optional
+
+from zenml.client import Client
+from zenml.io import fileio
+from zenml.logger import get_logger
+from zenml.utils import source_utils, string_utils
+from zenml.utils.archivable import Archivable, ArchiveType
+
+if TYPE_CHECKING:
+    from git.repo.base import Repo
+
+    from zenml.artifact_stores import BaseArtifactStore
+
+
+logger = get_logger(__name__)
+
+
+class CodeArchive(Archivable):
+    """Code archive class.
+
+    This class is used to archive user code before uploading it to the artifact
+    store. If the user code is stored in a Git repository, only files not
+    excluded by gitignores will be included in the archive.
+    """
+
+    def __init__(self, root: Optional[str] = None) -> None:
+        """Initialize the object.
+
+        Args:
+            root: Root directory of the archive.
+        """
+        super().__init__()
+        self._root = root
+
+    @property
+    def git_repo(self) -> Optional["Repo"]:
+        """Git repository active at the code archive root.
+
+        Returns:
+            The git repository if available.
+        """
+        try:
+            # These imports fail when git is not installed on the machine
+            from git.exc import InvalidGitRepositoryError
+            from git.repo.base import Repo
+        except ImportError:
+            return None
+
+        try:
+            git_repo = Repo(path=self._root, search_parent_directories=True)
+        except InvalidGitRepositoryError:
+            return None
+
+        return git_repo
+
+    def _get_all_files(self, archive_root: str) -> Dict[str, str]:
+        """Get all files inside the archive root.
+
+        Args:
+            archive_root: The root directory from which to get all files.
+
+        Returns:
+            All files inside the archive root.
+        """
+        all_files = {}
+        for root, _, files in os.walk(archive_root):
+            for file in files:
+                file_path = os.path.join(root, file)
+                path_in_archive = os.path.relpath(file_path, archive_root)
+                all_files[path_in_archive] = file_path
+
+        return all_files
+
+    def get_files(self) -> Dict[str, str]:
+        """Gets all regular files that should be included in the archive.
+
+        Raises:
+            RuntimeError: If the code archive would not include any files.
+
+        Returns:
+            A dict {path_in_archive: path_on_filesystem} for all regular files
+            in the archive.
+        """
+        if not self._root:
+            return {}
+
+        all_files = {}
+
+        if repo := self.git_repo:
+            try:
+                result = repo.git.ls_files(
+                    "--cached",
+                    "--others",
+                    "--modified",
+                    "--exclude-standard",
+                    self._root,
+                )
+            except Exception as e:
+                logger.warning(
+                    "Failed to get non-ignored files from git: %s", str(e)
+                )
+                all_files = self._get_all_files(archive_root=self._root)
+            else:
+                for file in result.split():
+                    file_path = os.path.join(repo.working_dir, file)
+                    path_in_archive = os.path.relpath(file_path, self._root)
+
+                    if os.path.exists(file_path):
+                        all_files[path_in_archive] = file_path
+        else:
+            all_files = self._get_all_files(archive_root=self._root)
+
+        if not all_files:
+            raise RuntimeError(
+                "The code archive to be uploaded does not contain any files. "
+                "This is probably because all files in your source root "
+                f"`{self._root}` are ignored by a .gitignore file."
+            )
+
+        # Explicitly remove .zen directories as we write an updated version
+        # to disk everytime ZenML is called. This updates the mtime of the
+        # file, which invalidates the code upload caching. The values in
+        # the .zen directory are not needed anyway as we set them as
+        # environment variables.
+        all_files = {
+            path_in_archive: file_path
+            for path_in_archive, file_path in sorted(all_files.items())
+            if ".zen" not in Path(path_in_archive).parts[:-1]
+        }
+
+        return all_files
+
+    def write_archive(
+        self,
+        output_file: IO[bytes],
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
+    ) -> None:
+        """Writes an archive of the build context to the given file.
+
+        Args:
+            output_file: The file to write the archive to.
+            archive_type: The type of archive to create.
+        """
+        super().write_archive(
+            output_file=output_file, archive_type=archive_type
+        )
+        archive_size = os.path.getsize(output_file.name)
+        if archive_size > 20 * 1024 * 1024:
+            logger.warning(
+                "Code archive size: `%s`. If you believe this is "
+                "unreasonably large, make sure to version your code in git and "
+                "ignore unnecessary files using a `.gitignore` file.",
+                string_utils.get_human_readable_filesize(archive_size),
+            )
+
+
+def compute_file_hash(file: IO[bytes]) -> str:
+    """Compute a hash of the content of a file.
+
+    This function will not seek the file before or after the hash computation.
+    This means that the content will be computed based on the current cursor
+    until the end of the file.
+
+    Args:
+        file: The file for which to compute the hash.
+
+    Returns:
+        A hash of the file content.
+    """
+    hash_ = hashlib.sha1()  # nosec
+
+    while True:
+        data = file.read(64 * 1024)
+        if not data:
+            break
+        hash_.update(data)
+
+    return hash_.hexdigest()
+
+
+def upload_code_if_necessary(code_archive: CodeArchive) -> str:
+    """Upload code to the artifact store if necessary.
+
+    This function computes a hash of the code to be uploaded, and if an archive
+    with the same hash already exists it will not re-upload but instead return
+    the path to the existing archive.
+
+    Args:
+        code_archive: The code archive to upload.
+
+    Returns:
+        The path where the archived code is uploaded.
+    """
+    artifact_store = Client().active_stack.artifact_store
+
+    with tempfile.NamedTemporaryFile(
+        mode="w+b", delete=False, suffix=".tar.gz"
+    ) as f:
+        code_archive.write_archive(f)
+        archive_path = f.name
+        archive_hash = compute_file_hash(f)
+
+    upload_dir = os.path.join(artifact_store.path, "code_uploads")
+    fileio.makedirs(upload_dir)
+    upload_path = os.path.join(upload_dir, f"{archive_hash}.tar.gz")
+
+    if not fileio.exists(upload_path):
+        archive_size = string_utils.get_human_readable_filesize(
+            os.path.getsize(archive_path)
+        )
+        logger.info(
+            "Uploading code to `%s` (Size: %s).", upload_path, archive_size
+        )
+        fileio.copy(archive_path, upload_path)
+        logger.info("Code upload finished.")
+    else:
+        logger.info("Code already exists in artifact store, skipping upload.")
+
+    if os.path.exists(archive_path):
+        os.remove(archive_path)
+
+    return upload_path
+
+
+def download_and_extract_code(code_path: str, extract_dir: str) -> None:
+    """Download and extract code.
+
+    Args:
+        code_path: Path where the code is uploaded.
+        extract_dir: Directory where to code should be extracted to.
+
+    Raises:
+        RuntimeError: If the code is stored in an artifact store which is
+            not active.
+    """
+    artifact_store = Client().active_stack.artifact_store
+
+    if not code_path.startswith(artifact_store.path):
+        raise RuntimeError("Code stored in different artifact store.")
+
+    download_path = os.path.basename(code_path)
+    fileio.copy(code_path, download_path)
+
+    shutil.unpack_archive(filename=download_path, extract_dir=extract_dir)
+    os.remove(download_path)
+
+
+def download_code_from_artifact_store(code_path: str) -> None:
+    """Download code from the artifact store.
+
+    Args:
+        code_path: Path where the code is stored.
+    """
+    logger.info("Downloading code from artifact store path `%s`.", code_path)
+
+    # Do not remove this line, we need to instantiate the artifact store to
+    # register the filesystem needed for the file download
+    _ = Client().active_stack.artifact_store
+
+    extract_dir = os.path.abspath("code")
+    os.makedirs(extract_dir)
+
+    download_and_extract_code(code_path=code_path, extract_dir=extract_dir)
+
+    source_utils.set_custom_source_root(extract_dir)
+    sys.path.insert(0, extract_dir)
+    os.chdir(extract_dir)
+
+
+def _get_notebook_upload_dir(artifact_store: "BaseArtifactStore") -> str:
+    """Get the upload directory for code extracted from notebook cells.
+
+    Args:
+        artifact_store: The artifact store in which the directory should be.
+
+    Returns:
+        The upload directory for code extracted from notebook cells.
+    """
+    return os.path.join(artifact_store.path, "notebook_code")
+
+
+def upload_notebook_code(
+    artifact_store: "BaseArtifactStore", cell_code: str, file_name: str
+) -> None:
+    """Upload code extracted from a notebook cell.
+
+    Args:
+        artifact_store: The artifact store in which to upload the code.
+        cell_code: The notebook cell code.
+        file_name: The filename to use for storing the cell code.
+    """
+    upload_dir = _get_notebook_upload_dir(artifact_store=artifact_store)
+    fileio.makedirs(upload_dir)
+    upload_path = os.path.join(upload_dir, file_name)
+
+    if not fileio.exists(upload_path):
+        with fileio.open(upload_path, "wb") as f:
+            f.write(cell_code.encode())
+
+        logger.info("Uploaded notebook cell code to %s.", upload_path)
+
+
+def download_notebook_code(
+    artifact_store: "BaseArtifactStore", file_name: str, download_path: str
+) -> None:
+    """Download code extracted from a notebook cell.
+
+    Args:
+        artifact_store: The artifact store from which to download the code.
+        file_name: The name of the code file.
+        download_path: The local path where the file should be downloaded to.
+
+    Raises:
+        FileNotFoundError: If no file with the given filename exists in this
+            artifact store.
+    """
+    code_dir = _get_notebook_upload_dir(artifact_store=artifact_store)
+    code_path = os.path.join(code_dir, file_name)
+
+    if not fileio.exists(code_path):
+        raise FileNotFoundError(
+            f"Notebook code at path {code_path} not found."
+        )
+
+    fileio.copy(code_path, download_path)

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()