Upload notebook code to artifact store instead

Author: schustmi

src/zenml/config/source.py

@@ -234,17 +234,20 @@ class NotebookSource(Source):
     """Source representing an object defined in a notebook.
 
     Attributes:
-        notebook_path: Path of the notebook (relative to the source root) in
-            which the object is defined.
-        cell_id: ID of the notebook cell in which the object is defined. This
-            will only be set for objects which explicitly store this by calling
-            `zenml.utils.notebook_utils.save_notebook_cell_id()`.
+        code_path: Path where the notebook cell code for this source is
+            uploaded.
+        replacement_module: Name of the module from which this source should
+            be loaded in case the code is not running in a notebook.
     """
 
-    notebook_path: Optional[str] = None
-    cell_id: Optional[str] = None
+    code_path: Optional[str] = None
+    replacement_module: Optional[str] = None
     type: SourceType = SourceType.NOTEBOOK
 
+    # Private attribute that is used to store the code but should not be
+    # serialized
+    _cell_code: Optional[str] = None
+
     @field_validator("type")
     @classmethod
     def _validate_type(cls, value: SourceType) -> SourceType:

src/zenml/constants.py

@@ -173,7 +173,6 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 )
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
-ENV_ZENML_NOTEBOOK_PATH = "ZENML_NOTEBOOK_PATH"
 
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"

src/zenml/entrypoints/base_entrypoint_configuration.py

@@ -15,7 +15,6 @@
 
 import argparse
 import os
-import shutil
 import sys
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Set
@@ -27,9 +26,13 @@
     ENV_ZENML_REQUIRES_CODE_DOWNLOAD,
     handle_bool_env_var,
 )
-from zenml.io import fileio
 from zenml.logger import get_logger
-from zenml.utils import code_repository_utils, source_utils, uuid_utils
+from zenml.utils import (
+    code_repository_utils,
+    code_utils,
+    source_utils,
+    uuid_utils,
+)
 
 if TYPE_CHECKING:
     from zenml.models import CodeReferenceResponse, PipelineDeploymentResponse
@@ -272,19 +275,14 @@ def download_code_from_artifact_store(self, code_path: str) -> None:
 
         # Do not remove this line, we need to instantiate the artifact store to
         # register the filesystem needed for the file download
-        artifact_store = Client().active_stack.artifact_store
-
-        if not code_path.startswith(artifact_store.path):
-            raise RuntimeError("Code stored in different artifact store.")
+        _ = Client().active_stack.artifact_store
 
         extract_dir = os.path.abspath("code")
         os.makedirs(extract_dir)
 
-        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)
+        code_utils.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)

src/zenml/materializers/base_materializer.py

@@ -100,7 +100,7 @@ def __new__(
 
             from zenml.utils import notebook_utils
 
-            notebook_utils.try_to_save_notebook_cell_id(cls)
+            notebook_utils.try_to_save_notebook_cell_code(cls)
 
         return cls
 

src/zenml/new/pipelines/build_utils.py

@@ -14,9 +14,7 @@
 """Pipeline build utilities."""
 
 import hashlib
-import os
 import platform
-import tempfile
 from typing import (
     TYPE_CHECKING,
     Dict,
@@ -29,7 +27,6 @@
 import zenml
 from zenml.client import Client
 from zenml.code_repositories import BaseCodeRepository
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.models import (
     BuildItem,
@@ -40,9 +37,8 @@
     PipelineDeploymentBase,
     StackResponse,
 )
-from zenml.new.pipelines.code_archive import CodeArchive
 from zenml.stack import Stack
-from zenml.utils import source_utils, string_utils
+from zenml.utils import source_utils
 from zenml.utils.pipeline_docker_image_builder import (
     PipelineDockerImageBuilder,
 )
@@ -724,56 +720,3 @@ def should_upload_code(
             return True
 
     return False
-
-
-def upload_code_if_necessary() -> 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.
-
-    Returns:
-        The path where to archived code is uploaded.
-    """
-    logger.info("Archiving code...")
-
-    code_archive = CodeArchive(root=source_utils.get_source_root())
-    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

src/zenml/new/pipelines/pipeline.py

@@ -72,8 +72,8 @@
 from zenml.new.pipelines.run_utils import (
     create_placeholder_run,
     deploy_pipeline,
-    fail_if_running_remotely_with_notebook_not_possible,
     prepare_model_versions,
+    upload_notebook_cell_code_if_necessary,
 )
 from zenml.stack import Stack
 from zenml.steps import BaseStep
@@ -83,6 +83,7 @@
 from zenml.steps.step_invocation import StepInvocation
 from zenml.utils import (
     code_repository_utils,
+    code_utils,
     dashboard_utils,
     dict_utils,
     pydantic_utils,
@@ -669,7 +670,7 @@ def _run(
 
             stack = Client().active_stack
             stack.validate()
-            fail_if_running_remotely_with_notebook_not_possible(
+            upload_notebook_cell_code_if_necessary(
                 deployment=deployment, stack=stack
             )
 
@@ -719,7 +720,11 @@ def _run(
                 build=build_model,
                 code_reference=code_reference,
             ):
-                code_path = build_utils.upload_code_if_necessary()
+                code_archive = code_utils.CodeArchive(
+                    root=source_utils.get_source_root()
+                )
+                logger.info("Archiving pipeline code...")
+                code_path = code_utils.upload_code_if_necessary(code_archive)
 
             deployment_request = PipelineDeploymentRequest(
                 user=Client().active_user.id,

src/zenml/new/pipelines/run_utils.py

@@ -1,5 +1,6 @@
 """Utility functions for running pipelines."""
 
+import hashlib
 import time
 from collections import defaultdict
 from datetime import datetime
@@ -24,7 +25,7 @@
 from zenml.new.pipelines.model_utils import NewModelRequest
 from zenml.orchestrators.utils import get_run_name
 from zenml.stack import Flavor, Stack
-from zenml.utils import cloud_utils, notebook_utils
+from zenml.utils import cloud_utils, code_utils, notebook_utils
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
@@ -364,42 +365,63 @@ def validate_run_config_is_runnable_from_server(
             )
 
 
-def fail_if_running_remotely_with_notebook_not_possible(
+def upload_notebook_cell_code_if_necessary(
     deployment: "PipelineDeploymentBase", stack: "Stack"
 ) -> None:
-    """Fail if running the deployment on the stack is not possible.
+    """Upload notebook cell code if necessary.
 
     This function checks if any of the steps of the pipeline that will be
     executed in a different process are defined in a notebook. If that is the
-    case, it will raise an exception if the active notebook path can't be
-    determined.
+    case, it will extract that notebook cell code into python files and upload
+    an archive of all the necessary files to the artifact store.
 
     Args:
         deployment: The deployment.
         stack: The stack on which the deployment will happen.
 
     Raises:
-        RuntimeError: If the active notebook can't be determined and steps that
-            are defined in that notebook should be executed out of process.
+        RuntimeError: If the code for one of the steps that will run out of
+            process cannot be extracted into a python file.
     """
+    code_archive = code_utils.CodeArchive(root=None)
+    should_upload = False
+    sources_that_require_upload = []
+
     for step in deployment.step_configurations.values():
         if step.spec.source.type == SourceType.NOTEBOOK:
             if (
                 stack.orchestrator.flavor != "local"
                 or step.config.step_operator
             ):
+                should_upload = True
+                cell_code = getattr(step.spec.source, "_cell_code", None)
+
                 # Code does not run in-process, which means we need to
-                # extract it from the notebook in the execution
-                # environment -> verify that we're able to detect the
-                # active notebook
-                if not notebook_utils.get_active_notebook_path():
+                # extract the step code into a python file
+                if not cell_code:
                     raise RuntimeError(
                         f"Unable to run step {step.config.name}. This step is "
                         "defined in a notebook and you're trying to run it "
                         "in a remote environment, but ZenML was not able to "
-                        "detect the notebook that you're running in. To fix "
-                        "this error, set the "
-                        f"{constants.ENV_ZENML_NOTEBOOK_PATH} environment "
-                        "variable to the path of the active notebook or define "
-                        "your step in a python file instead of a notebook."
+                        "detect the step code in the notebook. To fix "
+                        "this error, define your step in a python file instead "
+                        "of a notebook."
                     )
+
+                notebook_utils.warn_about_notebook_cell_magic_commands(
+                    cell_code=cell_code
+                )
+
+                code_hash = hashlib.sha1(cell_code.encode()).hexdigest()  # nosec
+                module_name = f"extracted_notebook_code_{code_hash}"
+                file_name = f"{module_name}.py"
+                code_archive.add_file(source=cell_code, destination=file_name)
+
+                setattr(step.spec.source, "replacement_module", module_name)
+                sources_that_require_upload.append(step.spec.source)
+
+    if should_upload:
+        logger.info("Archiving notebook code...")
+        code_path = code_utils.upload_code_if_necessary(code_archive)
+        for source in sources_that_require_upload:
+            setattr(source, "code_path", code_path)

src/zenml/steps/base_step.py

@@ -250,7 +250,7 @@ def __init__(
         )
         self._verify_and_apply_init_params(*args, **kwargs)
 
-        notebook_utils.try_to_save_notebook_cell_id(self.source_object)
+        notebook_utils.try_to_save_notebook_cell_code(self.source_object)
 
     @abstractmethod
     def entrypoint(self, *args: Any, **kwargs: Any) -> Any: