Add support for experiments in deployment bind/unbind commands (#2434)

## Changes
1. Changed `FindResourceByConfigKey` to return experiment resources

## Why
This PR adds support for experiment resources in deployment operations,
enabling users to:
- Bind experiments using `databricks bundle deployment bind
<myexperiment_key> <experiment_id>`
- Unbind experiments using `databricks bundle deployment unbind
<myexperiment_key>`

Where:
- `myexperiment_key` is a resource key defined in the bundle's .yml file
- `experiment_id` references an existing experiment in the Databricks
workspace

These capabilities allow for more flexible resource management of
experiments within bundles.

## Tests
Added a new acceptance test that tests bind and unbind methods together
with bundle deployment and destruction.
This commit is contained in:
Anton Nekipelov 2025-03-17 16:13:40 +01:00 committed by GitHub
parent 0b9958fc35
commit e5f39b5916
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 94 additions and 0 deletions

View File

@ -414,6 +414,7 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
require.NoError(t, err)
cmd.Env = append(cmd.Env, "TESTDIR="+absDir)
cmd.Env = append(cmd.Env, "CLOUD_ENV="+cloudEnv)
cmd.Env = append(cmd.Env, "CURRENT_USER_NAME="+user.UserName)
cmd.Dir = tmpDir
outputPath := filepath.Join(tmpDir, "output.txt")

View File

@ -0,0 +1,8 @@
bundle:
name: bind-ml-experiment-test-$BUNDLE_NAME_SUFFIX
resources:
experiments:
experiment1:
name: $EXPERIMENT_NAME

View File

@ -0,0 +1,31 @@
=== Bind experiment test:
=== Substitute variables in the template
=== Create a pre-defined experiment
=== Bind experiment: Updating deployment state...
Successfully bound databricks_mlflow_experiment with an id '[NUMID]'. Run 'bundle deploy' to deploy changes to your workspace
=== Deploy bundle: Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/bind-ml-experiment-test-[UUID]/default/files...
Deploying resources...
Updating deployment state...
Deployment complete!
=== Read the pre-defined experiment: {
"name": "/Users/[USERNAME]/test-experiment[UUID]",
"lifecycle_stage": "active"
}
=== Unbind the experiment: Updating deployment state...
=== Destroy the bundle: All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/bind-ml-experiment-test-[UUID]/default
Deleting files...
Destroy complete!
=== Read the pre-defined experiment again (expecting it still exists and is not deleted): {
"name": "/Users/[USERNAME]/test-experiment[UUID]",
"lifecycle_stage": "active"
}
=== Test cleanup:
=== Delete the pre-defined experiment: 0

View File

@ -0,0 +1,42 @@
title "Bind experiment test: "
title "Substitute variables in the template"
BUNDLE_NAME_SUFFIX=$(uuid)
export BUNDLE_NAME_SUFFIX
# double slash at the start prevents Windows to apply replacements to the path
EXPERIMENT_NAME="//Workspace/Users/${CURRENT_USER_NAME}/test-experiment$(uuid)"
if [ -z "$CLOUD_ENV" ]; then
EXPERIMENT_NAME="//Workspace/Users/${CURRENT_USER_NAME}/test-experiment6260d50f-e8ff-4905-8f28-812345678903" # use hard-coded uuid when running locally
fi
export EXPERIMENT_NAME
envsubst < databricks.yml > out.yml && mv out.yml databricks.yml
title "Create a pre-defined experiment"
EXPERIMENT_ID=$($CLI experiments create-experiment "${EXPERIMENT_NAME}" | jq -r '.experiment_id')
cleanupRemoveExperiment() {
title "Test cleanup: "
title "Delete the pre-defined experiment: "
$CLI experiments delete-experiment ${EXPERIMENT_ID}
echo $?
}
trap cleanupRemoveExperiment EXIT
title "Bind experiment: "
$CLI bundle deployment bind experiment1 ${EXPERIMENT_ID} --auto-approve
title "Deploy bundle: "
$CLI bundle deploy --force-lock --auto-approve
title "Read the pre-defined experiment: "
$CLI experiments get-experiment ${EXPERIMENT_ID} | jq '{name: .experiment.name, lifecycle_stage: .experiment.lifecycle_stage}'
title "Unbind the experiment: "
$CLI bundle deployment unbind experiment1
title "Destroy the bundle: "
$CLI bundle destroy --auto-approve
title "Read the pre-defined experiment again (expecting it still exists and is not deleted): "
$CLI experiments get-experiment ${EXPERIMENT_ID} | jq '{name: .experiment.name, lifecycle_stage: .experiment.lifecycle_stage}'

View File

@ -0,0 +1,6 @@
Local = false
Cloud = true
[[Repls]]
Old = "[0-9]{3,}"
New = "[NUMID]"

View File

@ -118,6 +118,12 @@ func (r *Resources) FindResourceByConfigKey(key string) (ConfigResource, error)
}
}
for k := range r.Experiments {
if k == key {
found = append(found, r.Experiments[k])
}
}
if len(found) == 0 {
return nil, fmt.Errorf("no such resource: %s", key)
}