Added Workspace property ProvisionNetworkNow. (#38404)

* Added Workspace property `ProvisionNetworkNow`.

 - Updated `CHANGELOG.md` to add this feature
 - Updated ARM templated
   - Patching for this property is not supported as defined in its requirement
 - Updated Workspace logic to support `ProvisionNetworkNow`
 - Updated sample code
 - Updated test cases

 - Please note that the Swagger (API spec) and the REST client for `v2024-10-01-preview` have updated earlier

* Added Workspace property `ProvisionNetworkNow`.

 - Updated `CHANGELOG.md` to add this feature
 - Updated ARM templated
   - Patching for this property is not supported as defined in its requirement
 - Updated Workspace logic to support `ProvisionNetworkNow`
 - Updated sample code
 - Updated test cases

 - Please note that the Swagger (API spec) and the REST client for `v2024-10-01-preview` have updated earlier

---------

Co-authored-by: Jing Li <[email protected]>

Author: u9009

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -11,7 +11,9 @@
   - 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.model_customization`
+  - Added support for Distillation jobs. Can be created by importing `disillation` from `azure.ai.ml.distillation`
+  - Added Workspace property `ProvisionNetworkNow` to trigger the provisioning of the managed VNet with the default
+    Options when creating a Workspace with the managed VNet enabled, or else it does nothing
 ### Bugs Fixed
   - Fix send email notification issue in model monitoring
 

sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json

@@ -432,6 +432,17 @@
                 "description": "Managed network settings to be used for the workspace. If not specified, isolation mode Disabled is the default"
             }
         },
+        "provisionNetworkNow": {
+            "type": "string",
+            "defaultValue": "false",
+            "allowedValues": [
+                "true",
+                "false"
+            ],
+            "metadata": {
+                "description": "Set to provision the managed VNet with the default Options when creating a Workspace with the managed VNet enabled"
+            }
+        },
         "allowRoleAssignmentOnRG": {
             "type": "string",
             "defaultValue": "true",
@@ -885,6 +896,7 @@
                 "systemDatastoresAuthMode": "[parameters('systemDatastoresAuthMode')]",
                 "allowRoleAssignmentOnRG": "[parameters('allowRoleAssignmentOnRG')]",
                 "managedNetwork": "[parameters('managedNetwork')]",
+                "provisionNetworkNow": "[parameters('provisionNetworkNow')]",
                 "featureStoreSettings": {
                     "computeruntime": {
                         "SparkRuntimeVersion": "[parameters('spark_runtime_version')]"
@@ -1001,6 +1013,7 @@
                                 "systemDatastoresAuthMode": "[parameters('systemDatastoresAuthMode')]",
                                 "allowRoleAssignmentOnRG": "[parameters('allowRoleAssignmentOnRG')]",
                                 "managedNetwork": "[parameters('managedNetwork')]",
+                                "provisionNetworkNow": "[parameters('provisionNetworkNow')]",
                                 "featureStoreSettings": {
                                     "computeruntime": {
                                         "SparkRuntimeVersion": "[parameters('spark_runtime_version')]"
@@ -1178,4 +1191,4 @@
             }
         }
     ]
-}
+}

sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json

@@ -134,6 +134,9 @@
     "managedNetwork": {
         "value": {}
     },
+    "provisionNetworkNow": {
+        "value": "false"
+    },
     "spark_runtime_version" : {
         "value": null
     },

sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_project.json

@@ -432,6 +432,17 @@
                 "description": "Managed network settings to be used for the workspace. If not specified, isolation mode Disabled is the default"
             }
         },
+        "provisionNetworkNow": {
+            "type": "string",
+            "defaultValue": "false",
+            "allowedValues": [
+                "true",
+                "false"
+            ],
+            "metadata": {
+                "description": "Set to provision the managed VNet with the default Options when creating a Workspace with the managed VNet enabled"
+            }
+        },
         "enable_data_isolation": {
             "type": "string",
             "defaultValue": "false",
@@ -717,4 +728,4 @@
             }
         }
     ]
-}
+}

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

@@ -41,6 +41,7 @@ class WorkspaceSchema(PathAwareSchema):
     primary_user_assigned_identity = fields.Str()
     workspace_hub = fields.Str(validate=validate_arm_str)
     managed_network = NestedField(ManagedNetworkSchema, unknown=EXCLUDE)
+    provision_network_now = fields.Bool()
     enable_data_isolation = fields.Bool()
     allow_roleassignment_on_rg = fields.Bool()
     serverless_compute = NestedField(ServerlessComputeSettingsSchema)

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

@@ -81,6 +81,9 @@ class Workspace(Resource):
     :type primary_user_assigned_identity: str
     :param managed_network: workspace's Managed Network configuration
     :type managed_network: ~azure.ai.ml.entities.ManagedNetwork
+    :param provision_network_now: Set to trigger the provisioning of the managed vnet with the default options when
+        creating a workspace with the managed vnet enable, or else it does nothing
+    :type provision_network_now: Optional[bool]
     :param system_datastores_auth_mode: The authentication mode for system datastores.
     :type system_datastores_auth_mode: str
     :param enable_data_isolation: A flag to determine if workspace has data isolation enabled.
@@ -104,6 +107,7 @@ class Workspace(Resource):
             :caption: Creating a Workspace object.
     """
 
+    # pylint: disable=too-many-locals
     def __init__(
         self,
         *,
@@ -124,6 +128,7 @@ def __init__(
         identity: Optional[IdentityConfiguration] = None,
         primary_user_assigned_identity: Optional[str] = None,
         managed_network: Optional[ManagedNetwork] = None,
+        provision_network_now: Optional[bool] = None,
         system_datastores_auth_mode: Optional[str] = None,
         enable_data_isolation: bool = False,
         allow_roleassignment_on_rg: Optional[bool] = None,
@@ -164,6 +169,7 @@ def __init__(
         self.identity = identity
         self.primary_user_assigned_identity = primary_user_assigned_identity
         self.managed_network = managed_network
+        self.provision_network_now = provision_network_now
         self.system_datastores_auth_mode = system_datastores_auth_mode
         self.enable_data_isolation = enable_data_isolation
         self.allow_roleassignment_on_rg = allow_roleassignment_on_rg
@@ -363,6 +369,11 @@ def _from_rest_object(
                     rest_obj.managed_network
                 )
 
+        # TODO: Remove once it's included in response
+        provision_network_now = None
+        if hasattr(rest_obj, "provision_network_now"):
+            provision_network_now = rest_obj.provision_network_now
+
         armid_parts = str(rest_obj.id).split("/")
         group = None if len(armid_parts) < 4 else armid_parts[4]
         identity = None
@@ -407,6 +418,7 @@ def _from_rest_object(
             identity=identity,
             primary_user_assigned_identity=rest_obj.primary_user_assigned_identity,
             managed_network=managed_network,
+            provision_network_now=provision_network_now,
             system_datastores_auth_mode=system_datastores_auth_mode,
             feature_store_settings=feature_store_settings,
             enable_data_isolation=rest_obj.enable_data_isolation,
@@ -454,6 +466,7 @@ def _to_rest_object(self) -> RestWorkspace:
                 if self.managed_network
                 else None
             ),  # pylint: disable=protected-access
+            provision_network_now=self.provision_network_now,
             system_datastores_auth_mode=self.system_datastores_auth_mode,
             feature_store_settings=feature_store_settings,
             enable_data_isolation=self.enable_data_isolation,

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

@@ -781,6 +781,9 @@ def _populate_arm_parameters(self, workspace: Workspace, **kwargs: Any) -> Tuple
             if not kwargs.get("grant_materialization_permissions", None):
                 _set_val(param["grant_materialization_permissions"], "false")
 
+        if workspace.provision_network_now:
+            _set_val(param["provisionNetworkNow"], "true")
+
         managed_network = None
         if workspace.managed_network:
             managed_network = workspace.managed_network._to_rest_object()

sdk/ml/azure-ai-ml/samples/ml_samples_workspace.py

@@ -116,6 +116,13 @@ def ml_workspace_config_sample_snippets_entities(self):
         ws = Workspace(name="ws-name", location="eastus", managed_network=network)
         # [END workspace_managed_network]
 
+        # [START workspace_managed_network_provision_now]
+        from azure.ai.ml.entities import IsolationMode, ManagedNetwork, Workspace
+
+        managed_net = ManagedNetwork(isolation_mode=IsolationMode.ALLOW_INTERNET_OUTBOUND)
+        ws = Workspace(name="ws-name", location="eastus", managed_network=managed_net, provision_network_now=True)
+        # [END workspace_managed_network_provision_now]
+
         # [START fqdn_outboundrule]
         from azure.ai.ml.entities import FqdnDestination
 

sdk/ml/azure-ai-ml/tests/conftest.py

@@ -308,6 +308,11 @@ def mock_aml_services_2024_07_01_preview(mocker: MockFixture) -> Mock:
     return mocker.patch("azure.ai.ml._restclient.v2024_07_01_preview")
 
 
+@pytest.fixture
+def mock_aml_services_2024_10_01_preview(mocker: MockFixture) -> Mock:
+    return mocker.patch("azure.ai.ml._restclient.v2024_10_01_preview")
+
+
 @pytest.fixture
 def mock_aml_services_run_history(mocker: MockFixture) -> Mock:
     return mocker.patch("azure.ai.ml._restclient.runhistory")

sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py

@@ -136,6 +136,7 @@ def outgoing_get_call(rg, name):
                     ),
                 ],
             )
+            ws.provision_network_now = True
             ws.system_datastores_auth_mode = "identity"
             ws.allow_roleassignment_on_rg = True
             return ws._to_rest_object()
@@ -163,6 +164,8 @@ def outgoing_get_call(rg, name):
         assert rules[2].protocol == "*"
         assert rules[2].port_ranges == "1,2"
 
+        assert ws.provision_network_now == True
+
     def test_create_get_exception_swallow(
         self,
         mock_workspace_operation_base: WorkspaceOperationsBase,