zenml-nightly 0.64.0.dev20240811__py3-none-any.whl → 0.66.0.dev20240910__py3-none-any.whl

zenml/integrations/azure/flavors/azureml.py

@@ -0,0 +1,139 @@
+#  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.
+"""AzureML definitions."""
+
+from typing import Optional
+
+from pydantic import model_validator
+
+from zenml.config.base_settings import BaseSettings
+from zenml.logger import get_logger
+from zenml.utils.enum_utils import StrEnum
+
+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 AzureMLComputeSettings(BaseSettings):
+    """Settings for the AzureML compute.
+
+    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 behavior):
+            - The `mode` is set to `serverless` (default behavior).
+            - 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 set 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) -> "AzureMLComputeSettings":
+        """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.
+        """
+        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 in AzureMLComputeSettings.model_fields
+            ):
+                logger.warning(
+                    f"In the {self.__class__.__name__} 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

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.