From f6390463b8319d01f25e7eb046e0ed57cdd78bc0 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Mon, 12 Oct 2026 16:18:06 +0000 Subject: [PATCH 1/4] Exit air run --watch when the command succeeds on every node Jobs reports an AI Runtime run terminal tens of seconds after the user command exits, so `air run --watch` kept running after the work was done. The stream now watches for the entry script's success summary. Once the streamed node prints it, every node's newest lines are checked for the same block, pinned to the current attempt, and the command exits 0 with a SUCCESS closing envelope. A failure summary, a missing summary, an unknown node count, or an API error keeps the existing wait for Jobs, since a failed attempt may be retried. Co-authored-by: Isaac --- .nextchanges/air/watch-exit-on-completion.md | 1 + cmd/air/logbricklens.go | 1 + cmd/air/logcompletion.go | 57 ++++ cmd/air/logcompletion_test.go | 295 +++++++++++++++++++ cmd/air/logstream.go | 61 ++++ cmd/air/run.go | 30 +- cmd/air/run_watch_test.go | 40 +++ 7 files changed, 474 insertions(+), 11 deletions(-) create mode 100644 .nextchanges/air/watch-exit-on-completion.md create mode 100644 cmd/air/logcompletion.go create mode 100644 cmd/air/logcompletion_test.go diff --git a/.nextchanges/air/watch-exit-on-completion.md b/.nextchanges/air/watch-exit-on-completion.md new file mode 100644 index 00000000000..b1895ae1bce --- /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 every node, instead of waiting for the job run to be marked terminal. diff --git a/cmd/air/logbricklens.go b/cmd/air/logbricklens.go index ca3b8170130..e230a8c5ca2 100644 --- a/cmd/air/logbricklens.go +++ b/cmd/air/logbricklens.go @@ -35,6 +35,7 @@ func (r logRecord) nano() int64 { type bricklensLogsResponse struct { LogRecords []logRecord `json:"log_records"` NextPageToken string `json:"next_page_token"` + NodeCount int `json:"node_count"` } // bricklensLogsQuery is the request-field surface of the log endpoint. diff --git a/cmd/air/logcompletion.go b/cmd/air/logcompletion.go new file mode 100644 index 00000000000..cc96ef7b999 --- /dev/null +++ b/cmd/air/logcompletion.go @@ -0,0 +1,57 @@ +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 every node has succeeded. +const ( + summaryHeader = "Execution Summary" + summaryRulePrefix = "====" + summaryTotalPrefix = "Total time: " + summaryUserCodePrefix = "User code execution time: " + summarySuccessPrefix = "Completed at " +) + +// completionTailLines is how many of a node's newest log lines are searched for +// the summary block. +const completionTailLines = 20 + +// completionMatcher recognizes the entry script's success summary in one node's +// log lines, fed oldest-first. 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..0054ee91bca --- /dev/null +++ b/cmd/air/logcompletion_test.go @@ -0,0 +1,295 @@ +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) + }) + } +} + +// withNode prefixes a node's training output to a summary. +func withNode(node int, summary []string) []string { + return append([]string{"training node " + strconv.Itoa(node)}, summary...) +} + +// logsBody renders a log endpoint response for one node's attempt. Records are +// timestamped one second apart from 1700000000s and keyed by attempt, position +// and content, so a poll that re-serves earlier lines dedups them while new or +// repeated lines (the summary's rules) stay distinct; newestFirst matches an +// ascending=false query. +func logsBody(t *testing.T, nodeCount, node int, attempt string, lines []string, newestFirst bool) string { + t.Helper() + type record struct { + RecordID string `json:"record_id"` + TimeUnixNano int64 `json:"time_unix_nano"` + Body string `json:"body"` + NodeIndex int `json:"node_index"` + } + records := make([]record, len(lines)) + for i, line := range lines { + records[i] = record{ + RecordID: strconv.Itoa(node) + ":" + attempt + ":" + strconv.Itoa(i) + ":" + line, + TimeUnixNano: (1_700_000_000 + int64(i)) * 1_000_000_000, + Body: line, + NodeIndex: node, + } + } + if newestFirst { + slices.Reverse(records) + } + b, err := json.Marshal(map[string]any{"log_records": records, "node_count": nodeCount}) + 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}]}` + retrySucceededRun = `{"run_id":1,"start_time":1700000000000,"end_time":1700000100000,"state":{"life_cycle_state":"TERMINATED","result_state":"SUCCESS"},"tasks":[{"run_id":2,"attempt_number":0},{"run_id":3,"attempt_number":1}]}` +) + +// logQuery is one log request the completion server received. attempt is the raw +// ref.attempt_number ("" for latest) and resolved the attempt it was served for; +// tail is true for a newest-first (completion check) query. +type logQuery struct { + node int + attempt string + resolved string + tail bool +} + +// completionServer serves an active run for the completion tests. runBodies are +// returned by successive runs/get calls (the last repeats). latestAttempt, when +// set, resolves a query without an attempt number, as the endpoint does; +// otherwise it resolves to attempt 0. +type completionServer struct { + runGets int + queries []logQuery + latestAttempt func() string +} + +// tailAttempts returns the raw attempt numbers of the completion checks for a node. +func (cs *completionServer) tailAttempts(node int) []string { + var out []string + for _, q := range cs.queries { + if q.node == node && q.tail { + out = append(out, q.attempt) + } + } + return out +} + +// newCompletionServer starts the server; nodeLines returns the lines for a log +// request, given the requests received so far including it. +func newCompletionServer(t *testing.T, nodeCount int, runBodies []string, nodeLines func(q logQuery, cs *completionServer) []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"): + node, err := strconv.Atoi(r.URL.Query().Get("filter.node_index")) + if !assert.NoError(t, err) { + return + } + q := logQuery{ + node: node, + attempt: r.URL.Query().Get("ref.attempt_number"), + resolved: r.URL.Query().Get("ref.attempt_number"), + tail: r.URL.Query().Get("ascending") == "false", + } + if q.resolved == "" { + q.resolved = "0" + if cs.latestAttempt != nil { + q.resolved = cs.latestAttempt() + } + } + cs.queries = append(cs.queries, q) + _, _ = w.Write([]byte(logsBody(t, nodeCount, node, q.resolved, nodeLines(q, cs), q.tail))) + 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 TestStreamBricklensExitsWhenAllNodesComplete(t *testing.T) { + cs, srv := newCompletionServer(t, 2, []string{runningRunBody}, func(q logQuery, _ *completionServer) []string { + return withNode(q.node, 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) + // Both nodes, the streamed one included, were checked for attempt 0. + assert.Equal(t, []string{"0"}, cs.tailAttempts(0)) + assert.Equal(t, []string{"0"}, cs.tailAttempts(1)) + assert.Contains(t, out, `"line":"training node 0"`) +} + +func TestStreamBricklensWaitsForEveryNode(t *testing.T) { + cs, srv := newCompletionServer(t, 2, []string{runningRunBody}, func(q logQuery, cs *completionServer) []string { + if q.node == 1 && len(cs.tailAttempts(1)) < 3 { + return []string{"still training"} + } + return successSummary + }) + + ok, _, err := streamForCompletion(t, srv.URL, true) + require.NoError(t, err) + assert.True(t, ok) + assert.Len(t, cs.tailAttempts(1), 3) +} + +func TestStreamBricklensIgnoresFailureSummary(t *testing.T) { + cs, srv := newCompletionServer(t, 1, []string{runningRunBody, failedRunBody}, func(q logQuery, _ *completionServer) []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) + assert.Empty(t, cs.tailAttempts(0)) +} + +func TestStreamBricklensCompletionRequiresOptIn(t *testing.T) { + cs, srv := newCompletionServer(t, 1, []string{runningRunBody, succeededRunBody}, func(q logQuery, _ *completionServer) []string { + return successSummary + }) + + ok, _, err := streamForCompletion(t, srv.URL, false) + require.NoError(t, err) + assert.True(t, ok) + assert.GreaterOrEqual(t, cs.runGets, 2) + assert.Empty(t, cs.tailAttempts(0)) +} + +func TestStreamBricklensBindsCompletionToLatestAttempt(t *testing.T) { + // Attempt 0 succeeds on node 0 but fails on node 1, and is retried. On attempt + // 1, node 0 succeeds while node 1 is still running when Jobs reports the run + // terminal. Every check must target the attempt that triggered it, so the + // stream ends on Jobs' terminal state. + attemptOneStreamPolls := 0 + cs, srv := newCompletionServer(t, 2, []string{retryingRunBody, retryingRunBody, retrySucceededRun}, func(q logQuery, _ *completionServer) []string { + switch { + case q.resolved == "0" && q.node == 0: + return successSummary + case q.resolved == "0" && q.node == 1: + return failureSummary + case q.node == 1: + return []string{"still training"} + case q.tail: + return successSummary + default: + // Node 0's stream on attempt 1: the summary arrives on the second poll. + attemptOneStreamPolls++ + if attemptOneStreamPolls >= 2 { + return append([]string{"attempt 1 node 0"}, successSummary...) + } + return []string{"attempt 1 node 0"} + } + }) + cs.latestAttempt = func() string { + if cs.runGets == 0 { + return "0" + } + return "1" + } + + ok, _, err := streamForCompletion(t, srv.URL, true) + require.NoError(t, err) + // Ended on Jobs' terminal SUCCESS, not on summaries from different attempts. + assert.True(t, ok) + assert.GreaterOrEqual(t, cs.runGets, 3) + assert.Equal(t, []string{"0", "1"}, cs.tailAttempts(1)) +} diff --git a/cmd/air/logstream.go b/cmd/air/logstream.go index 11c03d86d47..cefe07dc483 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 every node of the latest attempt has + // printed 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,14 @@ type bricklensStreamer struct { onFirstLog func() // updateSpinner, when set, refreshes the waiting-spinner text each poll. updateSpinner func(string) + // nodeCount is the run's node count as reported by the log endpoint; 0 until + // a response carries it. + nodeCount int + // completion watches the streamed node for the success summary of + // completionAttempt. It only triggers the per-attempt check in + // allNodesCompleted, and a retry starts a fresh matcher. + completion completionMatcher + completionAttempt int } // waitingSpinnerText returns the waiting-spinner text: the server-set STATUS @@ -460,6 +472,10 @@ func (st *bricklensStreamer) run() (bool, error) { st.reportStatusChange() st.req.reportRetry(st.ctx, st.out, st.status, st.onFirstLog) + if st.status.latestAttempt != st.completionAttempt { + st.completion = completionMatcher{} + st.completionAttempt = st.status.latestAttempt + } terminal := st.status.terminal() toSec := st.req.toSeconds(st.status) @@ -518,6 +534,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.allNodesCompleted() { + log.Infof(st.ctx, "air logs: run %d completed on all %d node(s); not waiting for Jobs to report it terminal", st.req.runID, st.nodeCount) + return true, nil + } if !st.firstLogSeen { emptyStreamPolls++ if emptyStreamPolls%bricklensEmptyMLflowProbeEveryNPolls == 0 && mlflowLogsExist(st.ctx, st.w, st.req) { @@ -533,6 +553,45 @@ func (st *bricklensStreamer) run() (bool, error) { } } +// allNodesCompleted reports whether every node of the run's latest attempt has +// printed the entry script's success summary. The streamed node's summary only +// triggers the check: every node, the streamed one included, is then read with an +// explicit attempt number, so all summaries come from the same attempt. Best +// effort: an unknown node count or a failed fetch reads as not completed, so the +// stream keeps following the run. +func (st *bricklensStreamer) allNodesCompleted() bool { + if !st.completion.completed || st.nodeCount < 1 { + return false + } + for node := range st.nodeCount { + if !st.nodeCompleted(node, st.completionAttempt) { + return false + } + } + return true +} + +// nodeCompleted searches the newest lines of a node's log for an attempt for the +// success summary. From the stream start rather than the live cursor, since the +// node may have finished before the lookback window. +func (st *bricklensStreamer) nodeCompleted(node, attempt int) bool { + resp, err := getBricklensLogs(st.ctx, st.apiClient, st.headers, st.req.runID, bricklensLogsQuery{ + fromSeconds: st.streamStartSec, + pageSize: completionTailLines, + attemptNumber: attempt, + nodeIndex: node, + }) + if err != nil { + log.Debugf(st.ctx, "air logs: completion check for node %d failed: %v", node, err) + return false + } + var m completionMatcher + for _, rec := range slices.Backward(resp.LogRecords) { + m.observe(rec.Body) + } + return m.completed +} + func waitForNextBricklensPoll(ctx context.Context, pollStarted time.Time) error { return sleepOrCancel(ctx, max(time.Duration(0), bricklensPollInterval-time.Since(pollStarted))) } @@ -634,6 +693,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++ @@ -678,6 +738,7 @@ func (st *bricklensStreamer) requestPage(pageToken string, toSec int64, pageSize for { resp, err := getBricklensLogs(st.ctx, st.apiClient, st.headers, st.req.runID, q) if err == nil { + st.nodeCount = max(st.nodeCount, resp.NodeCount) return resp, nil } diff --git a/cmd/air/run.go b/cmd/air/run.go index c8027b09e9a..7873e21b3cb 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 + // every 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 every + // 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..be71c21d92e 100644 --- a/cmd/air/run_watch_test.go +++ b/cmd/air/run_watch_test.go @@ -185,6 +185,46 @@ 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/"): + newestFirst := r.URL.Query().Get("ascending") == "false" + _, _ = w.Write([]byte(logsBody(t, 1, 0, "0", append([]string{"step 1"}, successSummary...), newestFirst))) + 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)) From 21ace16d3616dccab3db2793ce40e473c27eab43 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Mon, 12 Oct 2026 16:20:14 +0000 Subject: [PATCH 2/4] Link the changelog fragment to its PR Co-authored-by: Isaac --- .nextchanges/air/watch-exit-on-completion.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.nextchanges/air/watch-exit-on-completion.md b/.nextchanges/air/watch-exit-on-completion.md index b1895ae1bce..31615ecda6e 100644 --- a/.nextchanges/air/watch-exit-on-completion.md +++ b/.nextchanges/air/watch-exit-on-completion.md @@ -1 +1 @@ -* `air run --watch` now exits as soon as the command succeeds on every node, instead of waiting for the job run to be marked terminal. +* `air run --watch` now exits as soon as the command succeeds on every node, instead of waiting for the job run to be marked terminal. ([#7031](https://github.com/databricks/cli/pull/7031)) From 93a2fa27f1b4dc4e358b695e8ce42dc2de701534 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Mon, 12 Oct 2026 16:48:04 +0000 Subject: [PATCH 3/4] Resume the all-nodes completion check at the first unconfirmed node Once the streamed node's summary triggered the check, every poll re-fetched every node up to the first unfinished one, so a failing node late in a large run cost one request per node per second until Jobs reported the run terminal. A node that succeeded stays succeeded, so the check now resumes at the first unconfirmed node, and a retry resets that progress. Also pin with tests that a missing node count blocks the early exit and that each node is confirmed again on a new attempt. Co-authored-by: Isaac --- cmd/air/logcompletion_test.go | 24 ++++++++++++++++++------ cmd/air/logstream.go | 33 ++++++++++++++++++--------------- 2 files changed, 36 insertions(+), 21 deletions(-) diff --git a/cmd/air/logcompletion_test.go b/cmd/air/logcompletion_test.go index 0054ee91bca..9e6469fbc1a 100644 --- a/cmd/air/logcompletion_test.go +++ b/cmd/air/logcompletion_test.go @@ -72,11 +72,6 @@ func TestCompletionMatcher(t *testing.T) { } } -// withNode prefixes a node's training output to a summary. -func withNode(node int, summary []string) []string { - return append([]string{"training node " + strconv.Itoa(node)}, summary...) -} - // logsBody renders a log endpoint response for one node's attempt. Records are // timestamped one second apart from 1700000000s and keyed by attempt, position // and content, so a poll that re-serves earlier lines dedups them while new or @@ -201,7 +196,7 @@ func streamForCompletion(t *testing.T, srvURL string, exitOnCompletion bool) (bo func TestStreamBricklensExitsWhenAllNodesComplete(t *testing.T) { cs, srv := newCompletionServer(t, 2, []string{runningRunBody}, func(q logQuery, _ *completionServer) []string { - return withNode(q.node, successSummary) + return append([]string{"training node " + strconv.Itoa(q.node)}, successSummary...) }) ok, out, err := streamForCompletion(t, srv.URL, true) @@ -227,6 +222,21 @@ func TestStreamBricklensWaitsForEveryNode(t *testing.T) { require.NoError(t, err) assert.True(t, ok) assert.Len(t, cs.tailAttempts(1), 3) + // Node 0 was confirmed once and not re-fetched while node 1 caught up. + assert.Len(t, cs.tailAttempts(0), 1) +} + +func TestStreamBricklensCompletionRequiresNodeCount(t *testing.T) { + cs, srv := newCompletionServer(t, 0, []string{runningRunBody, succeededRunBody}, func(q logQuery, _ *completionServer) []string { + return successSummary + }) + + ok, _, err := streamForCompletion(t, srv.URL, true) + require.NoError(t, err) + assert.True(t, ok) + // Without a node count there is nothing to confirm, so it waited for Jobs. + assert.GreaterOrEqual(t, cs.runGets, 2) + assert.Empty(t, cs.tailAttempts(0)) } func TestStreamBricklensIgnoresFailureSummary(t *testing.T) { @@ -292,4 +302,6 @@ func TestStreamBricklensBindsCompletionToLatestAttempt(t *testing.T) { assert.True(t, ok) assert.GreaterOrEqual(t, cs.runGets, 3) assert.Equal(t, []string{"0", "1"}, cs.tailAttempts(1)) + // Node 0, confirmed on attempt 0, had to be confirmed again on attempt 1. + assert.Equal(t, []string{"0", "1"}, cs.tailAttempts(0)) } diff --git a/cmd/air/logstream.go b/cmd/air/logstream.go index cefe07dc483..330cf1af0ad 100644 --- a/cmd/air/logstream.go +++ b/cmd/air/logstream.go @@ -363,10 +363,12 @@ type bricklensStreamer struct { // a response carries it. nodeCount int // completion watches the streamed node for the success summary of - // completionAttempt. It only triggers the per-attempt check in - // allNodesCompleted, and a retry starts a fresh matcher. + // completionAttempt and only triggers allNodesCompleted. completedNodes counts + // the leading nodes confirmed to have succeeded on that attempt, so a check + // resumes where the last one stopped. A retry resets all three. completion completionMatcher completionAttempt int + completedNodes int } // waitingSpinnerText returns the waiting-spinner text: the server-set STATUS @@ -475,6 +477,7 @@ func (st *bricklensStreamer) run() (bool, error) { if st.status.latestAttempt != st.completionAttempt { st.completion = completionMatcher{} st.completionAttempt = st.status.latestAttempt + st.completedNodes = 0 } terminal := st.status.terminal() @@ -553,32 +556,32 @@ func (st *bricklensStreamer) run() (bool, error) { } } -// allNodesCompleted reports whether every node of the run's latest attempt has -// printed the entry script's success summary. The streamed node's summary only -// triggers the check: every node, the streamed one included, is then read with an -// explicit attempt number, so all summaries come from the same attempt. Best -// effort: an unknown node count or a failed fetch reads as not completed, so the -// stream keeps following the run. +// allNodesCompleted reports whether every node of completionAttempt has printed +// the entry script's success summary. Every node, the streamed one included, is +// read with an explicit attempt number, so all summaries come from the same +// attempt. A node that succeeded stays succeeded, so later calls resume at the +// first unconfirmed node. Best effort: an unknown node count or a failed fetch +// reads as not completed, so the stream keeps following the run. func (st *bricklensStreamer) allNodesCompleted() bool { if !st.completion.completed || st.nodeCount < 1 { return false } - for node := range st.nodeCount { - if !st.nodeCompleted(node, st.completionAttempt) { + for ; st.completedNodes < st.nodeCount; st.completedNodes++ { + if !st.nodeCompleted(st.completedNodes) { return false } } return true } -// nodeCompleted searches the newest lines of a node's log for an attempt for the -// success summary. From the stream start rather than the live cursor, since the -// node may have finished before the lookback window. -func (st *bricklensStreamer) nodeCompleted(node, attempt int) bool { +// nodeCompleted searches the newest lines of a node's log for completionAttempt +// for the success summary. From the stream start rather than the live cursor, +// since the node may have finished before the lookback window. +func (st *bricklensStreamer) nodeCompleted(node int) bool { resp, err := getBricklensLogs(st.ctx, st.apiClient, st.headers, st.req.runID, bricklensLogsQuery{ fromSeconds: st.streamStartSec, pageSize: completionTailLines, - attemptNumber: attempt, + attemptNumber: st.completionAttempt, nodeIndex: node, }) if err != nil { From e6328c83531db26c376f0567f38db644ca2dbfa8 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Tue, 13 Oct 2026 03:06:35 +0000 Subject: [PATCH 4/4] Exit air run --watch on the streamed node's success only --watch follows one node, so stop when that node's command succeeds instead of first confirming every node through extra per-node log requests. This drops the node count, the per-attempt checks, and the confirmed-node tracking. A failure summary still waits for Jobs, so a retried attempt keeps being followed and a failed run still exits non-zero. Co-authored-by: Isaac --- .nextchanges/air/watch-exit-on-completion.md | 2 +- cmd/air/logbricklens.go | 1 - cmd/air/logcompletion.go | 10 +- cmd/air/logcompletion_test.go | 190 +++++-------------- cmd/air/logstream.go | 67 +------ cmd/air/run.go | 8 +- cmd/air/run_watch_test.go | 3 +- 7 files changed, 67 insertions(+), 214 deletions(-) diff --git a/.nextchanges/air/watch-exit-on-completion.md b/.nextchanges/air/watch-exit-on-completion.md index 31615ecda6e..879b13e4475 100644 --- a/.nextchanges/air/watch-exit-on-completion.md +++ b/.nextchanges/air/watch-exit-on-completion.md @@ -1 +1 @@ -* `air run --watch` now exits as soon as the command succeeds on every node, instead of waiting for the job run to be marked terminal. ([#7031](https://github.com/databricks/cli/pull/7031)) +* `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/logbricklens.go b/cmd/air/logbricklens.go index e230a8c5ca2..ca3b8170130 100644 --- a/cmd/air/logbricklens.go +++ b/cmd/air/logbricklens.go @@ -35,7 +35,6 @@ func (r logRecord) nano() int64 { type bricklensLogsResponse struct { LogRecords []logRecord `json:"log_records"` NextPageToken string `json:"next_page_token"` - NodeCount int `json:"node_count"` } // bricklensLogsQuery is the request-field surface of the log endpoint. diff --git a/cmd/air/logcompletion.go b/cmd/air/logcompletion.go index cc96ef7b999..19237ac0b1a 100644 --- a/cmd/air/logcompletion.go +++ b/cmd/air/logcompletion.go @@ -16,7 +16,7 @@ import "strings" // "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 every node has succeeded. +// `air run --watch` uses it to stop as soon as the streamed node has succeeded. const ( summaryHeader = "Execution Summary" summaryRulePrefix = "====" @@ -25,13 +25,9 @@ const ( summarySuccessPrefix = "Completed at " ) -// completionTailLines is how many of a node's newest log lines are searched for -// the summary block. -const completionTailLines = 20 - // completionMatcher recognizes the entry script's success summary in one node's -// log lines, fed oldest-first. Only a consecutive block matches, so user output -// that happens to contain one of the lines does not. +// 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 diff --git a/cmd/air/logcompletion_test.go b/cmd/air/logcompletion_test.go index 9e6469fbc1a..112936763e7 100644 --- a/cmd/air/logcompletion_test.go +++ b/cmd/air/logcompletion_test.go @@ -72,78 +72,47 @@ func TestCompletionMatcher(t *testing.T) { } } -// logsBody renders a log endpoint response for one node's attempt. Records are -// timestamped one second apart from 1700000000s and keyed by attempt, position -// and content, so a poll that re-serves earlier lines dedups them while new or -// repeated lines (the summary's rules) stay distinct; newestFirst matches an -// ascending=false query. -func logsBody(t *testing.T, nodeCount, node int, attempt string, lines []string, newestFirst bool) string { +// 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"` - NodeIndex int `json:"node_index"` } records := make([]record, len(lines)) for i, line := range lines { records[i] = record{ - RecordID: strconv.Itoa(node) + ":" + attempt + ":" + strconv.Itoa(i) + ":" + line, + RecordID: strconv.Itoa(i) + ":" + line, TimeUnixNano: (1_700_000_000 + int64(i)) * 1_000_000_000, Body: line, - NodeIndex: node, } } - if newestFirst { - slices.Reverse(records) - } - b, err := json.Marshal(map[string]any{"log_records": records, "node_count": nodeCount}) + 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}]}` - retrySucceededRun = `{"run_id":1,"start_time":1700000000000,"end_time":1700000100000,"state":{"life_cycle_state":"TERMINATED","result_state":"SUCCESS"},"tasks":[{"run_id":2,"attempt_number":0},{"run_id":3,"attempt_number":1}]}` + 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}]}` ) -// logQuery is one log request the completion server received. attempt is the raw -// ref.attempt_number ("" for latest) and resolved the attempt it was served for; -// tail is true for a newest-first (completion check) query. -type logQuery struct { - node int - attempt string - resolved string - tail bool -} - -// completionServer serves an active run for the completion tests. runBodies are -// returned by successive runs/get calls (the last repeats). latestAttempt, when -// set, resolves a query without an attempt number, as the endpoint does; -// otherwise it resolves to attempt 0. +// completionServer counts the requests the streamer made to the fake server. type completionServer struct { - runGets int - queries []logQuery - latestAttempt func() string -} - -// tailAttempts returns the raw attempt numbers of the completion checks for a node. -func (cs *completionServer) tailAttempts(node int) []string { - var out []string - for _, q := range cs.queries { - if q.node == node && q.tail { - out = append(out, q.attempt) - } - } - return out + runGets int + logRequests int } -// newCompletionServer starts the server; nodeLines returns the lines for a log -// request, given the requests received so far including it. -func newCompletionServer(t *testing.T, nodeCount int, runBodies []string, nodeLines func(q logQuery, cs *completionServer) []string) (*completionServer, *httptest.Server) { +// 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 @@ -153,24 +122,9 @@ func newCompletionServer(t *testing.T, nodeCount int, runBodies []string, nodeLi srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch { case strings.HasSuffix(r.URL.Path, "/logs"): - node, err := strconv.Atoi(r.URL.Query().Get("filter.node_index")) - if !assert.NoError(t, err) { - return - } - q := logQuery{ - node: node, - attempt: r.URL.Query().Get("ref.attempt_number"), - resolved: r.URL.Query().Get("ref.attempt_number"), - tail: r.URL.Query().Get("ascending") == "false", - } - if q.resolved == "" { - q.resolved = "0" - if cs.latestAttempt != nil { - q.resolved = cs.latestAttempt() - } - } - cs.queries = append(cs.queries, q) - _, _ = w.Write([]byte(logsBody(t, nodeCount, node, q.resolved, nodeLines(q, cs), q.tail))) + 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++ @@ -194,9 +148,9 @@ func streamForCompletion(t *testing.T, srvURL string, exitOnCompletion bool) (bo return ok, buf.String(), err } -func TestStreamBricklensExitsWhenAllNodesComplete(t *testing.T) { - cs, srv := newCompletionServer(t, 2, []string{runningRunBody}, func(q logQuery, _ *completionServer) []string { - return append([]string{"training node " + strconv.Itoa(q.node)}, successSummary...) +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) @@ -204,43 +158,29 @@ func TestStreamBricklensExitsWhenAllNodesComplete(t *testing.T) { assert.True(t, ok) // Exited on the first poll, without Jobs ever reporting the run terminal. assert.Equal(t, 0, cs.runGets) - // Both nodes, the streamed one included, were checked for attempt 0. - assert.Equal(t, []string{"0"}, cs.tailAttempts(0)) - assert.Equal(t, []string{"0"}, cs.tailAttempts(1)) - assert.Contains(t, out, `"line":"training node 0"`) + assert.Equal(t, 1, cs.logRequests) + assert.Contains(t, out, `"line":"training"`) + assert.Contains(t, out, `"line":"Completed at `) } -func TestStreamBricklensWaitsForEveryNode(t *testing.T) { - cs, srv := newCompletionServer(t, 2, []string{runningRunBody}, func(q logQuery, cs *completionServer) []string { - if q.node == 1 && len(cs.tailAttempts(1)) < 3 { - return []string{"still training"} +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 successSummary - }) - - ok, _, err := streamForCompletion(t, srv.URL, true) - require.NoError(t, err) - assert.True(t, ok) - assert.Len(t, cs.tailAttempts(1), 3) - // Node 0 was confirmed once and not re-fetched while node 1 caught up. - assert.Len(t, cs.tailAttempts(0), 1) -} - -func TestStreamBricklensCompletionRequiresNodeCount(t *testing.T) { - cs, srv := newCompletionServer(t, 0, []string{runningRunBody, succeededRunBody}, func(q logQuery, _ *completionServer) []string { - return successSummary + return lines }) ok, _, err := streamForCompletion(t, srv.URL, true) require.NoError(t, err) assert.True(t, ok) - // Without a node count there is nothing to confirm, so it waited for Jobs. - assert.GreaterOrEqual(t, cs.runGets, 2) - assert.Empty(t, cs.tailAttempts(0)) + assert.Equal(t, 2, cs.logRequests) + assert.Equal(t, 1, cs.runGets) } func TestStreamBricklensIgnoresFailureSummary(t *testing.T) { - cs, srv := newCompletionServer(t, 1, []string{runningRunBody, failedRunBody}, func(q logQuery, _ *completionServer) []string { + cs, srv := newCompletionServer(t, []string{runningRunBody, failedRunBody}, func(int) []string { return failureSummary }) @@ -249,59 +189,31 @@ func TestStreamBricklensIgnoresFailureSummary(t *testing.T) { // Waited for Jobs: a failed attempt may still be retried. assert.False(t, ok) assert.GreaterOrEqual(t, cs.runGets, 2) - assert.Empty(t, cs.tailAttempts(0)) } -func TestStreamBricklensCompletionRequiresOptIn(t *testing.T) { - cs, srv := newCompletionServer(t, 1, []string{runningRunBody, succeededRunBody}, func(q logQuery, _ *completionServer) []string { - return successSummary +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, false) + ok, _, err := streamForCompletion(t, srv.URL, true) require.NoError(t, err) assert.True(t, ok) - assert.GreaterOrEqual(t, cs.runGets, 2) - assert.Empty(t, cs.tailAttempts(0)) + assert.Equal(t, 2, cs.logRequests) } -func TestStreamBricklensBindsCompletionToLatestAttempt(t *testing.T) { - // Attempt 0 succeeds on node 0 but fails on node 1, and is retried. On attempt - // 1, node 0 succeeds while node 1 is still running when Jobs reports the run - // terminal. Every check must target the attempt that triggered it, so the - // stream ends on Jobs' terminal state. - attemptOneStreamPolls := 0 - cs, srv := newCompletionServer(t, 2, []string{retryingRunBody, retryingRunBody, retrySucceededRun}, func(q logQuery, _ *completionServer) []string { - switch { - case q.resolved == "0" && q.node == 0: - return successSummary - case q.resolved == "0" && q.node == 1: - return failureSummary - case q.node == 1: - return []string{"still training"} - case q.tail: - return successSummary - default: - // Node 0's stream on attempt 1: the summary arrives on the second poll. - attemptOneStreamPolls++ - if attemptOneStreamPolls >= 2 { - return append([]string{"attempt 1 node 0"}, successSummary...) - } - return []string{"attempt 1 node 0"} - } +func TestStreamBricklensCompletionRequiresOptIn(t *testing.T) { + cs, srv := newCompletionServer(t, []string{runningRunBody, succeededRunBody}, func(int) []string { + return successSummary }) - cs.latestAttempt = func() string { - if cs.runGets == 0 { - return "0" - } - return "1" - } - ok, _, err := streamForCompletion(t, srv.URL, true) + ok, _, err := streamForCompletion(t, srv.URL, false) require.NoError(t, err) - // Ended on Jobs' terminal SUCCESS, not on summaries from different attempts. assert.True(t, ok) - assert.GreaterOrEqual(t, cs.runGets, 3) - assert.Equal(t, []string{"0", "1"}, cs.tailAttempts(1)) - // Node 0, confirmed on attempt 0, had to be confirmed again on attempt 1. - assert.Equal(t, []string{"0", "1"}, cs.tailAttempts(0)) + assert.GreaterOrEqual(t, cs.runGets, 2) } diff --git a/cmd/air/logstream.go b/cmd/air/logstream.go index 330cf1af0ad..7870175a2ae 100644 --- a/cmd/air/logstream.go +++ b/cmd/air/logstream.go @@ -114,9 +114,9 @@ type logRequest struct { // `air run --watch -o json` to emit STATUS events. onStatusChange func(current, previous string) retryTracker *retryTracker - // exitOnCompletion ends the stream once every node of the latest attempt has - // printed the entry script's success summary, without waiting for Jobs to - // report the run terminal. Used by `air run --watch`. + // 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 } @@ -359,16 +359,8 @@ type bricklensStreamer struct { onFirstLog func() // updateSpinner, when set, refreshes the waiting-spinner text each poll. updateSpinner func(string) - // nodeCount is the run's node count as reported by the log endpoint; 0 until - // a response carries it. - nodeCount int - // completion watches the streamed node for the success summary of - // completionAttempt and only triggers allNodesCompleted. completedNodes counts - // the leading nodes confirmed to have succeeded on that attempt, so a check - // resumes where the last one stopped. A retry resets all three. - completion completionMatcher - completionAttempt int - completedNodes int + // completion watches the streamed node's new lines for the success summary. + completion completionMatcher } // waitingSpinnerText returns the waiting-spinner text: the server-set STATUS @@ -474,11 +466,6 @@ func (st *bricklensStreamer) run() (bool, error) { st.reportStatusChange() st.req.reportRetry(st.ctx, st.out, st.status, st.onFirstLog) - if st.status.latestAttempt != st.completionAttempt { - st.completion = completionMatcher{} - st.completionAttempt = st.status.latestAttempt - st.completedNodes = 0 - } terminal := st.status.terminal() toSec := st.req.toSeconds(st.status) @@ -537,8 +524,8 @@ 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.allNodesCompleted() { - log.Infof(st.ctx, "air logs: run %d completed on all %d node(s); not waiting for Jobs to report it terminal", st.req.runID, st.nodeCount) + 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 { @@ -556,45 +543,6 @@ func (st *bricklensStreamer) run() (bool, error) { } } -// allNodesCompleted reports whether every node of completionAttempt has printed -// the entry script's success summary. Every node, the streamed one included, is -// read with an explicit attempt number, so all summaries come from the same -// attempt. A node that succeeded stays succeeded, so later calls resume at the -// first unconfirmed node. Best effort: an unknown node count or a failed fetch -// reads as not completed, so the stream keeps following the run. -func (st *bricklensStreamer) allNodesCompleted() bool { - if !st.completion.completed || st.nodeCount < 1 { - return false - } - for ; st.completedNodes < st.nodeCount; st.completedNodes++ { - if !st.nodeCompleted(st.completedNodes) { - return false - } - } - return true -} - -// nodeCompleted searches the newest lines of a node's log for completionAttempt -// for the success summary. From the stream start rather than the live cursor, -// since the node may have finished before the lookback window. -func (st *bricklensStreamer) nodeCompleted(node int) bool { - resp, err := getBricklensLogs(st.ctx, st.apiClient, st.headers, st.req.runID, bricklensLogsQuery{ - fromSeconds: st.streamStartSec, - pageSize: completionTailLines, - attemptNumber: st.completionAttempt, - nodeIndex: node, - }) - if err != nil { - log.Debugf(st.ctx, "air logs: completion check for node %d failed: %v", node, err) - return false - } - var m completionMatcher - for _, rec := range slices.Backward(resp.LogRecords) { - m.observe(rec.Body) - } - return m.completed -} - func waitForNextBricklensPoll(ctx context.Context, pollStarted time.Time) error { return sleepOrCancel(ctx, max(time.Duration(0), bricklensPollInterval-time.Since(pollStarted))) } @@ -741,7 +689,6 @@ func (st *bricklensStreamer) requestPage(pageToken string, toSec int64, pageSize for { resp, err := getBricklensLogs(st.ctx, st.apiClient, st.headers, st.req.runID, q) if err == nil { - st.nodeCount = max(st.nodeCount, resp.NodeCount) return resp, nil } diff --git a/cmd/air/run.go b/cmd/air/run.go index 7873e21b3cb..8ff2bcc5c4d 100644 --- a/cmd/air/run.go +++ b/cmd/air/run.go @@ -203,8 +203,8 @@ 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). It exits early once - // every node's entry script reports success, since Jobs lags the user - // command's exit by tens of seconds. + // 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, @@ -370,8 +370,8 @@ func watchTerminalStatus(ctx context.Context, w *databricks.WorkspaceClient, run if err != nil { return "UNKNOWN" } - // A stream that returned success on a still-active run stopped early on every - // node's success summary; Jobs reports the run terminal only later. + // 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" } diff --git a/cmd/air/run_watch_test.go b/cmd/air/run_watch_test.go index be71c21d92e..c21b51ec4f2 100644 --- a/cmd/air/run_watch_test.go +++ b/cmd/air/run_watch_test.go @@ -196,8 +196,7 @@ func watchServerCompleted(t *testing.T) *httptest.Server { 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/"): - newestFirst := r.URL.Query().Get("ascending") == "false" - _, _ = w.Write([]byte(logsBody(t, 1, 0, "0", append([]string{"step 1"}, successSummary...), newestFirst))) + _, _ = w.Write([]byte(logsBody(t, append([]string{"step 1"}, successSummary...)))) default: _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) }