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 index bb00aeff5a12..620eb8cfccab 100644 --- 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 @@ -444,6 +444,13 @@ "description": "Feature store offline store connection target" } }, + "online_store_connection_target": { + "type": "string", + "defaultValue": "", + "metadata": { + "description": "Feature store online store connection target" + } + }, "materialization_identity_client_id": { "type": "string", "defaultValue": "", @@ -499,7 +506,8 @@ }, "defaultPEConnections": "[array(variables('privateEndpointSettings'))]", "privateEndpointDeploymentName": "[concat('DeployPrivateEndpoint-', uniqueString(parameters('privateEndpointName')))]", - "offlineStoreConnectionName": "[if(equals(parameters('offline_store_connection_name'), ''), 'OfflineStoreConnectionName', parameters('offline_store_connection_name'))]" + "offlineStoreConnectionName": "[if(equals(parameters('offline_store_connection_name'), ''), 'OfflineStoreConnectionName', parameters('offline_store_connection_name'))]", + "onlineStoreConnectionName": "[if(equals(parameters('online_store_connection_name'), ''), 'OnlineStoreConnectionName', parameters('online_store_connection_name'))]" }, "resources": [ { @@ -685,35 +693,13 @@ "offlinestoreconnectionname": "[parameters('offline_store_connection_name')]", "onlinestoreconnectionname": "[parameters('online_store_connection_name')]" } - }, - "resources": [{ - "condition": "[equals(parameters('setup_materialization_store'), 'true')]", - "type": "connections", - "apiVersion": "2022-05-01", - "name": "[variables('offlineStoreConnectionName')]", - "location": "[parameters('location')]", - "dependsOn": [ - "[resourceId('Microsoft.MachineLearningServices/workspaces', parameters('workspaceName'))]" - ], - "identity": { - "type": "SystemAssigned" - }, - "properties": { - "category": "AzureDataLakeGen2", - "target": "[parameters('offline_store_connection_target')]", - "authType": "ManagedIdentity", - "credentials": { - "clientid": "[parameters('materialization_identity_client_id')]", - "resourceid": "[parameters('materialization_identity_resource_id')]" - } - } - }] + } }, { - "condition": "[equals(parameters('setup_materialization_store'), 'true')]", + "condition":"[and(equals(parameters('setup_materialization_store'), 'true'),not(equals(parameters('offline_store_connection_target'), '')))]", "type": "Microsoft.Resources/deployments", "apiVersion": "2022-05-01", - "name": "[concat(parameters('workspaceName'), '-update-materialization-identity')]", + "name": "[concat(parameters('workspaceName'), '-update-offline-connection')]", "dependsOn": [ "[parameters('workspaceName')]" ], @@ -724,7 +710,7 @@ "$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#", "contentVersion": "1.0.0.1", "resources": [{ - "apiVersion": "2022-01-01-preview", + "apiVersion": "2022-12-01-preview", "name": "[parameters('workspaceName')]", "location": "[parameters('location')]", "kind": "featurestore", @@ -763,7 +749,114 @@ "offlinestoreconnectionname": "[parameters('offline_store_connection_name')]", "onlinestoreconnectionname": "[parameters('online_store_connection_name')]" } - } + }, + "resources": [ + { + "condition": "[and(equals(parameters('setup_materialization_store'), 'true'),not(equals(parameters('offline_store_connection_target'), '')))]", + "type": "connections", + "apiVersion": "2022-05-01", + "name": "[variables('offlineStoreConnectionName')]", + "location": "[parameters('location')]", + "dependsOn": [ + "[resourceId('Microsoft.MachineLearningServices/workspaces', parameters('workspaceName'))]" + ], + "identity": { + "type": "SystemAssigned" + }, + "properties": { + "category": "AzureDataLakeGen2", + "target": "[parameters('offline_store_connection_target')]", + "authType": "ManagedIdentity", + "credentials": { + "clientid": "[parameters('materialization_identity_client_id')]", + "resourceid": "[parameters('materialization_identity_resource_id')]" + } + } + } + ] + }] + } + } + }, + { + "condition":"[and(equals(parameters('setup_materialization_store'), 'true'),not(equals(parameters('online_store_connection_target'), '')))]", + "type": "Microsoft.Resources/deployments", + "apiVersion": "2022-05-01", + "name": "[concat(parameters('workspaceName'), '-update-online-connection')]", + "dependsOn": [ + "[parameters('workspaceName')]" + ], + "properties": { + "mode": "Incremental", + "parameters": {}, + "template": { + "$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#", + "contentVersion": "1.0.0.1", + "resources": [{ + "apiVersion": "2022-12-01-preview", + "name": "[parameters('workspaceName')]", + "location": "[parameters('location')]", + "kind": "featurestore", + "type": "Microsoft.MachineLearningServices/workspaces", + "identity": { + "type": "SystemAssigned,UserAssigned", + "userAssignedIdentities": { + "[parameters('materialization_identity_resource_id')]": {} + } + }, + "properties": { + "friendlyName": "[parameters('friendlyName')]", + "description": "[parameters('description')]", + "storageAccount": "[variables('storageAccount')]", + "keyVault": "[variables('keyVault')]", + "applicationInsights": "[variables('applicationInsights')]", + "containerRegistry": "[if(not(equals(parameters('containerRegistryOption'), 'none')), variables('containerRegistry'), json('null'))]", + "hbiWorkspace": "[parameters('confidential_data')]", + "imageBuildCompute": "[parameters('imageBuildCompute')]", + "publicNetworkAccess": "[parameters('publicNetworkAccess')]", + "encryption": { + "status": "[parameters('encryption_status')]", + "keyVaultProperties": { + "keyVaultArmId": "[parameters('cmk_keyvault')]", + "keyIdentifier": "[parameters('resource_cmk_uri')]" + }, + "cosmosDbArmId": "[parameters('encryption_cosmosdb_resourceid')]", + "storageAccountArmId": "[parameters('encryption_storage_resourceid')]", + "SearchAccountArmId": "[parameters('encryption_search_resourceid')]" + }, + "primaryUserAssignedIdentity": "[parameters('primaryUserAssignedIdentity')]", + "featureStoreSettings": { + "computeruntime": { + "SparkRuntimeVersion": "[parameters('spark_runtime_version')]" + }, + "offlinestoreconnectionname": "[parameters('offline_store_connection_name')]", + "onlinestoreconnectionname": "[parameters('online_store_connection_name')]" + } + }, + "resources": [ + { + "condition": "[and(equals(parameters('setup_materialization_store'), 'true'),not(equals(parameters('online_store_connection_target'), '')))]", + "type": "connections", + "apiVersion": "2022-05-01", + "name": "[variables('onlineStoreConnectionName')]", + "location": "[parameters('location')]", + "dependsOn": [ + "[resourceId('Microsoft.MachineLearningServices/workspaces', parameters('workspaceName'))]" + ], + "identity": { + "type": "SystemAssigned" + }, + "properties": { + "category": "Redis", + "target": "[parameters('online_store_connection_target')]", + "authType": "ManagedIdentity", + "credentials": { + "clientid": "[parameters('materialization_identity_client_id')]", + "resourceid": "[parameters('materialization_identity_resource_id')]" + } + } + } + ] }] } } 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 index bda9a3ede13e..5ab9f65573b4 100644 --- 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 @@ -148,5 +148,8 @@ }, "online_store_connection_name" : { "value": "" + }, + "online_store_connection_target" : { + "value": "" } } \ No newline at end of file diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_store/feature_store_schema.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_store/feature_store_schema.py index 80a7e2454360..402c8aea7e78 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_store/feature_store_schema.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_feature_store/feature_store_schema.py @@ -19,6 +19,7 @@ class FeatureStoreSchema(PathAwareSchema): name = fields.Str(required=True) compute_runtime = NestedField(ComputeRuntimeSchema) offline_store = NestedField(MaterializationStoreSchema) + online_store = NestedField(MaterializationStoreSchema) materialization_identity = NestedField(UserAssignedIdentitySchema) description = fields.Str() tags = fields.Dict(keys=fields.Str(), values=fields.Str()) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/_constants.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/_constants.py index 646a94ba2f25..8063c74e4363 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/_constants.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/_constants.py @@ -5,5 +5,8 @@ OFFLINE_STORE_CONNECTION_NAME = "OfflineStoreConnectionName" OFFLINE_MATERIALIZATION_STORE_TYPE = "azure_data_lake_gen2" OFFLINE_STORE_CONNECTION_CATEGORY = "ADLSGen2" -DEFAULT_SPARK_RUNTIME_VERSION = "3.1.0" +ONLINE_STORE_CONNECTION_NAME = "OnlineStoreConnectionName" +ONLINE_MATERIALIZATION_STORE_TYPE = "redis" +ONLINE_STORE_CONNECTION_CATEGORY = "Redis" +DEFAULT_SPARK_RUNTIME_VERSION = "3.2.0" FEATURE_STORE_KIND = "featurestore" diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py index 7731f6dea952..2ba2cd928eef 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py @@ -22,7 +22,12 @@ from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY from .materialization_store import _MaterializationStore -from ._constants import OFFLINE_STORE_CONNECTION_NAME, DEFAULT_SPARK_RUNTIME_VERSION, FEATURE_STORE_KIND +from ._constants import ( + OFFLINE_STORE_CONNECTION_NAME, + ONLINE_STORE_CONNECTION_NAME, + DEFAULT_SPARK_RUNTIME_VERSION, + FEATURE_STORE_KIND, +) @experimental @@ -33,6 +38,7 @@ def __init__( name: str, compute_runtime: Optional[_ComputeRuntime] = None, offline_store: Optional[_MaterializationStore] = None, + online_store: Optional[_MaterializationStore] = None, materialization_identity: Optional[ManagedIdentityConfiguration] = None, description: Optional[str] = None, tags: Optional[Dict[str, str]] = None, @@ -61,6 +67,9 @@ def __init__( :param offline_store: Offline store for feature store. materialization_identity is required when offline_store is passed. :type offline_store: ~azure.ai.ml.entities._MaterializationStore + :param online_store: Online store for feature store. + materialization_identity is required when online_store is passed. + :type online_store: ~azure.ai.ml.entities._MaterializationStore :param materialization_identity: Identity used for materialization. :type materialization_identity: ~azure.ai.ml.entities.ManagedIdentityConfiguration :param description: Description of the feature store. @@ -109,6 +118,9 @@ def __init__( if offline_store and not materialization_identity: raise ValidationError("materialization_identity is required to setup offline store") + if online_store and not materialization_identity: + raise ValidationError("materialization_identity is required to setup online store") + feature_store_settings = _FeatureStoreSettings( compute_runtime=compute_runtime if compute_runtime @@ -116,6 +128,9 @@ def __init__( offline_store_connection_name=( OFFLINE_STORE_CONNECTION_NAME if materialization_identity and offline_store else None ), + online_store_connection_name=( + ONLINE_STORE_CONNECTION_NAME if materialization_identity and online_store else None + ), ) self._workspace_id = kwargs.pop("workspace_id", "") super().__init__( @@ -140,6 +155,7 @@ def __init__( **kwargs, ) self.offline_store = offline_store + self.online_store = online_store self.materialization_identity = materialization_identity self.identity = identity diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py index 535c5eeefd26..49fc7c863979 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py @@ -15,22 +15,29 @@ @experimental class _FeatureStoreSettings(RestTranslatableMixin): def __init__( - self, *, compute_runtime: Optional[_ComputeRuntime] = None, offline_store_connection_name: Optional[str] = None + self, + *, + compute_runtime: Optional[_ComputeRuntime] = None, + offline_store_connection_name: Optional[str] = None, + online_store_connection_name: Optional[str] = None, ): """ :keyword compute_runtime: :paramtype compute_runtime: ~azure.ai.ml.entities.ComputeRuntime :keyword offline_store_connection_name: :paramtype offline_store_connection_name: str + :keyword online_store_connection_name: + :paramtype online_store_connection_name: str """ - self.compute_runtime = compute_runtime if compute_runtime else _ComputeRuntime(spark_runtime_version="3.1.0") + self.compute_runtime = compute_runtime if compute_runtime else _ComputeRuntime(spark_runtime_version="3.2.0") self.offline_store_connection_name = offline_store_connection_name + self.online_store_connection_name = online_store_connection_name def _to_rest_object(self) -> RestFeatureStoreSettings: return RestFeatureStoreSettings( compute_runtime=_ComputeRuntime._to_rest_object(self.compute_runtime), offline_store_connection_name=self.offline_store_connection_name, - online_store_connection_name=None, + online_store_connection_name=self.online_store_connection_name, ) @classmethod @@ -40,4 +47,5 @@ def _from_rest_object(cls, obj: RestFeatureStoreSettings) -> "_FeatureStoreSetti return _FeatureStoreSettings( compute_runtime=_ComputeRuntime._from_rest_object(obj.compute_runtime), offline_store_connection_name=obj.offline_store_connection_name, + online_store_connection_name=obj.online_store_connection_name, ) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py index faff0da32877..b06d0dce7ca8 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py @@ -30,6 +30,9 @@ OFFLINE_STORE_CONNECTION_NAME, OFFLINE_MATERIALIZATION_STORE_TYPE, OFFLINE_STORE_CONNECTION_CATEGORY, + ONLINE_STORE_CONNECTION_NAME, + ONLINE_MATERIALIZATION_STORE_TYPE, + ONLINE_STORE_CONNECTION_CATEGORY, FEATURE_STORE_KIND, ) from azure.ai.ml.constants import ManagedServiceIdentityType @@ -111,30 +114,54 @@ def get(self, name: str, **kwargs: Dict) -> _FeatureStore: feature_store = _FeatureStore._from_rest_object(rest_workspace_obj) if feature_store: - offline_Store_connection = None + offline_store_connection = None if ( rest_workspace_obj.feature_store_settings and rest_workspace_obj.feature_store_settings.offline_store_connection_name ): try: - offline_Store_connection = self._workspace_connection_operation.get( + offline_store_connection = self._workspace_connection_operation.get( resource_group, name, rest_workspace_obj.feature_store_settings.offline_store_connection_name ) except ResourceNotFoundError: pass - if offline_Store_connection: + if offline_store_connection: if ( - offline_Store_connection.properties - and offline_Store_connection.properties.category == OFFLINE_STORE_CONNECTION_CATEGORY + offline_store_connection.properties + and offline_store_connection.properties.category == OFFLINE_STORE_CONNECTION_CATEGORY ): feature_store.offline_store = _MaterializationStore( - type=OFFLINE_MATERIALIZATION_STORE_TYPE, target=offline_Store_connection.properties.target + type=OFFLINE_MATERIALIZATION_STORE_TYPE, target=offline_store_connection.properties.target ) - # materialization identity = identity when created through feature store operations + + online_store_connection = None + if ( + rest_workspace_obj.feature_store_settings + and rest_workspace_obj.feature_store_settings.online_store_connection_name + ): + try: + online_store_connection = self._workspace_connection_operation.get( + resource_group, name, rest_workspace_obj.feature_store_settings.online_store_connection_name + ) + except ResourceNotFoundError: + pass + + if online_store_connection: if ( - offline_Store_connection.name == OFFLINE_STORE_CONNECTION_NAME - and feature_store.identity + online_store_connection.properties + and online_store_connection.properties.category == ONLINE_STORE_CONNECTION_CATEGORY + ): + feature_store.online_store = _MaterializationStore( + type=ONLINE_MATERIALIZATION_STORE_TYPE, target=online_store_connection.properties.target + ) + + # materialization identity = identity when created through feature store operations + if (offline_store_connection and offline_store_connection.name == OFFLINE_STORE_CONNECTION_NAME) or ( + online_store_connection and online_store_connection.name == ONLINE_STORE_CONNECTION_NAME + ): + if ( + feature_store.identity and feature_store.identity.user_assigned_identities and isinstance(feature_store.identity.user_assigned_identities[0], ManagedIdentityConfiguration) ): @@ -166,6 +193,11 @@ def begin_create( if feature_store.offline_store and not feature_store.materialization_identity: raise ValidationError("materialization_identity is required to setup offline store") + if feature_store.online_store and feature_store.online_store.type != ONLINE_MATERIALIZATION_STORE_TYPE: + raise ValidationError("online store type should be redis") + if feature_store.online_store and not feature_store.materialization_identity: + raise ValidationError("materialization_identity is required to setup online store") + def get_callback(): return self.get(feature_store.name) @@ -174,6 +206,7 @@ def get_callback(): update_dependent_resources=update_dependent_resources, get_callback=get_callback, offline_store_target=feature_store.offline_store.target if feature_store.offline_store else None, + online_store_target=feature_store.online_store.target if feature_store.online_store else None, materialization_identity=feature_store.materialization_identity, **kwargs, ) @@ -213,28 +246,49 @@ def begin_update( resource_group = kwargs.get("resource_group") or self._resource_group_name offline_store = kwargs.get("offline_store", feature_store.offline_store) + online_store = kwargs.get("online_store", feature_store.online_store) materialization_identity = kwargs.get("materialization_identity", feature_store.materialization_identity) if offline_store and offline_store.type != OFFLINE_MATERIALIZATION_STORE_TYPE: raise ValidationError("offline store type should be azure_data_lake_gen2") if offline_store and rest_workspace_obj.feature_store_settings.offline_store_connection_name: - existing_offline_Store_connection = self._workspace_connection_operation.get( + existing_offline_store_connection = self._workspace_connection_operation.get( resource_group, feature_store.name, rest_workspace_obj.feature_store_settings.offline_store_connection_name, ) - if existing_offline_Store_connection: + if existing_offline_store_connection: if ( - not existing_offline_Store_connection.properties - or existing_offline_Store_connection.properties.target != offline_store.target + not existing_offline_store_connection.properties + or existing_offline_store_connection.properties.target != offline_store.target ): raise ValidationError("Cannot update the offline store target") else: if not materialization_identity: raise ValidationError("Materialization identity is required to setup offline store connection") + if online_store and online_store.type != ONLINE_MATERIALIZATION_STORE_TYPE: + raise ValidationError("online store type should be redis") + + if online_store and rest_workspace_obj.feature_store_settings.online_store_connection_name: + existing_online_store_connection = self._workspace_connection_operation.get( + resource_group, + feature_store.name, + rest_workspace_obj.feature_store_settings.online_store_connection_name, + ) + + if existing_online_store_connection: + if ( + not existing_online_store_connection.properties + or existing_online_store_connection.properties.target != online_store.target + ): + raise ValidationError("Cannot update the online store target") + else: + if not materialization_identity: + raise ValidationError("Materialization identity is required to setup online store connection") + feature_store_settings = _FeatureStoreSettings._from_rest_object(rest_workspace_obj.feature_store_settings) if offline_store and materialization_identity: @@ -258,6 +312,27 @@ def begin_update( ) feature_store_settings.offline_store_connection_name = offline_store_connection_name + if online_store and materialization_identity: + online_store_connection_name = ( + feature_store_settings.online_store_connection_name + if feature_store_settings.online_store_connection_name + else ONLINE_STORE_CONNECTION_NAME + ) + online_store_connection = WorkspaceConnection( + name=online_store_connection_name, + type=online_store.type, + target=online_store.target, + credentials=materialization_identity, + ) + rest_online_store_connection = online_store_connection._to_rest_object() + self._workspace_connection_operation.create( + resource_group_name=resource_group, + workspace_name=feature_store.name, + connection_name=online_store_connection_name, + parameters=rest_online_store_connection, + ) + feature_store_settings.online_store_connection_name = online_store_connection_name + identity = kwargs.get("identity", feature_store.identity) if materialization_identity: identity = IdentityConfiguration( 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 index d598a0dcc01e..e19eee074b8f 100644 --- 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 @@ -430,18 +430,27 @@ def _populate_arm_paramaters(self, workspace: Workspace, **kwargs: Dict) -> Tupl if workspace._feature_store_settings.offline_store_connection_name else "", ) - _set_val(param["online_store_connection_name"], "") + _set_val( + param["online_store_connection_name"], + workspace._feature_store_settings.online_store_connection_name + if workspace._feature_store_settings.online_store_connection_name + else "", + ) setup_materialization_store = False if workspace._kind and workspace._kind.lower() == "featurestore": materialization_identity = kwargs.get("materialization_identity", None) offline_store_target = kwargs.get("offline_store_target", None) + online_store_target = kwargs.get("online_store_target", None) - setup_materialization_store = offline_store_target and materialization_identity + setup_materialization_store = (offline_store_target or online_store_target) and materialization_identity _set_val(param["setup_materialization_store"], "true" if setup_materialization_store else "false") if setup_materialization_store: - _set_val(param["offline_store_connection_target"], offline_store_target) + if offline_store_target is not None: + _set_val(param["offline_store_connection_target"], offline_store_target) + if online_store_target is not None: + _set_val(param["online_store_connection_target"], online_store_target) _set_val(param["materialization_identity_client_id"], materialization_identity.client_id) _set_val(param["materialization_identity_resource_id"], materialization_identity.resource_id)