zenml-nightly 0.63.0.dev20240802__py3-none-any.whl → 0.64.0.dev20240810__py3-none-any.whl

zenml/new/pipelines/build_utils.py

@@ -30,6 +30,7 @@ from zenml.code_repositories import BaseCodeRepository
 from zenml.logger import get_logger
 from zenml.models import (
     BuildItem,
+    CodeReferenceRequest,
     PipelineBuildBase,
     PipelineBuildRequest,
     PipelineBuildResponse,
@@ -37,9 +38,7 @@ from zenml.models import (
     StackResponse,
 )
 from zenml.stack import Stack
-from zenml.utils import (
-    source_utils,
-)
+from zenml.utils import source_utils
 from zenml.utils.pipeline_docker_image_builder import (
     PipelineDockerImageBuilder,
 )
@@ -64,6 +63,93 @@ def build_required(deployment: "PipelineDeploymentBase") -> bool:
     return bool(stack.get_docker_builds(deployment=deployment))
 
 
+def requires_included_code(
+    deployment: "PipelineDeploymentBase",
+    code_repository: Optional["BaseCodeRepository"] = None,
+) -> bool:
+    """Checks whether the deployment requires included code.
+
+    Args:
+        deployment: The deployment.
+        code_repository: If provided, this code repository can be used to
+            download the code inside the container images.
+
+    Returns:
+        If the deployment requires code included in the container images.
+    """
+    for step in deployment.step_configurations.values():
+        docker_settings = step.config.docker_settings
+
+        if docker_settings.allow_download_from_artifact_store:
+            return False
+
+        if docker_settings.allow_download_from_code_repository:
+            if code_repository:
+                continue
+
+        if docker_settings.allow_including_files_in_images:
+            return True
+
+    return False
+
+
+def requires_download_from_code_repository(
+    deployment: "PipelineDeploymentBase",
+) -> bool:
+    """Checks whether the deployment needs to download code from a repository.
+
+    Args:
+        deployment: The deployment.
+
+    Returns:
+        If the deployment needs to download code from a code repository.
+    """
+    for step in deployment.step_configurations.values():
+        docker_settings = step.config.docker_settings
+
+        if docker_settings.allow_download_from_artifact_store:
+            return False
+
+        if docker_settings.allow_including_files_in_images:
+            return False
+
+        if docker_settings.allow_download_from_code_repository:
+            # The other two options are false, which means download from a
+            # code repo is required.
+            return True
+
+    return False
+
+
+def code_download_possible(
+    deployment: "PipelineDeploymentBase",
+    code_repository: Optional["BaseCodeRepository"] = None,
+) -> bool:
+    """Checks whether code download is possible for the deployment.
+
+    Args:
+        deployment: The deployment.
+        code_repository: If provided, this code repository can be used to
+            download the code inside the container images.
+
+    Returns:
+        Whether code download is possible for the deployment.
+    """
+    for step in deployment.step_configurations.values():
+        if step.config.docker_settings.allow_download_from_artifact_store:
+            continue
+
+        if (
+            step.config.docker_settings.allow_download_from_code_repository
+            and code_repository
+        ):
+            continue
+
+        return False
+
+    return True
+
+
 def reuse_or_create_pipeline_build(
     deployment: "PipelineDeploymentBase",
     allow_build_reuse: bool,
@@ -82,8 +168,8 @@ def reuse_or_create_pipeline_build(
         build: Optional existing build. If given, the build will be fetched
             (or registered) in the database. If not given, a new build will
             be created.
-        code_repository: If provided, this code repository will be used to
-            download inside the build images.
+        code_repository: If provided, this code repository can be used to
+            download code inside the container images.
 
     Returns:
         The build response.
@@ -91,8 +177,10 @@ def reuse_or_create_pipeline_build(
     if not build:
         if (
             allow_build_reuse
-            and code_repository
-            and not deployment.requires_included_files
+            and not deployment.should_prevent_build_reuse
+            and not requires_included_code(
+                deployment=deployment, code_repository=code_repository
+            )
             and build_required(deployment=deployment)
         ):
             existing_build = find_existing_build(
@@ -108,17 +196,13 @@ def reuse_or_create_pipeline_build(
                 return existing_build
             else:
                 logger.info(
-                    "Unable to find a build to reuse. When using a code "
-                    "repository, a previous build can be reused when the "
-                    "following conditions are met:\n"
+                    "Unable to find a build to reuse. A previous build can be "
+                    "reused when the following conditions are met:\n"
                     "  * The existing build was created for the same stack, "
                     "ZenML version and Python version\n"
                     "  * The stack contains a container registry\n"
                     "  * The Docker settings of the pipeline and all its steps "
-                    "are the same as for the existing build\n"
-                    "  * The build does not include code. This will only be "
-                    "the case if the existing build was created with a clean "
-                    "code repository."
+                    "are the same as for the existing build."
                 )
 
         return create_pipeline_build(
@@ -150,7 +234,7 @@ def reuse_or_create_pipeline_build(
 
 def find_existing_build(
     deployment: "PipelineDeploymentBase",
-    code_repository: "BaseCodeRepository",
+    code_repository: Optional["BaseCodeRepository"] = None,
 ) -> Optional["PipelineBuildResponse"]:
     """Find an existing build for a deployment.
 
@@ -280,6 +364,11 @@ def create_pipeline_build(
             download_files = build_config.should_download_files(
                 code_repository=code_repository,
             )
+            pass_code_repo = (
+                build_config.should_download_files_from_code_repository(
+                    code_repository=code_repository
+                )
+            )
 
             (
                 image_name_or_digest,
@@ -293,7 +382,7 @@ def create_pipeline_build(
                 download_files=download_files,
                 entrypoint=build_config.entrypoint,
                 extra_files=build_config.extra_files,
-                code_repository=code_repository,
+                code_repository=code_repository if pass_code_repo else None,
             )
             contains_code = include_files
 
@@ -389,30 +478,30 @@ def verify_local_repository_context(
         deployment, or None if code download is not possible.
     """
     if build_required(deployment=deployment):
-        if deployment.requires_code_download:
+        if requires_download_from_code_repository(deployment=deployment):
             if not local_repo_context:
                 raise RuntimeError(
                     "The `DockerSettings` of the pipeline or one of its "
-                    "steps specify that code should be included in the "
-                    "Docker image (`source_files='download'`), but there is no "
-                    "code repository active at your current source root "
-                    f"`{source_utils.get_source_root()}`."
+                    "steps specify that code should be downloaded from a "
+                    "code repository, but "
+                    "there is no code repository active at your current source "
+                    f"root `{source_utils.get_source_root()}`."
                 )
             elif local_repo_context.is_dirty:
                 raise RuntimeError(
                     "The `DockerSettings` of the pipeline or one of its "
-                    "steps specify that code should be included in the "
-                    "Docker image (`source_files='download'`), but the code "
-                    "repository active at your current source root "
+                    "steps specify that code should be downloaded from a "
+                    "code repository, but "
+                    "the code repository active at your current source root "
                     f"`{source_utils.get_source_root()}` has uncommitted "
                     "changes."
                 )
             elif local_repo_context.has_local_changes:
                 raise RuntimeError(
                     "The `DockerSettings` of the pipeline or one of its "
-                    "steps specify that code should be included in the "
-                    "Docker image (`source_files='download'`), but the code "
-                    "repository active at your current source root "
+                    "steps specify that code should be downloaded from a "
+                    "code repository, but "
+                    "the code repository active at your current source root "
                     f"`{source_utils.get_source_root()}` has unpushed "
                     "changes."
                 )
@@ -420,13 +509,13 @@ def verify_local_repository_context(
         if local_repo_context:
             if local_repo_context.is_dirty:
                 logger.warning(
-                    "Unable to use code repository to download code for this run "
-                    "as there are uncommitted changes."
+                    "Unable to use code repository to download code for this "
+                    "run as there are uncommitted changes."
                 )
             elif local_repo_context.has_local_changes:
                 logger.warning(
-                    "Unable to use code repository to download code for this run "
-                    "as there are unpushed changes."
+                    "Unable to use code repository to download code for this "
+                    "run as there are unpushed changes."
                 )
 
     code_repository = None
@@ -475,13 +564,41 @@ def verify_custom_build(
             "might differ from the local code in your client environment."
         )
 
-    if build.requires_code_download and not code_repository:
-        raise RuntimeError(
-            "The build you specified does not include code but code download "
-            "not possible. This might be because you don't have a code "
-            "repository registered or the code repository contains local "
-            "changes."
-        )
+    if build.requires_code_download:
+        if requires_included_code(
+            deployment=deployment, code_repository=code_repository
+        ):
+            raise RuntimeError(
+                "The `DockerSettings` of the pipeline or one of its "
+                "steps specify that code should be included in the Docker "
+                "image, but the build you "
+                "specified requires code download. Either update your "
+                "`DockerSettings` or specify a different build and try "
+                "again."
+            )
+
+        if (
+            requires_download_from_code_repository(deployment=deployment)
+            and not code_repository
+        ):
+            raise RuntimeError(
+                "The `DockerSettings` of the pipeline or one of its "
+                "steps specify that code should be downloaded from a "
+                "code repository but "
+                "there is no code repository active at your current source "
+                f"root `{source_utils.get_source_root()}`."
+            )
+
+        if not code_download_possible(
+            deployment=deployment, code_repository=code_repository
+        ):
+            raise RuntimeError(
+                "The `DockerSettings` of the pipeline or one of its "
+                "steps specify that code can not be downloaded from the "
+                "artifact store, but the build you specified requires code "
+                "download. Either update your `DockerSettings` or specify a "
+                "different build and try again."
+            )
 
     if build.checksum:
         build_checksum = compute_build_checksum(
@@ -561,3 +678,41 @@ def compute_stack_checksum(stack: StackResponse) -> str:
         hash_.update(integration.encode())
 
     return hash_.hexdigest()
+
+
+def should_upload_code(
+    deployment: PipelineDeploymentBase,
+    build: Optional[PipelineBuildResponse],
+    code_reference: Optional[CodeReferenceRequest],
+) -> bool:
+    """Checks whether the current code should be uploaded for the deployment.
+
+    Args:
+        deployment: The deployment.
+        build: The build for the deployment.
+        code_reference: The code reference for the deployment.
+
+    Returns:
+        Whether the current code should be uploaded for the deployment.
+    """
+    if not build:
+        # No build means we don't need to download code into a Docker container
+        # for step execution. In other remote orchestrators that don't use
+        # Docker containers but instead use e.g. Wheels to run, the code should
+        # already be included.
+        return False
+
+    for step in deployment.step_configurations.values():
+        docker_settings = step.config.docker_settings
+
+        if (
+            code_reference
+            and docker_settings.allow_download_from_code_repository
+        ):
+            # No upload needed for this step
+            continue
+
+        if docker_settings.allow_download_from_artifact_store:
+            return True
+
+    return False

zenml/new/pipelines/code_archive.py

@@ -0,0 +1,157 @@
+#  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 archive."""
+
+import os
+from pathlib import Path
+from typing import IO, TYPE_CHECKING, Dict, Optional
+
+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: str) -> 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) -> Dict[str, str]:
+        """Get all files inside the archive root.
+
+        Returns:
+            All files inside the archive root.
+        """
+        all_files = {}
+        for root, _, files in os.walk(self._root):
+            for file in files:
+                file_path = os.path.join(root, file)
+                path_in_archive = os.path.relpath(file_path, self._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.
+        """
+        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()
+            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()
+
+        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),
+            )

zenml/new/pipelines/pipeline.py

@@ -73,6 +73,7 @@ from zenml.new.pipelines.run_utils import (
     create_placeholder_run,
     deploy_pipeline,
     prepare_model_versions,
+    upload_notebook_cell_code_if_necessary,
 )
 from zenml.stack import Stack
 from zenml.steps import BaseStep
@@ -82,6 +83,7 @@ from zenml.steps.entrypoint_function_utils import (
 from zenml.steps.step_invocation import StepInvocation
 from zenml.utils import (
     code_repository_utils,
+    code_utils,
     dashboard_utils,
     dict_utils,
     pydantic_utils,
@@ -579,7 +581,8 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 method.
             unlisted: Whether the pipeline run should be unlisted (not assigned
                 to any pipeline).
-            prevent_build_reuse: Whether to prevent the reuse of a build.
+            prevent_build_reuse: DEPRECATED: Use
+                `DockerSettings.prevent_build_reuse` instead.
 
         Returns:
             Model of the pipeline run if running without a schedule, `None` if
@@ -667,6 +670,9 @@ To avoid this consider setting pipeline parameters only in one place (config or
 
             stack = Client().active_stack
             stack.validate()
+            upload_notebook_cell_code_if_necessary(
+                deployment=deployment, stack=stack
+            )
 
             prepare_model_versions(deployment)
 
@@ -677,6 +683,13 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 deployment=deployment, local_repo_context=local_repo_context
             )
 
+            if prevent_build_reuse:
+                logger.warning(
+                    "Passing `prevent_build_reuse=True` to "
+                    "`pipeline.with_opitions(...)` is deprecated. Use "
+                    "`DockerSettings.prevent_build_reuse` instead."
+                )
+
             build_model = build_utils.reuse_or_create_pipeline_build(
                 deployment=deployment,
                 pipeline_id=pipeline_id,
@@ -701,6 +714,18 @@ To avoid this consider setting pipeline parameters only in one place (config or
                     code_repository=local_repo_context.code_repository_id,
                 )
 
+            code_path = None
+            if build_utils.should_upload_code(
+                deployment=deployment,
+                build=build_model,
+                code_reference=code_reference,
+            ):
+                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,
                 workspace=Client().active_workspace.id,
@@ -709,6 +734,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 build=build_id,
                 schedule=schedule_id,
                 code_reference=code_reference,
+                code_path=code_path,
                 **deployment.model_dump(),
             )
             deployment_model = Client().zen_store.create_deployment(
@@ -1271,7 +1297,8 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 method.
             unlisted: Whether the pipeline run should be unlisted (not assigned
                 to any pipeline).
-            prevent_build_reuse: Whether to prevent the reuse of a build.
+            prevent_build_reuse: DEPRECATED: Use
+                `DockerSettings.prevent_build_reuse` instead.
             **kwargs: Pipeline configuration options. These will be passed
                 to the `pipeline.configure(...)` method.
 

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
@@ -9,6 +10,7 @@ from uuid import UUID
 from zenml import constants
 from zenml.client import Client
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
+from zenml.config.source import SourceType
 from zenml.config.step_configurations import StepConfigurationUpdate
 from zenml.enums import ExecutionStatus, ModelStages
 from zenml.logger import get_logger
@@ -23,7 +25,7 @@ from zenml.models import (
 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
+from zenml.utils import cloud_utils, code_utils, notebook_utils
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
@@ -361,3 +363,67 @@ def validate_run_config_is_runnable_from_server(
             raise ValueError(
                 "Can't set DockerSettings when running pipeline via Rest API."
             )
+
+
+def upload_notebook_cell_code_if_necessary(
+    deployment: "PipelineDeploymentBase", stack: "Stack"
+) -> None:
+    """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 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 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():
+        source = step.spec.source
+
+        if 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 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 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(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)

zenml/service_connectors/service_connector_utils.py

@@ -368,6 +368,20 @@ def get_resources_options_from_resource_model_for_full_stack(
                             flavor_display_name="Skypilot (VM)",
                         )
                     )
+                    orchestrators.append(
+                        _prepare_resource_info(
+                            connector_details=connector_details,
+                            resource_ids=each.resource_ids,
+                            stack_component_type=StackComponentType.ORCHESTRATOR,
+                            flavor="azureml",
+                            required_configuration={
+                                "subscription_id": "subscription ID",
+                                "resource_group": "resource group",
+                                "workspace": "workspace",
+                            },
+                            flavor_display_name="AzureML",
+                        )
+                    )
 
                 if each.resource_type == "kubernetes-cluster":
                     orchestrators.append(