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

zenml/stack_deployments/aws_stack_deployment.py

@@ -13,11 +13,14 @@
 #  permissions and limitations under the License.
 """Functionality to deploy a ZenML stack to AWS."""
 
-from typing import ClassVar, Dict, List
+from typing import ClassVar, Dict, List, Optional
 
 from zenml.enums import StackDeploymentProvider
 from zenml.models import StackDeploymentConfig
-from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
+from zenml.stack_deployments.stack_deployment import (
+    STACK_DEPLOYMENT_TERRAFORM,
+    ZenMLCloudStackDeployment,
+)
 from zenml.utils.string_utils import random_str
 
 AWS_DEPLOYMENT_TYPE = "cloud-formation"
@@ -217,6 +220,7 @@ console.
         deploy the ZenML stack. The URL should include as many pre-filled
         URL query parameters as possible.
         * a textual description of the URL
+        * a Terraform script used to deploy the ZenML stack
         * some deployment providers may require additional configuration
         parameters or scripts to be passed to the cloud provider in addition to
         the deployment URL query parameters. Where that is the case, this method
@@ -248,8 +252,27 @@ console.
             f"{region}#/stacks/create/review?{query_params}"
         )
 
+        config: Optional[str] = None
+        if self.deployment_type == STACK_DEPLOYMENT_TERRAFORM:
+            config = f"""module "zenml_stack" {{
+    source  = "zenml-io/zenml-stack/aws"
+
+    region = "{self.location or "eu-central-1"}"
+    zenml_server_url = "{self.zenml_server_url}"
+    zenml_api_key = ""
+    zenml_api_token = "{self.zenml_server_api_token}"
+    zenml_stack_name = "{self.stack_name}"
+    zenml_stack_deployment = "{self.deployment_type}"
+}}
+output "zenml_stack_id" {{
+    value = module.zenml_stack.zenml_stack_id
+}}
+output "zenml_stack_name" {{
+    value = module.zenml_stack.zenml_stack_name
+}}"""
+
         return StackDeploymentConfig(
             deployment_url=url,
             deployment_url_text="AWS CloudFormation Console",
-            configuration=None,
+            configuration=config,
         )

zenml/stack_deployments/azure_stack_deployment.py

@@ -20,12 +20,14 @@ from zenml.enums import StackDeploymentProvider
 from zenml.models import StackDeploymentConfig
 from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
 
+AZURE_DEPLOYMENT_TYPE = "azure-cloud-shell"
+
 
 class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
     """Azure ZenML Cloud Stack Deployment."""
 
     provider: ClassVar[StackDeploymentProvider] = StackDeploymentProvider.AZURE
-    deployment: ClassVar[str] = "azure-cloud-shell"
+    deployment: ClassVar[str] = AZURE_DEPLOYMENT_TYPE
 
     @classmethod
     def description(cls) -> str:
@@ -67,7 +69,8 @@ permissions and are aware of any potential costs:
 - An Azure Resource Group to contain all the resources required for the ZenML stack
 - An Azure Storage Account and Blob Storage Container registered as a [ZenML artifact store](https://docs.zenml.io/stack-components/artifact-stores/azure).
 - An Azure Container Registry registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/azure).
-- SkyPilot will be registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/skypilot-vm) and used to run pipelines in your Azure subscription.
+- An AzureML Workspace registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/azureml) and used to run pipelines.
+A Key Vault and Application Insights instance will also be created in the same Resource Group and used to construct the AzureML Workspace.
 - An Azure Service Principal with the minimum necessary permissions to access
 the above resources.
 - An Azure Service Principal client secret used to give access to ZenML to
@@ -117,7 +120,7 @@ ZenML's access to your Azure subscription.
             The list of ZenML integrations that need to be installed for the
             stack to be usable.
         """
-        return ["azure", "skypilot_azure"]
+        return ["azure"]
 
     @classmethod
     def permissions(cls) -> Dict[str, List[str]]:
@@ -136,8 +139,9 @@ ZenML's access to your Azure subscription.
                 "AcrPush",
                 "Contributor",
             ],
-            "Subscription": [
-                "Owner (required by SkyPilot)",
+            "AzureML Workspace": [
+                "AzureML Compute Operator",
+                "AzureML Data Scientist",
             ],
         }
 
@@ -258,11 +262,12 @@ ZenML's access to your Azure subscription.
     source  = "zenml-io/zenml-stack/azure"
 
     location = "{self.location or "eastus"}"
+    orchestrator = "azureml"
     zenml_server_url = "{self.zenml_server_url}"
     zenml_api_key = ""
     zenml_api_token = "{self.zenml_server_api_token}"
     zenml_stack_name = "{self.stack_name}"
-    zenml_stack_deployment = "{self.deployment}"
+    zenml_stack_deployment = "{self.deployment_type}"
 }}
 output "zenml_stack_id" {{
     value = module.zenml_stack.zenml_stack_id

zenml/stack_deployments/gcp_stack_deployment.py

@@ -18,7 +18,10 @@ from typing import ClassVar, Dict, List
 
 from zenml.enums import StackDeploymentProvider
 from zenml.models import StackDeploymentConfig
-from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
+from zenml.stack_deployments.stack_deployment import (
+    STACK_DEPLOYMENT_TERRAFORM,
+    ZenMLCloudStackDeployment,
+)
 
 GCP_DEPLOYMENT_TYPE = "deployment-manager"
 
@@ -255,7 +258,26 @@ GCP project and to clean up the resources created by the stack by using
             f"https://shell.cloud.google.com/cloudshell/editor?{query_params}"
         )
 
-        config = f"""
+        if self.deployment_type == STACK_DEPLOYMENT_TERRAFORM:
+            config = f"""module "zenml_stack" {{
+    source  = "zenml-io/zenml-stack/gcp"
+
+    project_id = "my-gcp-project"
+    region = "{self.location or "europe-west3"}"
+    zenml_server_url = "{self.zenml_server_url}"
+    zenml_api_key = ""
+    zenml_api_token = "{self.zenml_server_api_token}"
+    zenml_stack_name = "{self.stack_name}"
+    zenml_stack_deployment = "{self.deployment_type}"
+}}
+output "zenml_stack_id" {{
+    value = module.zenml_stack.zenml_stack_id
+}}
+output "zenml_stack_name" {{
+    value = module.zenml_stack.zenml_stack_name
+}}"""
+        else:
+            config = f"""
 ### BEGIN CONFIGURATION ###
 ZENML_STACK_NAME={self.stack_name}
 ZENML_STACK_REGION={self.location or "europe-west3"}

zenml/stack_deployments/stack_deployment.py

@@ -27,12 +27,15 @@ from zenml.models import (
     StackDeploymentInfo,
 )
 
+STACK_DEPLOYMENT_TERRAFORM = "terraform"
+
 
 class ZenMLCloudStackDeployment(BaseModel):
     """ZenML Cloud Stack CLI Deployment base class."""
 
     provider: ClassVar[StackDeploymentProvider]
     deployment: ClassVar[str]
+    terraform: bool = False
     stack_name: str
     zenml_server_url: str
     zenml_server_api_token: str
@@ -105,6 +108,17 @@ class ZenMLCloudStackDeployment(BaseModel):
             names to region descriptions.
         """
 
+    @property
+    def deployment_type(self) -> str:
+        """Return the type of deployment.
+
+        Returns:
+            The type of deployment.
+        """
+        if self.terraform:
+            return STACK_DEPLOYMENT_TERRAFORM
+        return self.deployment
+
     @classmethod
     def skypilot_default_regions(cls) -> Dict[str, str]:
         """Returns the regions supported by default for the Skypilot.
@@ -158,7 +172,8 @@ class ZenMLCloudStackDeployment(BaseModel):
         """
 
     def get_stack(
-        self, date_start: Optional[datetime.datetime] = None
+        self,
+        date_start: Optional[datetime.datetime] = None,
     ) -> Optional[DeployedStack]:
         """Return the ZenML stack that was deployed and registered.
 
@@ -201,7 +216,7 @@ class ZenMLCloudStackDeployment(BaseModel):
             if stack.labels.get("zenml:provider") != self.provider.value:
                 continue
 
-            if stack.labels.get("zenml:deployment") != self.deployment:
+            if stack.labels.get("zenml:deployment") != self.deployment_type:
                 continue
 
             artifact_store = stack.components[

zenml/steps/base_step.py

@@ -54,6 +54,7 @@ from zenml.steps.utils import (
 )
 from zenml.utils import (
     dict_utils,
+    notebook_utils,
     pydantic_utils,
     settings_utils,
     source_code_utils,
@@ -249,6 +250,8 @@ class BaseStep(metaclass=BaseStepMeta):
         )
         self._verify_and_apply_init_params(*args, **kwargs)
 
+        notebook_utils.try_to_save_notebook_cell_code(self.source_object)
+
     @abstractmethod
     def entrypoint(self, *args: Any, **kwargs: Any) -> Any:
         """Abstract method for core step logic.

zenml/utils/archivable.py

@@ -0,0 +1,149 @@
+#  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
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import IO, Any, Dict
+
+from zenml.io import fileio
+
+
+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], 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.
+        """
+        files = self.get_files()
+        extra_files = self.get_extra_files()
+
+        if use_gzip:
+            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.
+            fileobj: Any = GzipFile(
+                filename="", mode="wb", fileobj=output_file, mtime=0.0
+            )
+        else:
+            fileobj = output_file
+
+        with tarfile.open(mode="w", fileobj=fileobj) as tf:
+            for archive_path, file_path in files.items():
+                if archive_path in extra_files:
+                    continue
+
+                if info := tf.gettarinfo(file_path, arcname=archive_path):
+                    if info.isfile():
+                        with open(file_path, "rb") as f:
+                            tf.addfile(info, f)
+                    else:
+                        tf.addfile(info, None)
+
+            for archive_path, contents in extra_files.items():
+                info = tarfile.TarInfo(archive_path)
+                contents_encoded = contents.encode("utf-8")
+                info.size = len(contents_encoded)
+                tf.addfile(info, io.BytesIO(contents_encoded))
+
+        if use_gzip:
+            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/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)