Address comments from Azure SDK team (Azure#29379)

* address comments

* fix pylint and tests

* fix pylint

* fix pylint

* fix OutboundRule classes

* fix schema dump for managed network

* black, cspell, pylint fixes

---------

Co-authored-by: Josh Harrington <[email protected]>

Author: joshharrin

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py

@@ -11,7 +11,6 @@
 from azure.ai.ml._schema.core.fields import NestedField
 from azure.ai.ml.entities._workspace.networking import (
     ManagedNetwork,
-    OutboundRule,
     FqdnDestination,
     ServiceTagDestination,
     PrivateEndpointDestination,
@@ -66,20 +65,23 @@ def createdestobject(self, data, **kwargs):
         category = data.get("category", OutboundRuleCategory.USER_DEFINED)
         if dest:
             if isinstance(dest, str):
-                return FqdnDestination(dest, _snake_to_camel(category))
+                return FqdnDestination(rule_name=None, destination=dest, category=_snake_to_camel(category))
             else:
                 if dest.get("subresource_target", False):
                     return PrivateEndpointDestination(
-                        dest["service_resource_id"],
-                        dest["subresource_target"],
-                        dest["spark_enabled"],
-                        _snake_to_camel(category),
+                        rule_name=None,
+                        service_resource_id=dest["service_resource_id"],
+                        subresource_target=dest["subresource_target"],
+                        spark_enabled=dest["spark_enabled"],
+                        category=_snake_to_camel(category),
                     )
-                if dest.get("service_tag", False):
-                    return ServiceTagDestination(
-                        dest["service_tag"], dest["protocol"], dest["port_ranges"], _snake_to_camel(category)
-                    )
-        return OutboundRule(data)
+            return ServiceTagDestination(
+                rule_name=None,
+                service_tag=dest["service_tag"],
+                protocol=dest["protocol"],
+                port_ranges=dest["port_ranges"],
+                category=_snake_to_camel(category),
+            )
 
     def fqdn_dest2dict(self, fqdndest):
         res = fqdndest
@@ -118,7 +120,18 @@ class ManagedNetworkSchema(metaclass=PatchedSchemaMeta):
 
     @post_load
     def make(self, data, **kwargs):
-        if data.get("outbound_rules", False):
-            return ManagedNetwork(_snake_to_camel(data["isolation_mode"]), data["outbound_rules"])
+        rules_dict = data.get("outbound_rules", False)
+        if rules_dict:
+            rules_as_list = []
+            for rule_name in rules_dict:
+                rule = rules_dict[rule_name]
+                rule.rule_name = rule_name
+                rules_as_list.append(rule)
+            return ManagedNetwork(_snake_to_camel(data["isolation_mode"]), rules_as_list)
         else:
             return ManagedNetwork(_snake_to_camel(data["isolation_mode"]))
+
+    @pre_dump
+    def predump(self, data, **kwargs):
+        data.outbound_rules = {rule.rule_name: rule for rule in data.outbound_rules}
+        return data

sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py

@@ -133,6 +133,7 @@
     FqdnDestination,
     ServiceTagDestination,
     PrivateEndpointDestination,
+    IsolationMode,
 )
 from ._workspace.private_endpoint import EndpointConnection, PrivateEndpoint
 from ._workspace.workspace import Workspace
@@ -222,6 +223,7 @@
     "FqdnDestination",
     "ServiceTagDestination",
     "PrivateEndpointDestination",
+    "IsolationMode",
     "EndpointConnection",
     "CustomerManagedKey",
     "DataImport",

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/networking.py

@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, List
 
 from azure.ai.ml._restclient.v2022_12_01_preview.models import (
     ManagedNetworkSettings as RestManagedNetwork,
@@ -19,23 +19,36 @@
 
 @experimental
 class OutboundRule:
+    """Base class for Outbound Rules, should not be instantiated directly.
+
+    :param rule_name: Name of the outbound rule.
+    :type rule_name: str
+    :param type: Type of the outbound rule. Supported types are "FQDN", "PrivateEndpoint", "ServiceTag"
+    :type type: str
+    """
+
     def __init__(
-        self, type: str = None, category: str = OutboundRuleCategory.USER_DEFINED  # pylint: disable=redefined-builtin
+        self,
+        *,
+        rule_name: str = None,
+        **kwargs,
     ) -> None:
-        self.type = type
-        self.category = category
+        self.rule_name = rule_name
+        self.type = kwargs.pop("type", None)
+        self.category = kwargs.pop("category", OutboundRuleCategory.USER_DEFINED)
 
     @classmethod
-    def _from_rest_object(cls, rest_obj: Any) -> "OutboundRule":
+    def _from_rest_object(cls, rest_obj: Any, rule_name: str) -> "OutboundRule":
         if isinstance(rest_obj, RestFqdnOutboundRule):
-            rule = FqdnDestination(destination=rest_obj.destination)
+            rule = FqdnDestination(destination=rest_obj.destination, rule_name=rule_name)
             rule.category = rest_obj.category
             return rule
         if isinstance(rest_obj, RestPrivateEndpointOutboundRule):
             rule = PrivateEndpointDestination(
                 service_resource_id=rest_obj.destination.service_resource_id,
                 subresource_target=rest_obj.destination.subresource_target,
                 spark_enabled=rest_obj.destination.spark_enabled,
+                rule_name=rule_name,
             )
             rule.category = rest_obj.category
             return rule
@@ -44,37 +57,44 @@ def _from_rest_object(cls, rest_obj: Any) -> "OutboundRule":
                 service_tag=rest_obj.destination.service_tag,
                 protocol=rest_obj.destination.protocol,
                 port_ranges=rest_obj.destination.port_ranges,
+                rule_name=rule_name,
             )
             rule.category = rest_obj.category
             return rule
 
 
 @experimental
 class FqdnDestination(OutboundRule):
-    def __init__(self, destination: str, category: str = OutboundRuleCategory.USER_DEFINED) -> None:
+    def __init__(self, *, rule_name: str, destination: str, **kwargs) -> None:
         self.destination = destination
-        OutboundRule.__init__(self, type=OutboundRuleType.FQDN, category=category)
+        category = kwargs.pop("category", OutboundRuleCategory.USER_DEFINED)
+        OutboundRule.__init__(self, type=OutboundRuleType.FQDN, category=category, rule_name=rule_name)
 
     def _to_rest_object(self) -> RestFqdnOutboundRule:
         return RestFqdnOutboundRule(type=self.type, category=self.category, destination=self.destination)
 
     def _to_dict(self) -> Dict:
-        return {"type": OutboundRuleType.FQDN, "category": self.category, "destination": self.destination}
+        return {
+            self.rule_name: {"type": OutboundRuleType.FQDN, "category": self.category, "destination": self.destination}
+        }
 
 
 @experimental
 class PrivateEndpointDestination(OutboundRule):
     def __init__(
         self,
+        *,
+        rule_name: str,
         service_resource_id: str,
         subresource_target: str,
         spark_enabled: bool = False,
-        category: str = OutboundRuleCategory.USER_DEFINED,
+        **kwargs,
     ) -> None:
         self.service_resource_id = service_resource_id
         self.subresource_target = subresource_target
         self.spark_enabled = spark_enabled
-        OutboundRule.__init__(self, OutboundRuleType.PRIVATE_ENDPOINT, category=category)
+        category = kwargs.pop("category", OutboundRuleCategory.USER_DEFINED)
+        OutboundRule.__init__(self, type=OutboundRuleType.PRIVATE_ENDPOINT, category=category, rule_name=rule_name)
 
     def _to_rest_object(self) -> RestPrivateEndpointOutboundRule:
         return RestPrivateEndpointOutboundRule(
@@ -89,25 +109,34 @@ def _to_rest_object(self) -> RestPrivateEndpointOutboundRule:
 
     def _to_dict(self) -> Dict:
         return {
-            "type": OutboundRuleType.PRIVATE_ENDPOINT,
-            "category": self.category,
-            "destination": {
-                "service_resource_id": self.service_resource_id,
-                "subresource_target": self.subresource_target,
-                "spark_enabled": self.spark_enabled,
-            },
+            self.rule_name: {
+                "type": OutboundRuleType.PRIVATE_ENDPOINT,
+                "category": self.category,
+                "destination": {
+                    "service_resource_id": self.service_resource_id,
+                    "subresource_target": self.subresource_target,
+                    "spark_enabled": self.spark_enabled,
+                },
+            }
         }
 
 
 @experimental
 class ServiceTagDestination(OutboundRule):
     def __init__(
-        self, service_tag: str, protocol: str, port_ranges: str, category: str = OutboundRuleCategory.USER_DEFINED
+        self,
+        *,
+        rule_name: str,
+        service_tag: str,
+        protocol: str,
+        port_ranges: str,
+        **kwargs,
     ) -> None:
         self.service_tag = service_tag
         self.protocol = protocol
         self.port_ranges = port_ranges
-        OutboundRule.__init__(self, OutboundRuleType.SERVICE_TAG, category=category)
+        category = kwargs.pop("category", OutboundRuleCategory.USER_DEFINED)
+        OutboundRule.__init__(self, type=OutboundRuleType.SERVICE_TAG, category=category, rule_name=rule_name)
 
     def _to_rest_object(self) -> RestServiceTagOutboundRule:
         return RestServiceTagOutboundRule(
@@ -120,13 +149,15 @@ def _to_rest_object(self) -> RestServiceTagOutboundRule:
 
     def _to_dict(self) -> Dict:
         return {
-            "type": OutboundRuleType.SERVICE_TAG,
-            "category": self.category,
-            "destination": {
-                "service_tag": self.service_tag,
-                "protocol": self.protocol,
-                "port_ranges": self.port_ranges,
-            },
+            self.rule_name: {
+                "type": OutboundRuleType.SERVICE_TAG,
+                "category": self.category,
+                "destination": {
+                    "service_tag": self.service_tag,
+                    "protocol": self.protocol,
+                    "port_ranges": self.port_ranges,
+                },
+            }
         }
 
 
@@ -135,7 +166,7 @@ class ManagedNetwork:
     def __init__(
         self,
         isolation_mode: str = IsolationMode.DISABLED,
-        outbound_rules: Optional[Dict[str, OutboundRule]] = None,
+        outbound_rules: Optional[List[OutboundRule]] = None,
         network_id: Optional[str] = None,
     ) -> None:
         self.isolation_mode = isolation_mode
@@ -145,8 +176,8 @@ def __init__(
     def _to_rest_object(self) -> RestManagedNetwork:
         rest_outbound_rules = (
             {
-                rule_name: self.outbound_rules[rule_name]._to_rest_object()  # pylint: disable=protected-access
-                for rule_name in self.outbound_rules
+                outbound_rule.rule_name: outbound_rule._to_rest_object()  # pylint: disable=protected-access
+                for outbound_rule in self.outbound_rules
             }
             if self.outbound_rules
             else None
@@ -156,12 +187,12 @@ def _to_rest_object(self) -> RestManagedNetwork:
     @classmethod
     def _from_rest_object(cls, obj: RestManagedNetwork) -> "ManagedNetwork":
         from_rest_outbound_rules = (
-            {
-                rule_name: OutboundRule._from_rest_object(  # pylint: disable=protected-access
-                    obj.outbound_rules[rule_name]
+            [
+                OutboundRule._from_rest_object(  # pylint: disable=protected-access
+                    obj.outbound_rules[rule_name], rule_name=rule_name
                 )
                 for rule_name in obj.outbound_rules
-            }
+            ]
             if obj.outbound_rules
             else {}
         )

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py

@@ -260,7 +260,7 @@ def callback(_, deserialized, args):
         poller = self._operation.begin_update(resource_group, workspace_name, update_param, polling=True, cls=callback)
         return poller
 
-    def begin_delete(self, name: str, *, delete_dependent_resources: bool, **kwargs: Dict) -> LROPoller:
+    def begin_delete(self, name: str, *, delete_dependent_resources: bool, **kwargs: Dict) -> LROPoller[None]:
         workspace = self.get(name, **kwargs)
         resource_group = kwargs.get("resource_group") or self._resource_group_name
         if delete_dependent_resources:

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_outbound_rule_operations.py

@@ -2,7 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-from typing import Dict
+from typing import Dict, Iterable
 from azure.ai.ml._restclient.v2022_12_01_preview import AzureMachineLearningWorkspaces as ServiceClient122022Preview
 from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope
 
@@ -36,29 +36,31 @@ def __init__(
         self._credentials = credentials
         self._init_kwargs = kwargs
 
-    @monitor_with_activity(logger, "WorkspaceOutboundRule.Show", ActivityType.PUBLICAPI)
-    def show(self, resource_group: str, ws_name: str, outbound_rule_name: str, **kwargs) -> OutboundRule:
+    @monitor_with_activity(logger, "WorkspaceOutboundRule.Get", ActivityType.PUBLICAPI)
+    def get(self, resource_group: str, ws_name: str, outbound_rule_name: str, **kwargs) -> OutboundRule:
         workspace_name = self._check_workspace_name(ws_name)
         resource_group = kwargs.get("resource_group") or self._resource_group_name
 
         obj = self._rule_operation.get(resource_group, workspace_name, outbound_rule_name)
-        return OutboundRule._from_rest_object(obj)  # pylint: disable=protected-access
+        return OutboundRule._from_rest_object(obj, rule_name=outbound_rule_name)  # pylint: disable=protected-access
 
     @monitor_with_activity(logger, "WorkspaceOutboundRule.List", ActivityType.PUBLICAPI)
-    def list(self, resource_group: str, ws_name: str, **kwargs) -> Dict[str, OutboundRule]:
+    def list(self, resource_group: str, ws_name: str, **kwargs) -> Iterable[OutboundRule]:
         workspace_name = self._check_workspace_name(ws_name)
         resource_group = kwargs.get("resource_group") or self._resource_group_name
 
         rest_rules = self._rule_operation.list(resource_group, workspace_name)
 
-        result = {
-            rule_name: OutboundRule._from_rest_object(rest_rules[rule_name])  # pylint: disable=protected-access
+        result = [
+            OutboundRule._from_rest_object(  # pylint: disable=protected-access
+                rest_obj=rest_rules[rule_name], rule_name=rule_name
+            )
             for rule_name in rest_rules.keys()
-        }
+        ]
         return result
 
     @monitor_with_activity(logger, "WorkspaceOutboundRule.Remove", ActivityType.PUBLICAPI)
-    def remove(self, resource_group: str, ws_name: str, outbound_rule_name: str, **kwargs) -> LROPoller:
+    def begin_remove(self, resource_group: str, ws_name: str, outbound_rule_name: str, **kwargs) -> LROPoller[None]:
         workspace_name = self._check_workspace_name(ws_name)
         resource_group = kwargs.get("resource_group") or self._resource_group_name