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

zenml/stack_deployments/azure_stack_deployment.py

@@ -17,15 +17,17 @@ import re
 from typing import ClassVar, Dict, List
 
 from zenml.enums import StackDeploymentProvider
+from zenml.models import StackDeploymentConfig
 from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
 
+AZURE_DEPLOYMENT_TYPE = "azure-cloud-shell"
+
 
-# TODO: this class just implements the regions list, and is not suitable for other
-# deployment tasks.
 class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
     """Azure ZenML Cloud Stack Deployment."""
 
     provider: ClassVar[StackDeploymentProvider] = StackDeploymentProvider.AZURE
+    deployment: ClassVar[str] = AZURE_DEPLOYMENT_TYPE
 
     @classmethod
     def description(cls) -> str:
@@ -37,8 +39,11 @@ class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
         Returns:
             A MarkDown description of the ZenML Cloud Stack Deployment.
         """
-        # TODO: Implement this
-        return ""
+        return """
+Provision and register a basic Azure ZenML stack authenticated and connected to
+all the necessary cloud infrastructure resources required to run pipelines in
+Azure.
+"""
 
     @classmethod
     def instructions(cls) -> str:
@@ -51,8 +56,45 @@ class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
             MarkDown instructions on how to deploy the ZenML stack to the
             specified cloud provider.
         """
-        # TODO: Implement this
-        return ""
+        return """
+You will be redirected to an Azure Cloud Shell console in your browser where
+you'll be asked to log into your Azure project and then use
+[the Azure ZenML Stack Terraform module](https://registry.terraform.io/modules/zenml-io/zenml-stack/azure)
+to provision the necessary cloud resources for ZenML.
+
+**NOTE**: The Azure ZenML Stack Terraform module will create the following new
+resources in your Azure subscription. Please ensure you have the necessary
+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).
+- 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
+connect to the above resources through a [ZenML service connector](https://docs.zenml.io/how-to/auth-management/azure-service-connector).
+
+The Azure ZenML Stack Terraform module will automatically create an Azure
+Service Principal client secret and will share it with ZenML to give it
+permission to access the resources created by the stack. You can revoke these
+permissions at any time by deleting the Service Principal in your Azure
+subscription.
+
+**Estimated costs**
+
+A small training job would cost around: $0.60
+
+These are rough estimates and actual costs may vary based on your usage and specific Azure pricing. 
+Some services may be eligible for the Azure Free Tier. Use [the Azure Pricing Calculator](https://azure.microsoft.com/en-us/pricing/calculator)
+for a detailed estimate based on your usage.
+
+
+💡 **After the Terraform deployment is complete, you can close the Cloud
+Shell session and return to the CLI to view details about the associated ZenML
+stack automatically registered with ZenML.**
+"""
 
     @classmethod
     def post_deploy_instructions(cls) -> str:
@@ -64,8 +106,11 @@ class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
             MarkDown instructions on what to do after the deployment is
             complete.
         """
-        # TODO: Implement this
-        return ""
+        return """
+The ZenML stack has been successfully deployed and registered. You can delete
+the provisioned Service Principal and Resource Group at any time to revoke
+ZenML's access to your Azure subscription.
+"""
 
     @classmethod
     def integrations(cls) -> List[str]:
@@ -85,8 +130,20 @@ class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
             The permissions granted to ZenML to access the cloud resources, as
             a dictionary grouping permissions by resource.
         """
-        # TODO: Implement this
-        return {}
+        return {
+            "Storage Account": [
+                "Storage Blob Data Contributor",
+            ],
+            "Container Registry": [
+                "AcrPull",
+                "AcrPush",
+                "Contributor",
+            ],
+            "AzureML Workspace": [
+                "AzureML Compute Operator",
+                "AzureML Data Scientist",
+            ],
+        }
 
     @classmethod
     def locations(cls) -> Dict[str, str]:
@@ -177,3 +234,58 @@ class AZUREZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
             for k, v in cls.locations().items()
             if "(US)" in k and matcher.match(v)
         }
+
+    def get_deployment_config(
+        self,
+    ) -> StackDeploymentConfig:
+        """Return the configuration to deploy the ZenML stack to the specified cloud provider.
+
+        The configuration should include:
+
+        * a cloud provider console URL where the user will be redirected to
+        deploy the ZenML stack. The URL should include as many pre-filled
+        URL query parameters as possible.
+        * a textual description of the URL
+        * 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
+        should also return a string that the user can copy and paste into the
+        cloud provider console to deploy the ZenML stack (e.g. a set of
+        environment variables, YAML configuration snippet, bash or Terraform
+        script etc.).
+
+        Returns:
+            The configuration or script to deploy the ZenML stack to the
+            specified cloud provider.
+        """
+        config = f"""module "zenml_stack" {{
+    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_type}"
+}}
+output "zenml_stack_id" {{
+    value = module.zenml_stack.zenml_stack_id
+}}
+output "zenml_stack_name" {{
+    value = module.zenml_stack.zenml_stack_name
+}}"""
+        instructions = """
+1. The Azure Cloud Shell console will open in your browser.
+2. Create a file named `main.tf` in the Cloud Shell and copy and paste the
+Terraform configuration below into it.
+3. Run `terraform init --upgrade` to initialize the Terraform configuration.
+4. Run `terraform apply` to deploy the ZenML stack to Azure.
+"""
+
+        return StackDeploymentConfig(
+            deployment_url="https://shell.azure.com",
+            deployment_url_text="Azure Cloud Shell Console",
+            configuration=config,
+            instructions=instructions,
+        )

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"
 
@@ -231,15 +234,16 @@ GCP project and to clean up the resources created by the stack by using
         URL query parameters as possible.
         * a textual description of the URL
         * some deployment providers may require additional configuration
-        parameters to be passed to the cloud provider in addition to the
-        deployment URL query parameters. Where that is the case, this method
+        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
         should also return a string that the user can copy and paste into the
         cloud provider console to deploy the ZenML stack (e.g. a set of
-        environment variables, or YAML configuration snippet etc.).
+        environment variables, YAML configuration snippet, bash or Terraform
+        script etc.).
 
         Returns:
-            The configuration to deploy the ZenML stack to the specified cloud
-            provider.
+            The configuration or script to deploy the ZenML stack to the
+            specified cloud provider.
         """
         params = dict(
             cloudshell_git_repo="https://github.com/zenml-io/zenml",
@@ -254,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"}
@@ -262,8 +285,14 @@ ZENML_SERVER_URL={self.zenml_server_url}
 ZENML_SERVER_API_TOKEN={self.zenml_server_api_token}
 ### END CONFIGURATION ###"""
 
+        instructions = (
+            "You will be asked to provide the following configuration values "
+            "during the deployment process:"
+        )
+
         return StackDeploymentConfig(
             deployment_url=url,
             deployment_url_text="GCP Cloud Shell Console",
             configuration=config,
+            instructions=instructions,
         )

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.
@@ -145,19 +159,21 @@ class ZenMLCloudStackDeployment(BaseModel):
         URL query parameters as possible.
         * a textual description of the URL
         * some deployment providers may require additional configuration
-        parameters to be passed to the cloud provider in addition to the
-        deployment URL query parameters. Where that is the case, this method
+        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
         should also return a string that the user can copy and paste into the
         cloud provider console to deploy the ZenML stack (e.g. a set of
-        environment variables, or YAML configuration snippet etc.).
+        environment variables, YAML configuration snippet, bash or Terraform
+        script etc.).
 
         Returns:
-            The configuration to deploy the ZenML stack to the specified cloud
-            provider.
+            The configuration or script to deploy the ZenML stack to the
+            specified cloud provider.
         """
 
     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.
 
@@ -200,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/steps/utils.py

@@ -413,7 +413,6 @@ def log_step_metadata(
     metadata: Dict[str, "MetadataType"],
     step_name: Optional[str] = None,
     pipeline_name_id_or_prefix: Optional[Union[str, UUID]] = None,
-    pipeline_version: Optional[str] = None,
     run_id: Optional[str] = None,
 ) -> None:
     """Logs step metadata.
@@ -424,8 +423,6 @@ def log_step_metadata(
             when being called inside a step.
         pipeline_name_id_or_prefix: The name of the pipeline to log metadata
             for. Can be omitted when being called inside a step.
-        pipeline_version: The version of the pipeline to log metadata for.
-            Can be omitted when being called inside a step.
         run_id: The ID of the run to log metadata for. Can be omitted when
             being called inside a step.
 
@@ -460,7 +457,6 @@ def log_step_metadata(
             )
         pipeline_run = client.get_pipeline(
             name_id_or_prefix=pipeline_name_id_or_prefix,
-            version=pipeline_version,
         ).last_run
         step_run_id = pipeline_run.steps[step_name].id
     client.create_run_metadata(

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()