zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.65.0.dev20240906__py3-none-any.whl

zenml/integrations/azure/flavors/azureml_orchestrator_flavor.py

@@ -15,20 +15,19 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
-from pydantic import Field, model_validator
+from pydantic import Field
 
-from zenml.config.base_settings import BaseSettings
 from zenml.integrations.azure import (
     AZURE_RESOURCE_TYPE,
     AZUREML_ORCHESTRATOR_FLAVOR,
 )
+from zenml.integrations.azure.flavors.azureml import AzureMLComputeSettings
 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
@@ -36,118 +35,8 @@ if TYPE_CHECKING:
 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 AzureMLOrchestratorSettings(AzureMLComputeSettings):
+    """Settings for the AzureML orchestrator."""
 
 
 class AzureMLOrchestratorConfig(
@@ -164,6 +53,10 @@ class AzureMLOrchestratorConfig(
     workspace: str = Field(
         description="Name of the workspace that AzureML is running on."
     )
+    synchronous: bool = Field(
+        default=True,
+        description="Whether the orchestrator runs synchronously or not.",
+    )
 
     @property
     def is_remote(self) -> bool:
@@ -180,12 +73,21 @@ class AzureMLOrchestratorConfig(
 
     @property
     def is_synchronous(self) -> bool:
-        """Whether the orchestrator runs synchronous or not.
+        """Whether the orchestrator runs synchronously or not.
 
         Returns:
-            Whether the orchestrator runs synchronous or not.
+            Whether the orchestrator runs synchronously or not.
         """
-        return False
+        return self.synchronous
+
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return True
 
 
 class AzureMLOrchestratorFlavor(BaseOrchestratorFlavor):

zenml/integrations/azure/flavors/azureml_step_operator_flavor.py

@@ -13,29 +13,65 @@
 #  permissions and limitations under the License.
 """AzureML step operator flavor."""
 
-from typing import TYPE_CHECKING, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
-from zenml.config.base_settings import BaseSettings
-from zenml.integrations.azure import AZUREML_STEP_OPERATOR_FLAVOR
+from pydantic import Field, model_validator
+
+from zenml.integrations.azure import (
+    AZURE_RESOURCE_TYPE,
+    AZUREML_STEP_OPERATOR_FLAVOR,
+)
+from zenml.integrations.azure.flavors.azureml import (
+    AzureMLComputeSettings,
+    AzureMLComputeTypes,
+)
+from zenml.models import ServiceConnectorRequirements
 from zenml.step_operators.base_step_operator import (
     BaseStepOperatorConfig,
     BaseStepOperatorFlavor,
 )
+from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.utils.secret_utils import SecretField
 
 if TYPE_CHECKING:
     from zenml.integrations.azure.step_operators import AzureMLStepOperator
 
 
-class AzureMLStepOperatorSettings(BaseSettings):
+class AzureMLStepOperatorSettings(AzureMLComputeSettings):
     """Settings for the AzureML step operator.
 
     Attributes:
-        environment_name: The name of the environment if there
-            already exists one.
+        compute_target_name: The name of the configured ComputeTarget.
+            Deprecated in favor of `compute_name`.
     """
 
-    environment_name: Optional[str] = None
+    compute_target_name: Optional[str] = Field(
+        default=None,
+        description="Name of the configured ComputeTarget. Deprecated in favor "
+        "of `compute_name`.",
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _migrate_compute_name(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Backward compatibility for compute_target_name.
+
+        Args:
+            data: The model data.
+
+        Returns:
+            The migrated data.
+        """
+        if (
+            "compute_target_name" in data
+            and "compute_name" not in data
+            and "mode" not in data
+        ):
+            data["compute_name"] = data.pop("compute_target_name")
+            data["mode"] = AzureMLComputeTypes.COMPUTE_INSTANCE
+
+        return data
 
 
 class AzureMLStepOperatorConfig(
@@ -48,19 +84,21 @@ class AzureMLStepOperatorConfig(
         resource_group: The resource group to which the AzureML workspace
             is deployed.
         workspace_name: The name of the AzureML Workspace.
-        compute_target_name: The name of the configured ComputeTarget.
-            An instance of it has to be created on the portal if it doesn't
-            exist already.
         tenant_id: The Azure Tenant ID.
         service_principal_id: The ID for the service principal that is created
             to allow apps to access secure resources.
         service_principal_password: Password for the service principal.
     """
 
-    subscription_id: str
-    resource_group: str
-    workspace_name: str
-    compute_target_name: str
+    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_name: str = Field(
+        description="Name of the workspace that AzureML is running on."
+    )
 
     # Service principal authentication
     # https://docs.microsoft.com/en-us/azure/machine-learning/how-to-setup-authentication#configure-a-service-principal
@@ -94,6 +132,21 @@ class AzureMLStepOperatorFlavor(BaseStepOperatorFlavor):
         """
         return AZUREML_STEP_OPERATOR_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.

zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -26,6 +26,7 @@ from typing import (
     Union,
     cast,
 )
+from uuid import UUID
 
 from azure.ai.ml import Input, MLClient, Output
 from azure.ai.ml.constants import TimeZone
@@ -40,15 +41,16 @@ from azure.ai.ml.entities import (
 from azure.core.exceptions import (
     HttpResponseError,
     ResourceExistsError,
-    ResourceNotFoundError,
 )
 from azure.identity import DefaultAzureCredential
 
 from zenml.config.base_settings import BaseSettings
 from zenml.config.step_configurations import Step
+from zenml.constants import METADATA_ORCHESTRATOR_URL
 from zenml.enums import StackComponentType
+from zenml.integrations.azure.azureml_utils import create_or_get_compute
+from zenml.integrations.azure.flavors.azureml import AzureMLComputeTypes
 from zenml.integrations.azure.flavors.azureml_orchestrator_flavor import (
-    AzureMLComputeTypes,
     AzureMLOrchestratorConfig,
     AzureMLOrchestratorSettings,
 )
@@ -56,14 +58,13 @@ from zenml.integrations.azure.orchestrators.azureml_orchestrator_entrypoint_conf
     AzureMLEntrypointConfiguration,
 )
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType, Uri
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 from zenml.utils.string_utils import b64_encode
 
 if TYPE_CHECKING:
-    from azure.ai.ml.entities import AmlCompute, ComputeInstance
-
     from zenml.models import PipelineDeploymentResponse
     from zenml.stack import Stack
 
@@ -193,172 +194,6 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
             command=" ".join(command + arguments),
         )
 
-    @staticmethod
-    def _check_settings_and_compute_configuration(
-        parameter: str,
-        settings: AzureMLOrchestratorSettings,
-        compute: Union["ComputeInstance", "AmlCompute"],
-    ) -> None:
-        """Utility function comparing a parameter between settings and compute.
-
-        Args:
-            parameter: the name of the parameter.
-            settings: The AzureML orchestrator settings.
-            compute: The compute instance or cluster from AzureML.
-        """
-        # Check the compute size
-        compute_value = getattr(compute, parameter)
-        settings_value = getattr(settings, parameter)
-
-        if settings_value is not None and settings_value != compute_value:
-            logger.warning(
-                f"The '{parameter}' defined in the settings '{settings_value}' "
-                "does not match the actual parameter of the instance: "
-                f"'{compute_value}'. Will ignore this setting for now."
-            )
-
-    def _create_or_get_compute(
-        self, client: MLClient, settings: AzureMLOrchestratorSettings
-    ) -> Optional[str]:
-        """Creates or fetches the compute target if defined in the settings.
-
-        Args:
-            client: the AzureML client.
-            settings: the settings for the orchestrator.
-
-        Returns:
-            None, if the orchestrator is using serverless compute or
-            str, the name of the compute target (instance or cluster).
-
-        Raises:
-            RuntimeError: if the fetched compute target is unsupported or the
-                mode defined in the setting does not match the type of the
-                compute target.
-        """
-        # If the mode is serverless, we can not fetch anything anyhow
-        if settings.mode == AzureMLComputeTypes.SERVERLESS:
-            return None
-
-        # If a name is not provided, generate one based on the orchestrator id
-        compute_name = settings.compute_name or f"compute_{self.id}"
-        # Try to fetch the compute target
-        try:
-            compute = client.compute.get(compute_name)
-
-            logger.info(f"Using existing compute target: '{compute_name}'.")
-
-            # Check if compute size matches with the settings
-            self._check_settings_and_compute_configuration(
-                parameter="size", settings=settings, compute=compute
-            )
-
-            compute_type = compute.type
-
-            # Check the type and matches the settings
-            if compute_type == "computeinstance":  # Compute Instance
-                if settings.mode != AzureMLComputeTypes.COMPUTE_INSTANCE:
-                    raise RuntimeError(
-                        "The mode of operation for the compute target defined"
-                        f"in the settings '{settings.mode}' does not match "
-                        f"the type of the compute target: `{compute_name}` "
-                        "which is a 'compute-instance'. Please make sure that "
-                        "the settings are adjusted properly."
-                    )
-
-                if compute.state != "Running":
-                    raise RuntimeError(
-                        f"The compute instance `{compute_name}` is not in a "
-                        "running state at the moment. Please make sure that "
-                        "the compute target is running, before executing the "
-                        "pipeline."
-                    )
-
-                # Idle time before shutdown
-                self._check_settings_and_compute_configuration(
-                    parameter="idle_time_before_shutdown_minutes",
-                    settings=settings,
-                    compute=compute,
-                )
-
-            elif compute_type == "amIcompute":  # Compute Cluster
-                if settings.mode != AzureMLComputeTypes.COMPUTE_CLUSTER:
-                    raise RuntimeError(
-                        "The mode of operation for the compute target defined "
-                        f"in the settings '{settings.mode}' does not match "
-                        f"the type of the compute target: `{compute_name}` "
-                        "which is a 'compute-cluster'. Please make sure that "
-                        "the settings are adjusted properly."
-                    )
-
-                if compute.provisioning_state != "Succeeded":
-                    raise RuntimeError(
-                        f"The provisioning state '{compute.provisioning_state}'"
-                        f"of the compute cluster `{compute_name}` is not "
-                        "successful. Please make sure that the compute cluster "
-                        "is provisioned properly, before executing the "
-                        "pipeline."
-                    )
-
-                for parameter in [
-                    "idle_time_before_scale_down",
-                    "max_instances",
-                    "min_instances",
-                    "tier",
-                    "location",
-                ]:
-                    # Check all possible configurations
-                    self._check_settings_and_compute_configuration(
-                        parameter=parameter, settings=settings, compute=compute
-                    )
-            else:
-                raise RuntimeError(f"Unsupported compute type: {compute_type}")
-            return compute_name
-
-        # If the compute target does not exist create it
-        except ResourceNotFoundError:
-            logger.info(
-                "Can not find the compute target with name: "
-                f"'{compute_name}':"
-            )
-
-            if settings.mode == AzureMLComputeTypes.COMPUTE_INSTANCE:
-                logger.info(
-                    "Creating a new compute instance. This might take a "
-                    "few minutes."
-                )
-
-                from azure.ai.ml.entities import ComputeInstance
-
-                compute_instance = ComputeInstance(
-                    name=compute_name,
-                    size=settings.size,
-                    idle_time_before_shutdown_minutes=settings.idle_time_before_shutdown_minutes,
-                )
-                client.begin_create_or_update(compute_instance).result()
-                return compute_name
-
-            elif settings.mode == AzureMLComputeTypes.COMPUTE_CLUSTER:
-                logger.info(
-                    "Creating a new compute cluster. This might take a "
-                    "few minutes."
-                )
-
-                from azure.ai.ml.entities import AmlCompute
-
-                compute_cluster = AmlCompute(
-                    name=compute_name,
-                    size=settings.size,
-                    location=settings.location,
-                    min_instances=settings.min_instances,
-                    max_instances=settings.max_instances,
-                    idle_time_before_scale_down=settings.idle_time_before_scaledown_down,
-                    tier=settings.tier,
-                )
-                client.begin_create_or_update(compute_cluster).result()
-                return compute_name
-
-        return None
-
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -429,7 +264,9 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
         )
         pipeline_args["name"] = run_name
 
-        if compute_target := self._create_or_get_compute(ml_client, settings):
+        if compute_target := create_or_get_compute(
+            ml_client, settings, default_compute_name=f"zenml_{self.id}"
+        ):
             pipeline_args["compute"] = compute_target
 
         @pipeline(force_rerun=True, **pipeline_args)  # type: ignore[call-overload, misc]
@@ -540,5 +377,54 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
                 )
 
         else:
-            ml_client.jobs.create_or_update(pipeline_job)
+            job = ml_client.jobs.create_or_update(pipeline_job)
             logger.info(f"Pipeline {run_name} has been started.")
+            assert job.services is not None
+            assert job.name is not None
+
+            logger.info(
+                f"Pipeline {run_name} is running. "
+                "You can view the pipeline in the AzureML portal at "
+                f"{job.services['Studio'].endpoint}"
+            )
+
+            if self.config.synchronous:
+                logger.info("Waiting for pipeline to finish...")
+                ml_client.jobs.stream(job.name)
+
+    def get_pipeline_run_metadata(
+        self, run_id: UUID
+    ) -> Dict[str, "MetadataType"]:
+        """Get general component-specific metadata for a pipeline run.
+
+        Args:
+            run_id: The ID of the pipeline run.
+
+        Returns:
+            A dictionary of metadata.
+        """
+        try:
+            if connector := self.get_connector():
+                credentials = connector.connect()
+            else:
+                credentials = DefaultAzureCredential()
+
+            ml_client = MLClient(
+                credential=credentials,
+                subscription_id=self.config.subscription_id,
+                resource_group_name=self.config.resource_group,
+                workspace_name=self.config.workspace,
+            )
+
+            azureml_root_run_id = os.environ[ENV_ZENML_AZUREML_RUN_ID]
+            azureml_job = ml_client.jobs.get(azureml_root_run_id)
+
+            return {
+                METADATA_ORCHESTRATOR_URL: Uri(azureml_job.studio_url),
+            }
+        except Exception as e:
+            logger.warning(
+                f"Failed to fetch the Studio URL of the AzureML pipeline "
+                f"job: {e}"
+            )
+            return {}

zenml/integrations/azure/orchestrators/azureml_orchestrator_entrypoint_config.py

@@ -69,6 +69,7 @@ class AzureMLEntrypointConfiguration(StepEntrypointConfiguration):
 
         # Azure automatically changes the working directory, we have to set it
         # back to /app before running the step.
+        os.makedirs("/app", exist_ok=True)
         os.chdir("/app")
 
         # Run the step

zenml/integrations/azure/service_connectors/azure_service_connector.py

@@ -14,6 +14,7 @@
 """Azure Service Connector."""
 
 import datetime
+import logging
 import re
 import subprocess
 from typing import Any, Dict, List, Optional, Tuple
@@ -58,6 +59,9 @@ from zenml.service_connectors.service_connector import (
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.secret_utils import PlainSerializedSecretStr
 
+# Configure the logging level for azure.identity
+logging.getLogger("azure.identity").setLevel(logging.WARNING)
+
 logger = get_logger(__name__)