diff --git a/.nextchanges/air/7637.md b/.nextchanges/air/7637.md new file mode 100644 index 00000000000..44ff9dda9ea --- /dev/null +++ b/.nextchanges/air/7637.md @@ -0,0 +1 @@ +* `air get` and the `air list` info pane now show the environment of runs on the Databricks AI environment (`databricks_ai_v*`) instead of `N/A`. ([#7637](https://github.com/databricks/cli/pull/7637)) diff --git a/acceptance/air/get-databricks-ai/out.test.toml b/acceptance/air/get-databricks-ai/out.test.toml new file mode 100644 index 00000000000..18b1a88417e --- /dev/null +++ b/acceptance/air/get-databricks-ai/out.test.toml @@ -0,0 +1 @@ +Cloud = false diff --git a/acceptance/air/get-databricks-ai/output.txt b/acceptance/air/get-databricks-ai/output.txt new file mode 100644 index 00000000000..97133d0f25f --- /dev/null +++ b/acceptance/air/get-databricks-ai/output.txt @@ -0,0 +1,20 @@ + +>>> [CLI] air get 123 + +╭─ Metadata ─────────────────────────────────────────────────────╮ +│ │ +│ Run ID 123 │ +│ Status ● SUCCESS │ +│ Submitted 2023-11-14 22:13 UTC │ +│ Retries 0 │ +│ Max Retries 3 │ +│ Duration 12s │ +│ Experiment my-exp │ +│ MLflow Run N/A │ +│ User user@example.com │ +│ Accelerators 1x A10 │ +│ Environment databricks_ai_v6 │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +Run URL: [DATABRICKS_URL]/jobs/runs/123?o=[WORKSPACE_ID] diff --git a/acceptance/air/get-databricks-ai/script b/acceptance/air/get-databricks-ai/script new file mode 100644 index 00000000000..0733c4e5f53 --- /dev/null +++ b/acceptance/air/get-databricks-ai/script @@ -0,0 +1 @@ +trace $CLI air get 123 diff --git a/acceptance/air/get-databricks-ai/test.toml b/acceptance/air/get-databricks-ai/test.toml new file mode 100644 index 00000000000..6a03a7dd207 --- /dev/null +++ b/acceptance/air/get-databricks-ai/test.toml @@ -0,0 +1,38 @@ +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# The GetRun response for a run submitted with +# `environment.version: databricks_ai_v6`. `air run` submits such a version as a +# base_environment rather than an environment_version, and the Environment cell +# shows it as the config wrote it. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "run_page_url": "https://my-workspace.cloud.databricks.test/jobs/runs/123", + "creator_user_name": "user@example.com", + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "environments": [ + {"environment_key": "default", "spec": {"base_environment": "workspace-base-environments/databricks_ai_v6"}} + ], + "tasks": [ + { + "task_key": "train", + "run_id": 456, + "attempt_number": 0, + "max_retries": 3, + "environment_key": "default", + "ai_runtime_task": { + "experiment": "my-exp", + "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}] + } + } + ] +} +''' diff --git a/cmd/air/convert_to_dabs.go b/cmd/air/convert_to_dabs.go index 363bc07e4d7..b20d539481b 100644 --- a/cmd/air/convert_to_dabs.go +++ b/cmd/air/convert_to_dabs.go @@ -370,7 +370,7 @@ func buildBundleValue(ctx context.Context, cfg *runConfig, configPath, codeSourc envVersion, deps := bundleEnvironmentDeps(ctx, cfg) envSpec := map[string]dyn.Value{} if strings.HasPrefix(envVersion, databricksAIPrefix) { - envSpec["base_environment"] = nv("workspace-base-environments/"+envVersion, 1) + envSpec["base_environment"] = nv(workspaceBaseEnvironmentPrefix+envVersion, 1) } else { envSpec["environment_version"] = nv(envVersion, 1) } diff --git a/cmd/air/render.go b/cmd/air/render.go index 6e2f3f5b1a3..09475905bc8 100644 --- a/cmd/air/render.go +++ b/cmd/air/render.go @@ -180,6 +180,7 @@ func aiRuntimeEnvironmentVersion(ctx context.Context, w *databricks.WorkspaceCli EnvironmentKey string `json:"environment_key"` Spec struct { EnvironmentVersion string `json:"environment_version"` + BaseEnvironment string `json:"base_environment"` } `json:"spec"` } `json:"environments"` } @@ -192,6 +193,11 @@ func aiRuntimeEnvironmentVersion(ctx context.Context, w *databricks.WorkspaceCli } for _, e := range resp.Environments { if e.EnvironmentKey == aiRuntimeEnvironmentKey { + // buildRunEnvironment submits a databricks_ai_v* image as a base + // environment instead of an environment_version. + if version, ok := strings.CutPrefix(e.Spec.BaseEnvironment, workspaceBaseEnvironmentPrefix); ok { + return version + } return e.Spec.EnvironmentVersion } } diff --git a/cmd/air/render_test.go b/cmd/air/render_test.go index 29ad7cc6913..1c0d81a121a 100644 --- a/cmd/air/render_test.go +++ b/cmd/air/render_test.go @@ -55,6 +55,31 @@ func TestConfigYAML(t *testing.T) { }) } +func TestAiRuntimeEnvironmentVersion(t *testing.T) { + tests := []struct { + name string + spec string + want string + }{ + { + name: "numeric version is read from environment_version", + spec: `{"environment_version": "6"}`, + want: "6", + }, + { + name: "databricks_ai image is read from base_environment", + spec: `{"base_environment": "workspace-base-environments/databricks_ai_v6"}`, + want: "databricks_ai_v6", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + srv := runGetServer(t, `{"run_id": 5, "environments": [{"environment_key": "default", "spec": `+tt.spec+`}]}`) + assert.Equal(t, tt.want, aiRuntimeEnvironmentVersion(t.Context(), newTestWorkspaceClient(t, srv.URL), 5)) + }) + } +} + func TestResolveConfigYAML(t *testing.T) { ctx := t.Context() diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index c06165dd2b4..c8f8606b3c3 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -34,6 +34,10 @@ const defaultDlRuntimeImage = "CLIENT-GPU-6" // carries the runtime channel. const aiRuntimeEnvironmentKey = "default" +// workspaceBaseEnvironmentPrefix turns a databricks_ai_v* image into the name of +// the Databricks-provided base environment it runs on. +const workspaceBaseEnvironmentPrefix = "workspace-base-environments/" + // dlRuntimeImage resolves the bare runtime channel (config version, else env, // else default), always stripping the CLIENT-GPU- prefix. func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { @@ -50,7 +54,7 @@ func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { func buildRunEnvironment(dlImage string, deps []string) *compute.Environment { environment := &compute.Environment{} if strings.HasPrefix(dlImage, databricksAIPrefix) { - environment.BaseEnvironment = "workspace-base-environments/" + dlImage + environment.BaseEnvironment = workspaceBaseEnvironmentPrefix + dlImage } else { environment.EnvironmentVersion = dlImage }