diff --git a/.nextchanges/cli/air-status-parity.md b/.nextchanges/cli/air-status-parity.md new file mode 100644 index 00000000000..7afabc29af9 --- /dev/null +++ b/.nextchanges/cli/air-status-parity.md @@ -0,0 +1 @@ +Added AIR failure reasons, timeout handling, and pending task display status parity. diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index 1dfb7e11847..fbcc7b44aa5 100644 --- a/experimental/air/cmd/format.go +++ b/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 { diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go index 1063e20ca1f..3df5c799948 100644 --- a/experimental/air/cmd/format_test.go +++ b/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"}})) +} diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 94af41f35c4..3fb32887611 100644 --- a/experimental/air/cmd/get.go +++ b/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 diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go index aa23e58159e..dfd9fbb38de 100644 --- a/experimental/air/cmd/get_test.go +++ b/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{{ diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 7fc0839ea76..e0c3c3da483 100644 --- a/experimental/air/cmd/list.go +++ b/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`. diff --git a/experimental/air/cmd/list_format.go b/experimental/air/cmd/list_format.go index 393b48a768c..8f5776134f2 100644 --- a/experimental/air/cmd/list_format.go +++ b/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, } } diff --git a/experimental/air/cmd/list_tui_render.go b/experimental/air/cmd/list_tui_render.go index e8fcc4fa476..3f23535a8e2 100644 --- a/experimental/air/cmd/list_tui_render.go +++ b/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. diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index 51e7e54a768..9d5c736c35e 100644 --- a/experimental/air/cmd/logstream.go +++ b/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 { diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go index be4a9d0c452..07fb16f9390 100644 --- a/experimental/air/cmd/logstream_test.go +++ b/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 { diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go index 44ed7eedd55..6bcb82f5bf7 100644 --- a/experimental/air/cmd/render.go +++ b/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.