Add support for experiments in deployment bind/unbind commands

This commit is contained in:
Anton Nekipelov 2025-03-05 15:52:42 +01:00
parent a24a7f5738
commit eb1a16e27b
4 changed files with 92 additions and 0 deletions

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

View File

@ -17,8 +17,70 @@ import (
"github.com/stretchr/testify/require"
"github.com/databricks/databricks-sdk-go/service/catalog"
"github.com/databricks/databricks-sdk-go/service/ml"
)
func TestBindExperimentToExistingExperiment(t *testing.T) {
ctx, wt := acc.UcWorkspaceTest(t)
currentUser, err := wt.W.CurrentUser.Me(ctx)
require.NoError(t, err)
// create a pre-defined experiment:
uniqueId := uuid.New().String()
experimentName := "/Workspace/Users/" + currentUser.UserName + "/test-experiment" + uniqueId
predefinedExperiment, err := wt.W.Experiments.CreateExperiment(ctx, ml.CreateExperiment{
Name: experimentName,
ArtifactLocation: "s3://test-location",
})
require.NoError(t, err)
t.Cleanup(func() {
err := wt.W.Experiments.DeleteExperiment(ctx, ml.DeleteExperiment{ExperimentId: predefinedExperiment.ExperimentId})
require.NoError(t, err)
})
// setup the bundle:
bundleRoot := initTestTemplate(t, ctx, "ml_experiment", map[string]any{
"unique_id": uniqueId,
"experiment_name": experimentName,
})
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
// run the bind command:
c := testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "experiment1", predefinedExperiment.ExperimentId, "--auto-approve")
_, _, err = c.Run()
require.NoError(t, err)
// deploy the bundle:
deployBundle(t, ctx, bundleRoot)
// check that the predefinedExperiment was not re-created / deleted (it is still active):
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)
updatedExperiment, err := w.Experiments.GetExperiment(ctx, ml.GetExperimentRequest{
ExperimentId: predefinedExperiment.ExperimentId,
})
require.NoError(t, err)
require.Equal(t, "active", updatedExperiment.Experiment.LifecycleStage)
// unbind the experiment:
c = testcli.NewRunner(t, ctx, "bundle", "deployment", "unbind", "experiment1")
_, _, err = c.Run()
require.NoError(t, err)
// destroy the bundle:
destroyBundle(t, ctx, bundleRoot)
// Check that schema is unbound and exists after bundle is destroyed
postDestroyExperiment, err := w.Experiments.GetExperiment(ctx, ml.GetExperimentRequest{
ExperimentId: predefinedExperiment.ExperimentId,
})
require.NoError(t, err)
require.Equal(t, "active", postDestroyExperiment.Experiment.LifecycleStage)
}
func TestBindSchemaToExistingSchema(t *testing.T) {
ctx, wt := acc.UcWorkspaceTest(t)

View File

@ -0,0 +1,12 @@
{
"properties": {
"unique_id": {
"type": "string",
"description": "Unique ID for the experiment name"
},
"experiment_name": {
"type": "string",
"description": "Experiment name. An experiment name must be an absolute path within the Databricks workspace, e.g. '/Users/<some-username>/my-experiment'"
}
}
}

View File

@ -0,0 +1,12 @@
bundle:
name: ml-experiment
workspace:
root_path: "~/.bundle/{{.unique_id}}"
resources:
experiments:
experiment1:
name: {{.experiment_name}}
artifact_location: s3://test-location