diff --git a/acceptance/experimental/air/get-ai-runtime/output.txt b/acceptance/experimental/air/get-ai-runtime/output.txt index c47790eaa31..21719cb8d78 100644 --- a/acceptance/experimental/air/get-ai-runtime/output.txt +++ b/acceptance/experimental/air/get-ai-runtime/output.txt @@ -27,7 +27,7 @@ │ MLflow Run my-run │ │ User user@example.com │ │ Accelerators 1x A10 │ -│ Environment N/A │ +│ Environment 4 │ │ │ ╰────────────────────────────────────────────────────────────────╯ diff --git a/acceptance/experimental/air/get-ai-runtime/test.toml b/acceptance/experimental/air/get-ai-runtime/test.toml index 07d75ae83c8..061238e5e37 100644 --- a/acceptance/experimental/air/get-ai-runtime/test.toml +++ b/acceptance/experimental/air/get-ai-runtime/test.toml @@ -9,8 +9,10 @@ MSYS_NO_PATHCONV = "1" Pattern = "HEAD /" Response.Body = '' -# The typed SDK GetRun response: an ai_runtime_task run has no gen_ai_compute_task, -# so the task comes back empty (the SDK has no field for ai_runtime_task). +# The GetRun response for an ai_runtime_task run. It has no gen_ai_compute_task, +# so the Environment cell instead comes from the run's environments[].spec (keyed +# by the task's environment_key), which the typed SDK Run drops — `air get` reads +# it with a raw request, so this route is hit twice (typed GetRun + that request). [[Server]] Pattern = "GET /api/2.2/jobs/runs/get" Response.Body = ''' @@ -21,12 +23,16 @@ Response.Body = ''' "start_time": 1700000000000, "end_time": 1700000012000, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "environments": [ + {"environment_key": "default", "spec": {"environment_version": "4"}} + ], "tasks": [ { "task_key": "train", "run_id": 456, "attempt_number": 0, "max_retries": 3, + "environment_key": "default", "ai_runtime_task": { "experiment": "my-exp", "deployments": [ diff --git a/acceptance/experimental/air/get/output.txt b/acceptance/experimental/air/get/output.txt index 6e51d7debf9..ff5cbd2ab04 100644 --- a/acceptance/experimental/air/get/output.txt +++ b/acceptance/experimental/air/get/output.txt @@ -27,7 +27,7 @@ │ MLflow Run my-run │ │ User user@example.com │ │ Accelerators 1x A10 │ -│ Environment ml-runtime-gpu:1.0 │ +│ Environment N/A │ │ │ ╰────────────────────────────────────────────────────────────────╯ diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index 8e46c70d48c..de7046de3f8 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -285,19 +285,6 @@ func gpuDisplayName(gpuType string) string { return gpuType } -// environment returns the run's runtime image (the training environment), or an -// empty string if the run has no GenAI-compute task. -func environment(run *jobs.Run) string { - if len(run.Tasks) == 0 { - return "" - } - task := run.Tasks[0].GenAiComputeTask - if task == nil { - return "" - } - return task.DlRuntimeImage -} - // maxRetries returns the configured retry limit for the run's latest task as a // display string: "unlimited" for the backend's -1, otherwise the count. func maxRetries(run *jobs.Run) string { diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 2f8bd8dd09e..dcde5473d19 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -236,7 +236,10 @@ func buildGetData(run *jobs.Run) getData { } data.UserDisplay = orNA(run.CreatorUserName) data.AcceleratorsDisplay = orNA(accelerators(run)) - data.EnvironmentDisplay = orNA(environment(run)) + // EnvironmentDisplay is resolved at render time: the serverless environment + // version needs a raw GetRun read (the typed SDK Run drops it), so it is not + // filled here alongside the fields read straight off the run. + data.EnvironmentDisplay = na data.MaxRetriesDisplay = maxRetries(run) return data } diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go index f3fb63db807..5d4cf9bb6d0 100644 --- a/experimental/air/cmd/render.go +++ b/experimental/air/cmd/render.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "io" + "net/http" "strconv" "strings" @@ -11,6 +12,7 @@ import ( "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/muesli/termenv" "go.yaml.in/yaml/v3" @@ -22,6 +24,10 @@ const ( metadataBoxTitle = "Metadata" ) +// jobsRunsGetPath is the Jobs GetRun endpoint, called with a raw request to read +// the run's environments[] block, which the typed SDK Run does not expose. +const jobsRunsGetPath = "/api/2.2/jobs/runs/get" + // minBoxInnerWidth keeps all boxes a uniform, comfortable width; boxHPad and // boxVPad are the horizontal and vertical padding inside each box. const ( @@ -91,6 +97,11 @@ func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceCl renderer, colorOn := cmdio.NewRenderer(ctx, out) p := newPalette(renderer) + // The serverless environment version lives on the run's environments[].spec, + // which the typed SDK Run drops, so read it with a raw request. Empty (a run + // with no serverless environment) stays "N/A". + data.EnvironmentDisplay = orNA(aiRuntimeEnvironmentVersion(ctx, w, run.RunId)) + view := runView{ runID: data.RunID, dashboardURL: data.DashboardURL, @@ -143,6 +154,40 @@ func genAIComputeTask(run *jobs.Run) *jobs.GenAiComputeTask { return run.Tasks[0].GenAiComputeTask } +// aiRuntimeEnvironmentVersion returns the serverless environment version an +// ai_runtime run used (e.g. "4"), read from the Jobs GetRun response's +// environments[] entry keyed by aiRuntimeEnvironmentKey. The typed SDK Run has +// no environments field, so the value is fetched with a raw request. Best-effort: +// returns "" (logged) on any error or when the run declares no environment. +func aiRuntimeEnvironmentVersion(ctx context.Context, w *databricks.WorkspaceClient, runID int64) string { + apiClient, err := client.New(w.Config) + if err != nil { + log.Warnf(ctx, "air get: could not create client to read environment: %v", err) + return "" + } + var resp struct { + Environments []struct { + EnvironmentKey string `json:"environment_key"` + Spec struct { + EnvironmentVersion string `json:"environment_version"` + } `json:"spec"` + } `json:"environments"` + } + // For a GET the SDK serializes the request value into query parameters, so + // run_id is passed as the request, mirroring the other raw calls in this package. + query := map[string]any{"run_id": runID} + if err := apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &resp); err != nil { + log.Warnf(ctx, "air get: could not read environment for run %d: %v", runID, err) + return "" + } + for _, e := range resp.Environments { + if e.EnvironmentKey == aiRuntimeEnvironmentKey { + return e.Spec.EnvironmentVersion + } + } + return "" +} + // resolveConfigYAML returns the config box body: from the downloaded config file // when we have its path, else from the legacy task. func resolveConfigYAML(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run, data *getData) string { diff --git a/experimental/air/cmd/render_test.go b/experimental/air/cmd/render_test.go index 067f5817fb1..66e5b09deb7 100644 --- a/experimental/air/cmd/render_test.go +++ b/experimental/air/cmd/render_test.go @@ -153,7 +153,7 @@ func TestRenderFields(t *testing.T) { mlflowURL: "https://h.test/ml/experiments/E1/runs/R1", user: "user@example.com", accelerators: "1x A10", - environment: "ml-runtime-gpu:1.0", + environment: "4", }) // Labels are padded to the longest ("Accelerators"), so values align. @@ -161,7 +161,7 @@ func TestRenderFields(t *testing.T) { assert.Contains(t, out, "Accelerators 1x A10") // Max retries and environment show alongside the other fields. assert.Contains(t, out, "Max Retries 3") - assert.Contains(t, out, "Environment ml-runtime-gpu:1.0") + assert.Contains(t, out, "Environment 4") // The status carries its dot prefix. assert.Contains(t, out, "● SUCCESS") // Off a terminal, links render as the bare label (URLs live in JSON output).