air: show the Databricks AI environment in air get

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

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                                │
 │                                                                │
 ╰────────────────────────────────────────────────────────────────╯
 

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,

acceptance/air/get-databricks-ai/out.test.toml

@@ -0,0 +1 @@
+Cloud = false

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]

acceptance/air/get-databricks-ai/script

@@ -0,0 +1 @@
+trace $CLI air get 123

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}}]
+      }
+    }
+  ]
+}
+'''

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

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_<uuid>, 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)),

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

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_<uuid>.
+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
 	}