[ml] NMS workspace managed network outbound rule operations updates (Azure#31663)

joshharrin · Josh Harrington · web-flow · commit ac06e81ece10 · 2023-08-23T16:23:51.000-07:00
* docstrings, set rules, expose rule operations

* updates

* remove duplicate import

* remove duplicate import

* fix format complaints

* more pylint

* fix another pylint

---------

Co-authored-by: Josh Harrington &lt;joharrington@microsoft.com&gt;
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
@@ -665,6 +665,15 @@ def workspaces(self) -> WorkspaceOperations:
         """
         return self._workspaces
 
+    @property
+    def workspace_outbound_rules(self) -> WorkspaceOutboundRuleOperations:
+        """A collection of workspace outbound rule related operations.
+
+        :return: Workspace outbound rule operations
+        :rtype: ~azure.ai.ml.operations.WorkspaceOutboundRuleOperations
+        """
+        return self._workspace_outbound_rules
+
     @property
     @experimental
     def registries(self) -> RegistryOperations:
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
@@ -4,6 +4,7 @@
 
 from typing import Any, Dict, Optional, List
 
+from abc import ABC
 from azure.ai.ml._restclient.v2023_06_01_preview.models import (
     ManagedNetworkSettings as RestManagedNetwork,
     FqdnOutboundRule as RestFqdnOutboundRule,
@@ -18,9 +19,8 @@
 from azure.ai.ml._utils._experimental import experimental
 
 
-@experimental
-class OutboundRule:
-    """Base class for Outbound Rules, should not be instantiated directly.
+class OutboundRule(ABC):
+    """Base class for Outbound Rules, cannot be instantiated directly.
 
     :param name: Name of the outbound rule.
     :type name: str
@@ -74,6 +74,15 @@ def _from_rest_object(cls, rest_obj: Any, name: str) -> Optional["OutboundRule"]
 
 @experimental
 class FqdnDestination(OutboundRule):
+    """Class representing a FQDN outbound rule.
+
+    :param name: Name of the outbound rule.
+    :type name: str
+    :param destination: Fully qualified domain name to which outbound connections are allowed.
+        For example: “*.contoso.com”.
+    :type destination: str
+    """
+
     def __init__(self, *, name: str, destination: str, **kwargs) -> None:
         self.destination = destination
         OutboundRule.__init__(self, type=OutboundRuleType.FQDN, name=name, **kwargs)
@@ -93,6 +102,18 @@ def _to_dict(self) -> Dict:
 
 @experimental
 class PrivateEndpointDestination(OutboundRule):
+    """Class representing a Private Endpoint outbound rule.
+
+    :param name: Name of the outbound rule.
+    :type name: str
+    :param service_resource_id: The resource URI of the root service that supports creation of the private link.
+    :type service_resource_id: str
+    :param subresource_target: The target endpoint of the subresource of the service.
+    :type subresource_target: str
+    :param spark_enabled: Indicates if the private endpoint can be used for Spark jobs, default is “false”.
+    :type spark_enabled: bool
+    """
+
     def __init__(
         self,
         *,
@@ -134,6 +155,19 @@ def _to_dict(self) -> Dict:
 
 @experimental
 class ServiceTagDestination(OutboundRule):
+    """Class representing a Service Tag outbound rule.
+
+    :param name: Name of the outbound rule.
+    :type name: str
+    :param service_tag: Service Tag of an Azure service, maps to predefined IP addresses for its service endpoints.
+    :type service_tag: str
+    :param protocol: Allowed transport protocol, can be "TCP", "UDP", "ICMP" or "*" for all supported protocols.
+    :type protocol: str
+    :param port_ranges: A comma-separated list of single ports and/or range of ports, such as "80,1024-65535".
+        Traffics should be allowed to these port ranges.
+    :type port_ranges: str
+    """
+
     def __init__(
         self,
         *,
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
@@ -31,7 +31,7 @@
 from azure.ai.ml._version import VERSION
 from azure.ai.ml.constants import ManagedServiceIdentityType
 from azure.ai.ml.constants._common import ArmConstants, LROConfigurations, WorkspaceResourceConstants
-from azure.ai.ml.constants._workspace import IsolationMode
+from azure.ai.ml.constants._workspace import IsolationMode, OutboundRuleCategory
 from azure.ai.ml.entities import Workspace
 from azure.ai.ml.entities._credentials import IdentityConfiguration
 from azure.ai.ml.entities._workspace.networking import ManagedNetwork
@@ -264,6 +264,14 @@ def begin_update(
                 )
             )
 
+        if workspace.managed_network is not None and workspace.managed_network.outbound_rules is not None:
+            # drop recommended and required rules from the update request since it would result in bad request
+            workspace.managed_network.outbound_rules = [
+                rule
+                for rule in workspace.managed_network.outbound_rules
+                if rule.category not in (OutboundRuleCategory.REQUIRED, OutboundRuleCategory.RECOMMENDED)
+            ]
+
         update_role_assignment = (
             kwargs.get("update_workspace_role_assignment", None)
             or kwargs.get("update_offline_store_role_assignment", None)
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_outbound_rule_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_outbound_rule_operations.py
@@ -3,7 +3,8 @@
 # ---------------------------------------------------------
 
 from typing import Dict, Iterable
-from azure.ai.ml._restclient.v2023_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient122022Preview
+from azure.ai.ml._restclient.v2023_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042023Preview
+from azure.ai.ml._restclient.v2023_04_01_preview.models import OutboundRuleBasicResource
 from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope
 
 from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
@@ -18,10 +19,16 @@
 
 
 class WorkspaceOutboundRuleOperations:
+    """WorkspaceOutboundRuleOperations.
+
+    You should not instantiate this class directly. Instead, you should create an MLClient instance that instantiates it
+    for you and attaches it as an attribute.
+    """
+
     def __init__(
         self,
         operation_scope: OperationScope,
-        service_client: ServiceClient122022Preview,
+        service_client: ServiceClient042023Preview,
         all_operations: OperationsContainer,
         credentials: TokenCredential = None,
         **kwargs: Dict,
@@ -36,16 +43,92 @@ def __init__(
         self._init_kwargs = kwargs
 
     @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)
+    def get(self, workspace_name: str, outbound_rule_name: str, **kwargs) -> OutboundRule:
+        """Get a workspace OutboundRule by name.
+
+        :param workspace_name: Name of the workspace.
+        :type workspace_name: str
+        :param outbound_rule_name: Name of the outbound rule.
+        :type outbound_rule_name: str
+        :return: The OutboundRule with the provided name for the workspace.
+        :rtype: OutboundRule
+        """
+
+        workspace_name = self._check_workspace_name(workspace_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.properties, name=obj.name)  # pylint: disable=protected-access
 
+    @monitor_with_activity(logger, "WorkspaceOutboundRule.BeginCreate", ActivityType.PUBLICAPI)
+    def begin_create(self, workspace_name: str, rule: OutboundRule, **kwargs) -> LROPoller[OutboundRule]:
+        """Create a Workspace OutboundRule.
+
+        :param workspace_name: Name of the workspace.
+        :type workspace_name: str
+        :param rule: OutboundRule definition (FqdnDestination, PrivateEndpointDestination, or ServiceTagDestination).
+        :type rule: OutboundRule
+        :return: An instance of LROPoller that returns an OutboundRule.
+        :rtype: ~azure.core.polling.LROPoller[~azure.ai.ml.entities.OutboundRule]
+        """
+
+        workspace_name = self._check_workspace_name(workspace_name)
+        resource_group = kwargs.get("resource_group") or self._resource_group_name
+
+        rule_params = OutboundRuleBasicResource(properties=rule._to_rest_object())  # pylint: disable=protected-access
+
+        # pylint: disable=unused-argument
+        def callback(_, deserialized, args):
+            properties = deserialized.properties
+            name = deserialized.name
+            return OutboundRule._from_rest_object(properties, name=name)  # pylint: disable=protected-access
+
+        poller = self._rule_operation.begin_create_or_update(
+            resource_group, workspace_name, rule.name, rule_params, polling=True, cls=callback
+        )
+        module_logger.info("Create request initiated for outbound rule with name: %s\n", rule.name)
+        return poller
+
+    @monitor_with_activity(logger, "WorkspaceOutboundRule.BeginUpdate", ActivityType.PUBLICAPI)
+    def begin_update(self, workspace_name: str, rule: OutboundRule, **kwargs) -> LROPoller[OutboundRule]:
+        """Update a Workspace OutboundRule.
+
+        :param workspace_name: Name of the workspace.
+        :type workspace_name: str
+        :param rule: OutboundRule definition (FqdnDestination, PrivateEndpointDestination, or ServiceTagDestination).
+        :type rule: OutboundRule
+        :return: An instance of LROPoller that returns an OutboundRule.
+        :rtype: ~azure.core.polling.LROPoller[~azure.ai.ml.entities.OutboundRule]
+        """
+
+        workspace_name = self._check_workspace_name(workspace_name)
+        resource_group = kwargs.get("resource_group") or self._resource_group_name
+
+        rule_params = OutboundRuleBasicResource(properties=rule._to_rest_object())  # pylint: disable=protected-access
+
+        # pylint: disable=unused-argument
+        def callback(_, deserialized, args):
+            properties = deserialized.properties
+            name = deserialized.name
+            return OutboundRule._from_rest_object(properties, name=name)  # pylint: disable=protected-access
+
+        poller = self._rule_operation.begin_create_or_update(
+            resource_group, workspace_name, rule.name, rule_params, polling=True, cls=callback
+        )
+        module_logger.info("Update request initiated for outbound rule with name: %s\n", rule.name)
+        return poller
+
     @monitor_with_activity(logger, "WorkspaceOutboundRule.List", ActivityType.PUBLICAPI)
-    def list(self, resource_group: str, ws_name: str, **kwargs) -> Iterable[OutboundRule]:
-        workspace_name = self._check_workspace_name(ws_name)
+    def list(self, workspace_name: str, **kwargs) -> Iterable[OutboundRule]:
+        """List Workspace OutboundRules.
+
+        :param workspace_name: Name of the workspace.
+        :type workspace_name: str
+        :return: An Iterable of OutboundRule.
+        :rtype: Iterable[OutboundRule]
+        """
+
+        workspace_name = self._check_workspace_name(workspace_name)
         resource_group = kwargs.get("resource_group") or self._resource_group_name
 
         rest_rules = self._rule_operation.list(resource_group, workspace_name)
@@ -57,8 +140,18 @@ def list(self, resource_group: str, ws_name: str, **kwargs) -> Iterable[Outbound
         return result
 
     @monitor_with_activity(logger, "WorkspaceOutboundRule.Remove", ActivityType.PUBLICAPI)
-    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)
+    def begin_remove(self, workspace_name: str, outbound_rule_name: str, **kwargs) -> LROPoller[None]:
+        """Remove a Workspace OutboundRule.
+
+        :param workspace_name: Name of the workspace.
+        :type workspace_name: str
+        :param outbound_rule_name: Name of the outbound rule to remove.
+        :type outbound_rule_name: str
+        :return: An Iterable of OutboundRule.
+        :rtype: Iterable[OutboundRule]
+        """
+
+        workspace_name = self._check_workspace_name(workspace_name)
         resource_group = kwargs.get("resource_group") or self._resource_group_name
 
         poller = self._rule_operation.begin_delete(
diff --git a/sdk/ml/azure-ai-ml/tests/recordings/workspace/e2etests/test_workspace_outbound_rule_operations.pyTestWorkspaceOutboundRulestest_workspace_create_with_managed_network_list_show_remove_rules.json b/sdk/ml/azure-ai-ml/tests/recordings/workspace/e2etests/test_workspace_outbound_rule_operations.pyTestWorkspaceOutboundRulestest_workspace_create_with_managed_network_list_show_remove_rules.json
diff --git a/sdk/ml/azure-ai-ml/tests/workspace/e2etests/test_workspace_outbound_rule_operations.py b/sdk/ml/azure-ai-ml/tests/workspace/e2etests/test_workspace_outbound_rule_operations.py
