diff --git a/.nextchanges/air/watch-exit-on-completion.md b/.nextchanges/air/watch-exit-on-completion.md new file mode 100644 index 00000000000..879b13e4475 --- /dev/null +++ b/.nextchanges/air/watch-exit-on-completion.md @@ -0,0 +1 @@ +* `air run --watch` now exits as soon as the command succeeds on the node it streams, instead of waiting for the job run to be marked terminal. ([#7031](https://github.com/databricks/cli/pull/7031)) diff --git a/cmd/air/logcompletion.go b/cmd/air/logcompletion.go new file mode 100644 index 00000000000..19237ac0b1a --- /dev/null +++ b/cmd/air/logcompletion.go @@ -0,0 +1,53 @@ +package aircmd + +import "strings" + +// The remote entry script prints this block from its exit handler, as the last +// thing it does on each node: +// +// ================================ +// Execution Summary +// ================================ +// Total time: 32s +// User code execution time: 32s +// Completed at Fri Oct 9 09:46:18 PM UTC 2026 +// ================================ +// +// "Completed at" is printed only for exit code 0; a failure prints "ERROR: Script +// failed ..." instead, which is not treated as completion because the run may be +// retried. Jobs reports the run terminal tens of seconds after the block, so +// `air run --watch` uses it to stop as soon as the streamed node has succeeded. +const ( + summaryHeader = "Execution Summary" + summaryRulePrefix = "====" + summaryTotalPrefix = "Total time: " + summaryUserCodePrefix = "User code execution time: " + summarySuccessPrefix = "Completed at " +) + +// completionMatcher recognizes the entry script's success summary in one node's +// log lines, fed oldest-first across polls. Only a consecutive block matches, so +// user output that happens to contain one of the lines does not. +type completionMatcher struct { + step int + completed bool +} + +func (m *completionMatcher) observe(line string) { + line = strings.TrimSpace(line) + switch { + case line == summaryHeader: + m.step = 1 + case m.step == 1 && strings.HasPrefix(line, summaryRulePrefix): + m.step = 2 + case m.step == 2 && strings.HasPrefix(line, summaryTotalPrefix): + m.step = 3 + case m.step == 3 && strings.HasPrefix(line, summaryUserCodePrefix): + // Optional line between the total time and the outcome. + case m.step == 3 && strings.HasPrefix(line, summarySuccessPrefix): + m.completed = true + m.step = 0 + default: + m.step = 0 + } +} diff --git a/cmd/air/logcompletion_test.go b/cmd/air/logcompletion_test.go new file mode 100644 index 00000000000..112936763e7 --- /dev/null +++ b/cmd/air/logcompletion_test.go @@ -0,0 +1,219 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "slices" + "strconv" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +var successSummary = []string{ + "================================", + "Execution Summary", + "================================", + "Total time: 32s", + "User code execution time: 32s", + "Completed at Fri Oct 9 09:46:18 PM UTC 2026", + "================================", +} + +var failureSummary = []string{ + "================================", + "Execution Summary", + "================================", + "Total time: 32s", + "User code execution time: 32s", + "ERROR: Script failed with exit code 1 after 32s", + "================================", +} + +func TestCompletionMatcher(t *testing.T) { + tests := []struct { + name string + lines []string + want bool + }{ + {name: "success summary", lines: successSummary, want: true}, + { + name: "success summary without user code time", + lines: slices.Delete(slices.Clone(successSummary), 4, 5), + want: true, + }, + {name: "failure summary", lines: failureSummary, want: false}, + { + name: "trailing whitespace", + lines: []string{"Execution Summary\r", "====\r", "Total time: 1s", "Completed at now\r"}, + want: true, + }, + { + name: "interrupted block", + lines: []string{"Execution Summary", "====", "step 5", "Total time: 1s", "Completed at now"}, + want: false, + }, + {name: "completed line alone", lines: []string{"Completed at epoch 3"}, want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var m completionMatcher + for _, line := range tt.lines { + m.observe(line) + } + assert.Equal(t, tt.want, m.completed) + }) + } +} + +// logsBody renders a log endpoint response for the streamed node. Records are +// timestamped one second apart from 1700000000s and keyed by position and +// content, so a poll that re-serves earlier lines dedups them while new lines, +// including the summary's repeated rules, stay distinct. +func logsBody(t *testing.T, lines []string) string { + t.Helper() + type record struct { + RecordID string `json:"record_id"` + TimeUnixNano int64 `json:"time_unix_nano"` + Body string `json:"body"` + } + records := make([]record, len(lines)) + for i, line := range lines { + records[i] = record{ + RecordID: strconv.Itoa(i) + ":" + line, + TimeUnixNano: (1_700_000_000 + int64(i)) * 1_000_000_000, + Body: line, + } + } + b, err := json.Marshal(map[string]any{"log_records": records}) + assert.NoError(t, err) + return string(b) +} + +const ( + runningRunBody = `{"run_id":1,"start_time":1700000000000,"state":{"life_cycle_state":"RUNNING"},"tasks":[{"run_id":2,"attempt_number":0}]}` + succeededRunBody = `{"run_id":1,"start_time":1700000000000,"end_time":1700000100000,"state":{"life_cycle_state":"TERMINATED","result_state":"SUCCESS"},"tasks":[{"run_id":2,"attempt_number":0}]}` + failedRunBody = `{"run_id":1,"start_time":1700000000000,"end_time":1700000100000,"state":{"life_cycle_state":"TERMINATED","result_state":"FAILED"},"tasks":[{"run_id":2,"attempt_number":0}]}` + retryingRunBody = `{"run_id":1,"start_time":1700000000000,"state":{"life_cycle_state":"RUNNING"},"tasks":[{"run_id":2,"attempt_number":0},{"run_id":3,"attempt_number":1}]}` +) + +// completionServer counts the requests the streamer made to the fake server. +type completionServer struct { + runGets int + logRequests int +} + +// newCompletionServer serves an active run. runBodies are returned by successive +// runs/get calls (the last repeats); lines returns node 0's log lines for the +// n-th log request (1-based). +func newCompletionServer(t *testing.T, runBodies []string, lines func(n int) []string) (*completionServer, *httptest.Server) { + t.Helper() + original := bricklensPollInterval + bricklensPollInterval = time.Millisecond + t.Cleanup(func() { bricklensPollInterval = original }) + + cs := &completionServer{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/logs"): + assert.Equal(t, "0", r.URL.Query().Get("filter.node_index")) + cs.logRequests++ + _, _ = w.Write([]byte(logsBody(t, lines(cs.logRequests)))) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runBodies[min(cs.runGets, len(runBodies)-1)])) + cs.runGets++ + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return cs, srv +} + +func streamForCompletion(t *testing.T, srvURL string, exitOnCompletion bool) (bool, string, error) { + t.Helper() + // Every scenario ends within a few polls; fail rather than hang if one does not. + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + var buf bytes.Buffer + ok, err := streamBricklensLogs(ctx, newTestWorkspaceClient(t, srvURL), &buf, + logRequest{runID: 1, attempt: -1, tailLines: -1, jsonOutput: true, exitOnCompletion: exitOnCompletion}, + logRunStatus{lifeCycleState: "RUNNING", startTimeMs: 1_700_000_000_000}) + return ok, buf.String(), err +} + +func TestStreamBricklensExitsOnSuccessSummary(t *testing.T) { + cs, srv := newCompletionServer(t, []string{runningRunBody}, func(int) []string { + return append([]string{"training"}, successSummary...) + }) + + ok, out, err := streamForCompletion(t, srv.URL, true) + require.NoError(t, err) + assert.True(t, ok) + // Exited on the first poll, without Jobs ever reporting the run terminal. + assert.Equal(t, 0, cs.runGets) + assert.Equal(t, 1, cs.logRequests) + assert.Contains(t, out, `"line":"training"`) + assert.Contains(t, out, `"line":"Completed at `) +} + +func TestStreamBricklensMatchesSummarySplitAcrossPolls(t *testing.T) { + cs, srv := newCompletionServer(t, []string{runningRunBody}, func(n int) []string { + lines := append([]string{"training"}, successSummary...) + if n == 1 { + return lines[:4] + } + return lines + }) + + ok, _, err := streamForCompletion(t, srv.URL, true) + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, 2, cs.logRequests) + assert.Equal(t, 1, cs.runGets) +} + +func TestStreamBricklensIgnoresFailureSummary(t *testing.T) { + cs, srv := newCompletionServer(t, []string{runningRunBody, failedRunBody}, func(int) []string { + return failureSummary + }) + + ok, _, err := streamForCompletion(t, srv.URL, true) + require.NoError(t, err) + // Waited for Jobs: a failed attempt may still be retried. + assert.False(t, ok) + assert.GreaterOrEqual(t, cs.runGets, 2) +} + +func TestStreamBricklensExitsOnRetrySuccess(t *testing.T) { + // Attempt 0 fails, so the stream keeps following the run; the retry then + // succeeds and the stream exits without waiting for Jobs. + cs, srv := newCompletionServer(t, []string{retryingRunBody}, func(n int) []string { + if n == 1 { + return failureSummary + } + return append([]string{"attempt 1"}, successSummary...) + }) + + ok, _, err := streamForCompletion(t, srv.URL, true) + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, 2, cs.logRequests) +} + +func TestStreamBricklensCompletionRequiresOptIn(t *testing.T) { + cs, srv := newCompletionServer(t, []string{runningRunBody, succeededRunBody}, func(int) []string { + return successSummary + }) + + ok, _, err := streamForCompletion(t, srv.URL, false) + require.NoError(t, err) + assert.True(t, ok) + assert.GreaterOrEqual(t, cs.runGets, 2) +} diff --git a/cmd/air/logstream.go b/cmd/air/logstream.go index 11c03d86d47..7870175a2ae 100644 --- a/cmd/air/logstream.go +++ b/cmd/air/logstream.go @@ -114,6 +114,10 @@ type logRequest struct { // `air run --watch -o json` to emit STATUS events. onStatusChange func(current, previous string) retryTracker *retryTracker + // exitOnCompletion ends the stream once the streamed node prints the entry + // script's success summary, without waiting for Jobs to report the run + // terminal. Used by `air run --watch`. + exitOnCompletion bool } // logRunStatus is the subset of a run's state the log path needs, resolved once @@ -355,6 +359,8 @@ type bricklensStreamer struct { onFirstLog func() // updateSpinner, when set, refreshes the waiting-spinner text each poll. updateSpinner func(string) + // completion watches the streamed node's new lines for the success summary. + completion completionMatcher } // waitingSpinnerText returns the waiting-spinner text: the server-set STATUS @@ -518,6 +524,10 @@ func (st *bricklensStreamer) run() (bool, error) { log.Infof(st.ctx, "air logs: run %d finished in state %s", st.req.runID, st.status.displayState()) return st.status.succeeded(), nil } + if st.req.exitOnCompletion && st.completion.completed { + log.Infof(st.ctx, "air logs: run %d succeeded on node %d; not waiting for Jobs to report it terminal", st.req.runID, st.req.node) + return true, nil + } if !st.firstLogSeen { emptyStreamPolls++ if emptyStreamPolls%bricklensEmptyMLflowProbeEveryNPolls == 0 && mlflowLogsExist(st.ctx, st.w, st.req) { @@ -634,6 +644,7 @@ func (st *bricklensStreamer) drainPages(toSec int64, suppressInitialHistory bool if st.seen.has(rec) { continue } + st.completion.observe(rec.Body) if !suppressInitialHistory || (nano != 0 && nano > initialTailNano) { st.emit(rec.Body) emitted++ diff --git a/cmd/air/run.go b/cmd/air/run.go index c8027b09e9a..8ff2bcc5c4d 100644 --- a/cmd/air/run.go +++ b/cmd/air/run.go @@ -202,15 +202,18 @@ The path must be a separate argument: cobra reserves -h as a boolean, so // --watch: stream the submitted run's logs until it reaches a terminal // state, then exit with the run's outcome. This is the same pipeline as - // `air logs ` (Bricklens with MLflow fallback). + // `air logs ` (Bricklens with MLflow fallback). It exits early once + // the streamed node's entry script reports success, since Jobs lags the + // user command's exit by tens of seconds. maxRetries := cfg.maxRetries() req := logRequest{ - runID: runID, - attempt: -1, - tailLines: -1, - jsonOutput: jsonOut, - maxRetries: &maxRetries, - retryTracker: newRetryTracker(logRunStatus{}), + runID: runID, + attempt: -1, + tailLines: -1, + jsonOutput: jsonOut, + maxRetries: &maxRetries, + retryTracker: newRetryTracker(logRunStatus{}), + exitOnCompletion: true, } watchCtx, stop := notifyInterrupt(ctx) @@ -258,7 +261,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so // the Bricklens path, so the terminal status must come from the run's // actual state — correct whether Bricklens or the MLflow fallback served // the logs. - printTerminalEvent(out, runIDStr, watchTerminalStatus(ctx, w, runID), dashboardURL) + printTerminalEvent(out, runIDStr, watchTerminalStatus(ctx, w, runID, err), dashboardURL) return err } @@ -360,12 +363,17 @@ func printLogsDivider(ctx context.Context, out io.Writer) { } // watchTerminalStatus resolves a watched run's final display state for the -// closing --watch envelope. The run is terminal once streaming returns; if the -// status can't be re-fetched, "UNKNOWN" is reported rather than guessing. -func watchTerminalStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) string { +// closing --watch envelope. If the status can't be re-fetched, "UNKNOWN" is +// reported rather than guessing. +func watchTerminalStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64, streamErr error) string { status, err := resolveRunStatus(ctx, w, runID) if err != nil { return "UNKNOWN" } + // A stream that returned success on a still-active run stopped early on the + // streamed node's success summary; Jobs reports the run terminal only later. + if streamErr == nil && !status.terminal() { + return "SUCCESS" + } return status.displayState() } diff --git a/cmd/air/run_watch_test.go b/cmd/air/run_watch_test.go index e1843e78ed1..c21b51ec4f2 100644 --- a/cmd/air/run_watch_test.go +++ b/cmd/air/run_watch_test.go @@ -185,6 +185,45 @@ func TestRunWatchFailedRunExitsNonZero(t *testing.T) { assert.Contains(t, buf.String(), "step 1\nstep 2") } +// watchServerCompleted serves a run that Jobs still reports RUNNING but whose +// single node has printed the entry script's success summary. +func watchServerCompleted(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/jobs/runs/submit"): + _, _ = w.Write([]byte(`{"run_id": 777}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runningRunBody)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + _, _ = w.Write([]byte(logsBody(t, append([]string{"step 1"}, successSummary...)))) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestRunWatchExitsOnCompletionSummary(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputText, &buf, watchServerCompleted(t).URL) + require.NoError(t, err) + assert.Contains(t, buf.String(), "step 1\n") + assert.Contains(t, buf.String(), "Completed at ") +} + +func TestRunWatchJSONCompletionSummaryTerminalEnvelope(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServerCompleted(t).URL) + require.NoError(t, err) + // Jobs still reports RUNNING, but the closing envelope reflects the + // completed user command. + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + assert.Contains(t, lines[len(lines)-1], `"status":"SUCCESS"`) + assert.Contains(t, lines[len(lines)-1], `"run_id":"777"`) +} + func TestHandleWatchResultPrintsProfileAwareResumeCommand(t *testing.T) { var buf bytes.Buffer err := handleWatchResult(&buf, "team profile", "777", fmt.Errorf("stream logs: %w", context.Canceled))