Improve AIR run log guidance

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

@@ -5,7 +5,9 @@ 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.
+Tip: use --watch when submitting a run to stream logs to your terminal.
+Stream logs after submission using:
+  databricks experimental air logs 555
 
 === only config + command are uploaded; no requirements.yaml
 >>> print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep

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

@@ -6,7 +6,9 @@ Uploading [SNAPSHOT_TARBALL]...
 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.
+Tip: use --watch when submitting a run to stream logs to your terminal.
+Stream logs after submission using:
+  databricks experimental air logs 555
 
 === the ai_runtime_task carries the code_source_path
 >>> print_requests.py //api/2.2/jobs/runs/submit

experimental/air/cmd/format.go

@@ -104,16 +104,39 @@ func runStatus(state *jobs.RunState) string {
 
 func displayRunStatus(run *jobs.Run) string {
 	status := runStatus(run.State)
-	if status != string(jobs.RunLifeCycleStateRunning) || len(run.Tasks) == 0 || run.Tasks[0].State == nil {
-		return status
+	if runWaitingForCompute(run) {
+		return "PENDING"
+	}
+	return status
+}
+
+func runWaitingForCompute(run *jobs.Run) bool {
+	if run.State == nil || run.State.ResultState != "" {
+		return false
+	}
+	if run.State.LifeCycleState == jobs.RunLifeCycleStatePending {
+		return true
+	}
+	if run.State.LifeCycleState != jobs.RunLifeCycleStateRunning || len(run.Tasks) == 0 || run.Tasks[0].State == nil {
+		return false
 	}
 	switch run.Tasks[0].State.LifeCycleState {
 	case jobs.RunLifeCycleStatePending, jobs.RunLifeCycleStateQueued,
 		jobs.RunLifeCycleStateWaitingForRetry, jobs.RunLifeCycleStateBlocked:
-		return "PENDING"
+		return true
 	default:
-		return status
+		return false
+	}
+}
+
+func detailedDisplayRunStatus(run *jobs.Run, statusMessage string) string {
+	if run.State != nil && run.State.ResultState == "" && statusMessage != "" {
+		return statusMessage
+	}
+	if runWaitingForCompute(run) {
+		return waitingForComputeStatus
 	}
+	return displayRunStatus(run)
 }
 
 func terminationReason(run *jobs.Run) *string {

experimental/air/cmd/format_test.go

@@ -262,6 +262,27 @@ func TestDisplayRunStatus(t *testing.T) {
 	}))
 }
 
+func TestDetailedDisplayRunStatus(t *testing.T) {
+	waiting := &jobs.Run{
+		State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning},
+		Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateQueued}}},
+	}
+	assert.Equal(t, "Waiting for GPU capacity...", detailedDisplayRunStatus(waiting, "Waiting for GPU capacity..."))
+	assert.Equal(t, waitingForComputeStatus, detailedDisplayRunStatus(waiting, ""))
+
+	started := &jobs.Run{
+		State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning},
+		Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}}},
+	}
+	assert.Equal(t, "RUNNING", detailedDisplayRunStatus(started, ""))
+
+	failed := &jobs.Run{
+		State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated, ResultState: jobs.RunResultStateFailed},
+		Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStatePending}}},
+	}
+	assert.Equal(t, "FAILED", detailedDisplayRunStatus(failed, "stale status..."))
+}
+
 func TestTerminationReason(t *testing.T) {
 	reason := terminationReason(&jobs.Run{
 		State:  &jobs.RunState{ResultState: jobs.RunResultStateFailed, StateMessage: "parent reason"},

experimental/air/cmd/get.go

@@ -161,7 +161,8 @@ func newGetCommand() *cobra.Command {
 
 		data := buildGetData(run)
 		data.DashboardURL = dashboardURL(w.Config.Host, runID, workspaceID)
-		ids := mlflowIDs(ctx, w, run)
+		taskOutput := aiRuntimeTaskOutput(ctx, w, run)
+		ids := mlflowIDsFromOutput(taskOutput)
 		if ids != nil {
 			url := mlflowLogsURL(w.Config.Host, ids)
 			data.MLflowURL = &url
@@ -194,6 +195,11 @@ func newGetCommand() *cobra.Command {
 			fmt.Fprintf(out, "Job Link: %s\n\n", hyperlink(ctx, out, data.DashboardURL, data.DashboardURL))
 			return renderEnvelope(ctx, data)
 		}
+		statusMessage := ""
+		if taskOutput != nil {
+			statusMessage = normalizeStatusMessage(taskOutput.StatusMessage)
+		}
+		data.DisplayStatus = detailedDisplayRunStatus(run, statusMessage)
 
 		renderRunText(ctx, out, w, run, &data, ids)
 		return nil

experimental/air/cmd/logs.go

@@ -105,7 +105,7 @@ func newLogsCommand() *cobra.Command {
 			tailLines = lines
 		}
 
-		return runLogs(ctx, cmd, logRequest{
+		err = runLogs(ctx, cmd, logRequest{
 			runID:         runID,
 			node:          node,
 			nodeSet:       cmd.Flags().Changed("node"),
@@ -115,6 +115,10 @@ func newLogsCommand() *cobra.Command {
 			downloadTo:    downloadTo,
 			jsonOutput:    root.OutputType(cmd) == flags.OutputJSON,
 		})
+		if downloadTo != "" || root.OutputType(cmd) == flags.OutputJSON {
+			return err
+		}
+		return handleWatchResult(cmd.OutOrStdout(), cmdctx.WorkspaceClient(ctx).Config.Profile, args[0], err)
 	}
 
 	return cmd

experimental/air/cmd/logs_test.go

@@ -2,11 +2,13 @@ package aircmd
 
 import (
 	"bytes"
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/flags"
@@ -194,6 +196,37 @@ func TestLogsFallsBackToMLflow(t *testing.T) {
 	assert.Equal(t, "line one\nline two\n", buf.String())
 }
 
+func TestLogsCommandPrintsGuidanceWhenStreamingIsInterrupted(t *testing.T) {
+	ctx, cancel := context.WithCancel(t.Context())
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case r.URL.Path == "/api/2.2/jobs/runs/get":
+			_, _ = w.Write([]byte(`{"run_id":5,"state":{"life_cycle_state":"RUNNING"},"tasks":[{"run_id":456}]}`))
+		case strings.HasSuffix(r.URL.Path, "/logs"):
+			cancel()
+			_, _ = w.Write([]byte(`{"log_records":[]}`))
+		default:
+			_, _ = w.Write([]byte(`{"userName":"u@example.com"}`))
+		}
+	}))
+	t.Cleanup(srv.Close)
+
+	var buf bytes.Buffer
+	client := newTestWorkspaceClient(t, srv.URL)
+	client.Config.Profile = "team profile"
+	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, nil, &buf, &buf, "", ""))
+	ctx = cmdctx.SetWorkspaceClient(ctx, client)
+	cmd := withOutput(newLogsCommand(), flags.OutputText)
+	cmd.SetContext(ctx)
+	cmd.SetOut(&buf)
+
+	err := cmd.RunE(cmd, []string{"5"})
+	require.ErrorIs(t, err, root.ErrAlreadyPrinted)
+	assert.Contains(t, buf.String(), "Streaming logs interrupted.")
+	assert.Contains(t, buf.String(), "To check status:\ndatabricks experimental air get 5 -p 'team profile'")
+	assert.Contains(t, buf.String(), "To resume streaming logs:\ndatabricks experimental air logs 5 -p 'team profile'")
+}
+
 // activeRunPastRetryServer serves a still-RUNNING run with two attempts and a
 // single page of Bricklens logs. runs/get always returns RUNNING; a test that
 // follows the run would poll forever, so it also asserts the static path never

experimental/air/cmd/logstream.go

@@ -58,6 +58,9 @@ func normalizeStatusMessage(raw string) string {
 	if payload == "" {
 		return ""
 	}
+	if strings.EqualFold(payload, "Waiting for GPU compute capacity to become available") {
+		return waitingForComputeStatus
+	}
 	return payload + "..."
 }
 
@@ -104,11 +107,12 @@ type logRequest struct {
 // logRunStatus is the subset of a run's state the log path needs, resolved once
 // and reused.
 type logRunStatus struct {
-	lifeCycleState string
-	resultState    string
-	stateMessage   string
-	startTimeMs    int64
-	endTimeMs      int64
+	lifeCycleState          string
+	firstTaskLifeCycleState string
+	resultState             string
+	stateMessage            string
+	startTimeMs             int64
+	endTimeMs               int64
 	// latestAttempt is the highest attempt_number across the run's tasks.
 	latestAttempt int
 }
@@ -128,6 +132,25 @@ func (s logRunStatus) succeeded() bool {
 	return s.resultState == "SUCCESS"
 }
 
+func (s logRunStatus) waitingForCompute() bool {
+	if s.resultState != "" {
+		return false
+	}
+	if s.lifeCycleState == string(jobs.RunLifeCycleStatePending) {
+		return true
+	}
+	if s.lifeCycleState != string(jobs.RunLifeCycleStateRunning) {
+		return false
+	}
+	switch jobs.RunLifeCycleState(s.firstTaskLifeCycleState) {
+	case jobs.RunLifeCycleStatePending, jobs.RunLifeCycleStateQueued,
+		jobs.RunLifeCycleStateWaitingForRetry, jobs.RunLifeCycleStateBlocked:
+		return true
+	default:
+		return false
+	}
+}
+
 // downloadOutcome is the exit status for a one-shot fetch, which unlike streaming
 // can run against an active run. An active run has no result state yet, and
 // not-yet-finished is not a failure, so only a terminal run decides the exit code.
@@ -157,6 +180,9 @@ func projectRunStatus(run *jobs.Run) logRunStatus {
 		s.resultState = string(run.State.ResultState)
 		s.stateMessage = run.State.StateMessage
 	}
+	if len(run.Tasks) > 0 && run.Tasks[0].State != nil {
+		s.firstTaskLifeCycleState = string(run.Tasks[0].State.LifeCycleState)
+	}
 	for i := range run.Tasks {
 		s.latestAttempt = max(s.latestAttempt, run.Tasks[i].AttemptNumber)
 	}
@@ -260,7 +286,7 @@ func (st *bricklensStreamer) waitingSpinnerText() string {
 	if msg := st.serverStatusMessage(); msg != "" {
 		return msg
 	}
-	if st.status.lifeCycleState == "PENDING" {
+	if st.status.waitingForCompute() {
 		return waitingForComputeStatus
 	}
 	return fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node)

experimental/air/cmd/logstream_test.go

@@ -82,7 +82,7 @@ func TestProjectRunStatus(t *testing.T) {
 			StateMessage:   "done",
 		},
 		Tasks: []jobs.RunTask{
-			{AttemptNumber: 0},
+			{AttemptNumber: 0, State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateQueued}},
 			{AttemptNumber: 2},
 			{AttemptNumber: 1},
 		},
@@ -92,6 +92,7 @@ func TestProjectRunStatus(t *testing.T) {
 	assert.Equal(t, "TERMINATED", s.lifeCycleState)
 	assert.Equal(t, "SUCCESS", s.resultState)
 	assert.Equal(t, "done", s.stateMessage)
+	assert.Equal(t, "QUEUED", s.firstTaskLifeCycleState)
 	assert.Equal(t, int64(1000), s.startTimeMs)
 	assert.Equal(t, int64(2000), s.endTimeMs)
 	assert.Equal(t, 2, s.latestAttempt)
@@ -213,6 +214,7 @@ func TestNormalizeStatusMessage(t *testing.T) {
 	}{
 		{"STATUS: Waiting for GPU capacity.", "Waiting for GPU capacity..."},
 		{"STATUS:Waiting for GPU capacity", "Waiting for GPU capacity..."},
+		{"STATUS: Waiting for GPU compute capacity to become available", waitingForComputeStatus},
 		{"status: provisioning", "provisioning..."}, // type match is case-insensitive
 		{"STATUS: done...", "done..."},              // trailing dots collapse to one "..."
 		{"INFO: not a status", ""},                  // other type ignored
@@ -228,7 +230,7 @@ func TestNormalizeStatusMessage(t *testing.T) {
 
 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 {
+	newStreamer := func(t *testing.T, statusMessage, lifeCycle, firstTaskLifeCycle string) *bricklensStreamer {
 		t.Helper()
 		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			switch r.URL.Path {
@@ -245,21 +247,30 @@ func TestWaitingSpinnerText(t *testing.T) {
 			ctx:    t.Context(),
 			w:      newTestWorkspaceClient(t, srv.URL),
 			req:    logRequest{runID: 1, node: 0},
-			status: logRunStatus{lifeCycleState: lifeCycle},
+			status: logRunStatus{lifeCycleState: lifeCycle, firstTaskLifeCycleState: firstTaskLifeCycle},
 		}
 	}
 
 	// Server STATUS message wins.
 	assert.Equal(t, "Waiting for GPU capacity...",
-		newStreamer(t, "STATUS: Waiting for GPU capacity", "PENDING").waitingSpinnerText())
+		newStreamer(t, "STATUS: Waiting for GPU capacity", "PENDING", "").waitingSpinnerText())
 
 	// No status message + PENDING -> compute-capacity fallback.
 	assert.Equal(t, waitingForComputeStatus,
-		newStreamer(t, "", "PENDING").waitingSpinnerText())
+		newStreamer(t, "", "PENDING", "").waitingSpinnerText())
+
+	for _, state := range []string{"PENDING", "QUEUED", "WAITING_FOR_RETRY", "BLOCKED"} {
+		assert.Equal(t, waitingForComputeStatus,
+			newStreamer(t, "", "RUNNING", state).waitingSpinnerText(), state)
+	}
 
 	// 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())
+		newStreamer(t, "", "RUNNING", "RUNNING").waitingSpinnerText())
+	assert.Equal(t, "Waiting for run to start (node 0)...",
+		newStreamer(t, "", "RUNNING", "").waitingSpinnerText())
+	assert.Equal(t, "Waiting for run to start (node 0)...",
+		newStreamer(t, "", "TERMINATED", "PENDING").waitingSpinnerText())
 }
 
 func TestEmitLogLineJSON(t *testing.T) {

experimental/air/cmd/mlflow.go

@@ -30,18 +30,25 @@ type mlflowIdentifiers struct {
 // mlflowIDs fetches the MLflow IDs for a run via its latest task. Returns nil if
 // they can't be obtained.
 func mlflowIDs(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run) *mlflowIdentifiers {
+	return mlflowIDsFromOutput(aiRuntimeTaskOutput(ctx, w, run))
+}
+
+func aiRuntimeTaskOutput(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run) *jobs.AiRuntimeTaskOutput {
 	if len(run.Tasks) == 0 {
 		return nil
 	}
-	// The MLflow output is attached to the task run, not the parent job run.
-	return mlflowIDsForTask(ctx, w, run.Tasks[len(run.Tasks)-1].RunId)
+	return aiRuntimeTaskOutputForTask(ctx, w, run.Tasks[len(run.Tasks)-1].RunId)
 }
 
 // mlflowIDsForTask fetches a task run's MLflow experiment and run IDs from
 // runs/get-output, or nil if they can't be obtained. They drive a convenience
 // link, so any failure (endpoint error, run not yet started, no MLflow output)
 // is logged and treated as "no link" rather than failing the command.
 func mlflowIDsForTask(ctx context.Context, w *databricks.WorkspaceClient, taskRunID int64) *mlflowIdentifiers {
+	return mlflowIDsFromOutput(aiRuntimeTaskOutputForTask(ctx, w, taskRunID))
+}
+
+func aiRuntimeTaskOutputForTask(ctx context.Context, w *databricks.WorkspaceClient, taskRunID int64) *jobs.AiRuntimeTaskOutput {
 	if taskRunID == 0 {
 		return nil
 	}
@@ -52,10 +59,14 @@ func mlflowIDsForTask(ctx context.Context, w *databricks.WorkspaceClient, taskRu
 		return nil
 	}
 
-	if o := out.AiRuntimeTaskOutput; o != nil && o.MlflowExperimentId != "" && o.MlflowRunId != "" {
-		return &mlflowIdentifiers{ExperimentID: o.MlflowExperimentId, RunID: o.MlflowRunId}
+	return out.AiRuntimeTaskOutput
+}
+
+func mlflowIDsFromOutput(output *jobs.AiRuntimeTaskOutput) *mlflowIdentifiers {
+	if output == nil || output.MlflowExperimentId == "" || output.MlflowRunId == "" {
+		return nil
 	}
-	return nil
+	return &mlflowIdentifiers{ExperimentID: output.MlflowExperimentId, RunID: output.MlflowRunId}
 }
 
 // mlflowLogsURL is the deep link to a run's node-0 logs. It is the value of the

experimental/air/cmd/mlflow_test.go

@@ -78,6 +78,17 @@ func TestMLflowIDs(t *testing.T) {
 	})
 }
 
+func TestAiRuntimeTaskOutput(t *testing.T) {
+	var hit bool
+	srv := runOutputServer(t, `{"ai_runtime_task_output":{"status_message":"STATUS: Waiting for GPU capacity"}}`, &hit)
+	run := &jobs.Run{Tasks: []jobs.RunTask{{RunId: 99}}}
+
+	got := aiRuntimeTaskOutput(t.Context(), newTestWorkspaceClient(t, srv.URL), run)
+	require.NotNil(t, got)
+	assert.True(t, hit)
+	assert.Equal(t, "STATUS: Waiting for GPU capacity", got.StatusMessage)
+}
+
 func TestMLflowIDsForTask(t *testing.T) {
 	ctx := t.Context()