add experimental decorator for firewallsku (Azure#38657)

* add experimental decorator * mark firewallsku as experimental field * remove duplicate import * isort * format file --------- Co-authored-by: Nethra Sashikar <[email protected]>
kingernupur · Dec 19, 2024 · 8c18f91 · 8c18f91
1 parent 57de74a
commit 8c18f91
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
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
@@ -7,6 +7,7 @@
 from marshmallow import EXCLUDE, fields
 from marshmallow.decorators import post_load, pre_dump
 
+from azure.ai.ml._schema import ExperimentalField
 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._utils.utils import _snake_to_camel, camel_to_snake
@@ -47,7 +48,12 @@ def createdestobject(self, data, **kwargs):
         category = data.get("category", OutboundRuleCategory.USER_DEFINED)
         name = data.get("name")
         status = data.get("status", None)
-        return FqdnDestination(name=name, destination=dest, category=_snake_to_camel(category), status=status)
+        return FqdnDestination(
+            name=name,
+            destination=dest,
+            category=_snake_to_camel(category),
+            status=status,
+        )
 
 
 class ServiceTagDestinationSchema(metaclass=PatchedSchemaMeta):
@@ -185,13 +191,15 @@ 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"},
+    firewall_sku = ExperimentalField(
+        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)
@@ -211,5 +219,6 @@ def make(self, data, **kwargs):
             )
         else:
             return ManagedNetwork(
-                isolation_mode=_snake_to_camel(data["isolation_mode"]), firewall_sku=firewall_sku_value
+                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
@@ -3,8 +3,8 @@
 # ---------------------------------------------------------
 from enum import Enum
 
-from azure.core import CaseInsensitiveEnumMeta
 from azure.ai.ml._utils._experimental import experimental
+from azure.core import CaseInsensitiveEnumMeta
 
 
 class ManagedServiceIdentityType(str, Enum, metaclass=CaseInsensitiveEnumMeta):
@@ -24,6 +24,7 @@ class IsolationMode:
     ALLOW_ONLY_APPROVED_OUTBOUND = "AllowOnlyApprovedOutbound"
 
 
+@experimental
 class FirewallSku:
     """Firewall Sku for FQDN rules in AllowOnlyApprovedOutbound."""