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/control-chars/output.txt b/acceptance/air/control-chars/output.txt index b320060ce29..e47edb09c90 100644 --- a/acceptance/air/control-chars/output.txt +++ b/acceptance/air/control-chars/output.txt @@ -15,7 +15,7 @@ │ MLflow Run my-run │ │ User user@example.com │ │ Accelerators 1x A10 │ -│ Environment N/A │ +│ Environment databricks_ai_v6 │ │ │ ╰────────────────────────────────────────────────────────────────╯ diff --git a/acceptance/air/control-chars/test.toml b/acceptance/air/control-chars/test.toml index 770413a86a4..34cba10837a 100644 --- a/acceptance/air/control-chars/test.toml +++ b/acceptance/air/control-chars/test.toml @@ -8,7 +8,8 @@ Env.DATABRICKS_CACHE_ENABLED = "false" Pattern = "HEAD /" Response.Body = '' -# Experiment and creator names with embedded escapes, for `air get` and `air list`. +# Experiment, creator and base environment names with embedded escapes, for +# `air get` and `air list`. [[Server]] Pattern = "GET /api/2.2/jobs/runs/get" Response.Body = ''' @@ -19,6 +20,9 @@ Response.Body = ''' "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\u001b]0;pwned\u0007_v6"}} + ], "tasks": [{ "task_key": "train", "run_id": 456, 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..5786dfc20e2 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,12 @@ 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. Any other workspace + // base environment, such as a custom dbe_, shows its ID too. + if version, ok := strings.CutPrefix(e.Spec.BaseEnvironment, workspaceBaseEnvironmentPrefix); ok { + return version + } return e.Spec.EnvironmentVersion } } @@ -376,13 +383,14 @@ func renderBox(p palette, title, body string) string { // renderFields draws the two-column summary: muted labels right-padded to the // longest one, neutral values, a status-colored Status, and blue Run ID / MLflow -// Run hyperlinks. The experiment, MLflow run and user values are server-supplied -// names, so they are stripped of control characters before they reach the -// terminal. +// Run hyperlinks. The experiment, MLflow run, user and environment values are +// server-supplied, so they are stripped of control characters before they reach +// the terminal. func renderFields(p palette, colorOn bool, v runView) string { v.experiment = stripControl(v.experiment) v.mlflowLabel = stripControl(v.mlflowLabel) v.user = stripControl(v.user) + v.environment = stripControl(v.environment) status := statusStyle(p, v.status).Render("● " + v.status) rows := []string{ field(p, "Run ID", link(colorOn, p.blue, v.runID, v.dashboardURL)), diff --git a/cmd/air/render_test.go b/cmd/air/render_test.go index 29ad7cc6913..be872174fe8 100644 --- a/cmd/air/render_test.go +++ b/cmd/air/render_test.go @@ -55,6 +55,36 @@ 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", + }, + { + name: "custom workspace base environment is read from base_environment", + spec: `{"base_environment": "workspace-base-environments/dbe_b849b66e-b31a-4cb5-b161-1f2b10877fb7"}`, + want: "dbe_b849b66e-b31a-4cb5-b161-1f2b10877fb7", + }, + } + 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() @@ -184,12 +214,13 @@ func TestRenderFieldsStripsControlCharacters(t *testing.T) { mlflowURL: "https://h.test/ml/experiments/E1/runs/R1", user: "user\x1b[31m@example.com", accelerators: "1x A10", - environment: "4", + environment: "\x1b]8;;https://attacker.test/\x07databricks_ai_v6\x1b]8;;\x07", }) assert.Contains(t, out, "Experiment experiment") assert.Contains(t, out, "MLflow Run runname") assert.Contains(t, out, "User user@example.com") + assert.Contains(t, out, "Environment databricks_ai_v6") assert.NotContains(t, out, "\x1b") assert.NotContains(t, out, "\r") } diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index c06165dd2b4..20591275a62 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -34,6 +34,11 @@ const defaultDlRuntimeImage = "CLIENT-GPU-6" // carries the runtime channel. const aiRuntimeEnvironmentKey = "default" +// workspaceBaseEnvironmentPrefix is the name prefix of every workspace base +// environment, for example the Databricks-provided databricks_ai_v6 or a custom +// dbe_. +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 +55,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 }