Cross subscription storage account support for workspace (#37288)

* fixing sweep job service response parsing issue * fixing sweep job service response parsing issue * raising a work around for cmd job and spark job name lowering case * fixing black issue * fixing serialization issue in sweep job * include black fix * found and fix the issue in compute validation * fixing a corrupt wokspace arm template * fixing the failed test case * updating the date for 1.18.0 release * removing empty section * adding in progress version details * adding IDC code reviewers * correcting the directory path * lint fix * adding change log entry * addressing review comment * fixing ut * addressing api review comment * adding a fix for subnet validation in compute instance * my py fix * fixing autoML serialization issue * formatting * moving auto ML job create API to 2024 * subscription spell correction * adding change log entry * fixing UT * fix default workspace creation
Azure · Sep 19, 2024 · 3b725b4 · 3b725b4
1 parent bddc491
commit 3b725b4
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 3 deletions.
diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ## 1.21.0 (unreleased)
 
+### Features Added
+  - Cross subscription storage account support for workspace and feature store. Developer can provide a storage account from another subscription while creating a workspace or storage account.
+
 ## 1.20.0 (2024-09-10)
 
 ### Features Added

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json
@@ -83,6 +83,10 @@
                 "description": "Determines whether or not to put the storage account behind VNet"
             }
         },
+        "storageAccountSubscriptionId": {
+            "type": "string",
+            "defaultValue": "[subscription().subscriptionId]"
+        },
         "storageAccountResourceGroupName": {
             "type": "string",
             "defaultValue": "[parameters('resourceGroupName')]"
@@ -635,7 +639,7 @@
     },
     "variables": {
         "tenantId": "[subscription().tenantId]",
-        "storageAccount": "[resourceId(parameters('storageAccountResourceGroupName'), 'Microsoft.Storage/storageAccounts', parameters('storageAccountName'))]",
+        "storageAccount": "[resourceId(parameters('storageAccountSubscriptionId'), parameters('storageAccountResourceGroupName'), 'Microsoft.Storage/storageAccounts', parameters('storageAccountName'))]",
         "keyVault": "[resourceId(parameters('keyVaultResourceGroupName'), 'Microsoft.KeyVault/vaults', parameters('keyVaultName'))]",
         "containerRegistry": "[resourceId(parameters('containerRegistryResourceGroupName'), 'Microsoft.ContainerRegistry/registries', parameters('containerRegistryName'))]",
         "applicationInsights": "[resourceId(parameters('applicationInsightsResourceGroupName'), 'Microsoft.Insights/components', parameters('applicationInsightsName'))]",

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json
@@ -38,6 +38,9 @@
     "storageAccountName": {
         "value": ""
     },
+    "storageAccountSubscriptionId": {
+        "value": ""
+    },
     "storageAccountResourceGroupName": {
         "value": ""
     },

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_workspace_utils.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_workspace_utils.py
@@ -84,6 +84,10 @@ def get_resource_and_group_name(armstr: str) -> str:
     return armstr.split("/")[-1], armstr.split("/")[-5]
 
 
+def get_sub_id_resource_and_group_name(armstr: str) -> str:
+    return armstr.split("/")[-7], armstr.split("/")[-1], armstr.split("/")[-5]
+
+
 def get_endpoint_parts(arm_id: str, subnet_arm_id: str) -> ():
     arm_id_parts = arm_id.split("/")
     subnet_id_parts = subnet_arm_id.split("/")

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
@@ -29,16 +29,17 @@
     get_name_for_dependent_resource,
     get_resource_and_group_name,
     get_resource_group_location,
+    get_sub_id_resource_and_group_name,
 )
 from azure.ai.ml._utils.utils import camel_to_snake, from_iso_duration_format_min_sec
 from azure.ai.ml._version import VERSION
 from azure.ai.ml.constants import ManagedServiceIdentityType
 from azure.ai.ml.constants._common import (
+    WORKSPACE_PATCH_REJECTED_KEYS,
     ArmConstants,
     LROConfigurations,
     WorkspaceKind,
     WorkspaceResourceConstants,
-    WORKSPACE_PATCH_REJECTED_KEYS,
 )
 from azure.ai.ml.constants._workspace import IsolationMode, OutboundRuleCategory
 from azure.ai.ml.entities import Hub, Project, Workspace
@@ -582,17 +583,22 @@ def _populate_arm_parameters(self, workspace: Workspace, **kwargs: Any) -> Tuple
             )
 
         if workspace.storage_account:
-            resource_name, group_name = get_resource_and_group_name(workspace.storage_account)
+            subscription_id, resource_name, group_name = get_sub_id_resource_and_group_name(workspace.storage_account)
             _set_val(param["storageAccountName"], resource_name)
             _set_val(param["storageAccountOption"], "existing")
             _set_val(param["storageAccountResourceGroupName"], group_name)
+            _set_val(param["storageAccountSubscriptionId"], subscription_id)
         else:
             storage = _generate_storage(workspace.name, resources_being_deployed)
             _set_val(param["storageAccountName"], storage)
             _set_val(
                 param["storageAccountResourceGroupName"],
                 workspace.resource_group,
             )
+            _set_val(
+                param["storageAccountSubscriptionId"],
+                self._subscription_id,
+            )
 
         if workspace.application_insights:
             resource_name, group_name = get_resource_and_group_name(workspace.application_insights)

diff --git a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py
@@ -8,6 +8,8 @@
 from azure.ai.ml._restclient.v2024_07_01_preview.models import (
     EncryptionKeyVaultUpdateProperties,
     EncryptionUpdateProperties,
+)
+from azure.ai.ml._restclient.v2024_07_01_preview.models import (
     ServerlessComputeSettings as RestServerlessComputeSettings,
 )
 from azure.ai.ml._scope_dependent_operations import OperationScope
@@ -305,6 +307,10 @@ def _populate_arm_parameters(
     def test_populate_arm_parameters_other_branches(
         self, mock_workspace_operation_base: WorkspaceOperationsBase, mocker: MockFixture
     ) -> None:
+        mocker.patch(
+            "azure.ai.ml.operations._workspace_operations_base.get_sub_id_resource_and_group_name",
+            return_value=["random_name", "random_grp", "random_storage"],
+        )
         mocker.patch(
             "azure.ai.ml.operations._workspace_operations_base.get_resource_group_location", return_value="random_name"
         )
@@ -330,6 +336,10 @@ def test_populate_arm_parameters_other_branches(
     def test_populate_feature_store_arm_parameters(
         self, mock_workspace_operation_base: WorkspaceOperationsBase, mocker: MockFixture
     ) -> None:
+        mocker.patch(
+            "azure.ai.ml.operations._workspace_operations_base.get_sub_id_resource_and_group_name",
+            return_value=["random_sub", "random_grp", "random_storage"],
+        )
         mocker.patch(
             "azure.ai.ml.operations._workspace_operations_base.get_resource_group_location", return_value="random_name"
         )