diff --git a/.nextchanges/bundles/mlflow-experiment-tags.md b/.nextchanges/bundles/mlflow-experiment-tags.md new file mode 100644 index 00000000000..b6639a2e6da --- /dev/null +++ b/.nextchanges/bundles/mlflow-experiment-tags.md @@ -0,0 +1 @@ +* Fix bundle deployment to add or change configured MLflow experiment tags without recreating the experiment. Removing a tag now reports that deletion is unsupported. ([#7302](https://github.com/databricks/cli/pull/7302)) diff --git a/acceptance/bundle/resources/experiments/basic/out.plan.txt b/acceptance/bundle/resources/experiments/basic/out.plan.txt index c54c9d511c0..986d3cf68cd 100644 --- a/acceptance/bundle/resources/experiments/basic/out.plan.txt +++ b/acceptance/bundle/resources/experiments/basic/out.plan.txt @@ -1 +1,3 @@ -Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged +update experiments.my_experiment + +Plan: 0 to add, 1 to change, 0 to delete, 0 unchanged diff --git a/acceptance/bundle/resources/experiments/basic/output.txt b/acceptance/bundle/resources/experiments/basic/output.txt index be16a709a17..d1081fbf592 100644 --- a/acceptance/bundle/resources/experiments/basic/output.txt +++ b/acceptance/bundle/resources/experiments/basic/output.txt @@ -136,7 +136,7 @@ Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged ] } -=== adding a new tag should be a no-op +=== adding a new tag updates the existing experiment >>> envsubst @@ -154,24 +154,80 @@ Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged "value": "v1" }, { - "key": "mlflow.ownerId", - "value": "[USERID]" + "key": "k2", + "value": "v2" }, { "key": "mlflow.experiment.sourceName", "value": "/Users/[USERNAME]/new-name-[UNIQUE_NAME]" }, { - "key": "mlflow.ownerId", - "value": "[USERID]" + "key": "mlflow.experimentType", + "value": "MLFLOW_EXPERIMENT" }, { "key": "mlflow.ownerEmail", "value": "[USERNAME]" }, + { + "key": "mlflow.ownerId", + "value": "[USERID]" + }, + { + "key": "mlflow.ownerId", + "value": "[USERID]" + } + ] +} + +=== removing all configured tags is rejected + +>>> envsubst + +>>> musterr [CLI] bundle deploy -qq +Error: cannot update resources.experiments.my_experiment: updating id=[MY_EXPERIMENT_ID_3]: cannot remove experiment tag "k1": the MLflow API does not support deleting experiment tags + + +=== removing a configured tag is rejected + +>>> envsubst + +>>> musterr [CLI] bundle deploy -qq +Error: cannot update resources.experiments.my_experiment: updating id=[MY_EXPERIMENT_ID_3]: cannot remove experiment tag "k2": the MLflow API does not support deleting experiment tags + + +>>> [CLI] experiments get-experiment [MY_EXPERIMENT_ID_3] +{ + "name": "/Users/[USERNAME]/new-name-[UNIQUE_NAME]", + "artifact_location": "s3://new-[UNIQUE_NAME]", + "tags": [ + { + "key": "k1", + "value": "v1" + }, + { + "key": "k2", + "value": "v2" + }, + { + "key": "mlflow.experiment.sourceName", + "value": "/Users/[USERNAME]/new-name-[UNIQUE_NAME]" + }, { "key": "mlflow.experimentType", "value": "MLFLOW_EXPERIMENT" + }, + { + "key": "mlflow.ownerEmail", + "value": "[USERNAME]" + }, + { + "key": "mlflow.ownerId", + "value": "[USERID]" + }, + { + "key": "mlflow.ownerId", + "value": "[USERID]" } ] } diff --git a/acceptance/bundle/resources/experiments/basic/script b/acceptance/bundle/resources/experiments/basic/script index ab6ca3de78d..f4e3c93d33e 100755 --- a/acceptance/bundle/resources/experiments/basic/script +++ b/acceptance/bundle/resources/experiments/basic/script @@ -29,9 +29,22 @@ trace $CLI bundle deploy -qq experiment_id=$(read_id.py my_experiment) trace $CLI experiments get-experiment $experiment_id | jq '.experiment | {name, artifact_location, tags}' -title "adding a new tag should be a no-op" +title "adding a new tag updates the existing experiment" +existing_experiment_id=$experiment_id trace envsubst < templates/two_tag.tmpl > databricks.yml trace $CLI bundle plan > out.plan.txt trace $CLI bundle deploy -qq experiment_id=$(read_id.py my_experiment) -trace $CLI experiments get-experiment $experiment_id | jq '.experiment | {name, artifact_location, tags}' +test "$experiment_id" = "$existing_experiment_id" +trace $CLI experiments get-experiment $experiment_id | jq '.experiment | {name, artifact_location, tags: (.tags | sort_by(.key))}' + +title "removing all configured tags is rejected" +trace envsubst < templates/no_tag.tmpl > databricks.yml +trace musterr $CLI bundle deploy -qq +test "$(read_id.py my_experiment)" = "$existing_experiment_id" + +title "removing a configured tag is rejected" +trace envsubst < templates/one_tag.tmpl > databricks.yml +trace musterr $CLI bundle deploy -qq +test "$(read_id.py my_experiment)" = "$existing_experiment_id" +trace $CLI experiments get-experiment $experiment_id | jq '.experiment | {name, artifact_location, tags: (.tags | sort_by(.key))}' diff --git a/acceptance/bundle/resources/experiments/basic/templates/no_tag.tmpl b/acceptance/bundle/resources/experiments/basic/templates/no_tag.tmpl new file mode 100644 index 00000000000..4e6e89d5c2a --- /dev/null +++ b/acceptance/bundle/resources/experiments/basic/templates/no_tag.tmpl @@ -0,0 +1,8 @@ +bundle: + name: experiment-basic-$UNIQUE_NAME + +resources: + experiments: + my_experiment: + artifact_location: $EXPERIMENT_ARTIFACT_LOCATION + name: /Users/${CURRENT_USER_NAME}/${EXPERIMENT_NAME_PREFIX}-${UNIQUE_NAME} diff --git a/bundle/direct/dresources/configs/experiments.yml b/bundle/direct/dresources/configs/experiments.yml index 6a261514379..e7eb2da82f3 100644 --- a/bundle/direct/dresources/configs/experiments.yml +++ b/bundle/direct/dresources/configs/experiments.yml @@ -3,9 +3,6 @@ recreate_on_changes: - field: artifact_location reason: immutable ignore_remote_changes: - # Tags updates are not supported by TF. This mirrors that behaviour. + # The backend adds tags that the bundle does not manage. - field: tags - reason: terraform_compat -ignore_local_changes: - - field: tags - reason: terraform_compat + reason: backend_generated diff --git a/bundle/direct/dresources/experiment.go b/bundle/direct/dresources/experiment.go index e4f2e8ebbd7..77aea8d50f4 100644 --- a/bundle/direct/dresources/experiment.go +++ b/bundle/direct/dresources/experiment.go @@ -2,6 +2,9 @@ package dresources import ( "context" + "encoding/json" + "fmt" + "slices" "strings" "github.com/databricks/cli/bundle/config/resources" @@ -42,6 +45,12 @@ func (*ResourceExperiment) RemapState(experiment *ml.Experiment) *ml.CreateExper } } +func (*ResourceExperiment) KeyedSlices() map[string]any { + return map[string]any{ + "tags": func(tag ml.ExperimentTag) (string, string) { return "key", tag.Key }, + } +} + func (r *ResourceExperiment) DoRead(ctx context.Context, id string) (*ml.Experiment, error) { result, err := r.client.Experiments.GetExperiment(ctx, ml.GetExperimentRequest{ ExperimentId: id, @@ -60,14 +69,97 @@ func (r *ResourceExperiment) DoCreate(ctx context.Context, config *ml.CreateExpe return result.ExperimentId, nil, nil } -func (r *ResourceExperiment) DoUpdate(ctx context.Context, id string, config *ml.CreateExperiment, _ *PlanEntry) (*ml.Experiment, error) { - updateReq := ml.UpdateExperiment{ - ExperimentId: id, - NewName: config.Name, - ForceSendFields: utils.FilterFields[ml.UpdateExperiment](config.ForceSendFields), +func (r *ResourceExperiment) DoUpdate(ctx context.Context, id string, config *ml.CreateExperiment, entry *PlanEntry) (*ml.Experiment, error) { + tagChanges := entry.Changes.HasChange(structpath.MustParsePath("tags")) + var remoteTags map[string]string + if tagChanges { + remoteTags = make(map[string]string) + for _, tag := range entry.RemoteState.(*ml.Experiment).Tags { + remoteTags[tag.Key] = tag.Value + } + removedKey, err := removedExperimentTag(entry.Changes, config.Tags, remoteTags) + if err != nil { + return nil, err + } + if removedKey != "" { + return nil, fmt.Errorf("cannot remove experiment tag %q: the MLflow API does not support deleting experiment tags", removedKey) + } + } + + if entry.Changes.HasChange(structpath.MustParsePath("name")) { + updateReq := ml.UpdateExperiment{ + ExperimentId: id, + NewName: config.Name, + ForceSendFields: utils.FilterFields[ml.UpdateExperiment](config.ForceSendFields), + } + if err := r.client.Experiments.UpdateExperiment(ctx, updateReq); err != nil { + return nil, err + } + } + + if tagChanges { + for _, tag := range config.Tags { + if remoteValue, exists := remoteTags[tag.Key]; exists && remoteValue == tag.Value { + continue + } + if err := r.client.Experiments.SetExperimentTag(ctx, ml.SetExperimentTag{ + ExperimentId: id, + Key: tag.Key, + Value: tag.Value, + }); err != nil { + return nil, fmt.Errorf("failed to set experiment %s tag %q: %w", id, tag.Key, err) + } + } + } + + return nil, nil +} + +func removedExperimentTag(changes Changes, desired []ml.ExperimentTag, remote map[string]string) (string, error) { + desiredKeys := make(map[string]bool, len(desired)) + for _, tag := range desired { + desiredKeys[tag.Key] = true } - return nil, r.client.Experiments.UpdateExperiment(ctx, updateReq) + var removedKeys []string + for field, change := range changes { + if change.Action == deployplan.Skip || change.Old == nil { + continue + } + path, err := structpath.ParsePath(field) + if err != nil || !path.HasPrefix(structpath.MustParsePath("tags")) { + continue + } + if field == "tags" { + // A loaded plan stores change values as untyped JSON. + data, err := json.Marshal(change.Old) + if err != nil { + return "", fmt.Errorf("encoding previously configured experiment tags: %w", err) + } + var oldTags []ml.ExperimentTag + if err := json.Unmarshal(data, &oldTags); err != nil { + return "", fmt.Errorf("decoding previously configured experiment tags: %w", err) + } + for _, tag := range oldTags { + removedKeys = append(removedKeys, tag.Key) + } + } else if change.New == nil { + keyField, key, ok := path.KeyValue() + if ok && keyField == "key" && path.Parent().String() == "tags" { + removedKeys = append(removedKeys, key) + } + } + } + + slices.Sort(removedKeys) + for _, key := range removedKeys { + if !desiredKeys[key] { + if _, exists := remote[key]; exists { + return key, nil + } + } + } + return "", nil } func (r *ResourceExperiment) DoDelete(ctx context.Context, id string, _ *ml.CreateExperiment) error { diff --git a/bundle/direct/dresources/experiment_test.go b/bundle/direct/dresources/experiment_test.go new file mode 100644 index 00000000000..8c4f337bc6d --- /dev/null +++ b/bundle/direct/dresources/experiment_test.go @@ -0,0 +1,95 @@ +package dresources + +import ( + "testing" + + "github.com/databricks/cli/bundle/deployplan" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/ml" + "github.com/stretchr/testify/require" +) + +func TestExperimentDoUpdateSetsConfiguredTags(t *testing.T) { + server := testserver.New(t) + testserver.AddDefaultHandlers(server) + client, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: server.URL, + Token: "testtoken", + }) + require.NoError(t, err) + + resource := (&ResourceExperiment{}).New(client) + ctx := t.Context() + id, _, err := resource.DoCreate(ctx, &ml.CreateExperiment{ + Name: "experiment-tags", + Tags: []ml.ExperimentTag{{Key: "mlflow.workload_creation_policy_id", Value: "old-policy"}}, + }) + require.NoError(t, err) + before, err := resource.DoRead(ctx, id) + require.NoError(t, err) + + _, err = resource.DoUpdate(ctx, id, &ml.CreateExperiment{ + Name: "experiment-tags", + Tags: []ml.ExperimentTag{ + {Key: "mlflow.workload_creation_policy_id", Value: "new-policy"}, + {Key: "added-tag", Value: "value"}, + }, + }, &deployplan.PlanEntry{ + RemoteState: before, + Changes: deployplan.Changes{ + "tags": {Action: deployplan.Update}, + }, + }) + require.NoError(t, err) + + after, err := resource.DoRead(ctx, id) + require.NoError(t, err) + require.Equal(t, id, after.ExperimentId) + tags := make(map[string]string, len(after.Tags)) + for _, tag := range after.Tags { + tags[tag.Key] = tag.Value + } + require.Equal(t, "new-policy", tags["mlflow.workload_creation_policy_id"]) + require.Equal(t, "value", tags["added-tag"]) +} + +func TestExperimentDoUpdateRejectsTagRemovalBeforeRename(t *testing.T) { + server := testserver.New(t) + testserver.AddDefaultHandlers(server) + client, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: server.URL, + Token: "testtoken", + }) + require.NoError(t, err) + + resource := (&ResourceExperiment{}).New(client) + ctx := t.Context() + id, _, err := resource.DoCreate(ctx, &ml.CreateExperiment{ + Name: "original-name", + Tags: []ml.ExperimentTag{ + {Key: "keep", Value: "value"}, + {Key: "remove", Value: "value"}, + }, + }) + require.NoError(t, err) + before, err := resource.DoRead(ctx, id) + require.NoError(t, err) + + _, err = resource.DoUpdate(ctx, id, &ml.CreateExperiment{ + Name: "new-name", + Tags: []ml.ExperimentTag{{Key: "keep", Value: "value"}}, + }, &deployplan.PlanEntry{ + RemoteState: before, + Changes: deployplan.Changes{ + "name": {Action: deployplan.Update}, + "tags[key='remove']": {Action: deployplan.Update, Old: ml.ExperimentTag{Key: "remove", Value: "value"}}, + }, + }) + require.ErrorContains(t, err, `cannot remove experiment tag "remove"`) + + after, err := resource.DoRead(ctx, id) + require.NoError(t, err) + require.Equal(t, before.Name, after.Name) + require.Equal(t, before.Tags, after.Tags) +} diff --git a/libs/testserver/experiments.go b/libs/testserver/experiments.go index c7c3fa6c77d..bc4b6ad31d1 100644 --- a/libs/testserver/experiments.go +++ b/libs/testserver/experiments.go @@ -123,6 +123,39 @@ func (s *FakeWorkspace) ExperimentUpdate(req Request) Response { return Response{} } +func (s *FakeWorkspace) ExperimentSetTag(req Request) Response { + defer s.LockUnlock()() + + var setReq ml.SetExperimentTag + if err := json.Unmarshal(req.Body, &setReq); err != nil { + return Response{ + Body: fmt.Sprintf("request parsing error: %s", err), + StatusCode: http.StatusBadRequest, + } + } + + experiment, exists := s.Experiments[setReq.ExperimentId] + if !exists { + return Response{ + StatusCode: http.StatusNotFound, + Body: map[string]string{"message": fmt.Sprintf("Experiment %s not found", setReq.ExperimentId)}, + } + } + + for i := range experiment.Experiment.Tags { + if experiment.Experiment.Tags[i].Key == setReq.Key { + experiment.Experiment.Tags[i].Value = setReq.Value + s.Experiments[setReq.ExperimentId] = experiment + return Response{} + } + } + experiment.Experiment.Tags = append(experiment.Experiment.Tags, ml.ExperimentTag{ + Key: setReq.Key, Value: setReq.Value, + }) + s.Experiments[setReq.ExperimentId] = experiment + return Response{} +} + func (s *FakeWorkspace) ExperimentDelete(req Request) Response { var deleteReq ml.DeleteExperiment if err := json.Unmarshal(req.Body, &deleteReq); err != nil { diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index a021fd8b53a..7866c4cb3dc 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -1435,6 +1435,10 @@ func AddDefaultHandlers(server *Server) { return req.Workspace.ExperimentUpdate(req) }) + server.Handle("POST", "/api/2.0/mlflow/experiments/set-experiment-tag", func(req Request) any { + return req.Workspace.ExperimentSetTag(req) + }) + server.Handle("POST", "/api/2.0/mlflow/experiments/delete", func(req Request) any { return req.Workspace.ExperimentDelete(req) })