zenml-nightly 0.62.0.dev20240729__py3-none-any.whl → 0.64.0.dev20240809__py3-none-any.whl

zenml/utils/code_utils.py

@@ -0,0 +1,244 @@
+#  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 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 string_utils
+from zenml.utils.archivable import Archivable
+
+if TYPE_CHECKING:
+    from git.repo.base import Repo
+
+
+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], use_gzip: bool = True
+    ) -> None:
+        """Writes an archive of the build context to the given file.
+
+        Args:
+            output_file: The file to write the archive to.
+            use_gzip: Whether to use `gzip` to compress the file.
+        """
+        super().write_archive(output_file=output_file, use_gzip=use_gzip)
+        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 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 to 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)
+
+        hash_ = hashlib.sha1()  # nosec
+
+        while True:
+            data = f.read(64 * 1024)
+            if not data:
+                break
+            hash_.update(data)
+
+        filename = f"{hash_.hexdigest()}.tar.gz"
+        upload_dir = os.path.join(artifact_store.path, "code_uploads")
+        fileio.makedirs(upload_dir)
+        upload_path = os.path.join(upload_dir, filename)
+
+        if not fileio.exists(upload_path):
+            archive_size = string_utils.get_human_readable_filesize(
+                os.path.getsize(f.name)
+            )
+            logger.info(
+                "Uploading code to `%s` (Size: %s).", upload_path, archive_size
+            )
+            fileio.copy(f.name, upload_path)
+            logger.info("Code upload finished.")
+        else:
+            logger.info(
+                "Code already exists in artifact store, skipping upload."
+            )
+
+    if os.path.exists(f.name):
+        os.remove(f.name)
+
+    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)

zenml/utils/notebook_utils.py

@@ -0,0 +1,122 @@
+#  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.
+"""Notebook utilities."""
+
+from typing import Any, Callable, Optional, TypeVar, Union
+
+from zenml.environment import Environment
+from zenml.logger import get_logger
+
+ZENML_NOTEBOOK_CELL_CODE_ATTRIBUTE_NAME = "__zenml_notebook_cell_code__"
+
+AnyObject = TypeVar("AnyObject", bound=Any)
+
+logger = get_logger(__name__)
+
+
+def is_defined_in_notebook_cell(obj: Any) -> bool:
+    """Check whether an object is defined in a notebook cell.
+
+    Args:
+        obj: The object to check.
+
+    Returns:
+        Whether the object is defined in a notebook cell.
+    """
+    if not Environment.in_notebook():
+        return False
+
+    module_name = getattr(obj, "__module__", None)
+    return module_name == "__main__"
+
+
+def enable_notebook_code_extraction(
+    _obj: Optional["AnyObject"] = None,
+) -> Union["AnyObject", Callable[["AnyObject"], "AnyObject"]]:
+    """Decorator to enable code extraction from notebooks.
+
+    Args:
+        _obj: The class or function for which to enable code extraction.
+
+    Returns:
+        The decorated class or function.
+    """
+
+    def inner_decorator(obj: "AnyObject") -> "AnyObject":
+        try_to_save_notebook_cell_code(obj)
+        return obj
+
+    if _obj is None:
+        return inner_decorator
+    else:
+        return inner_decorator(_obj)
+
+
+def get_active_notebook_cell_code() -> Optional[str]:
+    """Get the code of the currently active notebook cell.
+
+    Returns:
+        The code of the currently active notebook cell.
+    """
+    cell_code = None
+    try:
+        ipython = get_ipython()  # type: ignore[name-defined]
+        cell_code = ipython.get_parent()["content"]["code"]
+    except (NameError, KeyError) as e:
+        logger.warning("Unable to extract cell code: %s.", str(e))
+
+    return cell_code
+
+
+def try_to_save_notebook_cell_code(obj: Any) -> None:
+    """Try to save the notebook cell code for an object.
+
+    Args:
+        obj: The object for which to save the notebook cell code.
+    """
+    if is_defined_in_notebook_cell(obj):
+        if cell_code := get_active_notebook_cell_code():
+            setattr(
+                obj,
+                ZENML_NOTEBOOK_CELL_CODE_ATTRIBUTE_NAME,
+                cell_code,
+            )
+
+
+def load_notebook_cell_code(obj: Any) -> Optional[str]:
+    """Load the notebook cell code for an object.
+
+    Args:
+        obj: The object for which to load the cell code.
+
+    Returns:
+        The notebook cell code if it was saved.
+    """
+    return getattr(obj, ZENML_NOTEBOOK_CELL_CODE_ATTRIBUTE_NAME, None)
+
+
+def warn_about_notebook_cell_magic_commands(cell_code: str) -> None:
+    """Warn about magic commands in the cell code.
+
+    Args:
+        cell_code: The cell code.
+    """
+    if any(line.startswith(("%", "!")) for line in cell_code.splitlines()):
+        logger.warning(
+            "Some lines in your notebook cell start with a `!` or `%` "
+            "character. Running a ZenML step remotely from a notebook "
+            "only works if the cell only contains python code. If any "
+            "of these lines contain Jupyter notebook magic commands, "
+            "remove them and try again."
+        )

zenml/utils/package_utils.py

@@ -13,6 +13,8 @@
 #  permissions and limitations under the License.
 """Utility functions for the package."""
 
+from typing import List
+
 import requests
 from packaging import version
 
@@ -48,3 +50,40 @@ def is_latest_zenml_version() -> bool:
         return False
     else:
         return True
+
+
+def clean_requirements(requirements: List[str]) -> List[str]:
+    """Clean requirements list from redundant requirements.
+
+    Args:
+        requirements: List of requirements.
+
+    Returns:
+        Cleaned list of requirements
+
+    Raises:
+        TypeError: If input is not a list
+        ValueError: If any element in the list is not a string
+    """
+    if not isinstance(requirements, list):
+        raise TypeError("Input must be a list")
+
+    if not all(isinstance(req, str) for req in requirements):
+        raise ValueError("All elements in the list must be strings")
+
+    cleaned = {}
+    for req in requirements:
+        package = (
+            req.split(">=")[0]
+            .split("==")[0]
+            .split("<")[0]
+            .split("~=")[0]
+            .split("^=")[0]
+            .split("[")[0]
+            .strip()
+        )
+        if package not in cleaned or any(
+            op in req for op in ["=", ">", "<", "~", "^"]
+        ):
+            cleaned[package] = req
+    return sorted(cleaned.values())

zenml/utils/pipeline_docker_image_builder.py

@@ -17,11 +17,9 @@ import itertools
 import os
 import subprocess
 import sys
-from collections import defaultdict
 from typing import (
     TYPE_CHECKING,
     Any,
-    DefaultDict,
     Dict,
     List,
     Optional,
@@ -277,9 +275,7 @@ class PipelineDockerImageBuilder:
             requirements_files = self.gather_requirements_files(
                 docker_settings=docker_settings,
                 stack=stack,
-                # Only pass code repo to include its dependencies if we actually
-                # need to download code
-                code_repository=code_repository if download_files else None,
+                code_repository=code_repository,
             )
 
             self._add_requirements_files(
@@ -444,8 +440,9 @@ class PipelineDockerImageBuilder:
             requirements files.
             The files will be in the following order:
             - Packages installed in the local Python environment
+            - Requirements defined by stack integrations
+            - Requirements defined by user integrations
             - User-defined requirements
-            - Requirements defined by user-defined and/or stack integrations
         """
         requirements_files: List[Tuple[str, str, List[str]]] = []
 
@@ -481,43 +478,6 @@ class PipelineDockerImageBuilder:
                     "- Including python packages from local environment"
                 )
 
-        # Generate requirements files for all ZenML Hub plugins
-        if docker_settings.required_hub_plugins:
-            (
-                hub_internal_requirements,
-                hub_pypi_requirements,
-            ) = PipelineDockerImageBuilder._get_hub_requirements(
-                docker_settings.required_hub_plugins
-            )
-
-            # Plugin packages themselves
-            for i, (index, packages) in enumerate(
-                hub_internal_requirements.items()
-            ):
-                file_name = f".zenml_hub_internal_requirements_{i}"
-                file_lines = [f"-i {index}", *packages]
-                file_contents = "\n".join(file_lines)
-                requirements_files.append(
-                    (file_name, file_contents, ["--no-deps"])
-                )
-                if log:
-                    logger.info(
-                        "- Including internal hub packages from index `%s`: %s",
-                        index,
-                        ", ".join(f"`{r}`" for r in packages),
-                    )
-
-            # PyPI requirements of plugin packages
-            if hub_pypi_requirements:
-                file_name = ".zenml_hub_pypi_requirements"
-                file_contents = "\n".join(hub_pypi_requirements)
-                requirements_files.append((file_name, file_contents, []))
-                if log:
-                    logger.info(
-                        "- Including hub requirements from PyPI: %s",
-                        ", ".join(f"`{r}`" for r in hub_pypi_requirements),
-                    )
-
         if docker_settings.install_stack_requirements:
             stack_requirements = stack.requirements()
             if code_repository:
@@ -599,59 +559,6 @@ class PipelineDockerImageBuilder:
 
         return requirements_files
 
-    @staticmethod
-    def _get_hub_requirements(
-        required_hub_plugins: List[str],
-    ) -> Tuple[Dict[str, List[str]], List[str]]:
-        """Get package requirements for ZenML Hub plugins.
-
-        Args:
-            required_hub_plugins: List of hub plugin names in the format
-                `(<author_username>/)<plugin_name>(==<version>)`.
-
-        Returns:
-            - A dict of the hub plugin packages themselves (which need to be
-                installed from a custom index, mapping index URLs to lists of
-                package names.
-            - A list of all unique dependencies of the required hub plugins
-                (which can be installed from PyPI).
-        """
-        from zenml._hub.client import HubClient
-        from zenml._hub.utils import parse_plugin_name, plugin_display_name
-
-        client = HubClient()
-
-        internal_requirements: DefaultDict[str, List[str]] = defaultdict(list)
-        pypi_requirements: List[str] = []
-
-        for plugin_str in required_hub_plugins:
-            author, name, version = parse_plugin_name(
-                plugin_str, version_separator="=="
-            )
-
-            plugin = client.get_plugin(
-                name=name,
-                version=version,
-                author=author,
-            )
-
-            if plugin and plugin.index_url and plugin.package_name:
-                internal_requirements[plugin.index_url].append(
-                    plugin.package_name
-                )
-                if plugin.requirements:
-                    pypi_requirements.extend(plugin.requirements)
-            else:
-                display_name = plugin_display_name(name, version, author)
-                logger.warning(
-                    "Hub plugin `%s` does not exist or cannot be installed."
-                    "Skipping installation of this plugin.",
-                    display_name,
-                )
-
-        pypi_requirements = sorted(set(pypi_requirements))
-        return dict(internal_requirements), pypi_requirements
-
     @staticmethod
     def _generate_zenml_pipeline_dockerfile(
         parent_image: str,