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 fc5ad8ebb98..c248b9b1fd1 100644 --- a/acceptance/bundle/resources/experiments/basic/output.txt +++ b/acceptance/bundle/resources/experiments/basic/output.txt @@ -133,7 +133,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 >>> [CLI] bundle plan @@ -168,6 +168,10 @@ Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged { "key": "mlflow.experimentType", "value": "MLFLOW_EXPERIMENT" + }, + { + "key": "k2", + "value": "v2" } ] } diff --git a/acceptance/bundle/resources/experiments/basic/script b/acceptance/bundle/resources/experiments/basic/script index ab6ca3de78d..fb10fb13ed8 100755 --- a/acceptance/bundle/resources/experiments/basic/script +++ b/acceptance/bundle/resources/experiments/basic/script @@ -29,9 +29,11 @@ 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) +test "$experiment_id" = "$existing_experiment_id" trace $CLI experiments get-experiment $experiment_id | jq '.experiment | {name, artifact_location, tags}' 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..eeae626c707 100644 --- a/bundle/direct/dresources/experiment.go +++ b/bundle/direct/dresources/experiment.go @@ -2,6 +2,7 @@ package dresources import ( "context" + "fmt" "strings" "github.com/databricks/cli/bundle/config/resources" @@ -60,14 +61,38 @@ 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) { + 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 + } } - return nil, r.client.Experiments.UpdateExperiment(ctx, updateReq) + if entry.Changes.HasChange(structpath.MustParsePath("tags")) { + remoteTags := make(map[string]string) + for _, tag := range entry.RemoteState.(*ml.Experiment).Tags { + remoteTags[tag.Key] = tag.Value + } + 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 (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..924ac15d92d --- /dev/null +++ b/bundle/direct/dresources/experiment_test.go @@ -0,0 +1,55 @@ +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"]) +} 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 50ff904a573..6431795f1a2 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -1355,6 +1355,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) })