Featurestore/online store crud new pr

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')]"
+                                    }
+                                }
+                            }
+                        ]
 					}]
 				}
 			}

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": ""
     }
 }

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())

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"

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,13 +118,19 @@ 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
             else _ComputeRuntime(spark_runtime_version=DEFAULT_SPARK_RUNTIME_VERSION),
             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
 

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,
         )