Add AIR status and termination parity

.nextchanges/cli/air-status-parity.md

@@ -0,0 +1 @@
+Added AIR failure reasons, timeout handling, and pending task display status parity.

experimental/air/cmd/format.go

@@ -102,6 +102,43 @@ func runStatus(state *jobs.RunState) string {
 	return statusWord(string(state.LifeCycleState), string(state.ResultState))
 }
 
+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
+	}
+	switch run.Tasks[0].State.LifeCycleState {
+	case jobs.RunLifeCycleStatePending, jobs.RunLifeCycleStateQueued,
+		jobs.RunLifeCycleStateWaitingForRetry, jobs.RunLifeCycleStateBlocked:
+		return "PENDING"
+	default:
+		return status
+	}
+}
+
+func terminationReason(run *jobs.Run) *string {
+	status := runStatus(run.State)
+	if status != "FAILED" && status != "TIMEDOUT" && status != "INTERNAL_ERROR" {
+		return nil
+	}
+	if run.Status != nil && run.Status.TerminationDetails != nil {
+		if message := strings.TrimSpace(run.Status.TerminationDetails.Message); message != "" {
+			return &message
+		}
+	}
+	if run.State != nil {
+		if message := strings.TrimSpace(run.State.StateMessage); message != "" {
+			return &message
+		}
+	}
+	if len(run.Tasks) > 0 && run.Tasks[0].State != nil {
+		if message := strings.TrimSpace(run.Tasks[0].State.StateMessage); message != "" {
+			return &message
+		}
+	}
+	return nil
+}
+
 // statusWord picks the status word to show from a run's lifecycle and result
 // states: the result state is the more meaningful one, so it wins when set.
 func statusWord(lifeCycle, result string) string {

experimental/air/cmd/format_test.go

@@ -235,3 +235,48 @@ func TestStatusWord(t *testing.T) {
 	assert.Equal(t, "RUNNING", statusWord("RUNNING", ""))           // falls back to lifecycle
 	assert.Equal(t, "UNKNOWN", statusWord("", ""))
 }
+
+func TestDisplayRunStatus(t *testing.T) {
+	pendingStates := []jobs.RunLifeCycleState{
+		jobs.RunLifeCycleStatePending,
+		jobs.RunLifeCycleStateQueued,
+		jobs.RunLifeCycleStateWaitingForRetry,
+		jobs.RunLifeCycleStateBlocked,
+	}
+	for _, state := range pendingStates {
+		run := &jobs.Run{
+			State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning},
+			Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: state}}},
+		}
+		assert.Equal(t, "PENDING", displayRunStatus(run))
+	}
+
+	assert.Equal(t, "RUNNING", displayRunStatus(&jobs.Run{
+		State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning},
+		Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}}},
+	}))
+	assert.Equal(t, "RUNNING", displayRunStatus(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}}))
+	assert.Equal(t, "FAILED", displayRunStatus(&jobs.Run{
+		State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated, ResultState: jobs.RunResultStateFailed},
+		Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStatePending}}},
+	}))
+}
+
+func TestTerminationReason(t *testing.T) {
+	reason := terminationReason(&jobs.Run{
+		State:  &jobs.RunState{ResultState: jobs.RunResultStateFailed, StateMessage: "parent reason"},
+		Status: &jobs.RunStatus{TerminationDetails: &jobs.TerminationDetails{Message: " detailed reason "}},
+	})
+	require.NotNil(t, reason)
+	assert.Equal(t, "detailed reason", *reason)
+
+	reason = terminationReason(&jobs.Run{
+		State: &jobs.RunState{ResultState: jobs.RunResultStateTimedout},
+		Tasks: []jobs.RunTask{{State: &jobs.RunState{StateMessage: "task timed out"}}},
+	})
+	require.NotNil(t, reason)
+	assert.Equal(t, "task timed out", *reason)
+
+	assert.Nil(t, terminationReason(&jobs.Run{State: &jobs.RunState{ResultState: jobs.RunResultStateSuccess, StateMessage: "done"}}))
+	assert.Nil(t, terminationReason(&jobs.Run{State: &jobs.RunState{ResultState: jobs.RunResultStateCanceled, StateMessage: "cancelled"}}))
+}

experimental/air/cmd/get.go

@@ -21,14 +21,15 @@ import (
 // the machine-readable output; fields tagged `json:"-"` are shown only in the
 // human-readable text view.
 type getData struct {
-	RunID           string  `json:"run_id"`
-	Status          string  `json:"status"`
-	StartedAt       *string `json:"started_at"`
-	DurationSeconds *int64  `json:"duration_seconds"`
-	AttemptNumber   int     `json:"attempt_number"`
-	ExperimentName  *string `json:"experiment_name"`
-	DashboardURL    string  `json:"dashboard_url"`
-	MLflowURL       *string `json:"mlflow_url"`
+	RunID             string  `json:"run_id"`
+	Status            string  `json:"status"`
+	StartedAt         *string `json:"started_at"`
+	DurationSeconds   *int64  `json:"duration_seconds"`
+	AttemptNumber     int     `json:"attempt_number"`
+	ExperimentName    *string `json:"experiment_name"`
+	DashboardURL      string  `json:"dashboard_url"`
+	MLflowURL         *string `json:"mlflow_url"`
+	TerminationReason *string `json:"termination_reason,omitempty"`
 	// ESTRemainingSeconds and ESTPercentComplete are a best-effort progress
 	// estimate for a running run, or null when one can't be made (see
 	// estimateTrainingETA).
@@ -50,6 +51,7 @@ type getData struct {
 	// ProgressDisplay is the pre-rendered "Progress" cell ("45% · ~2h 15m left"),
 	// set only for a running run with an estimable remaining time.
 	ProgressDisplay string `json:"-"`
+	DisplayStatus   string `json:"-"`
 	// TrainingConfigPath is the run's config file, downloaded for the config box.
 	TrainingConfigPath string `json:"-"`
 	// Sweep replaces the single-run view for foreach runs.
@@ -61,7 +63,7 @@ type getData struct {
 // used only when .Data.Sweep is set. It reads from the JSON envelope, so every
 // field is reached through ".Data".
 const getTemplate = `Sweep Run ID: {{.Data.RunID}}
-Status:       {{.Data.Status}}
+Status:       {{if .Data.DisplayStatus}}{{.Data.DisplayStatus}}{{else}}{{.Data.Status}}{{end}}
 Total:        {{.Data.Sweep.Total}}
 Completed:    {{.Data.Sweep.Completed}}
 Succeeded:    {{.Data.Sweep.Succeeded}}
@@ -237,12 +239,14 @@ func aiRuntimeTaskOf(run *jobs.Run) *jobs.AiRuntimeTask {
 // hyperlinks and colors once the dashboard and MLflow identifiers are known.
 func buildGetData(run *jobs.Run) getData {
 	data := getData{
-		RunID:           strconv.FormatInt(run.RunId, 10),
-		Status:          runStatus(run.State),
-		StartedAt:       startedAt(run),
-		DurationSeconds: durationSeconds(run),
-		AttemptNumber:   latestAttemptNumber(run),
-		ExperimentName:  experimentName(run),
+		RunID:             strconv.FormatInt(run.RunId, 10),
+		Status:            runStatus(run.State),
+		DisplayStatus:     displayRunStatus(run),
+		TerminationReason: terminationReason(run),
+		StartedAt:         startedAt(run),
+		DurationSeconds:   durationSeconds(run),
+		AttemptNumber:     latestAttemptNumber(run),
+		ExperimentName:    experimentName(run),
 	}
 	data.SubmittedDisplay = submittedDisplay(run)
 	data.DurationDisplay = na

experimental/air/cmd/get_test.go

@@ -220,6 +220,38 @@ func TestBuildGetData(t *testing.T) {
 	assert.Equal(t, int64(12), *d.DurationSeconds)
 }
 
+func TestBuildGetDataTerminationReasonJSON(t *testing.T) {
+	failed := buildGetData(&jobs.Run{
+		RunId: 5,
+		State: &jobs.RunState{ResultState: jobs.RunResultStateFailed},
+		Status: &jobs.RunStatus{TerminationDetails: &jobs.TerminationDetails{
+			Message: "out of memory",
+		}},
+	})
+	body, err := json.Marshal(failed)
+	require.NoError(t, err)
+	assert.JSONEq(t, `{
+		"run_id":"5","status":"FAILED","started_at":null,"duration_seconds":null,
+		"attempt_number":0,"experiment_name":null,"dashboard_url":"","mlflow_url":null,
+		"termination_reason":"out of memory","est_remaining_seconds":null,"est_percent_complete":null
+	}`, string(body))
+
+	canceled := buildGetData(&jobs.Run{RunId: 6, State: &jobs.RunState{ResultState: jobs.RunResultStateCanceled, StateMessage: "user canceled"}})
+	body, err = json.Marshal(canceled)
+	require.NoError(t, err)
+	assert.NotContains(t, string(body), "termination_reason")
+}
+
+func TestGetTemplateUsesDisplayStatus(t *testing.T) {
+	out := renderGet(t, getData{
+		RunID:         "456",
+		Status:        "RUNNING",
+		DisplayStatus: "PENDING",
+		Sweep:         &sweepInfo{Total: 1, Active: 1},
+	})
+	assert.Contains(t, out, "Status:       PENDING")
+}
+
 func TestEnrichFromAiRuntimeTask(t *testing.T) {
 	t.Run("fills config path, experiment, and accelerators", func(t *testing.T) {
 		run := &jobs.Run{RunId: 5, Tasks: []jobs.RunTask{{

experimental/air/cmd/list.go

@@ -37,12 +37,13 @@ type listData struct {
 // machine-readable output; fields tagged `json:"-"` are shown only in the
 // human-readable table.
 type listRow struct {
-	RunID     string  `json:"run_id"`
-	RunName   string  `json:"run_name"`
-	User      string  `json:"user"`
-	Status    string  `json:"status"`
-	StartedAt *string `json:"started_at"`
-	IsSweep   bool    `json:"is_sweep"`
+	RunID         string  `json:"run_id"`
+	RunName       string  `json:"run_name"`
+	User          string  `json:"user"`
+	Status        string  `json:"status"`
+	DisplayStatus string  `json:"-"`
+	StartedAt     *string `json:"started_at"`
+	IsSweep       bool    `json:"is_sweep"`
 
 	// Experiment, Duration, Progress, MLflowURL and Accelerators are table-only
 	// columns, omitted from JSON to match `air list --json`.

experimental/air/cmd/list_format.go

@@ -34,17 +34,18 @@ func buildListRow(run *jobs.Run, host string, workspaceID int64) listRow {
 	}
 
 	return listRow{
-		RunID:        strconv.FormatInt(run.RunId, 10),
-		RunName:      run.RunName,
-		User:         run.CreatorUserName,
-		Status:       runStatus(run.State),
-		StartedAt:    startedAt,
-		IsSweep:      isSweep(run),
-		Experiment:   experiment,
-		Duration:     duration,
-		MLflowURL:    "-",
-		MLflowLabel:  "-",
-		RunURL:       dashboardURL(host, run.RunId, workspaceID),
-		Accelerators: accel,
+		RunID:         strconv.FormatInt(run.RunId, 10),
+		RunName:       run.RunName,
+		User:          run.CreatorUserName,
+		Status:        runStatus(run.State),
+		DisplayStatus: displayRunStatus(run),
+		StartedAt:     startedAt,
+		IsSweep:       isSweep(run),
+		Experiment:    experiment,
+		Duration:      duration,
+		MLflowURL:     "-",
+		MLflowLabel:   "-",
+		RunURL:        dashboardURL(host, run.RunId, workspaceID),
+		Accelerators:  accel,
 	}
 }

experimental/air/cmd/list_tui_render.go

@@ -58,7 +58,7 @@ func computeListCols(rows []listRow) listCols {
 	for _, r := range rows {
 		c.runID = max(c.runID, lipgloss.Width(r.RunID))
 		c.experiment = min(columnCap, max(c.experiment, lipgloss.Width(r.Experiment)))
-		c.status = max(c.status, lipgloss.Width("● "+r.Status))
+		c.status = max(c.status, lipgloss.Width("● "+listRowStatus(r)))
 		c.started = max(c.started, lipgloss.Width(startedDisplay(r)))
 		c.duration = max(c.duration, lipgloss.Width(r.Duration))
 		c.progress = max(c.progress, lipgloss.Width(progressDisplay(r)))
@@ -123,7 +123,7 @@ func (s listStyles) renderRow(cols listCols, r listRow, selected, links bool) st
 		s.cell(base, gutter, 1, fg(colN7), false, false, ""),
 		s.cell(base, r.RunID, cols.runID, fg(colRunID), false, runIDLink != "", runIDLink),
 		s.cell(base, r.Experiment, cols.experiment, fg(colN11), false, experimentLink != "", experimentLink),
-		s.cell(base, "● "+r.Status, cols.status, fg(statusColor(r.Status)), false, false, ""),
+		s.cell(base, "● "+listRowStatus(r), cols.status, fg(statusColor(listRowStatus(r))), false, false, ""),
 		s.cell(base, startedDisplay(r), cols.started, fg(colN9), false, false, ""),
 		s.cell(base, r.Duration, cols.duration, fg(colN9), true, false, ""),
 		s.cell(base, progressDisplay(r), cols.progress, fg(colAmber), true, false, ""),
@@ -135,6 +135,13 @@ func (s listStyles) renderRow(cols listCols, r listRow, selected, links bool) st
 	return strings.TrimRight(strings.Join(cells, base.Render("  ")), " ")
 }
 
+func listRowStatus(r listRow) string {
+	if r.DisplayStatus != "" {
+		return r.DisplayStatus
+	}
+	return r.Status
+}
+
 // cell renders one padded, colored cell. The text is truncated to width, then
 // padded with (background-only) spaces so columns align even when the text is
 // styled or hyperlinked.

experimental/air/cmd/logstream.go

@@ -114,7 +114,7 @@ type logRunStatus struct {
 // set (result states only appear on terminal runs).
 var (
 	terminalLifeCycleStates = map[string]bool{"TERMINATED": true, "SKIPPED": true, "INTERNAL_ERROR": true}
-	terminalResultStates    = map[string]bool{"SUCCESS": true, "FAILED": true, "CANCELED": true}
+	terminalResultStates    = map[string]bool{"SUCCESS": true, "FAILED": true, "TIMEDOUT": true, "CANCELED": true}
 )
 
 func (s logRunStatus) terminal() bool {

experimental/air/cmd/logstream_test.go

@@ -112,6 +112,7 @@ func TestLogRunStatusTerminal(t *testing.T) {
 		{"terminated lifecycle", "TERMINATED", "", true},
 		{"internal error lifecycle", "INTERNAL_ERROR", "", true},
 		{"failed result", "TERMINATING", "FAILED", true},
+		{"timed out result", "RUNNING", "TIMEDOUT", true},
 		{"canceled result", "RUNNING", "CANCELED", true},
 	}
 	for _, tt := range tests {

experimental/air/cmd/render.go

@@ -118,6 +118,9 @@ func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceCl
 		accelerators: data.AcceleratorsDisplay,
 		environment:  data.EnvironmentDisplay,
 	}
+	if data.DisplayStatus != "" {
+		view.status = data.DisplayStatus
+	}
 
 	if ids != nil {
 		view.mlflowLabel = mlflowRunLabel(fetchMLflowRunName(ctx, w, ids.RunID), ids.RunID)
@@ -129,6 +132,9 @@ func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceCl
 		sections = append(sections, renderBox(p, configBoxTitle, body))
 	}
 	sections = append(sections, renderBox(p, metadataBoxTitle, renderFields(p, colorOn, view)))
+	if data.TerminationReason != nil {
+		sections = append(sections, p.red.Render("Failure reason: ")+p.n12.Render(*data.TerminationReason))
+	}
 
 	// A single write: a blank line before the first box and after the last, and
 	// one between each box.