diff --git a/.nextchanges/bundles/model-serving-telemetry-config-converges.md b/.nextchanges/bundles/model-serving-telemetry-config-converges.md new file mode 100644 index 00000000000..e06f76c7610 --- /dev/null +++ b/.nextchanges/bundles/model-serving-telemetry-config-converges.md @@ -0,0 +1 @@ +direct: Fixed model serving `telemetry_config` drift and applied planned telemetry updates. Unsupported endpoint types now fail when telemetry is applied; create may still succeed because it drops the field ([#6106](https://github.com/databricks/cli/pull/6106)). diff --git a/acceptance/bundle/invariant/configs/model_serving_endpoint_telemetry.yml.tmpl b/acceptance/bundle/invariant/configs/model_serving_endpoint_telemetry.yml.tmpl new file mode 100644 index 00000000000..d8c0147e937 --- /dev/null +++ b/acceptance/bundle/invariant/configs/model_serving_endpoint_telemetry.yml.tmpl @@ -0,0 +1,20 @@ +bundle: + name: test-bundle-$UNIQUE_NAME + +resources: + model_serving_endpoints: + foo: + name: test-endpoint-$UNIQUE_NAME + config: + served_entities: + - name: prod + entity_name: main.default.test_model + entity_version: "1" + workload_size: Small + scale_to_zero_enabled: true + telemetry_config: + table_names: + logs_table: main.default.test_logs + traces_table: main.default.test_traces + inference_table_config: + sampling_fraction: 0.5 diff --git a/acceptance/bundle/invariant/continue_293/out.test.toml b/acceptance/bundle/invariant/continue_293/out.test.toml index ca367650f71..b06760dbfae 100644 --- a/acceptance/bundle/invariant/continue_293/out.test.toml +++ b/acceptance/bundle/invariant/continue_293/out.test.toml @@ -23,6 +23,7 @@ EnvMatrix.INPUT_CONFIG = [ "job_with_task.yml.tmpl", "model.yml.tmpl", "model_serving_endpoint.yml.tmpl", + "model_serving_endpoint_telemetry.yml.tmpl", "pipeline.yml.tmpl", "pipeline_allow_duplicate_names.yml.tmpl", "pipeline_apply_policy_default_values.yml.tmpl", diff --git a/acceptance/bundle/invariant/delete_idempotent/out.test.toml b/acceptance/bundle/invariant/delete_idempotent/out.test.toml index 9b8f7ea4435..04295a71d21 100644 --- a/acceptance/bundle/invariant/delete_idempotent/out.test.toml +++ b/acceptance/bundle/invariant/delete_idempotent/out.test.toml @@ -30,6 +30,7 @@ EnvMatrix.INPUT_CONFIG = [ "model.yml.tmpl", "model_with_permissions.yml.tmpl", "model_serving_endpoint.yml.tmpl", + "model_serving_endpoint_telemetry.yml.tmpl", "pipeline.yml.tmpl", "pipeline_allow_duplicate_names.yml.tmpl", "pipeline_apply_policy_default_values.yml.tmpl", diff --git a/acceptance/bundle/invariant/destroy_idempotent/out.test.toml b/acceptance/bundle/invariant/destroy_idempotent/out.test.toml index 9b8f7ea4435..04295a71d21 100644 --- a/acceptance/bundle/invariant/destroy_idempotent/out.test.toml +++ b/acceptance/bundle/invariant/destroy_idempotent/out.test.toml @@ -30,6 +30,7 @@ EnvMatrix.INPUT_CONFIG = [ "model.yml.tmpl", "model_with_permissions.yml.tmpl", "model_serving_endpoint.yml.tmpl", + "model_serving_endpoint_telemetry.yml.tmpl", "pipeline.yml.tmpl", "pipeline_allow_duplicate_names.yml.tmpl", "pipeline_apply_policy_default_values.yml.tmpl", diff --git a/acceptance/bundle/invariant/migrate/test.toml b/acceptance/bundle/invariant/migrate/test.toml index bb2337b32aa..9d79e11bef9 100644 --- a/acceptance/bundle/invariant/migrate/test.toml +++ b/acceptance/bundle/invariant/migrate/test.toml @@ -37,3 +37,6 @@ EnvMatrixExclude.no_pydabs_1000_tasks = ["INPUT_CONFIG=job_pydabs_1000_tasks.yml # migrate deploys via Terraform first, and the TF provider rejects an uppercase # volume schema_name ("inconsistent final plan"). Covered by no_drift on direct. EnvMatrixExclude.no_volume_uppercase = ["INPUT_CONFIG=volume_uppercase_name.yml.tmpl"] + +# Terraform types sampling_fraction as an integer and truncates 0.5; covered by no_drift. +EnvMatrixExclude.no_model_serving_endpoint_telemetry = ["INPUT_CONFIG=model_serving_endpoint_telemetry.yml.tmpl"] diff --git a/acceptance/bundle/invariant/no_drift/out.test.toml b/acceptance/bundle/invariant/no_drift/out.test.toml index 9b8f7ea4435..04295a71d21 100644 --- a/acceptance/bundle/invariant/no_drift/out.test.toml +++ b/acceptance/bundle/invariant/no_drift/out.test.toml @@ -30,6 +30,7 @@ EnvMatrix.INPUT_CONFIG = [ "model.yml.tmpl", "model_with_permissions.yml.tmpl", "model_serving_endpoint.yml.tmpl", + "model_serving_endpoint_telemetry.yml.tmpl", "pipeline.yml.tmpl", "pipeline_allow_duplicate_names.yml.tmpl", "pipeline_apply_policy_default_values.yml.tmpl", diff --git a/acceptance/bundle/invariant/test.toml b/acceptance/bundle/invariant/test.toml index 32a0808dbac..4d0777fdea3 100644 --- a/acceptance/bundle/invariant/test.toml +++ b/acceptance/bundle/invariant/test.toml @@ -48,6 +48,7 @@ EnvMatrix.INPUT_CONFIG = [ "model.yml.tmpl", "model_with_permissions.yml.tmpl", "model_serving_endpoint.yml.tmpl", + "model_serving_endpoint_telemetry.yml.tmpl", "pipeline.yml.tmpl", "pipeline_allow_duplicate_names.yml.tmpl", "pipeline_apply_policy_default_values.yml.tmpl", @@ -116,6 +117,9 @@ no_external_volume_on_cloud = ["CONFIG_Cloud=true", "INPUT_CONFIG=volume_externa # it here to keep the cloud invariant runs from timing out. Still exercised locally. no_vector_search_index_on_cloud = ["CONFIG_Cloud=true", "INPUT_CONFIG=vector_search_index.yml.tmpl"] +# Telemetry requires a real model, which cloud invariant tests do not provision. +no_model_serving_endpoint_telemetry_on_cloud = ["CONFIG_Cloud=true", "INPUT_CONFIG=model_serving_endpoint_telemetry.yml.tmpl"] + # Fake SQL endpoint for local tests [[Server]] Pattern = "POST /api/2.0/sql/statements/" diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/recreated_same_name/out.test.toml b/acceptance/bundle/resources/model_serving_endpoints/drift/recreated_same_name/out.test.toml index fe4076cdf9b..8c52d40aa2d 100644 --- a/acceptance/bundle/resources/model_serving_endpoints/drift/recreated_same_name/out.test.toml +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/recreated_same_name/out.test.toml @@ -1,4 +1,3 @@ -Local = true Cloud = true RequiresUnityCatalog = true EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/recreated_same_name/test.toml b/acceptance/bundle/resources/model_serving_endpoints/drift/recreated_same_name/test.toml index 83e36142b53..fbd03b35ba9 100644 --- a/acceptance/bundle/resources/model_serving_endpoints/drift/recreated_same_name/test.toml +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/recreated_same_name/test.toml @@ -1,6 +1,5 @@ Badness = "After deleting and recreating a model serving endpoint remotely with the same name but a different endpoint_id, bundle plan/deploy ends up with a permanent update on permissions because the V1 permissions API does not delete ACLs immediately when the parent is gone." -Local = true Cloud = true RequiresUnityCatalog = true RecordRequests = false diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/databricks.yml b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/databricks.yml new file mode 100644 index 00000000000..1291612ba80 --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/databricks.yml @@ -0,0 +1,14 @@ +bundle: + name: test-bundle + +resources: + model_serving_endpoints: + endpoint1: + name: test-endpoint + config: + served_entities: + - name: prod + entity_name: main.default.test_model + entity_version: "1" + workload_size: Small + scale_to_zero_enabled: true diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/out.test.toml b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/out.test.toml new file mode 100644 index 00000000000..0938e678987 --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/output.txt b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/output.txt new file mode 100644 index 00000000000..f6924643628 --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/output.txt @@ -0,0 +1,36 @@ + +=== Deploy without telemetry_config +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +=== Configure telemetry out of band +=== Remote-only telemetry_config is not drift +>>> [CLI] bundle plan --output json +{ + "telemetry_config": { + "action": "skip", + "reason": "backend_default", + "remote": { + "telemetry_profile_id": "[UUID]" + } + } +} + +=== Deploy does not remove it +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> [CLI] api get /api/2.0/serving-endpoints/test-endpoint +{ + "telemetry_profile_id": "[UUID]" +} + +=== Only the out-of-band telemetry PATCH was sent +>>> print_requests.py --method PATCH --oneline //serving-endpoints +{"method": "PATCH", "path": "/api/2.0/serving-endpoints/test-endpoint/telemetry-config", "body": {"telemetry_config": {"table_names": {"logs_table": "main.default.other_logs"}}}} diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/script b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/script new file mode 100644 index 00000000000..5e30339c21f --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/script @@ -0,0 +1,17 @@ +echo "*" > .gitignore + +title "Deploy without telemetry_config" +trace $CLI bundle deploy + +title "Configure telemetry out of band" +MSYS_NO_PATHCONV=1 $CLI api patch "/api/2.0/serving-endpoints/test-endpoint/telemetry-config" --json '{"telemetry_config": {"table_names": {"logs_table": "main.default.other_logs"}}}' > /dev/null + +title "Remote-only telemetry_config is not drift" +trace $CLI bundle plan --output json | jq '.plan[].changes | with_entries(select(.key | startswith("telemetry_config")))' + +title "Deploy does not remove it" +trace $CLI bundle deploy +MSYS_NO_PATHCONV=1 trace $CLI api get "/api/2.0/serving-endpoints/test-endpoint" | jq .telemetry_config + +title "Only the out-of-band telemetry PATCH was sent" +trace print_requests.py --method PATCH --oneline //serving-endpoints diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/test.toml b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/test.toml new file mode 100644 index 00000000000..ec0e9c5ca86 --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_unmanaged/test.toml @@ -0,0 +1,4 @@ +RecordRequests = true + +# Telemetry classification is direct-engine-specific. +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/databricks.yml b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/databricks.yml new file mode 100644 index 00000000000..889f5172c6a --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/databricks.yml @@ -0,0 +1,19 @@ +bundle: + name: test-bundle + +resources: + model_serving_endpoints: + endpoint1: + name: test-endpoint + config: + served_entities: + - name: prod + entity_name: main.default.test_model + entity_version: "1" + workload_size: Small + scale_to_zero_enabled: true + telemetry_config: # TELEMETRY + table_names: # TELEMETRY + logs_table: main.default.test_logs # TELEMETRY + inference_table_config: # TELEMETRY + sampling_fraction: 0.5 # TELEMETRY diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/out.test.toml b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/out.test.toml new file mode 100644 index 00000000000..0938e678987 --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/output.txt b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/output.txt new file mode 100644 index 00000000000..c0f25ba86f1 --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/output.txt @@ -0,0 +1,45 @@ + +=== Initial deployment +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +=== Update config and telemetry_config +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +=== Plan converges +>>> [CLI] bundle plan +Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged + +=== Request order: config PUT, poll, telemetry PATCH +>>> print_requests.py --get --oneline //serving-endpoints +1 {"method": "POST", "path": "/api/2.0/serving-endpoints", "body": {"config": {"served_entities": [{"entity_name": "main.default.test_model", "entity_version": "1", "name": "prod", "scale_to_zero_enabled": true, "workload_size": "Small"}]}, "name": "test-endpoint", "telemetry_config": {"inference_table_config": {"sampling_fraction": 0.5}, "table_names": {"logs_table": "main.default.test_logs"}}}} +2 {"method": "GET", "path": "/api/2.0/serving-endpoints/test-endpoint"} +1 {"method": "PUT", "path": "/api/2.0/serving-endpoints/test-endpoint/config", "body": {"served_entities": [{"entity_name": "main.default.test_model", "entity_version": "1", "name": "prod", "scale_to_zero_enabled": true, "workload_size": "Medium"}]}} +2 {"method": "GET", "path": "/api/2.0/serving-endpoints/test-endpoint"} +1 {"method": "PATCH", "path": "/api/2.0/serving-endpoints/test-endpoint/telemetry-config", "body": {"telemetry_config": {"inference_table_config": {"sampling_fraction": 0.9}, "table_names": {"logs_table": "main.default.test_logs"}}}} +2 {"method": "GET", "path": "/api/2.0/serving-endpoints/test-endpoint"} + +=== Remove telemetry_config +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> [CLI] api get /api/2.0/serving-endpoints/test-endpoint +null + +=== Plan converges after removal +>>> [CLI] bundle plan +Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged + +=== Removal sends an empty telemetry PATCH +>>> print_requests.py --method PATCH --oneline //serving-endpoints +{"method": "PATCH", "path": "/api/2.0/serving-endpoints/test-endpoint/telemetry-config", "body": {}} diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/script b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/script new file mode 100644 index 00000000000..3df97d41e8c --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/script @@ -0,0 +1,26 @@ +echo "*" > .gitignore + +title "Initial deployment" +trace $CLI bundle deploy + +title "Update config and telemetry_config" +update_file.py databricks.yml "workload_size: Small" "workload_size: Medium" +update_file.py databricks.yml "sampling_fraction: 0.5" "sampling_fraction: 0.9" +trace $CLI bundle deploy + +title "Plan converges" +trace $CLI bundle plan | contains.py "Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged" + +title "Request order: config PUT, poll, telemetry PATCH" +trace print_requests.py --get --oneline //serving-endpoints | uniq -c | sed 's/^ *//' + +title "Remove telemetry_config" +grep -v TELEMETRY databricks.yml > updated.yml && mv updated.yml databricks.yml +trace $CLI bundle deploy +MSYS_NO_PATHCONV=1 trace $CLI api get "/api/2.0/serving-endpoints/test-endpoint" | jq .telemetry_config + +title "Plan converges after removal" +trace $CLI bundle plan | contains.py "Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged" + +title "Removal sends an empty telemetry PATCH" +trace print_requests.py --method PATCH --oneline //serving-endpoints diff --git a/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/test.toml b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/test.toml new file mode 100644 index 00000000000..49a9063424b --- /dev/null +++ b/acceptance/bundle/resources/model_serving_endpoints/drift/telemetry_config_with_config_update/test.toml @@ -0,0 +1,4 @@ +RecordRequests = true + +# Telemetry updates are direct-engine-specific. +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/bundle/direct/dresources/model_serving_endpoint.go b/bundle/direct/dresources/model_serving_endpoint.go index e8a1917c2f2..c0cc8f82598 100644 --- a/bundle/direct/dresources/model_serving_endpoint.go +++ b/bundle/direct/dresources/model_serving_endpoint.go @@ -18,6 +18,7 @@ var ( pathAiGateway = structpath.MustParsePath("ai_gateway") pathConfig = structpath.MustParsePath("config") pathEmailNotifications = structpath.MustParsePath("email_notifications") + pathTelemetryConfig = structpath.MustParsePath("telemetry_config") ) type ResourceModelServingEndpoint struct { @@ -246,6 +247,19 @@ func (r *ResourceModelServingEndpoint) updateNotifications(ctx context.Context, return nil } +// updateTelemetryConfig removes telemetry when telemetryConfig is nil. +func (r *ResourceModelServingEndpoint) updateTelemetryConfig(ctx context.Context, id string, telemetryConfig *serving.TelemetryConfig) error { + req := serving.PatchTelemetryConfigRequest{ + Name: id, + TelemetryConfig: telemetryConfig, + } + _, err := r.client.ServingEndpoints.PatchTelemetryConfig(ctx, req) + if err != nil { + return fmt.Errorf("failed to update telemetry config: %w", err) + } + return nil +} + func diffTags(currentTags, desiredTags []serving.EndpointTag) (addTags []serving.EndpointTag, deleteTags []string) { addTags = make([]serving.EndpointTag, 0) @@ -343,6 +357,19 @@ func (r *ResourceModelServingEndpoint) DoUpdate(ctx context.Context, id string, } } + if entry.Changes.HasChange(pathTelemetryConfig) { + // The telemetry API rejects endpoints with an update in progress. + _, err = r.waitForEndpointReady(ctx, id) + if err != nil { + return nil, err + } + + err = r.updateTelemetryConfig(ctx, id, config.TelemetryConfig) + if err != nil { + return nil, err + } + } + return nil, nil } diff --git a/bundle/direct/dresources/resources.yml b/bundle/direct/dresources/resources.yml index 4099c536281..0c9c9bf5319 100644 --- a/bundle/direct/dresources/resources.yml +++ b/bundle/direct/dresources/resources.yml @@ -285,6 +285,9 @@ resources: # Accepted on write but not returned by GET. - field: config.served_entities[*].burst_scaling_enabled reason: input_only + # table_names creates a profile and is not returned; inference_table_config round-trips. + - field: telemetry_config.table_names + reason: input_only # Write-only secrets: the backend stores them and returns the reference field, not the plaintext. - field: config.served_entities[*].external_model.ai21labs_config.ai21labs_api_key_plaintext reason: input_only @@ -316,6 +319,9 @@ resources: - field: rate_limits reason: not_implemented backend_defaults: + # Remote-only telemetry is reported as a change at this parent path. + - field: telemetry_config + # https://github.com/databricks/terraform-provider-databricks/blob/4eba541abe1a9f50993ea7b9dd83874207e224a1/serving/resource_model_serving.go#L383 # common.CustomizeSchemaPath(m, "config", "served_entities", "name").SetComputed() - field: config.served_entities[*].name diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 2ef3ad0c7a2..4a5fc912350 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -957,7 +957,7 @@ func AddDefaultHandlers(server *Server) { // Serving Endpoints: server.Handle("GET", "/api/2.0/serving-endpoints/{name}", func(req Request) any { - return MapGet(req.Workspace, req.Workspace.ServingEndpoints, req.Vars["name"]) + return req.Workspace.ServingEndpointGet(req.Vars["name"]) }) server.Handle("POST", "/api/2.0/serving-endpoints", func(req Request) any { @@ -984,6 +984,10 @@ func AddDefaultHandlers(server *Server) { return req.Workspace.ServingEndpointPatchTags(req, req.Vars["name"]) }) + server.Handle("PATCH", "/api/2.0/serving-endpoints/{name}/telemetry-config", func(req Request) any { + return req.Workspace.ServingEndpointPatchTelemetryConfig(req, req.Vars["name"]) + }) + // Vector Search Endpoints: server.Handle("POST", "/api/2.0/vector-search/endpoints", func(req Request) any { diff --git a/libs/testserver/serving_endpoints.go b/libs/testserver/serving_endpoints.go index 41a6b9766b5..efa10d47463 100644 --- a/libs/testserver/serving_endpoints.go +++ b/libs/testserver/serving_endpoints.go @@ -35,6 +35,45 @@ func servedEntitiesInputToOutput(input []serving.ServedEntityInput) []serving.Se return entities } +// applyTelemetryConfig consumes table_names and discards configs that identify no profile. +func applyTelemetryConfig(previous, config *serving.TelemetryConfig) *serving.TelemetryConfig { + if config == nil { + return nil + } + if config.TableNames == nil && config.TelemetryProfileId == "" { + return previous + } + + applied := serving.TelemetryConfig{TelemetryProfileId: config.TelemetryProfileId} + if applied.TelemetryProfileId == "" { + // Do not reuse previous: table_names provisions a new profile. + applied.TelemetryProfileId = nextUUID() + } + if config.InferenceTableConfig != nil { + inferenceTable := *config.InferenceTableConfig + // The backend names the payload table after the logs table it was given. + if config.TableNames != nil { + inferenceTable.Name = config.TableNames.LogsTable + "_payload" + } + applied.InferenceTableConfig = &inferenceTable + } + + return &applied +} + +// telemetrySupported returns the unsupported endpoint type when telemetry cannot be applied. +func telemetrySupported(endpoint serving.ServingEndpointDetailed) (string, bool) { + if endpoint.Config == nil || len(endpoint.Config.ServedEntities) == 0 { + return "NO_CONFIG", false + } + for _, entity := range endpoint.Config.ServedEntities { + if entity.ExternalModel == nil && entity.EntityName != "" { + return "", true + } + } + return "EXTERNAL_MODELS", false +} + // clearExternalModelSecrets mirrors the backend, which persists the *_plaintext // API keys as secrets and never returns them on GET. func clearExternalModelSecrets(em *serving.ExternalModel) *serving.ExternalModel { @@ -225,6 +264,7 @@ func (s *FakeWorkspace) ServingEndpointCreate(req Request) Response { PermissionLevel: serving.ServingEndpointDetailedPermissionLevelCanManage, RouteOptimized: createReq.RouteOptimized, Tags: createReq.Tags, + TelemetryConfig: applyTelemetryConfig(nil, createReq.TelemetryConfig), State: &serving.EndpointState{ ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating, Ready: serving.EndpointStateReadyNotReady, @@ -234,6 +274,11 @@ func (s *FakeWorkspace) ServingEndpointCreate(req Request) Response { ForceSendFields: append(createReq.ForceSendFields, "PermissionLevel", "RouteOptimized", "Description"), } + // Create drops unsupported telemetry, while the telemetry API rejects it. + if _, ok := telemetrySupported(endpoint); !ok { + endpoint.TelemetryConfig = nil + } + s.ServingEndpoints[createReq.Name] = endpoint return Response{ @@ -241,6 +286,36 @@ func (s *FakeWorkspace) ServingEndpointCreate(req Request) Response { } } +// ServingEndpointGet reports an in-progress update once before settling it. +func (s *FakeWorkspace) ServingEndpointGet(name string) Response { + defer s.LockUnlock()() + + endpoint, exists := s.ServingEndpoints[name] + if !exists { + return Response{ + StatusCode: 404, + Body: map[string]string{"message": fmt.Sprintf("Resource %T not found: %v", endpoint, name)}, + } + } + + if endpointUpdating(endpoint) { + // This response stays IN_PROGRESS; settle the stored copy for the next read. + settled := endpoint + settled.State = &serving.EndpointState{ + ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating, + Ready: endpoint.State.Ready, + } + s.ServingEndpoints[name] = settled + } + + return Response{Body: endpoint} +} + +// endpointUpdating reports whether a config update is in progress. +func endpointUpdating(endpoint serving.ServingEndpointDetailed) bool { + return endpoint.State != nil && endpoint.State.ConfigUpdate == serving.EndpointStateConfigUpdateInProgress +} + func (s *FakeWorkspace) ServingEndpointUpdate(req Request, name string) Response { defer s.LockUnlock()() @@ -288,8 +363,9 @@ func (s *FakeWorkspace) ServingEndpointUpdate(req Request, name string) Response endpoint.Config = config endpoint.LastUpdatedTimestamp = nowMilli() + // Keep the update in progress until GET observes it. endpoint.State = &serving.EndpointState{ - ConfigUpdate: serving.EndpointStateConfigUpdateNotUpdating, + ConfigUpdate: serving.EndpointStateConfigUpdateInProgress, Ready: serving.EndpointStateReadyNotReady, } @@ -369,6 +445,58 @@ func (s *FakeWorkspace) ServingEndpointUpdateNotifications(req Request, name str } } +// ServingEndpointPatchTelemetryConfig applies telemetry after validating endpoint state. +func (s *FakeWorkspace) ServingEndpointPatchTelemetryConfig(req Request, name string) Response { + defer s.LockUnlock()() + + var patchReq serving.PatchTelemetryConfigRequest + err := json.Unmarshal(req.Body, &patchReq) + if err != nil { + return Response{ + Body: fmt.Sprintf("cannot unmarshal request body: %s", err), + StatusCode: 400, + } + } + + endpoint, exists := s.ServingEndpoints[name] + if !exists { + return Response{ + StatusCode: 404, + Body: map[string]string{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": fmt.Sprintf("Serving endpoint with name %s not found", name)}, + } + } + + if endpointType, ok := telemetrySupported(endpoint); !ok { + return Response{ + StatusCode: 400, + Body: map[string]string{ + "error_code": "INVALID_PARAMETER_VALUE", + "message": fmt.Sprintf("Telemetry configuration is not supported for endpoint type '%s'. This API only supports endpoints with custom served models.", endpointType), + }, + } + } + + // The telemetry API returns 409 while another update is in progress. + if endpointUpdating(endpoint) { + return Response{ + StatusCode: 409, + Body: map[string]string{ + "error_code": "RESOURCE_CONFLICT", + "message": fmt.Sprintf("Endpoint %s is currently updating. Wait for the update to complete before updating its telemetry configuration.", name), + }, + } + } + + // An omitted telemetry_config removes the configuration from the endpoint. + endpoint.TelemetryConfig = applyTelemetryConfig(endpoint.TelemetryConfig, patchReq.TelemetryConfig) + endpoint.LastUpdatedTimestamp = nowMilli() + s.ServingEndpoints[name] = endpoint + + return Response{ + Body: endpoint, + } +} + func (s *FakeWorkspace) ServingEndpointPatchTags(req Request, name string) Response { defer s.LockUnlock()()