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

zenml/image_builders/build_context.py

@@ -14,18 +14,18 @@
 """Image build context."""
 
 import os
-from pathlib import Path
-from typing import IO, Dict, List, Optional, Set, Tuple, cast
+from typing import IO, Dict, List, Optional, Set, cast
 
 from zenml.constants import REPOSITORY_DIRECTORY_NAME
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.utils import io_utils, string_utils
+from zenml.utils.archivable import Archivable
 
 logger = get_logger(__name__)
 
 
-class BuildContext:
+class BuildContext(Archivable):
     """Image build context.
 
     This class is responsible for creating an archive of the files needed to
@@ -45,9 +45,9 @@ class BuildContext:
                 given, a file called `.dockerignore` in the build context root
                 directory will be used instead if it exists.
         """
+        super().__init__()
         self._root = root
         self._dockerignore_file = dockerignore_file
-        self._extra_files: Dict[str, str] = {}
 
     @property
     def dockerignore_file(self) -> Optional[str]:
@@ -68,70 +68,26 @@ class BuildContext:
 
         return None
 
-    def add_file(self, source: str, destination: str) -> None:
-        """Adds a file to the build context.
-
-        Args:
-            source: The source of the file to add. This can either be a path
-                or the file content.
-            destination: The path inside the build context 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 build context.
-
-        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], gzip: bool = True) -> None:
+    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.
-            gzip: Whether to use `gzip` to compress the file.
+            use_gzip: Whether to use `gzip` to compress the file.
         """
         from docker.utils import build as docker_build_utils
 
-        files = self._get_files()
-        extra_files = self._get_extra_files()
+        files = self.get_files()
+        extra_files = self.get_extra_files()
 
         context_archive = docker_build_utils.create_archive(
             fileobj=output_file,
             root=self._root,
-            files=sorted(files),
-            gzip=gzip,
-            extra_files=extra_files,
+            files=sorted(files.keys()),
+            gzip=use_gzip,
+            extra_files=list(extra_files.items()),
         )
 
         build_context_size = os.path.getsize(context_archive.name)
@@ -151,33 +107,30 @@ class BuildContext:
                 os.path.join(self._root, ".dockerignore"),
             )
 
-    def _get_files(self) -> Set[str]:
-        """Gets all non-ignored files in the build context root directory.
+    def get_files(self) -> Dict[str, str]:
+        """Gets all regular files that should be included in the archive.
 
         Returns:
-            All build context files.
+            A dict {path_in_archive: path_on_filesystem} for all regular files
+            in the archive.
         """
         if self._root:
-            exclude_patterns = self._get_exclude_patterns()
             from docker.utils import build as docker_build_utils
 
-            return cast(
+            exclude_patterns = self._get_exclude_patterns()
+
+            archive_paths = cast(
                 Set[str],
                 docker_build_utils.exclude_paths(
                     self._root, patterns=exclude_patterns
                 ),
             )
+            return {
+                archive_path: os.path.join(self._root, archive_path)
+                for archive_path in archive_paths
+            }
         else:
-            return set()
-
-    def _get_extra_files(self) -> List[Tuple[str, str]]:
-        """Gets all extra files of the build context.
-
-        Returns:
-            A tuple (path, file_content) for all extra files in the build
-            context.
-        """
-        return list(self._extra_files.items())
+            return {}
 
     def _get_exclude_patterns(self) -> List[str]:
         """Gets all exclude patterns from the dockerignore file.

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -20,6 +20,7 @@ from uuid import UUID
 
 import boto3
 import sagemaker
+from botocore.exceptions import WaiterError
 from sagemaker.network import NetworkConfig
 from sagemaker.processing import ProcessingInput, ProcessingOutput
 from sagemaker.workflow.execution_variables import ExecutionVariables
@@ -373,10 +374,18 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             logger.info(
                 "Executing synchronously. Waiting for pipeline to finish..."
             )
-            pipeline_execution.wait(
-                delay=POLLING_DELAY, max_attempts=MAX_POLLING_ATTEMPTS
-            )
-            logger.info("Pipeline completed successfully.")
+            try:
+                pipeline_execution.wait(
+                    delay=POLLING_DELAY, max_attempts=MAX_POLLING_ATTEMPTS
+                )
+                logger.info("Pipeline completed successfully.")
+            except WaiterError:
+                raise RuntimeError(
+                    "Timed out while waiting for pipeline execution to finish. For long-running "
+                    "pipelines we recommend configuring your orchestrator for asynchronous execution. "
+                    "The following command does this for you: \n"
+                    f"`zenml orchestrator update {self.name} --synchronous=False`"
+                )
 
     def _get_region_name(self) -> str:
         """Returns the AWS region name.

zenml/integrations/azure/__init__.py

@@ -27,6 +27,7 @@ from zenml.stack import Flavor
 
 AZURE_ARTIFACT_STORE_FLAVOR = "azure"
 AZUREML_STEP_OPERATOR_FLAVOR = "azureml"
+AZUREML_ORCHESTRATOR_FLAVOR = "azureml"
 
 # Service connector constants
 AZURE_CONNECTOR_TYPE = "azure"
@@ -47,6 +48,7 @@ class AzureIntegration(Integration):
         "azure-mgmt-containerservice>=20.0.0",
         "azure-storage-blob==12.17.0",  # temporary fix for https://github.com/Azure/azure-sdk-for-python/issues/32056
         "kubernetes",
+        "azure-ai-ml==1.18.0"
     ]
     REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes"]
 
@@ -65,11 +67,13 @@ class AzureIntegration(Integration):
         from zenml.integrations.azure.flavors import (
             AzureArtifactStoreFlavor,
             AzureMLStepOperatorFlavor,
+            AzureMLOrchestratorFlavor,
         )
 
         return [
             AzureArtifactStoreFlavor,
             AzureMLStepOperatorFlavor,
+            AzureMLOrchestratorFlavor,
         ]
 
 

zenml/integrations/azure/flavors/__init__.py

@@ -16,10 +16,17 @@
 from zenml.integrations.azure.flavors.azure_artifact_store_flavor import (
     AzureArtifactStoreConfig,
     AzureArtifactStoreFlavor,
+
 )
 from zenml.integrations.azure.flavors.azureml_step_operator_flavor import (
     AzureMLStepOperatorConfig,
     AzureMLStepOperatorFlavor,
+    AzureMLStepOperatorSettings,
+)
+from zenml.integrations.azure.flavors.azureml_orchestrator_flavor import (
+    AzureMLOrchestratorConfig,
+    AzureMLOrchestratorFlavor,
+    AzureMLOrchestratorSettings,
 )
 
 __all__ = [
@@ -27,4 +34,8 @@ __all__ = [
     "AzureArtifactStoreConfig",
     "AzureMLStepOperatorFlavor",
     "AzureMLStepOperatorConfig",
+    "AzureMLStepOperatorSettings",
+    "AzureMLOrchestratorFlavor",
+    "AzureMLOrchestratorConfig",
+    "AzureMLOrchestratorSettings",
 ]

zenml/integrations/azure/flavors/azureml_orchestrator_flavor.py

@@ -0,0 +1,263 @@
+#  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.
+"""Implementation of the AzureML Orchestrator flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from pydantic import Field, model_validator
+
+from zenml.config.base_settings import BaseSettings
+from zenml.integrations.azure import (
+    AZURE_RESOURCE_TYPE,
+    AZUREML_ORCHESTRATOR_FLAVOR,
+)
+from zenml.logger import get_logger
+from zenml.models import ServiceConnectorRequirements
+from zenml.orchestrators.base_orchestrator import (
+    BaseOrchestratorConfig,
+    BaseOrchestratorFlavor,
+)
+from zenml.utils.enum_utils import StrEnum
+
+if TYPE_CHECKING:
+    from zenml.integrations.azure.orchestrators import AzureMLOrchestrator
+
+logger = get_logger(__name__)
+
+
+class AzureMLComputeTypes(StrEnum):
+    """Enum for different types of compute on AzureML."""
+
+    SERVERLESS = "serverless"
+    COMPUTE_INSTANCE = "compute-instance"
+    COMPUTE_CLUSTER = "compute-cluster"
+
+
+class AzureMLOrchestratorSettings(BaseSettings):
+    """Settings for the AzureML orchestrator.
+
+    These settings adjust the compute resources that will be used by the
+    pipeline execution.
+
+    There are three possible use cases for this implementation:
+
+        1. Serverless compute (default behaviour):
+            - The `mode` is set to `serverless` (default behaviour).
+            - All the other parameters become irrelevant and will throw a
+              warning if set.
+
+        2. Compute instance:
+            - The `mode` is set to `compute-instance`.
+            - In this case, users have to provide a `compute-name`.
+                - If a compute instance exists with this name, this instance
+                will be used and all the other parameters become irrelevant
+                and will throw a warning if set.
+                - If a compute instance does not already exist, ZenML will
+                create it. You can use the parameters `compute_size` and
+                `idle_type_before_shutdown_minutes` for this operation.
+
+        3. Compute cluster:
+            - The `mode` is set to `compute-cluster`.
+            - In this case, users have to provide a `compute-name`.
+                - If a compute cluster exists with this name, this instance
+                will be used and all the other parameters become irrelevant
+                and will throw a warning if set.
+                - If a compute cluster does not already exist, ZenML will
+                create it. You can all the additional parameters for this
+                operation.
+    """
+
+    # Mode for compute
+    mode: AzureMLComputeTypes = AzureMLComputeTypes.SERVERLESS
+
+    # Common Configuration for Compute Instances and Clusters
+    compute_name: Optional[str] = None
+    size: Optional[str] = None
+
+    # Additional configuration for a Compute Instance
+    idle_time_before_shutdown_minutes: Optional[int] = None
+
+    # Additional configuration for a Compute Cluster
+    idle_time_before_scaledown_down: Optional[int] = None
+    location: Optional[str] = None
+    min_instances: Optional[int] = None
+    max_instances: Optional[int] = None
+    tier: Optional[str] = None
+
+    @model_validator(mode="after")
+    def azureml_settings_validator(self) -> "AzureMLOrchestratorSettings":
+        """Checks whether the right configuration is set based on mode.
+
+        Returns:
+            the instance itself.
+
+        Raises:
+            AssertionError: if a name is not provided when working with
+                instances and clusters.
+        """
+        excluded_fields = {"subscription_id", "resource_group", "workspace"}
+
+        viable_configuration_fields = {
+            AzureMLComputeTypes.SERVERLESS: {"mode"},
+            AzureMLComputeTypes.COMPUTE_INSTANCE: {
+                "mode",
+                "compute_name",
+                "size",
+                "idle_time_before_shutdown_minutes",
+            },
+            AzureMLComputeTypes.COMPUTE_CLUSTER: {
+                "mode",
+                "compute_name",
+                "size",
+                "idle_time_before_scaledown_down",
+                "location",
+                "min_instances",
+                "max_instances",
+                "tier",
+            },
+        }
+        viable_fields = viable_configuration_fields[self.mode]
+
+        for field in self.model_fields_set:
+            if field not in viable_fields and field not in excluded_fields:
+                logger.warning(
+                    "In the AzureML Orchestrator Settings, the mode of "
+                    f"operation is set to {self.mode}. In this mode, you can "
+                    f"not configure the parameter '{field}'. This "
+                    "configuration will be ignored."
+                )
+
+        if (
+            self.mode == AzureMLComputeTypes.COMPUTE_INSTANCE
+            or self.mode == AzureMLComputeTypes.COMPUTE_CLUSTER
+        ):
+            assert self.compute_name is not None, (
+                "When you are working with compute instances and clusters, "
+                "please define a name for the compute target."
+            )
+
+        return self
+
+
+class AzureMLOrchestratorConfig(
+    BaseOrchestratorConfig, AzureMLOrchestratorSettings
+):
+    """Configuration for the AzureML orchestrator."""
+
+    subscription_id: str = Field(
+        description="Subscription ID that AzureML is running on."
+    )
+    resource_group: str = Field(
+        description="Name of the resource group that AzureML is running on.",
+    )
+    workspace: str = Field(
+        description="Name of the workspace that AzureML is running on."
+    )
+
+    @property
+    def is_remote(self) -> bool:
+        """Checks if this stack component is running remotely.
+
+        This designation is used to determine if the stack component can be
+        used with a local ZenML database or if it requires a remote ZenML
+        server.
+
+        Returns:
+            True if this config is for a remote component, False otherwise.
+        """
+        return True
+
+    @property
+    def is_synchronous(self) -> bool:
+        """Whether the orchestrator runs synchronous or not.
+
+        Returns:
+            Whether the orchestrator runs synchronous or not.
+        """
+        return False
+
+
+class AzureMLOrchestratorFlavor(BaseOrchestratorFlavor):
+    """Flavor for the AzureML orchestrator."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return AZUREML_ORCHESTRATOR_FLAVOR
+
+    @property
+    def service_connector_requirements(
+        self,
+    ) -> Optional[ServiceConnectorRequirements]:
+        """Service connector resource requirements for service connectors.
+
+        Specifies resource requirements that are used to filter the available
+        service connector types that are compatible with this flavor.
+
+        Returns:
+            Requirements for compatible service connectors, if a service
+            connector is required for this flavor.
+        """
+        return ServiceConnectorRequirements(resource_type=AZURE_RESOURCE_TYPE)
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A URL to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A URL to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A URL to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/orchestrator/azureml.png"
+
+    @property
+    def config_class(self) -> Type[AzureMLOrchestratorConfig]:
+        """Returns AzureMLOrchestratorConfig config class.
+
+        Returns:
+            The config class.
+        """
+        return AzureMLOrchestratorConfig
+
+    @property
+    def implementation_class(self) -> Type["AzureMLOrchestrator"]:
+        """Implementation class.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.azure.orchestrators import AzureMLOrchestrator
+
+        return AzureMLOrchestrator

zenml/{_hub → integrations/azure/orchestrators}/__init__.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
+#  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.
@@ -11,4 +11,9 @@
 #  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.
-"""ZenML Hub internal code."""
+"""AzureML orchestrator."""
+from zenml.integrations.azure.orchestrators.azureml_orchestrator import (
+    AzureMLOrchestrator,
+)
+
+__all__ = ["AzureMLOrchestrator"]