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/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          [email protected]                                │
+│  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": "[email protected]",
+  "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,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
 		}
 	}

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

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
 	}