Skip to content

Commit 1c05882

Browse files
maykumar-msftmaykumar-msft
authored
Support for Feature store settings for workspace creation (Azure#29041)
* Adding feature store settings and deployment steps in workspace arm template * Adding feature store settings in workspace
1 parent dc5bfcc commit 1c05882

File tree

10 files changed

+343
-7
lines changed

10 files changed

+343
-7
lines changed

‎sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_helper.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
ArmConstants.CONTAINER_REGISTRY_PARAMETER_NAME: "Creating Container Registry: ({0})",
4343
ArmConstants.STORAGE_ACCOUNT_PARAMETER_NAME: "Creating Storage Account: ({0})",
4444
AzureMLResourceType.WORKSPACE: "Creating workspace: ({0})",
45+
AzureMLResourceType.CONNECTIONS: "Creating connection: ({0})",
4546
}
4647

4748

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

Lines changed: 154 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,13 @@
2020
"description": "Friendly name."
2121
}
2222
},
23+
"kind": {
24+
"type": "string",
25+
"defaultValue": "Default",
26+
"metadata": {
27+
"description": "Specifies the Kind of the workspace."
28+
}
29+
},
2330
"location": {
2431
"type": "string",
2532
"metadata": {
@@ -395,6 +402,58 @@
395402
"metadata": {
396403
"description": "Managed network settings to be used for the workspace. If not specified, isolation mode Disabled is the default"
397404
}
405+
},
406+
"spark_runtime_version": {
407+
"type": "string",
408+
"metadata": {
409+
"description": "spark version to be used by all feature sets"
410+
}
411+
},
412+
"offline_store_connection_name": {
413+
"type": "string",
414+
"defaultValue": "",
415+
"metadata": {
416+
"description": "Feature store offline store config"
417+
}
418+
},
419+
"online_store_connection_name": {
420+
"type": "string",
421+
"defaultValue": "",
422+
"metadata": {
423+
"description": "Feature store online store config"
424+
}
425+
},
426+
"offline_store_connection_target": {
427+
"type": "string",
428+
"defaultValue": "",
429+
"metadata": {
430+
"description": "Feature store offline store connection target"
431+
}
432+
},
433+
"materialization_identity_client_id": {
434+
"type": "string",
435+
"defaultValue": "",
436+
"metadata": {
437+
"description": "Feature store offline store connection credential client id"
438+
}
439+
},
440+
"materialization_identity_resource_id": {
441+
"type": "string",
442+
"defaultValue": "",
443+
"metadata": {
444+
"description": "Feature store offline store connection credential resource id"
445+
}
446+
},
447+
"setup_materialization_store": {
448+
"type": "string",
449+
"defaultValue": "false",
450+
"allowedValues": [
451+
"false",
452+
"true"
453+
],
454+
"metadata": {
455+
"description": "Whether to set up materialization store"
456+
}
398457
}
399458
},
400459
"variables": {
@@ -425,7 +484,8 @@
425484
}
426485
},
427486
"defaultPEConnections": "[array(variables('privateEndpointSettings'))]",
428-
"privateEndpointDeploymentName": "[concat('DeployPrivateEndpoint-', uniqueString(parameters('privateEndpointName')))]"
487+
"privateEndpointDeploymentName": "[concat('DeployPrivateEndpoint-', uniqueString(parameters('privateEndpointName')))]",
488+
"offlineStoreConnectionName": "[if(equals(parameters('offline_store_connection_name'), ''), 'OfflineStoreConnectionName', parameters('offline_store_connection_name'))]"
429489
},
430490
"resources": [
431491
{
@@ -555,6 +615,7 @@
555615
"apiVersion": "2022-12-01-preview",
556616
"tags": "[parameters('tagValues')]",
557617
"name": "[parameters('workspaceName')]",
618+
"kind": "[parameters('kind')]",
558619
"location": "[parameters('location')]",
559620
"dependsOn": [
560621
"[resourceId('Microsoft.Storage/storageAccounts', parameters('storageAccountName'))]",
@@ -586,9 +647,99 @@
586647
"SearchAccountArmId": "[parameters('encryption_search_resourceid')]"
587648
},
588649
"primaryUserAssignedIdentity": "[parameters('primaryUserAssignedIdentity')]",
589-
"managedNetwork": "[parameters('managedNetwork')]"
590-
}
650+
"managedNetwork": "[parameters('managedNetwork')]",
651+
"featureStoreSettings": {
652+
"computeruntime": {
653+
"SparkRuntimeVersion": "[parameters('spark_runtime_version')]"
654+
},
655+
"offlinestoreconnectionname": "[parameters('offline_store_connection_name')]",
656+
"onlinestoreconnectionname": "[parameters('online_store_connection_name')]"
657+
}
658+
},
659+
"resources": [{
660+
"condition": "[equals(parameters('setup_materialization_store'), 'true')]",
661+
"type": "connections",
662+
"apiVersion": "2022-05-01",
663+
"name": "[variables('offlineStoreConnectionName')]",
664+
"location": "[parameters('location')]",
665+
"dependsOn": [
666+
"[resourceId('Microsoft.MachineLearningServices/workspaces', parameters('workspaceName'))]"
667+
],
668+
"identity": {
669+
"type": "SystemAssigned"
670+
},
671+
"properties": {
672+
"category": "AzureDataLakeGen2",
673+
"target": "[parameters('offline_store_connection_target')]",
674+
"authType": "ManagedIdentity",
675+
"credentials": {
676+
"clientid": "[parameters('materialization_identity_client_id')]",
677+
"resourceid": "[parameters('materialization_identity_resource_id')]"
678+
}
679+
}
680+
}]
591681
},
682+
{
683+
"condition": "[equals(parameters('setup_materialization_store'), 'true')]",
684+
"type": "Microsoft.Resources/deployments",
685+
"apiVersion": "2022-05-01",
686+
"name": "[concat(parameters('workspaceName'), '-update-UAI')]",
687+
"dependsOn": [
688+
"[parameters('workspaceName')]"
689+
],
690+
"properties": {
691+
"mode": "Incremental",
692+
"parameters": {},
693+
"template": {
694+
"$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
695+
"contentVersion": "1.0.0.1",
696+
"resources": [{
697+
"apiVersion": "2022-01-01-preview",
698+
"name": "[parameters('workspaceName')]",
699+
"location": "[parameters('location')]",
700+
"kind": "featurestore",
701+
"type": "Microsoft.MachineLearningServices/workspaces",
702+
"identity": {
703+
"type": "SystemAssigned,UserAssigned",
704+
"userAssignedIdentities": {
705+
"[parameters('materialization_identity_resource_id')]": {}
706+
}
707+
},
708+
"properties": {
709+
"friendlyName": "[parameters('friendlyName')]",
710+
"description": "[parameters('description')]",
711+
"storageAccount": "[variables('storageAccount')]",
712+
"keyVault": "[variables('keyVault')]",
713+
"applicationInsights": "[variables('applicationInsights')]",
714+
"containerRegistry": "[if(not(equals(parameters('containerRegistryOption'), 'none')), variables('containerRegistry'), json('null'))]",
715+
"hbiWorkspace": "[parameters('confidential_data')]",
716+
"imageBuildCompute": "[parameters('imageBuildCompute')]",
717+
"publicNetworkAccess": "[parameters('publicNetworkAccess')]",
718+
"softDeleteEnabled": "[parameters('soft_delete_enabled')]",
719+
"allowRecoverSoftDeletedWorkspace": "[parameters('allow_recover_softdeleted_workspace')]",
720+
"encryption": {
721+
"status": "[parameters('encryption_status')]",
722+
"keyVaultProperties": {
723+
"keyVaultArmId": "[parameters('cmk_keyvault')]",
724+
"keyIdentifier": "[parameters('resource_cmk_uri')]"
725+
},
726+
"cosmosDbArmId": "[parameters('encryption_cosmosdb_resourceid')]",
727+
"storageAccountArmId": "[parameters('encryption_storage_resourceid')]",
728+
"SearchAccountArmId": "[parameters('encryption_search_resourceid')]"
729+
},
730+
"primaryUserAssignedIdentity": "[parameters('primaryUserAssignedIdentity')]",
731+
"featureStoreSettings": {
732+
"computeruntime": {
733+
"SparkRuntimeVersion": "[parameters('spark_runtime_version')]"
734+
},
735+
"offlinestoreconnectionname": "[parameters('offline_store_connection_name')]",
736+
"onlinestoreconnectionname": "[parameters('online_store_connection_name')]"
737+
}
738+
}
739+
}]
740+
}
741+
}
742+
},
592743
{
593744
"condition": "[and(variables('enablePE'), not(equals(parameters('privateEndpointType'), 'none')))]",
594745
"type": "Microsoft.Resources/deployments",

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@
1111
"description": {
1212
"value": ""
1313
},
14+
"kind": {
15+
"value": "Default"
16+
},
1417
"friendlyName": {
1518
"value": ""
1619
},
@@ -118,5 +121,26 @@
118121
},
119122
"managedNetwork": {
120123
"value": {}
124+
},
125+
"setup_materialization_store" : {
126+
"value": "false"
127+
},
128+
"spark_runtime_version" : {
129+
"value": ""
130+
},
131+
"offline_store_connection_name" : {
132+
"value": ""
133+
},
134+
"offline_store_connection_target" : {
135+
"value": ""
136+
},
137+
"materialization_identity_client_id" : {
138+
"value": ""
139+
},
140+
"materialization_identity_resource_id" : {
141+
"value": ""
142+
},
143+
"online_store_connection_name" : {
144+
"value": ""
121145
}
122146
}

‎sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,7 @@ class AzureMLResourceType(object):
197197
COMPONENT = "components"
198198
SCHEDULE = "schedules"
199199
REGISTRY = "registries"
200+
CONNECTIONS = "connections"
200201

201202
NAMED_TYPES = {
202203
JOB,

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,8 @@
134134
from ._workspace.private_endpoint import EndpointConnection, PrivateEndpoint
135135
from ._workspace.workspace import Workspace
136136
from ._workspace.workspace_keys import ContainerRegistryCredential, NotebookAccessKeys, WorkspaceKeys
137+
from ._workspace.compute_runtime import ComputeRuntime
138+
from ._workspace.feature_store_settings import FeatureStoreSettings
137139

138140
# TODO: enable in PuP
139141
# from ._job.import_job import ImportJob
@@ -192,6 +194,8 @@
192194
"Workspace",
193195
"WorkspaceKeys",
194196
"WorkspaceConnection",
197+
"ComputeRuntime",
198+
"FeatureStoreSettings",
195199
"DiagnoseRequestProperties",
196200
"DiagnoseResult",
197201
"DiagnoseResponseResult",

‎sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/compute_runtime.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# ---------------------------------------------------------
2+
# Copyright (c) Microsoft Corporation. All rights reserved.
3+
# ---------------------------------------------------------
4+
5+
from typing import Optional
6+
7+
from azure.ai.ml._restclient.v2022_12_01_preview.models import ComputeRuntimeDto as RestComputeRuntimeDto
8+
from azure.ai.ml.entities._mixins import RestTranslatableMixin
9+
from azure.ai.ml._utils._experimental import experimental
10+
11+
12+
@experimental
13+
class ComputeRuntime(RestTranslatableMixin):
14+
def __init__(
15+
self,
16+
*,
17+
spark_runtime_version: Optional[str] = None,
18+
):
19+
"""
20+
:keyword spark_runtime_version:
21+
:paramtype spark_runtime_version: str
22+
"""
23+
self.spark_runtime_version = spark_runtime_version
24+
25+
def _to_rest_object(self) -> RestComputeRuntimeDto:
26+
return RestComputeRuntimeDto(spark_runtime_version=self.spark_runtime_version)
27+
28+
@classmethod
29+
def _from_rest_object(cls, obj: RestComputeRuntimeDto) -> "ComputeRuntime":
30+
if not obj:
31+
return None
32+
return ComputeRuntime(spark_runtime_version=obj.spark_runtime_version)

‎sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# ---------------------------------------------------------
2+
# Copyright (c) Microsoft Corporation. All rights reserved.
3+
# ---------------------------------------------------------
4+
5+
# pylint: disable=protected-access
6+
7+
from typing import Optional, Any
8+
9+
from azure.ai.ml._restclient.v2022_12_01_preview.models import FeatureStoreSettings as RestFeatureStoreSettings
10+
from azure.ai.ml.entities._mixins import RestTranslatableMixin
11+
from azure.ai.ml._utils._experimental import experimental
12+
from .compute_runtime import ComputeRuntime
13+
14+
15+
@experimental
16+
class FeatureStoreSettings(RestTranslatableMixin):
17+
def __init__(
18+
self, *, compute_runtime: Optional[ComputeRuntime] = None, offline_store_connection_name: Optional[str] = None
19+
):
20+
"""
21+
:keyword compute_runtime:
22+
:paramtype compute_runtime: ~azure.ai.ml.entities.ComputeRuntime
23+
:keyword offline_store_connection_name:
24+
:paramtype offline_store_connection_name: str
25+
"""
26+
self.compute_runtime = compute_runtime if compute_runtime else ComputeRuntime(spark_runtime_version="3.1.0")
27+
self.offline_store_connection_name = offline_store_connection_name
28+
29+
def _to_rest_object(self) -> RestFeatureStoreSettings:
30+
return RestFeatureStoreSettings(
31+
compute_runtime=ComputeRuntime._to_rest_object(self.compute_runtime),
32+
offline_store_connection_name=self.offline_store_connection_name,
33+
online_store_connection_name=None,
34+
)
35+
36+
@classmethod
37+
def _from_rest_object(cls, obj: RestFeatureStoreSettings) -> "FeatureStoreSettings":
38+
if not obj:
39+
return None
40+
return FeatureStoreSettings(
41+
compute_runtime=ComputeRuntime._from_rest_object(obj.compute_runtime),
42+
offline_store_connection_name=obj.offline_store_connection_name,
43+
)

0 commit comments

Comments
 (0)