Fix DiffSuppressFunc for databricks_model_serving resource (#2813)

* fix model serving diff func * add test * fix test * fix logic * add comment for future
databricks · Oct 19, 2023 · 4087b40 · 4087b40
1 parent 6f70477
commit 4087b40
Show file tree

Hide file tree

Showing 2 changed files with 134 additions and 5 deletions.
diff --git a/serving/resource_model_serving.go b/serving/resource_model_serving.go
@@ -23,12 +23,10 @@ func ResourceModelServing() *schema.Resource {
 			common.MustSchemaPath(m, "config", "served_models", "name").Computed = true
 
 			common.MustSchemaPath(m, "config", "served_models", "workload_type").Default = "CPU"
-			// if GPU serving is not enabled, workload_type will always be empty
+			// TODO: `config.served_models.workload_type` should be a `Optional+Computed` field. Also consider this for other similar fields.
+			// In this scenario, if a workspace does not have GPU serving, specifying `workload_type` = 'CPU' will get empty response from API.
 			common.MustSchemaPath(m, "config", "served_models", "workload_type").DiffSuppressFunc = func(k, old, new string, d *schema.ResourceData) bool {
-				if old == "" || new == "" {
-					return true
-				}
-				return false
+				return old == "" && new == "CPU"
 			}
 			common.MustSchemaPath(m, "config", "traffic_config").Computed = true
 

diff --git a/serving/resource_model_serving_test.go b/serving/resource_model_serving_test.go
@@ -140,6 +140,137 @@ func TestModelServingCreate(t *testing.T) {
 	}.ApplyNoError(t)
 }
 
+func TestModelServingCreateGPU(t *testing.T) {
+	qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:   http.MethodPost,
+				Resource: "/api/2.0/serving-endpoints",
+				ExpectedRequest: serving.CreateServingEndpoint{
+					Name: "test-endpoint",
+					Config: serving.EndpointCoreConfigInput{
+						ServedModels: []serving.ServedModelInput{
+							{
+								Name:               "prod_model",
+								ModelName:          "ads1",
+								ModelVersion:       "2",
+								WorkloadSize:       "Small",
+								WorkloadType:       "GPU_MEDIUM",
+								ScaleToZeroEnabled: true,
+							},
+							{
+								Name:               "candidate_model",
+								ModelName:          "ads1",
+								ModelVersion:       "4",
+								WorkloadSize:       "Small",
+								WorkloadType:       "GPU_MEDIUM",
+								ScaleToZeroEnabled: false,
+							},
+						},
+						TrafficConfig: &serving.TrafficConfig{
+							Routes: []serving.Route{
+								{
+									ServedModelName:   "prod_model",
+									TrafficPercentage: 90,
+								},
+								{
+									ServedModelName:   "candidate_model",
+									TrafficPercentage: 10,
+								},
+							},
+						},
+					},
+				},
+				Response: serving.ServingEndpointDetailed{
+					Name: "test-endpoint",
+				},
+			},
+			{
+				Method:   http.MethodGet,
+				Resource: "/api/2.0/serving-endpoints/test-endpoint?",
+				Response: serving.ServingEndpointDetailed{
+					Name: "test-endpoint",
+					State: &serving.EndpointState{
+						ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating,
+					},
+				},
+			},
+			{
+				Method:   http.MethodGet,
+				Resource: "/api/2.0/serving-endpoints/test-endpoint?",
+				Response: serving.ServingEndpointDetailed{
+					Id:   "test-endpoint",
+					Name: "test-endpoint",
+					State: &serving.EndpointState{
+						ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating,
+					},
+					Config: &serving.EndpointCoreConfigOutput{
+						ServedModels: []serving.ServedModelOutput{
+							{
+								Name:               "prod_model",
+								ModelName:          "ads1",
+								ModelVersion:       "2",
+								ScaleToZeroEnabled: true,
+							},
+							{
+								Name:               "candidate_model",
+								ModelName:          "ads1",
+								ModelVersion:       "4",
+								ScaleToZeroEnabled: false,
+							},
+						},
+						TrafficConfig: &serving.TrafficConfig{
+							Routes: []serving.Route{
+								{
+									ServedModelName:   "prod_model",
+									TrafficPercentage: 90,
+								},
+								{
+									ServedModelName:   "candidate_model",
+									TrafficPercentage: 10,
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		Resource: ResourceModelServing(),
+		HCL: `
+			name = "test-endpoint"
+			config {
+				served_models {
+					name = "prod_model"
+					model_name = "ads1"
+					model_version = "2"
+					workload_size = "Small"
+					workload_type = "GPU_MEDIUM"
+					scale_to_zero_enabled = true
+				}
+				served_models {
+					name = "candidate_model"
+					model_name = "ads1"
+					model_version = "4"
+					workload_size = "Small"
+					workload_type = "GPU_MEDIUM"
+					scale_to_zero_enabled = false
+				}
+				traffic_config {
+					routes {
+						served_model_name = "prod_model"
+						traffic_percentage = 90
+					}
+					routes {
+						served_model_name = "candidate_model"
+						traffic_percentage = 10
+					}
+				}
+			}
+			`,
+		Create: true,
+	}.ApplyNoError(t)
+}
+
 func TestModelServingCreate_Error(t *testing.T) {
 	qa.ResourceFixture{
 		Fixtures: []qa.HTTPFixture{