From d2848b53e1efd95589a66d989ef5b52ffa158759 Mon Sep 17 00:00:00 2001 From: xuzhang3 Date: Wed, 31 Jan 2024 11:00:13 +0800 Subject: [PATCH 1/2] machine learning workspace feature store --- .../machine_learning_workspace_resource.go | 107 +++++++++++++++ ...achine_learning_workspace_resource_test.go | 128 ++++++++++++++++++ .../machine_learning_workspace.html.markdown | 17 +++ 3 files changed, 252 insertions(+) diff --git a/internal/services/machinelearning/machine_learning_workspace_resource.go b/internal/services/machinelearning/machine_learning_workspace_resource.go index b04d29d79b9e..07aeee85d9af 100644 --- a/internal/services/machinelearning/machine_learning_workspace_resource.go +++ b/internal/services/machinelearning/machine_learning_workspace_resource.go @@ -5,6 +5,7 @@ package machinelearning import ( "fmt" + "strings" "time" "github.com/hashicorp/go-azure-helpers/lang/pointer" @@ -94,6 +95,39 @@ func resourceMachineLearningWorkspace() *pluginsdk.Resource { "identity": commonschema.SystemAssignedUserAssignedIdentityRequired(), + "kind": { + Type: pluginsdk.TypeString, + Optional: true, + ValidateFunc: validation.StringInSlice([]string{ + "Default", + "FeatureStore", + }, false), + Default: "Default", + }, + + "feature_store": { + Type: pluginsdk.TypeList, + Optional: true, + MaxItems: 1, + Elem: &pluginsdk.Resource{ + Schema: map[string]*pluginsdk.Schema{ + "computer_spark_runtime_version": { + Type: pluginsdk.TypeString, + Optional: true, + }, + "offline_connection_name": { + Type: pluginsdk.TypeString, + Optional: true, + }, + + "online_connection_name": { + Type: pluginsdk.TypeString, + Optional: true, + }, + }, + }, + }, + "primary_user_assigned_identity": { Type: pluginsdk.TypeString, Optional: true, @@ -250,6 +284,7 @@ func resourceMachineLearningWorkspaceCreateOrUpdate(d *pluginsdk.ResourceData, m Name: d.Get("sku_name").(string), Tier: pointer.To(workspaces.SkuTier(d.Get("sku_name").(string))), }, + Kind: utils.String(d.Get("kind").(string)), Identity: expandedIdentity, Properties: &workspaces.WorkspaceProperties{ @@ -290,6 +325,18 @@ func resourceMachineLearningWorkspaceCreateOrUpdate(d *pluginsdk.ResourceData, m workspace.Properties.PrimaryUserAssignedIdentity = pointer.To(v.(string)) } + featureStore := expandMachineLearningWorkspaceFeatureStore(d.Get("feature_store").([]interface{})) + if strings.EqualFold(*workspace.Kind, "Default") { + if featureStore != nil { + return fmt.Errorf("`feature_store` can only be set when `kind` is `FeatureStore`") + } + } else { + if featureStore == nil { + return fmt.Errorf("`feature_store` can not be empty when `kind` is `FeatureStore`") + } + workspace.Properties.FeatureStoreSettings = featureStore + } + future, err := client.CreateOrUpdate(ctx, id, workspace) if err != nil { return fmt.Errorf("creating/updating %s: %+v", id, err) @@ -333,6 +380,8 @@ func resourceMachineLearningWorkspaceRead(d *pluginsdk.ResourceData, meta interf d.Set("sku_name", sku.Name) } + d.Set("kind", resp.Model.Kind) + if props := resp.Model.Properties; props != nil { d.Set("application_insights_id", props.ApplicationInsights) d.Set("storage_account_id", props.StorageAccount) @@ -367,6 +416,11 @@ func resourceMachineLearningWorkspaceRead(d *pluginsdk.ResourceData, meta interf return fmt.Errorf("setting `identity`: %+v", err) } + featureStoreSettings := flattenMachineLearningWorkspaceFeatureStore(resp.Model.Properties.FeatureStoreSettings) + if err := d.Set("feature_store", featureStoreSettings); err != nil { + return fmt.Errorf("setting `feature_store`: %+v", err) + } + flattenedEncryption, err := flattenMachineLearningWorkspaceEncryption(resp.Model.Properties.Encryption) if err != nil { return fmt.Errorf("flattening `encryption`: %+v", err) @@ -511,3 +565,56 @@ func flattenMachineLearningWorkspaceEncryption(input *workspaces.EncryptionPrope }, }, nil } + +func expandMachineLearningWorkspaceFeatureStore(input []interface{}) *workspaces.FeatureStoreSettings { + if len(input) == 0 || input[0] == nil { + return nil + } + + raw := input[0].(map[string]interface{}) + out := workspaces.FeatureStoreSettings{} + + if raw["computer_spark_runtime_version"].(string) != "" { + out.ComputeRuntime = &workspaces.ComputeRuntimeDto{ + SparkRuntimeVersion: utils.String(raw["computer_spark_runtime_version"].(string)), + } + } + + if raw["offline_connection_name"].(string) != "" { + out.OfflineStoreConnectionName = utils.String(raw["offline_connection_name"].(string)) + } + + if raw["online_connection_name"].(string) != "" { + out.OnlineStoreConnectionName = utils.String(raw["online_connection_name"].(string)) + } + return &out +} + +func flattenMachineLearningWorkspaceFeatureStore(input *workspaces.FeatureStoreSettings) *[]interface{} { + if input == nil { + return &[]interface{}{} + } + + computerSparkRunTimeVersion := "" + offlineConnectionName := "" + onlineConnectionName := "" + + if input.ComputeRuntime != nil && input.ComputeRuntime.SparkRuntimeVersion != nil { + computerSparkRunTimeVersion = *input.ComputeRuntime.SparkRuntimeVersion + } + if input.OfflineStoreConnectionName != nil { + offlineConnectionName = *input.OfflineStoreConnectionName + } + + if input.OnlineStoreConnectionName != nil { + onlineConnectionName = *input.OnlineStoreConnectionName + } + + return &[]interface{}{ + map[string]interface{}{ + "computer_spark_runtime_version": computerSparkRunTimeVersion, + "offline_connection_name": offlineConnectionName, + "online_connection_name": onlineConnectionName, + }, + } +} diff --git a/internal/services/machinelearning/machine_learning_workspace_resource_test.go b/internal/services/machinelearning/machine_learning_workspace_resource_test.go index 895b0d6f9654..1833d315b32b 100644 --- a/internal/services/machinelearning/machine_learning_workspace_resource_test.go +++ b/internal/services/machinelearning/machine_learning_workspace_resource_test.go @@ -262,6 +262,78 @@ func TestAccMachineLearningWorkspace_systemAssignedAndCustomManagedKey(t *testin }) } +func TestAccMachineLearningWorkspace_featureStore(t *testing.T) { + data := acceptance.BuildTestData(t, "azurerm_machine_learning_workspace", "test") + r := WorkspaceResource{} + + data.ResourceTest(t, r, []acceptance.TestStep{ + { + Config: r.featureStore(data), + Check: acceptance.ComposeTestCheckFunc( + check.That(data.ResourceName).ExistsInAzure(r), + check.That(data.ResourceName).Key("kind").Exists(), + ), + }, + data.ImportStep(), + }) +} + +func TestAccMachineLearningWorkspace_featureStoreUpdate(t *testing.T) { + data := acceptance.BuildTestData(t, "azurerm_machine_learning_workspace", "test") + r := WorkspaceResource{} + + data.ResourceTest(t, r, []acceptance.TestStep{ + { + Config: r.featureStore(data), + Check: acceptance.ComposeTestCheckFunc( + check.That(data.ResourceName).ExistsInAzure(r), + check.That(data.ResourceName).Key("kind").Exists(), + ), + }, + data.ImportStep(), + { + Config: r.featureStoreUpdate(data), + Check: acceptance.ComposeTestCheckFunc( + check.That(data.ResourceName).ExistsInAzure(r), + check.That(data.ResourceName).Key("kind").Exists(), + ), + }, + data.ImportStep(), + }) +} + +func TestAccMachineLearningWorkspace_kindUpdate(t *testing.T) { + data := acceptance.BuildTestData(t, "azurerm_machine_learning_workspace", "test") + r := WorkspaceResource{} + + data.ResourceTest(t, r, []acceptance.TestStep{ + { + Config: r.featureStore(data), + Check: acceptance.ComposeTestCheckFunc( + check.That(data.ResourceName).ExistsInAzure(r), + check.That(data.ResourceName).Key("kind").Exists(), + ), + }, + data.ImportStep(), + { + Config: r.basic(data), + Check: acceptance.ComposeTestCheckFunc( + check.That(data.ResourceName).ExistsInAzure(r), + check.That(data.ResourceName).Key("kind").Exists(), + ), + }, + data.ImportStep(), + { + Config: r.featureStore(data), + Check: acceptance.ComposeTestCheckFunc( + check.That(data.ResourceName).ExistsInAzure(r), + check.That(data.ResourceName).Key("kind").Exists(), + ), + }, + data.ImportStep(), + }) +} + func (r WorkspaceResource) Exists(ctx context.Context, client *clients.Client, state *pluginsdk.InstanceState) (*bool, error) { workspacesClient := client.MachineLearning.Workspaces id, err := workspaces.ParseWorkspaceID(state.ID) @@ -851,3 +923,59 @@ resource "azurerm_machine_learning_workspace" "test" { } `, r.template(data), data.RandomInteger) } + +func (r WorkspaceResource) featureStore(data acceptance.TestData) string { + template := r.template(data) + return fmt.Sprintf(` +%s + +resource "azurerm_machine_learning_workspace" "test" { + name = "acctest-MLW-%d" + location = azurerm_resource_group.test.location + resource_group_name = azurerm_resource_group.test.name + application_insights_id = azurerm_application_insights.test.id + key_vault_id = azurerm_key_vault.test.id + storage_account_id = azurerm_storage_account.test.id + + kind = "FeatureStore" + + feature_store { + computer_spark_runtime_version = "3.1" + offline_connection_name = "offlineStoreConnectionName" + online_connection_name = "onlineStoreConnectionName" + } + + identity { + type = "SystemAssigned" + } +} +`, template, data.RandomInteger) +} + +func (r WorkspaceResource) featureStoreUpdate(data acceptance.TestData) string { + template := r.template(data) + return fmt.Sprintf(` +%s + +resource "azurerm_machine_learning_workspace" "test" { + name = "acctest-MLW-%d" + location = azurerm_resource_group.test.location + resource_group_name = azurerm_resource_group.test.name + application_insights_id = azurerm_application_insights.test.id + key_vault_id = azurerm_key_vault.test.id + storage_account_id = azurerm_storage_account.test.id + + kind = "FeatureStore" + + feature_store { + computer_spark_runtime_version = "3.5" + offline_connection_name = "offlineStoreConnectionNameUpdate" + online_connection_name = "onlineStoreConnectionNameUpdate" + } + + identity { + type = "SystemAssigned" + } +} +`, template, data.RandomInteger) +} diff --git a/website/docs/r/machine_learning_workspace.html.markdown b/website/docs/r/machine_learning_workspace.html.markdown index 10b442bd4709..3fae221ffeb7 100644 --- a/website/docs/r/machine_learning_workspace.html.markdown +++ b/website/docs/r/machine_learning_workspace.html.markdown @@ -358,6 +358,8 @@ The following arguments are supported: * `identity` - (Required) An `identity` block as defined below. +* `kind` - (Optional) The type of the Workspace. Possible values are `Default`, `FeatureStore`. Defaults to `Default` + * `container_registry_id` - (Optional) The ID of the container registry associated with this Machine Learning Workspace. Changing this forces a new resource to be created. -> **NOTE:** The `admin_enabled` should be `true` in order to associate the Container Registry to this Machine Learning Workspace. @@ -374,6 +376,8 @@ The following arguments are supported: * `encryption` - (Optional) An `encryption` block as defined below. Changing this forces a new resource to be created. +* `feature_store` - (Optional) A `feature_store` block as defined below. + * `friendly_name` - (Optional) Display name for this Machine Learning Workspace. * `high_business_impact` - (Optional) Flag to signal High Business Impact (HBI) data in the workspace and reduce diagnostic data collected by the service. Changing this forces a new resource to be created. @@ -408,6 +412,19 @@ An `encryption` block supports the following: ~> **Note:** `user_assigned_identity_id` must set when`identity.type` is `UserAssigned` or service won't be able to find the assigned permissions. +--- + +An `feature_store` block supports the following: + +* `computer_spark_runtime_version` - (Optional) The version of Spark runtime. + +* `offline_connection_name` - (Optional) The name of offline store connection. + +* `online_connection_name` - (Optional) The name of online store connection. + +~> **Note:** `feature_store` must be set when`kind` is `FeatureStore` + + ## Attributes Reference In addition to the Arguments listed above - the following Attributes are exported: From 17d6d0fdff73fa5f40c1314a2d9bf16092e0e692 Mon Sep 17 00:00:00 2001 From: xuzhang3 Date: Wed, 31 Jan 2024 11:07:08 +0800 Subject: [PATCH 2/2] format --- .../machine_learning_workspace_resource_test.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/internal/services/machinelearning/machine_learning_workspace_resource_test.go b/internal/services/machinelearning/machine_learning_workspace_resource_test.go index 1833d315b32b..53feebf9b992 100644 --- a/internal/services/machinelearning/machine_learning_workspace_resource_test.go +++ b/internal/services/machinelearning/machine_learning_workspace_resource_test.go @@ -941,8 +941,8 @@ resource "azurerm_machine_learning_workspace" "test" { feature_store { computer_spark_runtime_version = "3.1" - offline_connection_name = "offlineStoreConnectionName" - online_connection_name = "onlineStoreConnectionName" + offline_connection_name = "offlineStoreConnectionName" + online_connection_name = "onlineStoreConnectionName" } identity { @@ -969,8 +969,8 @@ resource "azurerm_machine_learning_workspace" "test" { feature_store { computer_spark_runtime_version = "3.5" - offline_connection_name = "offlineStoreConnectionNameUpdate" - online_connection_name = "onlineStoreConnectionNameUpdate" + offline_connection_name = "offlineStoreConnectionNameUpdate" + online_connection_name = "onlineStoreConnectionNameUpdate" } identity {