[draft] Apply configured MLflow experiment tags on bundle update

.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))

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

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]"
     }
   ]
 }

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))}'

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}

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

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 {

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)
+}