zenml-nightly 0.63.0.dev20240801__py3-none-any.whl → 0.64.0.dev20240811__py3-none-any.whl

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/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,

zenml/utils/source_utils.py

@@ -35,14 +35,19 @@ from typing import (
 from zenml.config.source import (
     CodeRepositorySource,
     DistributionPackageSource,
+    NotebookSource,
     Source,
     SourceType,
 )
 from zenml.constants import ENV_ZENML_CUSTOM_SOURCE_ROOT
 from zenml.environment import Environment
 from zenml.logger import get_logger
+from zenml.utils import notebook_utils
 
 logger = get_logger(__name__)
+
+ZENML_SOURCE_ATTRIBUTE_NAME = "__zenml_source__"
+
 NoneType = type(None)
 NoneTypeSource = Source(
     module=NoneType.__module__, attribute="NoneType", type=SourceType.BUILTIN
@@ -58,10 +63,13 @@ BuiltinFunctionTypeSource = Source(
     type=SourceType.BUILTIN,
 )
 
+
 _CUSTOM_SOURCE_ROOT: Optional[str] = os.getenv(
     ENV_ZENML_CUSTOM_SOURCE_ROOT, None
 )
 
+_SHARED_TEMPDIR: Optional[str] = None
+
 
 def load(source: Union[Source, str]) -> Any:
     """Load a source or import path.
@@ -105,6 +113,14 @@ def load(source: Union[Source, str]) -> Any:
                     source.version,
                     source.import_path,
                 )
+    elif source.type == SourceType.NOTEBOOK:
+        if Environment.in_notebook():
+            # If we're in a notebook, we don't need to do anything as the
+            # loading from the __main__ module should work just fine.
+            pass
+        else:
+            notebook_source = NotebookSource.model_validate(dict(source))
+            return _try_to_load_notebook_source(notebook_source)
     elif source.type in {SourceType.USER, SourceType.UNKNOWN}:
         # Unknown source might also refer to a user file, include source
         # root in python path just to be sure
@@ -152,6 +168,9 @@ def resolve(
         return FunctionTypeSource
     elif obj is BuiltinFunctionType:
         return BuiltinFunctionTypeSource
+    elif source := getattr(obj, ZENML_SOURCE_ATTRIBUTE_NAME, None):
+        assert isinstance(source, Source)
+        return source
     elif isinstance(obj, ModuleType):
         module = obj
         attribute_name = None
@@ -216,6 +235,16 @@ def resolve(
         else:
             # Fallback to an unknown source if we can't find the package
             source_type = SourceType.UNKNOWN
+    elif source_type == SourceType.NOTEBOOK:
+        source = NotebookSource(
+            module=module_name,
+            attribute=attribute_name,
+            type=source_type,
+        )
+        # Private attributes are ignored by pydantic if passed in the __init__
+        # method, so we set this afterwards
+        source._cell_code = notebook_utils.load_notebook_cell_code(obj)
+        return source
 
     return Source(
         module=module_name, attribute=attribute_name, type=source_type
@@ -362,7 +391,7 @@ def get_source_type(module: ModuleType) -> SourceType:
         file_path = inspect.getfile(module)
     except (TypeError, OSError):
         if module.__name__ == "__main__" and Environment.in_notebook():
-            return SourceType.USER
+            return SourceType.NOTEBOOK
 
         return SourceType.BUILTIN
 
@@ -529,6 +558,85 @@ def _load_module(
         return importlib.import_module(module_name)
 
 
+def _get_shared_temp_dir() -> str:
+    """Get path to a shared temporary directory.
+
+    Returns:
+        Path to a shared temporary directory.
+    """
+    global _SHARED_TEMPDIR
+
+    if not _SHARED_TEMPDIR:
+        import tempfile
+
+        _SHARED_TEMPDIR = tempfile.mkdtemp()
+
+    return _SHARED_TEMPDIR
+
+
+def _try_to_load_notebook_source(source: NotebookSource) -> Any:
+    """Helper function to load a notebook source outside of a notebook.
+
+    Args:
+        source: The source to load.
+
+    Raises:
+        RuntimeError: If the source can't be loaded.
+
+    Returns:
+        The loaded object.
+    """
+    if not source.code_path or not source.replacement_module:
+        raise RuntimeError(
+            f"Failed to load {source.import_path}. This object was defined in "
+            "a notebook and you're trying to load it outside of a notebook. "
+            "This is currently only enabled for ZenML steps."
+        )
+
+    extract_dir = _get_shared_temp_dir()
+    file_path = os.path.join(extract_dir, f"{source.replacement_module}.py")
+
+    if not os.path.exists(file_path):
+        from zenml.utils import code_utils
+
+        logger.info(
+            "Downloading notebook cell content from `%s` to load `%s`.",
+            source.code_path,
+            source.import_path,
+        )
+
+        code_utils.download_and_extract_code(
+            code_path=source.code_path, extract_dir=extract_dir
+        )
+
+    try:
+        module = _load_module(
+            module_name=source.replacement_module, import_root=extract_dir
+        )
+    except ImportError:
+        raise RuntimeError(
+            f"Unable to load {source.import_path}. This object was defined in "
+            "a notebook and you're trying to load it outside of a notebook. "
+            "To enable this, ZenML extracts the code of your cell into a "
+            "python file. This means your cell code needs to be "
+            "self-contained:\n"
+            "  * All required imports must be done in this cell, even if the "
+            "same imports already happen in previous notebook cells.\n"
+            "  * The cell can't use any code defined in other notebook cells."
+        )
+
+    if source.attribute:
+        obj = getattr(module, source.attribute)
+    else:
+        obj = module
+
+    # Store the original notebook source so resolving this object works as
+    # expected
+    setattr(obj, ZENML_SOURCE_ATTRIBUTE_NAME, source)
+
+    return obj
+
+
 def _get_package_for_module(module_name: str) -> Optional[str]:
     """Get the package name for a module.
 

zenml/zen_server/dashboard/assets/{404-CI13wQp4.js → 404-CRAA_Lew.js}

@@ -1 +1 @@
-import{j as e}from"./@radix-CFOkMR_E.js";import{h as s,r as t}from"./index-KsTz2dHG.js";import{E as r}from"./EmptyState-BMLnFVlB.js";import{H as a}from"./help-bu_DgLKI.js";import{L as o}from"./@react-router-CO-OsFwI.js";import"./@tanstack-k96lU_C-.js";import"./@reactflow-DIYUhKYX.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@radix-BXWm7HOa.js";import{h as s,r as t}from"./index-DhIZtpxB.js";import{E as r}from"./EmptyState-BHblM39I.js";import{H as a}from"./help-FuHlZwn0.js";import{L as o}from"./@react-router-l3lMcXA2.js";import"./@tanstack-FmcYZMuX.js";import"./@reactflow-CeVxyqYT.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};