Add firewallsku as ManagedNetwork property (#37885)

* Generated REST client for upcoming API version `v2024-10-01-preview`. - Created Swagger folder `2024-10-01-preview` by copying from the same one in GitHub project `azure-rest-api-specs` - https://github.com/Azure/azure-rest-api-specs/tree/release-machinelearningservices-Microsoft.MachineLearningServices-2024-10-01-preview/specification/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-10-01-preview - Up to Azure/azure-rest-api-specs#30776 - Safe to merge in for not bumping up callers' API versions with it - Cannot bump up before registering it for RPs in selected regions - Added tag `v2024-10-01-preview` to `readme.md` - Generated REST client for `v2024-10-01-preview` * Initial commit * Initial commit * add new test for firewallsku * add new test for firewallsku * remove irrelevant changes * add addressprefix changes * address review comments * address review comments * reset the file in rest client * address review comments --------- Co-authored-by: Jing Li <[email protected]> Co-authored-by: Nethra Sashikar <[email protected]>
Azure · Nov 5, 2024 · f425c3e · f425c3e
1 parent ed61b4c
commit f425c3e
Show file tree

Hide file tree

Showing 23 changed files with 204 additions and 95 deletions.
diff --git a/sdk/ml/azure-ai-ml/.pre-commit-config.yaml b/sdk/ml/azure-ai-ml/.pre-commit-config.yaml
@@ -41,7 +41,7 @@ repos:
                     "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple/",
                 ],
             ),
-            ("pylint", "3.0.3", []),
+            ("pylint", "3.2.7", []),
         ]
 
         # Make sure that correct versions are installed

diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md
@@ -3,6 +3,8 @@
 ## 1.22.0 (unreleased)
 
 ### Features Added
+  - Added support to select firewall sku to used for provisioning azure firewall when FQDN rules are added in
+  AllowOnlyApprovedOutbound mode. FirewallSku options are `Standard` or `Basic`, defaults to `Standard`
   - Update TLS version from 1.0 to 1.2
   - Added support for Distillation jobs. Can be created by importing `disillation` from `azure.ai.ml.distillation`
 ### Bugs Fixed

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py
@@ -39,6 +39,7 @@
 from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012024Preview
 from azure.ai.ml._restclient.v2024_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042024Preview
 from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview
+from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview
 from azure.ai.ml._restclient.workspace_dataplane import (
     AzureMachineLearningWorkspaces as ServiceClientWorkspaceDataplane,
 )
@@ -381,6 +382,17 @@ def __init__(
             **kwargs,
         )
 
+        self._service_client_10_2024_preview = ServiceClient102024Preview(
+            credential=self._credential,
+            subscription_id=(
+                self._ws_operation_scope._subscription_id
+                if registry_reference
+                else self._operation_scope._subscription_id
+            ),
+            base_url=base_url,
+            **kwargs,
+        )
+
         # A general purpose, user-configurable pipeline for making
         # http requests
         self._requests_pipeline = HttpPipeline(**kwargs)
@@ -478,7 +490,7 @@ def __init__(
 
         self._workspaces = WorkspaceOperations(
             self._ws_operation_scope if registry_reference else self._operation_scope,
-            self._service_client_07_2024_preview,
+            self._service_client_10_2024_preview,
             self._operation_container,
             self._credential,
             requests_pipeline=self._requests_pipeline,
@@ -489,7 +501,7 @@ def __init__(
 
         self._workspace_outbound_rules = WorkspaceOutboundRuleOperations(
             self._operation_scope,
-            self._service_client_07_2024_preview,
+            self._service_client_10_2024_preview,
             self._operation_container,
             self._credential,
             **kwargs,
@@ -706,7 +718,7 @@ def __init__(
 
         self._featurestores = FeatureStoreOperations(
             self._operation_scope,
-            self._service_client_07_2024_preview,
+            self._service_client_10_2024_preview,
             self._operation_container,
             self._credential,
             **app_insights_handler_kwargs,  # type: ignore[arg-type]

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py
@@ -4,18 +4,19 @@
 
 # pylint: disable=unused-argument,no-else-return
 
-from marshmallow import fields, EXCLUDE
+from marshmallow import EXCLUDE, fields
 from marshmallow.decorators import post_load, pre_dump
+
+from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum, UnionField
 from azure.ai.ml._schema.core.schema_meta import PatchedSchemaMeta
-from azure.ai.ml._schema.core.fields import StringTransformedEnum, NestedField, UnionField
+from azure.ai.ml._utils.utils import _snake_to_camel, camel_to_snake
+from azure.ai.ml.constants._workspace import FirewallSku, IsolationMode, OutboundRuleCategory
 from azure.ai.ml.entities._workspace.networking import (
-    ManagedNetwork,
     FqdnDestination,
-    ServiceTagDestination,
+    ManagedNetwork,
     PrivateEndpointDestination,
+    ServiceTagDestination,
 )
-from azure.ai.ml.constants._workspace import IsolationMode, OutboundRuleCategory
-from azure.ai.ml._utils.utils import camel_to_snake, _snake_to_camel
 
 
 class ManagedNetworkStatusSchema(metaclass=PatchedSchemaMeta):
@@ -184,13 +185,31 @@ class ManagedNetworkSchema(metaclass=PatchedSchemaMeta):
         ),
         allow_none=True,
     )
+    firewall_sku = StringTransformedEnum(
+        allowed_values=[
+            FirewallSku.STANDARD,
+            FirewallSku.BASIC,
+        ],
+        casing_transform=camel_to_snake,
+        metadata={"description": "Firewall sku for FQDN rules in AllowOnlyApprovedOutbound mode"},
+    )
     network_id = fields.Str(required=False, dump_only=True)
     status = NestedField(ManagedNetworkStatusSchema, allow_none=False, unknown=EXCLUDE)
 
     @post_load
     def make(self, data, **kwargs):
         outbound_rules = data.get("outbound_rules", False)
+
+        firewall_sku = data.get("firewall_sku", False)
+        firewall_sku_value = _snake_to_camel(data["firewall_sku"]) if firewall_sku else FirewallSku.STANDARD
+
         if outbound_rules:
-            return ManagedNetwork(isolation_mode=_snake_to_camel(data["isolation_mode"]), outbound_rules=outbound_rules)
+            return ManagedNetwork(
+                isolation_mode=_snake_to_camel(data["isolation_mode"]),
+                outbound_rules=outbound_rules,
+                firewall_sku=firewall_sku_value,
+            )
         else:
-            return ManagedNetwork(isolation_mode=_snake_to_camel(data["isolation_mode"]))
+            return ManagedNetwork(
+                isolation_mode=_snake_to_camel(data["isolation_mode"]), firewall_sku=firewall_sku_value
+            )
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py b/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py
@@ -23,6 +23,13 @@ class IsolationMode:
     ALLOW_ONLY_APPROVED_OUTBOUND = "AllowOnlyApprovedOutbound"
 
 
+class FirewallSku:
+    """Firewall Sku for FQDN rules in AllowOnlyApprovedOutbound."""
+
+    STANDARD = "Standard"
+    BASIC = "Basic"
+
+
 class OutboundRuleCategory:
     """Category for a managed network outbound rule."""
 

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py
@@ -9,17 +9,17 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
 
-from azure.ai.ml._restclient.v2024_07_01_preview.models import Workspace as RestWorkspace
+from azure.ai.ml._restclient.v2024_10_01_preview.models import Workspace as RestWorkspace
 from azure.ai.ml._schema._feature_store.feature_store_schema import FeatureStoreSchema
-from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY
+from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY, WorkspaceKind
 from azure.ai.ml.entities._credentials import IdentityConfiguration, ManagedIdentityConfiguration
 from azure.ai.ml.entities._util import load_from_dict
 from azure.ai.ml.entities._workspace.compute_runtime import ComputeRuntime
 from azure.ai.ml.entities._workspace.customer_managed_key import CustomerManagedKey
 from azure.ai.ml.entities._workspace.feature_store_settings import FeatureStoreSettings
 from azure.ai.ml.entities._workspace.networking import ManagedNetwork
 from azure.ai.ml.entities._workspace.workspace import Workspace
-from azure.ai.ml.constants._common import WorkspaceKind
+
 from ._constants import DEFAULT_SPARK_RUNTIME_VERSION
 from .materialization_store import MaterializationStore
 

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/_ai_workspaces/hub.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/_ai_workspaces/hub.py
@@ -5,10 +5,8 @@
 # pylint: disable=too-many-instance-attributes,protected-access
 from typing import Any, Dict, List, Optional
 
-from azure.ai.ml._restclient.v2024_07_01_preview.models import (
-    Workspace as RestWorkspace,
-    WorkspaceHubConfig as RestWorkspaceHubConfig,
-)
+from azure.ai.ml._restclient.v2024_10_01_preview.models import Workspace as RestWorkspace
+from azure.ai.ml._restclient.v2024_10_01_preview.models import WorkspaceHubConfig as RestWorkspaceHubConfig
 from azure.ai.ml._schema.workspace import HubSchema
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml.constants._common import WorkspaceKind

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/diagnose.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/diagnose.py
@@ -3,13 +3,17 @@
 # ---------------------------------------------------------
 
 import json
-from typing import Any, Dict, Optional, List
+from typing import Any, Dict, List, Optional
 
-from azure.ai.ml._restclient.v2024_07_01_preview.models import (
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     DiagnoseRequestProperties as RestDiagnoseRequestProperties,
-    DiagnoseResponseResult as RestDiagnoseResponseResult,
+)
+from azure.ai.ml._restclient.v2024_10_01_preview.models import DiagnoseResponseResult as RestDiagnoseResponseResult
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     DiagnoseResponseResultValue as RestDiagnoseResponseResultValue,
-    DiagnoseResult as RestDiagnoseResult,
+)
+from azure.ai.ml._restclient.v2024_10_01_preview.models import DiagnoseResult as RestDiagnoseResult
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     DiagnoseWorkspaceParameters as RestDiagnoseWorkspaceParameters,
 )
 

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py
@@ -6,7 +6,7 @@
 
 from typing import Optional
 
-from azure.ai.ml._restclient.v2024_07_01_preview.models import FeatureStoreSettings as RestFeatureStoreSettings
+from azure.ai.ml._restclient.v2024_10_01_preview.models import FeatureStoreSettings as RestFeatureStoreSettings
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
 
 from .compute_runtime import ComputeRuntime

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/networking.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/networking.py
@@ -5,15 +5,21 @@
 from abc import ABC
 from typing import Any, Dict, List, Optional
 
-from azure.ai.ml._restclient.v2024_07_01_preview.models import (
-    FqdnOutboundRule as RestFqdnOutboundRule,
+from azure.ai.ml._restclient.v2024_10_01_preview.models import FqdnOutboundRule as RestFqdnOutboundRule
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     ManagedNetworkProvisionStatus as RestManagedNetworkProvisionStatus,
-    ManagedNetworkSettings as RestManagedNetwork,
+)
+from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkSettings as RestManagedNetwork
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     PrivateEndpointDestination as RestPrivateEndpointOutboundRuleDestination,
+)
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     PrivateEndpointOutboundRule as RestPrivateEndpointOutboundRule,
+)
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     ServiceTagDestination as RestServiceTagOutboundRuleDestination,
-    ServiceTagOutboundRule as RestServiceTagOutboundRule,
 )
+from azure.ai.ml._restclient.v2024_10_01_preview.models import ServiceTagOutboundRule as RestServiceTagOutboundRule
 from azure.ai.ml.constants._workspace import IsolationMode, OutboundRuleCategory, OutboundRuleType
 
 
@@ -253,6 +259,8 @@ class ManagedNetwork:
 
     :param isolation_mode: Isolation of the managed network, defaults to Disabled.
     :type isolation_mode: str
+    :param firewall_sku: Firewall Sku for FQDN rules in AllowOnlyApprovedOutbound..
+    :type firewall_sku: str
     :param outbound_rules: List of outbound rules for the managed network.
     :type outbound_rules: List[~azure.ai.ml.entities.OutboundRule]
     :param network_id: Network id for the managed network, not meant to be set by user.
@@ -271,10 +279,12 @@ def __init__(
         *,
         isolation_mode: str = IsolationMode.DISABLED,
         outbound_rules: Optional[List[OutboundRule]] = None,
+        firewall_sku: Optional[str] = None,
         network_id: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
         self.isolation_mode = isolation_mode
+        self.firewall_sku = firewall_sku
         self.network_id = network_id
         self.outbound_rules = outbound_rules
         self.status = kwargs.pop("status", None)
@@ -289,7 +299,9 @@ def _to_rest_object(self) -> RestManagedNetwork:
             if self.outbound_rules
             else {}
         )
-        return RestManagedNetwork(isolation_mode=self.isolation_mode, outbound_rules=rest_outbound_rules)
+        return RestManagedNetwork(
+            isolation_mode=self.isolation_mode, outbound_rules=rest_outbound_rules, firewall_sku=self.firewall_sku
+        )
 
     @classmethod
     def _from_rest_object(cls, obj: RestManagedNetwork) -> "ManagedNetwork":
@@ -306,6 +318,7 @@ def _from_rest_object(cls, obj: RestManagedNetwork) -> "ManagedNetwork":
             outbound_rules=from_rest_outbound_rules,  # type: ignore[arg-type]
             network_id=obj.network_id,
             status=obj.status,
+            firewall_sku=obj.firewall_sku,
         )
 
 

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/serverless_compute.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/serverless_compute.py
@@ -5,7 +5,7 @@
 
 from marshmallow.exceptions import ValidationError
 
-from azure.ai.ml._restclient.v2024_07_01_preview.models import (
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     ServerlessComputeSettings as RestServerlessComputeSettings,
 )
 from azure.ai.ml._schema._utils.utils import ArmId

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py
@@ -8,13 +8,13 @@
 from pathlib import Path
 from typing import IO, Any, AnyStr, Dict, List, Optional, Tuple, Type, Union
 
-from azure.ai.ml._restclient.v2024_07_01_preview.models import (
-    FeatureStoreSettings as RestFeatureStoreSettings,
-    ManagedNetworkSettings as RestManagedNetwork,
-    ManagedServiceIdentity as RestManagedServiceIdentity,
+from azure.ai.ml._restclient.v2024_10_01_preview.models import FeatureStoreSettings as RestFeatureStoreSettings
+from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkSettings as RestManagedNetwork
+from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedServiceIdentity as RestManagedServiceIdentity
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     ServerlessComputeSettings as RestServerlessComputeSettings,
-    Workspace as RestWorkspace,
 )
+from azure.ai.ml._restclient.v2024_10_01_preview.models import Workspace as RestWorkspace
 from azure.ai.ml._schema.workspace.workspace import WorkspaceSchema
 from azure.ai.ml._utils.utils import dump_yaml_to_file
 from azure.ai.ml.constants._common import (

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace_keys.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace_keys.py
@@ -4,7 +4,7 @@
 
 from typing import List, Optional
 
-from azure.ai.ml._restclient.v2024_07_01_preview.models import ListWorkspaceKeysResult
+from azure.ai.ml._restclient.v2024_10_01_preview.models import ListWorkspaceKeysResult
 
 
 class ContainerRegistryCredential:

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py
@@ -10,19 +10,19 @@
 
 from marshmallow import ValidationError
 
-from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview
-from azure.ai.ml._restclient.v2024_07_01_preview.models import ManagedNetworkProvisionOptions
+from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview
+from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkProvisionOptions
 from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope
 from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
 from azure.ai.ml._utils._logger_utils import OpsLogger
 from azure.ai.ml._utils.utils import camel_to_snake
 from azure.ai.ml.constants import ManagedServiceIdentityType
 from azure.ai.ml.constants._common import Scope, WorkspaceKind
 from azure.ai.ml.entities import (
-    WorkspaceConnection,
     IdentityConfiguration,
     ManagedIdentityConfiguration,
     ManagedNetworkProvisionStatus,
+    WorkspaceConnection,
 )
 from azure.ai.ml.entities._feature_store._constants import (
     OFFLINE_MATERIALIZATION_STORE_TYPE,
@@ -58,7 +58,7 @@ class FeatureStoreOperations(WorkspaceOperationsBase):
     def __init__(
         self,
         operation_scope: OperationScope,
-        service_client: ServiceClient072024Preview,
+        service_client: ServiceClient102024Preview,
         all_operations: OperationsContainer,
         credentials: Optional[TokenCredential] = None,
         **kwargs: Dict,

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py
@@ -8,8 +8,8 @@
 
 from marshmallow import ValidationError
 
-from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview
-from azure.ai.ml._restclient.v2024_07_01_preview.models import ManagedNetworkProvisionOptions
+from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview
+from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkProvisionOptions
 from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope
 from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
 from azure.ai.ml._utils._http_utils import HttpPipeline
@@ -52,7 +52,7 @@ class WorkspaceOperations(WorkspaceOperationsBase):
     def __init__(
         self,
         operation_scope: OperationScope,
-        service_client: ServiceClient072024Preview,
+        service_client: ServiceClient102024Preview,
         all_operations: OperationsContainer,
         credentials: Optional[TokenCredential] = None,
         **kwargs: Any,

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py
@@ -11,8 +11,8 @@
 
 from azure.ai.ml._arm_deployments import ArmDeploymentExecutor
 from azure.ai.ml._arm_deployments.arm_helper import get_template
-from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview
-from azure.ai.ml._restclient.v2024_07_01_preview.models import (
+from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview
+from azure.ai.ml._restclient.v2024_10_01_preview.models import (
     EncryptionKeyVaultUpdateProperties,
     EncryptionUpdateProperties,
     WorkspaceUpdateParameters,
@@ -60,7 +60,7 @@ class WorkspaceOperationsBase(ABC):
     def __init__(
         self,
         operation_scope: OperationScope,
-        service_client: ServiceClient072024Preview,
+        service_client: ServiceClient102024Preview,
         all_operations: OperationsContainer,
         credentials: Optional[TokenCredential] = None,
         **kwargs: Dict,