AIR CLI: bring air run submit output to parity with the Python CLI

acceptance/experimental/air/run-submit-deps/output.txt

@@ -1,8 +1,9 @@
 
 === submit with inline dependencies
 >>> [CLI] experimental air run -f run.yaml
-Submitted run 555
-View at: [DATABRICKS_URL]/jobs/runs/555
+Submitting experiment: deps-smoke
+Submitted workload with Job Run ID: 555
+View job run at: [DATABRICKS_URL]/jobs/runs/555
 
 Tip: use --watch to stream logs until the run completes.
 
@@ -57,8 +58,9 @@ Tip: use --watch to stream logs until the run completes.
 
 === file-form deps: version comes from the requirements file
 >>> [CLI] experimental air run -f run-file.yaml
-Submitted run 555
-View at: [DATABRICKS_URL]/jobs/runs/555
+Submitting experiment: deps-file-smoke
+Submitted workload with Job Run ID: 555
+View job run at: [DATABRICKS_URL]/jobs/runs/555
 
 Tip: use --watch to stream logs until the run completes.
 

acceptance/experimental/air/run-submit/output.txt

@@ -1,9 +1,10 @@
 
 === submit with a git code_source
 >>> [CLI] experimental air run -f run.yaml
+Submitting experiment: submit-smoke
 Uploading [SNAPSHOT_TARBALL]...
-Submitted run 555
-View at: [DATABRICKS_URL]/jobs/runs/555
+Submitted workload with Job Run ID: 555
+View job run at: [DATABRICKS_URL]/jobs/runs/555
 
 Tip: use --watch to stream logs until the run completes.
 

experimental/air/cmd/logstream.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"slices"
+	"strings"
 	"time"
 
 	"github.com/databricks/cli/libs/cmdio"
@@ -26,8 +27,37 @@ const (
 	defaultCompletedRunTailLines = 10000
 	// seenRecordsCap bounds the dedup set, evicting oldest-inserted entries first.
 	seenRecordsCap = 100000
+	// statusMessageRefreshEveryNPolls throttles the status_message fetch so the
+	// waiting spinner doesn't issue a get-output on every poll tick.
+	statusMessageRefreshEveryNPolls = 5
 )
 
+// statusMessageType tags a client-facing message packed into
+// ai_runtime_task_output.status_message as "<TYPE>:<payload>"; only STATUS-typed
+// messages are surfaced.
+const statusMessageType = "STATUS"
+
+// waitingForComputeStatus is the fallback shown while a PENDING run waits for
+// accelerator compute.
+const waitingForComputeStatus = "Waiting for accelerator compute capacity to become available..."
+
+// normalizeStatusMessage returns the payload of a "STATUS:<payload>" message,
+// normalized for display (trailing "." stripped, "..." suffix added), or "" for
+// any other type or an empty/absent message.
+func normalizeStatusMessage(raw string) string {
+	messageType, payload, ok := strings.Cut(raw, ":")
+	if !ok || !strings.EqualFold(strings.TrimSpace(messageType), statusMessageType) {
+		return ""
+	}
+	payload = strings.TrimSpace(payload)
+	payload = strings.TrimRight(payload, ".")
+	payload = strings.TrimSpace(payload)
+	if payload == "" {
+		return ""
+	}
+	return payload + "..."
+}
+
 // retryCheckInterval is the wait between status/log polls. A var so tests can
 // shrink it.
 var retryCheckInterval = 3 * time.Second
@@ -206,6 +236,40 @@ type bricklensStreamer struct {
 	updateSpinner func(string)
 }
 
+// waitingSpinnerText returns the waiting-spinner text: the server-set STATUS
+// message if present, else the compute-capacity message for a PENDING run, else
+// the default "waiting for run to start".
+func (st *bricklensStreamer) waitingSpinnerText() string {
+	if msg := st.serverStatusMessage(); msg != "" {
+		return msg
+	}
+	if st.status.lifeCycleState == "PENDING" {
+		return waitingForComputeStatus
+	}
+	return fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node)
+}
+
+// serverStatusMessage returns the run's server-set STATUS message (normalized for
+// display), or "" if unavailable. The message lives on the latest task run's
+// output, re-resolved each call so a retry's new task run is picked up.
+// Best-effort: any fetch failure logs at debug and returns "".
+func (st *bricklensStreamer) serverStatusMessage() string {
+	run, err := st.w.Jobs.GetRun(st.ctx, jobs.GetRunRequest{RunId: st.req.runID})
+	if err != nil || len(run.Tasks) == 0 {
+		return ""
+	}
+	taskRunID := run.Tasks[len(run.Tasks)-1].RunId
+	out, err := st.w.Jobs.GetRunOutputByRunId(st.ctx, taskRunID)
+	if err != nil {
+		log.Debugf(st.ctx, "air logs: status_message fetch failed for run %d: %v", st.req.runID, err)
+		return ""
+	}
+	if out.AiRuntimeTaskOutput == nil {
+		return ""
+	}
+	return normalizeStatusMessage(out.AiRuntimeTaskOutput.StatusMessage)
+}
+
 // reportStatusChange fires onStatusChange when the run's display state differs
 // from the last reported one.
 func (st *bricklensStreamer) reportStatusChange() {
@@ -242,6 +306,11 @@ func (st *bricklensStreamer) run() (bool, error) {
 	}
 
 	firstIteration := true
+	// Throttled refresh of the waiting-spinner text: statusRefreshCounter gates the
+	// server status_message fetch to every Nth poll, and lastSpinnerText avoids
+	// redundant spinner updates.
+	statusRefreshCounter := 0
+	lastSpinnerText := ""
 	for {
 		if !firstIteration {
 			status, err := resolveRunStatus(st.ctx, st.w, st.req.runID)
@@ -269,9 +338,17 @@ func (st *bricklensStreamer) run() (bool, error) {
 		terminal := st.status.terminal()
 		toSec := st.req.toSeconds(st.status)
 
-		// While waiting on a still-active run with no logs yet, refresh the spinner.
+		// While waiting on a still-active run with no logs yet, refresh the spinner
+		// with the server-set status (throttled), so a run stuck waiting for compute
+		// shows why rather than a generic "waiting" message.
 		if !terminal && !st.firstLogSeen && st.updateSpinner != nil {
-			st.updateSpinner(fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node))
+			if statusRefreshCounter%statusMessageRefreshEveryNPolls == 0 {
+				if desired := st.waitingSpinnerText(); desired != lastSpinnerText {
+					st.updateSpinner(desired)
+					lastSpinnerText = desired
+				}
+			}
+			statusRefreshCounter++
 		}
 
 		// A run already terminal on the first iteration renders as a tail (most

experimental/air/cmd/logstream_test.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"net/http"
 	"net/http/httptest"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -204,6 +205,62 @@ func TestDisplayState(t *testing.T) {
 	assert.Equal(t, "UNKNOWN", logRunStatus{}.displayState())
 }
 
+func TestNormalizeStatusMessage(t *testing.T) {
+	tests := []struct {
+		raw  string
+		want string
+	}{
+		{"STATUS: Waiting for GPU capacity.", "Waiting for GPU capacity..."},
+		{"STATUS:Waiting for GPU capacity", "Waiting for GPU capacity..."},
+		{"status: provisioning", "provisioning..."}, // type match is case-insensitive
+		{"STATUS: done...", "done..."},              // trailing dots collapse to one "..."
+		{"INFO: not a status", ""},                  // other type ignored
+		{"no type prefix", ""},
+		{"STATUS:", ""},    // empty payload
+		{"STATUS:   ", ""}, // whitespace-only payload
+		{"", ""},
+	}
+	for _, tt := range tests {
+		assert.Equal(t, tt.want, normalizeStatusMessage(tt.raw), "raw=%q", tt.raw)
+	}
+}
+
+func TestWaitingSpinnerText(t *testing.T) {
+	// A server that returns the run (with a task) and a STATUS-typed status_message.
+	newStreamer := func(t *testing.T, statusMessage, lifeCycle string) *bricklensStreamer {
+		t.Helper()
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			switch r.URL.Path {
+			case "/api/2.2/jobs/runs/get":
+				_, _ = w.Write([]byte(`{"run_id": 1, "tasks": [{"run_id": 2}]}`))
+			case "/api/2.2/jobs/runs/get-output":
+				_, _ = w.Write([]byte(`{"ai_runtime_task_output": {"status_message": ` + strconv.Quote(statusMessage) + `}}`))
+			default:
+				_, _ = w.Write([]byte(`{}`))
+			}
+		}))
+		t.Cleanup(srv.Close)
+		return &bricklensStreamer{
+			ctx:    t.Context(),
+			w:      newTestWorkspaceClient(t, srv.URL),
+			req:    logRequest{runID: 1, node: 0},
+			status: logRunStatus{lifeCycleState: lifeCycle},
+		}
+	}
+
+	// Server STATUS message wins.
+	assert.Equal(t, "Waiting for GPU capacity...",
+		newStreamer(t, "STATUS: Waiting for GPU capacity", "PENDING").waitingSpinnerText())
+
+	// No status message + PENDING -> compute-capacity fallback.
+	assert.Equal(t, waitingForComputeStatus,
+		newStreamer(t, "", "PENDING").waitingSpinnerText())
+
+	// No status message + non-PENDING -> default "waiting for run to start".
+	assert.Equal(t, "Waiting for run to start (node 0)...",
+		newStreamer(t, "", "RUNNING").waitingSpinnerText())
+}
+
 func TestEmitLogLineJSON(t *testing.T) {
 	var buf bytes.Buffer
 	emitLogLine(&buf, logRequest{node: 2, jsonOutput: true}, "hello")

experimental/air/cmd/mlflow.go

@@ -4,13 +4,23 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/databricks/databricks-sdk-go/service/ml"
 )
 
+// mlflowLinkPollAttempts bounds the best-effort poll for a freshly-submitted
+// run's MLflow IDs (see resolveMLflowIDsForRun), kept short so a bare `air run`
+// returns promptly when the IDs aren't ready yet.
+const mlflowLinkPollAttempts = 3
+
+// mlflowLinkPollInterval is the delay between poll attempts. A var, not a const,
+// so tests can shrink it and avoid a real sleep.
+var mlflowLinkPollInterval = 500 * time.Millisecond
+
 // mlflowIdentifiers are the experiment and run IDs MLflow assigns to a run.
 type mlflowIdentifiers struct {
 	ExperimentID string
@@ -62,6 +72,46 @@ func mlflowRunURL(host string, ids *mlflowIdentifiers) string {
 		strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID)
 }
 
+// mlflowExperimentURL links to the MLflow experiment page. Omits the ?o= query
+// for consistency with mlflowRunURL and the run-submit dashboard URL.
+func mlflowExperimentURL(host string, ids *mlflowIdentifiers) string {
+	return fmt.Sprintf("%s/ml/experiments/%s", strings.TrimRight(host, "/"), ids.ExperimentID)
+}
+
+// resolveMLflowIDsForRun best-effort resolves a just-submitted run's MLflow IDs,
+// polling because they are assigned only once the task run starts. Returns nil
+// (treated as "no link", not an error) if they don't appear within the budget or
+// the context is cancelled.
+func resolveMLflowIDsForRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) *mlflowIdentifiers {
+	// The task run id is fixed at submit time, so resolve it once; only the MLflow
+	// output (runs/get-output) fills in later, so that is all we poll.
+	run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID})
+	if err != nil {
+		log.Debugf(ctx, "air run: could not fetch run %d for MLflow link: %v", runID, err)
+		return nil
+	}
+	if len(run.Tasks) == 0 {
+		return nil
+	}
+	// The MLflow output is attached to the task run, not the parent job run.
+	taskRunID := run.Tasks[len(run.Tasks)-1].RunId
+
+	for attempt := range mlflowLinkPollAttempts {
+		if ids := mlflowIDsForTask(ctx, w, taskRunID); ids != nil {
+			return ids
+		}
+		if attempt == mlflowLinkPollAttempts-1 {
+			break
+		}
+		select {
+		case <-ctx.Done():
+			return nil
+		case <-time.After(mlflowLinkPollInterval):
+		}
+	}
+	return nil
+}
+
 // fetchMLflowRunName fetches a run's MLflow run_name via the MLflow REST API,
 // returning "" if it can't be obtained. Best-effort, like the rest of the MLflow
 // enrichment.