Provision Network Status as public entity (#30306)

* prov network fix, obrule type fix, test added * remove creds from recording * address comments --------- Co-authored-by: Josh Harrington <[email protected]>
Azure · May 9, 2023 · afb4099 · afb4099
1 parent b2b8b74
commit afb4099
Show file tree

Hide file tree

Showing 9 changed files with 5,620 additions and 1,366 deletions.
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py
@@ -139,6 +139,7 @@
     ServiceTagDestination,
     PrivateEndpointDestination,
     IsolationMode,
+    ManagedNetworkProvisionStatus,
 )
 from ._workspace.private_endpoint import EndpointConnection, PrivateEndpoint
 from ._workspace.workspace import Workspace
@@ -268,6 +269,7 @@
     "ServiceTagDestination",
     "PrivateEndpointDestination",
     "IsolationMode",
+    "ManagedNetworkProvisionStatus",
     "EndpointConnection",
     "CustomerManagedKey",
     "DataImport",

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
@@ -11,6 +11,7 @@
     PrivateEndpointDestination as RestPrivateEndpointOutboundRuleDestination,
     ServiceTagOutboundRule as RestServiceTagOutboundRule,
     ServiceTagDestination as RestServiceTagOutboundRuleDestination,
+    ManagedNetworkProvisionStatus as RestManagedNetworkProvisionStatus,
 )
 from azure.ai.ml.constants._workspace import IsolationMode, OutboundRuleCategory, OutboundRuleType
 
@@ -25,6 +26,8 @@ class OutboundRule:
     :type name: str
     :param type: Type of the outbound rule. Supported types are "FQDN", "PrivateEndpoint", "ServiceTag"
     :type type: str
+    :ivar type: Type of the outbound rule. Supported types are "FQDN", "PrivateEndpoint", "ServiceTag"
+    :vartype type: str
     """
 
     def __init__(
@@ -207,3 +210,34 @@ def _from_rest_object(cls, obj: RestManagedNetwork) -> "ManagedNetwork":
             network_id=obj.network_id,
             status=obj.status,
         )
+
+
+@experimental
+class ManagedNetworkProvisionStatus:
+    """ManagedNetworkProvisionStatus.
+
+    :param status: Status for managed network provision.
+    :type status: str
+    :param spark_ready: Bool value indicating if managed network is spark ready
+    :type spark_ready: str
+    :ivar status: Status for managed network provision.
+    :vartype status: str
+    :ivar spark_ready: Bool value indicating if managed network is spark ready
+    :vartype spark_ready: str
+    """
+
+    def __init__(
+        self,
+        *,
+        status: Optional[str] = None,
+        spark_ready: Optional[bool] = None,
+    ):
+        self.status = status
+        self.spark_ready = spark_ready
+
+    @classmethod
+    def _from_rest_object(cls, rest_obj: RestManagedNetworkProvisionStatus) -> "ManagedNetworkProvisionStatus":
+        return cls(
+            status=rest_obj.status,
+            spark_ready=rest_obj.spark_ready,
+        )
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,7 +8,6 @@
 
 from azure.ai.ml._restclient.v2023_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042023Preview
 from azure.ai.ml._restclient.v2023_04_01_preview.models import ManagedNetworkProvisionOptions
-from azure.ai.ml._restclient.v2023_04_01_preview.models import ManagedNetworkProvisionStatus
 
 from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope
 
@@ -22,6 +21,7 @@
     DiagnoseWorkspaceParameters,
     Workspace,
     WorkspaceKeys,
+    ManagedNetworkProvisionStatus,
 )
 from azure.core.credentials import TokenCredential
 from azure.core.polling import LROPoller
@@ -137,12 +137,19 @@ def begin_provision_network(
         :param workspace_name: Name of the workspace.
         :type workspace_name: str
         :return: An instance of LROPoller.
-        :rtype: ~azure.core.polling.LROPoller[None]
+        :rtype: ~azure.core.polling.LROPoller[~azure.ai.ml.entities.ManagedNetworkProvisionStatus]
         """
         workspace_name = self._check_workspace_name(workspace_name)
-        return self._provision_network_operation.begin_provision_managed_network(
-            self._resource_group_name, workspace_name, ManagedNetworkProvisionOptions(include_spark=include_spark)
+
+        poller = self._provision_network_operation.begin_provision_managed_network(
+            self._resource_group_name,
+            workspace_name,
+            ManagedNetworkProvisionOptions(include_spark=include_spark),
+            polling=True,
+            cls=lambda response, deserialized, headers: ManagedNetworkProvisionStatus._from_rest_object(deserialized),
         )
+        module_logger.info("Provision network request initiated for workspace: %s\n", workspace_name)
+        return poller
 
     @monitor_with_activity(logger, "Workspace.BeginCreate", ActivityType.PUBLICAPI)
     @distributed_trace

diff --git a/...e/e2etests/test_workspace.pyTestWorkspacetest_uai_workspace_create_update_and_delete.json b/...e/e2etests/test_workspace.pyTestWorkspacetest_uai_workspace_create_update_and_delete.json
@@ -3867,7 +3867,7 @@
       "RequestHeaders": {
         "Accept": "*/*",
         "Accept-Encoding": "gzip, deflate",
-        "Authorization": "Bearer eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsIng1dCI6Ii1LSTNROW5OUjdiUm9meG1lWm9YcWJIWkdldyIsImtpZCI6Ii1LSTNROW5OUjdiUm9meG1lWm9YcWJIWkdldyJ9.eyJhdWQiOiJodHRwczovL21hbmFnZW1lbnQuYXp1cmUuY29tIiwiaXNzIjoiaHR0cHM6Ly9zdHMud2luZG93cy5uZXQvNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3LyIsImlhdCI6MTY3NjYxNTQ4NCwibmJmIjoxNjc2NjE1NDg0LCJleHAiOjE2NzY2MjAyNTIsIl9jbGFpbV9uYW1lcyI6eyJncm91cHMiOiJzcmMxIn0sIl9jbGFpbV9zb3VyY2VzIjp7InNyYzEiOnsiZW5kcG9pbnQiOiJodHRwczovL2dyYXBoLndpbmRvd3MubmV0LzcyZjk4OGJmLTg2ZjEtNDFhZi05MWFiLTJkN2NkMDExZGI0Ny91c2Vycy9jMWFmZTY2My05MmU3LTQ5M2MtOWI1Ny0xNDBlZTkxNzQ1NDkvZ2V0TWVtYmVyT2JqZWN0cyJ9fSwiYWNyIjoiMSIsImFpbyI6IkFWUUFxLzhUQUFBQWRLQnpkdzl3NWk2SnM4RDVEaFQwV3FTam5OL1pyL1l2VkJ0WWwyaDlNZGJIMUFjNzE2aGhnNytOVmdURkxmYTljd2plWWtaTDNOS2dIRSs3ZjVUVnJlekhKNitGSW8rR0RoYzVXZDhDdGZzPSIsImFtciI6WyJyc2EiLCJtZmEiXSwiYXBwaWQiOiIwNGIwNzc5NS04ZGRiLTQ2MWEtYmJlZS0wMmY5ZTFiZjdiNDYiLCJhcHBpZGFjciI6IjAiLCJkZXZpY2VpZCI6ImM0YzQ4YTY1LTQ4YjUtNDUyNS1hYjU3LWU0NjljYjcwYTU2ZSIsImZhbWlseV9uYW1lIjoiSGFycmluZ3RvbiIsImdpdmVuX25hbWUiOiJKb3NoIiwiaXBhZGRyIjoiNzYuMjguMTUzLjYxIiwibmFtZSI6Ikpvc2ggSGFycmluZ3RvbiIsIm9pZCI6ImMxYWZlNjYzLTkyZTctNDkzYy05YjU3LTE0MGVlOTE3NDU0OSIsIm9ucHJlbV9zaWQiOiJTLTEtNS0yMS0yMTI3NTIxMTg0LTE2MDQwMTI5MjAtMTg4NzkyNzUyNy01OTU5MTU4OCIsInB1aWQiOiIxMDAzMjAwMjE1RTM4QkJFIiwicmgiOiIwLkFSb0F2NGo1Y3ZHR3IwR1JxeTE4MEJIYlIwWklmM2tBdXRkUHVrUGF3ZmoyTUJNYUFPMC4iLCJzY3AiOiJ1c2VyX2ltcGVyc29uYXRpb24iLCJzdWIiOiJUMzVkNm1uai1qc1lET3J5UU54MDQ5VWw1eWtrNkh0RTRRT0FVRDhFemxZIiwidGlkIjoiNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3IiwidW5pcXVlX25hbWUiOiJqb2hhcnJpbmd0b25AbWljcm9zb2Z0LmNvbSIsInVwbiI6ImpvaGFycmluZ3RvbkBtaWNyb3NvZnQuY29tIiwidXRpIjoiQXRfVUQ0MEZRa0ttQkRES1Nfd1NBQSIsInZlciI6IjEuMCIsIndpZHMiOlsiYjc5ZmJmNGQtM2VmOS00Njg5LTgxNDMtNzZiMTk0ZTg1NTA5Il0sInhtc19jYyI6WyJDUDEiXSwieG1zX3RjZHQiOjEyODkyNDE1NDd9.E7T4jNI1OZXaiNg4S_9mm7Rf4HELm4HJ7ukJwcywrYBeW2qypcgt5BW3SMHIm7H8MqF5ibTKWUD1m1Am6udfyiAOphCH2gzCzGcFPjbHBdnCxPaozXGGdRzwTCvFJkqXjh3vBA4pCkilA3Zf8lw6wHdHCWTQFhiDfElNGEcfi51XNzV_yW0tlBZwWHb5Tqg96AbMmxNd68K6KmaV0qTIs2gCcxREhxDECZD7hmy1qGKZPFf5PuvOMTJLs7-__i51Jvc8Ur1o46e2LxZET19LtYCu0fTEsMkRdmFvDnEtXrFBk51adzI0lDTAWKLqW3_3YKA3oBlSwVQPCdExiWwt4A",
+        "Authorization": "Sanitized",
         "Connection": "keep-alive",
         "User-Agent": "azure-ai-ml/1.5.0 azsdk-python-mgmt-machinelearningservices/0.1.0 Python/3.8.10 (Windows-10-10.0.22621-SP0)",
         "x-ms-client-request-id": "361b1da9-ae8e-11ed-b940-f42679b37e91"

diff --git a/...ests/test_workspace.pyTestWorkspacetest_workspace_create_delete_with_managed_network.json b/...ests/test_workspace.pyTestWorkspacetest_workspace_create_delete_with_managed_network.json
@@ -5063,7 +5063,7 @@
       "RequestHeaders": {
         "Accept": "*/*",
         "Accept-Encoding": "gzip, deflate",
-        "Authorization": "Bearer eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsIng1dCI6Ii1LSTNROW5OUjdiUm9meG1lWm9YcWJIWkdldyIsImtpZCI6Ii1LSTNROW5OUjdiUm9meG1lWm9YcWJIWkdldyJ9.eyJhdWQiOiJodHRwczovL21hbmFnZW1lbnQuYXp1cmUuY29tIiwiaXNzIjoiaHR0cHM6Ly9zdHMud2luZG93cy5uZXQvNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3LyIsImlhdCI6MTY4MjQ1NjY0NiwibmJmIjoxNjgyNDU2NjQ2LCJleHAiOjE2ODI0NjA4MjAsIl9jbGFpbV9uYW1lcyI6eyJncm91cHMiOiJzcmMxIn0sIl9jbGFpbV9zb3VyY2VzIjp7InNyYzEiOnsiZW5kcG9pbnQiOiJodHRwczovL2dyYXBoLndpbmRvd3MubmV0LzcyZjk4OGJmLTg2ZjEtNDFhZi05MWFiLTJkN2NkMDExZGI0Ny91c2Vycy8yNTlkMGY1MS1iODBmLTQ2YjUtYWE2OC0wM2JiMzc1NDRkMjMvZ2V0TWVtYmVyT2JqZWN0cyJ9fSwiYWNyIjoiMSIsImFpbyI6IkFhUUFXLzhUQUFBQTNiNkVYUXhGd0lMRWhjK3hTQ0c4eVNldWVoaGU0TSsvdVRZWDd5VHNGbXN1OUw3REZhS3JBWWg2cnowa0QvcWZJb3BvMmVDSEk5c3pFVnNVRldHZTVIQTNVV2RmRXNPUHZxS3luc2crR2Z0L200c0V3QVNjazlQMVdwY3k1SmlmdTkyNmdZemVLeDd5SFVQUmhwNGQ3b0VCVU9pYnlObWJILzVibzNFSU9VbU9SVzN6b1pQdE1xVXZpQ1BFM2NYUUxrOUFPaWtrVjU4alRScjJXd2IwYkE9PSIsImFtciI6WyJwd2QiLCJyc2EiLCJ3aWEiLCJtZmEiXSwiYXBwaWQiOiIwNGIwNzc5NS04ZGRiLTQ2MWEtYmJlZS0wMmY5ZTFiZjdiNDYiLCJhcHBpZGFjciI6IjAiLCJkZXZpY2VpZCI6IjlmMjdmOGM5LWQ1YjQtNDM5Ni1hN2YyLTk2MGU0NWIwNjU4YSIsImZhbWlseV9uYW1lIjoiSGUiLCJnaXZlbl9uYW1lIjoiTWluZ3dlaSIsImlwYWRkciI6IjEzMS4xMDcuMS4xNzEiLCJuYW1lIjoiTWluZ3dlaSBIZSIsIm9pZCI6IjI1OWQwZjUxLWI4MGYtNDZiNS1hYTY4LTAzYmIzNzU0NGQyMyIsIm9ucHJlbV9zaWQiOiJTLTEtNS0yMS0yMTI3NTIxMTg0LTE2MDQwMTI5MjAtMTg4NzkyNzUyNy00OTcyNzY4OCIsInB1aWQiOiIxMDAzMjAwMTNBQUMzRTQyIiwicmgiOiIwLkFRRUF2NGo1Y3ZHR3IwR1JxeTE4MEJIYlIwWklmM2tBdXRkUHVrUGF3ZmoyTUJNYUFBUS4iLCJzY3AiOiJ1c2VyX2ltcGVyc29uYXRpb24iLCJzdWIiOiJQMERaZEZXMWgzNjdVN3k2cU5LbE0wbmVIYnkyZm1hc1pPS3F1VHpTT1FJIiwidGlkIjoiNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3IiwidW5pcXVlX25hbWUiOiJtaW5nd2VpaGVAbWljcm9zb2Z0LmNvbSIsInVwbiI6Im1pbmd3ZWloZUBtaWNyb3NvZnQuY29tIiwidXRpIjoieDlScTIyWjJTVWlvYTY0bEdCVWtBUSIsInZlciI6IjEuMCIsIndpZHMiOlsiYjc5ZmJmNGQtM2VmOS00Njg5LTgxNDMtNzZiMTk0ZTg1NTA5Il0sInhtc19jYyI6WyJDUDEiXSwieG1zX3RjZHQiOjEyODkyNDE1NDd9.adIOQMcO4bfoawy0tS4IJ0SlsoNxA1buhzwCfiLP3KQNzpqk9PjOx500pgHuVskz8u4VdVYwoQ9Xemz9C4Trj39lHa6o7MpkXYOPYRo11D_iTuV9w_RqDP8uZ51A3y0oMkBKMv0PNyMB3eZ5OJR9FdmI_geKlMKvCVdjDRiVvCJkoxcc6uliXppHLBugYrsIctZPH6fjPxtORP7qhJKQcbm6gqNichzcx8QqhnfcNxG97tB7R-ZNWCYfx0UayAvdE2NArxTZofxSHW-HqWgiowLhwK2TOeVqmf9eS3SirwP60MHY0mR_LXvAhRDM-RZaZgb3ifnyI5Sx42DeutVCJw",
+        "Authorization": "Sanitized",
         "Connection": "keep-alive",
         "User-Agent": "azure-ai-ml/1.6.0 azsdk-python-mgmt-machinelearningservices/0.1.0 Python/3.7.13 (Windows-10-10.0.19041-SP0)",
         "x-ms-client-request-id": "eb2ccf18-e3ae-11ed-a894-b07b251c73bc"

diff --git a/...st_workspace.pyTestWorkspacetest_workspace_create_update_delete_with_managed_network.json b/...st_workspace.pyTestWorkspacetest_workspace_create_update_delete_with_managed_network.json
@@ -3990,7 +3990,7 @@
       "RequestHeaders": {
         "Accept": "*/*",
         "Accept-Encoding": "gzip, deflate",
-        "Authorization": "Bearer eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsIng1dCI6Ii1LSTNROW5OUjdiUm9meG1lWm9YcWJIWkdldyIsImtpZCI6Ii1LSTNROW5OUjdiUm9meG1lWm9YcWJIWkdldyJ9.eyJhdWQiOiJodHRwczovL21hbmFnZW1lbnQuYXp1cmUuY29tIiwiaXNzIjoiaHR0cHM6Ly9zdHMud2luZG93cy5uZXQvNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3LyIsImlhdCI6MTY3NjYxNTQ4NCwibmJmIjoxNjc2NjE1NDg0LCJleHAiOjE2NzY2MjAyNTIsIl9jbGFpbV9uYW1lcyI6eyJncm91cHMiOiJzcmMxIn0sIl9jbGFpbV9zb3VyY2VzIjp7InNyYzEiOnsiZW5kcG9pbnQiOiJodHRwczovL2dyYXBoLndpbmRvd3MubmV0LzcyZjk4OGJmLTg2ZjEtNDFhZi05MWFiLTJkN2NkMDExZGI0Ny91c2Vycy9jMWFmZTY2My05MmU3LTQ5M2MtOWI1Ny0xNDBlZTkxNzQ1NDkvZ2V0TWVtYmVyT2JqZWN0cyJ9fSwiYWNyIjoiMSIsImFpbyI6IkFWUUFxLzhUQUFBQWRLQnpkdzl3NWk2SnM4RDVEaFQwV3FTam5OL1pyL1l2VkJ0WWwyaDlNZGJIMUFjNzE2aGhnNytOVmdURkxmYTljd2plWWtaTDNOS2dIRSs3ZjVUVnJlekhKNitGSW8rR0RoYzVXZDhDdGZzPSIsImFtciI6WyJyc2EiLCJtZmEiXSwiYXBwaWQiOiIwNGIwNzc5NS04ZGRiLTQ2MWEtYmJlZS0wMmY5ZTFiZjdiNDYiLCJhcHBpZGFjciI6IjAiLCJkZXZpY2VpZCI6ImM0YzQ4YTY1LTQ4YjUtNDUyNS1hYjU3LWU0NjljYjcwYTU2ZSIsImZhbWlseV9uYW1lIjoiSGFycmluZ3RvbiIsImdpdmVuX25hbWUiOiJKb3NoIiwiaXBhZGRyIjoiNzYuMjguMTUzLjYxIiwibmFtZSI6Ikpvc2ggSGFycmluZ3RvbiIsIm9pZCI6ImMxYWZlNjYzLTkyZTctNDkzYy05YjU3LTE0MGVlOTE3NDU0OSIsIm9ucHJlbV9zaWQiOiJTLTEtNS0yMS0yMTI3NTIxMTg0LTE2MDQwMTI5MjAtMTg4NzkyNzUyNy01OTU5MTU4OCIsInB1aWQiOiIxMDAzMjAwMjE1RTM4QkJFIiwicmgiOiIwLkFSb0F2NGo1Y3ZHR3IwR1JxeTE4MEJIYlIwWklmM2tBdXRkUHVrUGF3ZmoyTUJNYUFPMC4iLCJzY3AiOiJ1c2VyX2ltcGVyc29uYXRpb24iLCJzdWIiOiJUMzVkNm1uai1qc1lET3J5UU54MDQ5VWw1eWtrNkh0RTRRT0FVRDhFemxZIiwidGlkIjoiNzJmOTg4YmYtODZmMS00MWFmLTkxYWItMmQ3Y2QwMTFkYjQ3IiwidW5pcXVlX25hbWUiOiJqb2hhcnJpbmd0b25AbWljcm9zb2Z0LmNvbSIsInVwbiI6ImpvaGFycmluZ3RvbkBtaWNyb3NvZnQuY29tIiwidXRpIjoiQXRfVUQ0MEZRa0ttQkRES1Nfd1NBQSIsInZlciI6IjEuMCIsIndpZHMiOlsiYjc5ZmJmNGQtM2VmOS00Njg5LTgxNDMtNzZiMTk0ZTg1NTA5Il0sInhtc19jYyI6WyJDUDEiXSwieG1zX3RjZHQiOjEyODkyNDE1NDd9.E7T4jNI1OZXaiNg4S_9mm7Rf4HELm4HJ7ukJwcywrYBeW2qypcgt5BW3SMHIm7H8MqF5ibTKWUD1m1Am6udfyiAOphCH2gzCzGcFPjbHBdnCxPaozXGGdRzwTCvFJkqXjh3vBA4pCkilA3Zf8lw6wHdHCWTQFhiDfElNGEcfi51XNzV_yW0tlBZwWHb5Tqg96AbMmxNd68K6KmaV0qTIs2gCcxREhxDECZD7hmy1qGKZPFf5PuvOMTJLs7-__i51Jvc8Ur1o46e2LxZET19LtYCu0fTEsMkRdmFvDnEtXrFBk51adzI0lDTAWKLqW3_3YKA3oBlSwVQPCdExiWwt4A",
+        "Authorization": "Sanitized",
         "Connection": "keep-alive",
         "User-Agent": "azure-ai-ml/1.5.0 azsdk-python-mgmt-machinelearningservices/0.1.0 Python/3.8.10 (Windows-10-10.0.22621-SP0)",
         "x-ms-client-request-id": "a561b98c-ae8e-11ed-b4cb-f42679b37e91"