Skip to content

air: Report each --watch STATUS transition once and never after the run ends - #7621

Draft
artemy-kolesnikov wants to merge 2 commits into
mainfrom
CNXT-2521-watch-status-dedup
Draft

artemy-kolesnikov wants to merge 2 commits into
mainfrom
CNXT-2521-watch-status-dedup

Conversation

@artemy-kolesnikov

@artemy-kolesnikov artemy-kolesnikov commented Nov 13, 2026 •

Copy link
Copy Markdown

Title: air: Report each --watch STATUS transition once and never after the run ends

Changes

air run --watch -o json emits a STATUS event on each lifecycle transition. When the Bricklens stream hands off to the MLflow fallback, the fallback started its own transition tracking from the status resolved before streaming. A run cancelled while PENDING with no logs therefore emitted (current<-previous):

SUBMITTED, PENDING<-, CANCELED<-PENDING, PENDING<-, CANCELED<-PENDING   # empty Bricklens -> MLflow
SUBMITTED, PENDING<-, PENDING<-, CANCELED<-PENDING                      # FEATURE_DISABLED -> MLflow
SUBMITTED, PENDING<-, RUNNING<-PENDING, PENDING<-RUNNING, SUCCESS<-PENDING   # MLflow logs found while RUNNING

With this change all three emit each transition once, in order, ending at the terminal state:

SUBMITTED, PENDING<-, CANCELED<-PENDING
SUBMITTED, PENDING<-, CANCELED<-PENDING
SUBMITTED, PENDING<-, RUNNING<-PENDING, SUCCESS<-RUNNING

The change has two parts. The tests show that each one is needed.

  1. STATUS guard (newStatusTransitionReporter, logstream_support.go; wired at run.go where onStatusChange is set). It drops a repeat of the last reported state and anything after a terminal state, and it reports the last state it actually emitted as previous_status. Non-terminal order is not otherwise policed.
  2. Carry the latest status into the fallback (fetchLogs → mlflowLogFallbackReporting → followMLflowLogs). streamBricklensLogsWithLatestStatus returns the status the Bricklens loop last observed. The MLflow stream reports that status until its first refresh, so it never surfaces the stale pre-stream one. The fallback's control flow is unchanged: it still follows on the pre-stream status. A run that was being followed keeps the active MLflow stream, with its lenient first pass, 3 s grace wait, re-read, uncapped output, per-chunk offsets and partial-line hold-back. It does not switch to the 10,000-line completed-tail path. The two-pass grace and the three empty terminal Bricklens polls are untouched.

The existing streamBricklensLogs, mlflowLogFallback and streamMLflowLogs signatures are kept as thin wrappers, so their callers and tests don't change. Text mode is unaffected; it has no STATUS sink.

Why

A report against a workspace: air run --watch takes about 60 s to notice a cancel made from another terminal. Investigating it turned up this event-ordering bug. A JSON consumer that treats each STATUS as a state machine saw the run go back to PENDING after CANCELED.

This patch does not change how quickly a cancel is noticed. A live smoke against a workspace (below) showed the Jobs run state itself staying RUNNING for 67-87 s after a successful air cancel. Both the patched and the unpatched watch reported CANCELED within about 0.5 s of Jobs flipping, and exited about 2.4 s later. A proposal for the client-side blocking-call case is under "Follow-up design" and isn't implemented here.

Tests

  • cmd/air/run_watch_status_test.go (new). These are full-command air run --watch -o json tests against libs/testserver. They assert the exact SUBMITTED/STATUS sequence and the closing envelope:
    • TestRunWatchJSONStatusAfterEmptyBricklensHandoff: the run flips PENDING → CANCELED during the first empty workflow-logs poll, and the MLflow IDs appear only on the second fallback lookup. It asserts SUBMITTED, PENDING<-, CANCELED<-PENDING, at least 2 get-output calls (the grace re-read still happens), and all 10,002 lines across two chunks printed (uncapped). The envelope is CANCELED.
    • TestRunWatchJSONStatusAfterFeatureDisabledHandoff: a 403 FEATURE_DISABLED handoff. It asserts SUBMITTED, PENDING<-, CANCELED<-PENDING and a CANCELED envelope.
    • TestRunWatchJSONStatusAfterActiveMLflowProbeHandoff: MLflow logs are found while RUNNING. It asserts SUBMITTED, PENDING<-, RUNNING<-PENDING, SUCCESS<-RUNNING.
    • TestStatusTransitionReporterDropsRepeatsAndPostTerminalStates: a unit test of the guard.
  • On unpatched c7d55d3a6, the first two tests fail with the sequences shown above. With only the guard (status not carried), the third test fails with ... RUNNING<-PENDING, PENDING<-RUNNING, SUCCESS<-PENDING.
  • All existing tests pass unchanged: go test ./cmd/air/... -count=1 (194 s). go test -race on the touched watch, logstream and MLflow tests passes. A full -race ./cmd/air run fails only in TestIndexStrategy* (list_index_test.go:43). That failure reproduces on unpatched c7d55d3a6 and isn't touched here. gofmt, golangci-lint run ./cmd/air/... (0 issues) and tools/validate_nextchanges.py pass.

Live smoke (against a workspace, 2026-11-13)

One 1×GPU_1xA10 sleep 120 workload (max_retries: 0) per build, run with air run --watch -o json and cancelled from a second process about 5 s after the first STATUS event:

| build | STATUS events | cancel accepted → Jobs TERMINATED/CANCELED | → watch CANCELED event | → exit | exit code |
|---|---|---|---|---|---|---|
| patched | RUNNING<-, CANCELED<-RUNNING | 87.5 s | +0.2 s | +2.4 s | 1 |
| unpatched | RUNNING<-, CANCELED<-RUNNING | 68.2 s | -0.5 s (watch saw it first) | +2.4 s | 1 |

Both runs reached RUNNING immediately and printed 7 environment-setup log lines, so this exercised the normal path, not an MLflow handoff. An independent observer polling jobs get-run and air get -o json once per second saw Jobs RUNNING and air get PENDING for the whole interval. Both changed to CANCELED at the same moment.

Follow-up design (not implemented): an independent status monitor

Under a stalled log request, the watch can't see a cancel. Bricklens is polled serially: runs/get is refreshed only between iterations, after each workflow-logs GET and its SDK retries. A held request therefore delays noticing the cancel by up to the SDK's 60 s inactivity timeout (a ~68 s delay reproduced locally with an injected stall).

Proposal:

  • While a watch is following a run, run one Jobs status poll on its own goroutine (about 1 s cadence, one request in flight, normal SDK retry). It replaces the loop's inline refresh rather than adding to it.
  • When the poll sees a terminal state, it reports it at once and cancels only the in-flight log request: a child context per request, never the parent or Ctrl-C context. The loop keeps the pages already received, its dedup set and its cursor.
  • The terminal drain then runs exactly as today, unbounded: three empty terminal polls, the MLflow handoff with its grace re-read, uncapped output. The monitor only removes the wait on a request that started before the run ended.
  • Status-poll failures are tolerated as they are today ("a transient status blip should not abort a live stream"). The monitor never invents a terminal state, and never aborts or changes the exit code because status is unavailable.
  • The closing JSON envelope can reuse the terminal snapshot the monitor saw instead of re-fetching it.

We considered and rejected a fixed 8 s deadline on the whole terminal drain, plus an 8 s "status unknown → exit nonzero" rule:

  • The drain is unbounded on purpose. Bricklens resets its empty-poll count whenever a terminal poll returns records (late-searchable records, Fix AIR live log polling correctness #6694). Every MLflow pass re-downloads every chunk, so its cost grows with total log size. An 8 s cap would cut off the final chunk, which usually holds the traceback or final metrics. Exiting nonzero on an incomplete drain would also turn healthy SUCCESS runs into CI failures.
  • One transient page failure already costs a 3 s retry sleep. The existing "5 failures → MLflow" path needs at least 12 s, so it could never be reached under an 8 s cap.
  • An 8 s status-unknown abort with 2 s per-poll deadlines would kill multi-hour watches during Jobs throttling or deploys. A 429 with a multi-second Retry-After can't succeed within 2 s. That reverses the deliberate blip tolerance in logstream.go and logmlflow.go.

The smoke above suggests the delay users see may be dominated by Jobs cancel propagation, not by the client. It's worth confirming with request traces before investing in the monitor.

This pull request and its description were written by Isaac.

…un ends

When the Bricklens stream of `air run --watch -o json` hands off to the MLflow
fallback (an empty terminal stream, FEATURE_DISABLED, or MLflow logs found
while active), the fallback restarted transition tracking from the status
resolved before streaming. A run cancelled while PENDING therefore emitted
PENDING, CANCELED, PENDING, CANCELED; a FEATURE_DISABLED handoff emitted
PENDING twice; an active handoff regressed RUNNING to PENDING.

- The STATUS sink in `air run --watch -o json` drops repeats and anything after
  a terminal state, and reports the last state it actually emitted as
  previous_status.
- fetchLogs passes the status Bricklens last observed into the MLflow
  fallback, which reports it instead of the stale pre-stream status. The
  fallback still follows on the pre-stream status, so a followed run keeps
  the active MLflow stream (lenient first pass, grace wait, re-read, uncapped
  output, partial-line hold-back) rather than the capped completed tail.

Co-authored-by: Isaac <no-reply@databricks.com>
@github-actions github-actions Bot added the AIR Databricks AI Runtime CLI label Nov 13, 2026
@eng-dev-ecosystem-bot

eng-dev-ecosystem-bot commented Nov 13, 2026 •

Copy link
Copy Markdown
Collaborator

✅ Integration test report

Commit: 8783969

Run: 43529957338

Env 💚​RECOVERED ✅​pass 🙈​skip Time
💚​ aws linux-2core-8gb 1 270 15 6:14
💚​ aws-windows-latest-4core-16gb 1 272 13 4:37
💚​ azure linux-2core-8gb 1 269 15 5:25
💚​ azure-windows-latest-4core-16gb 1 271 13 3:40
💚​ gcp linux-2core-8gb 1 270 15 5:08
💚​ gcp-windows-latest-4core-16gb 1 272 13 3:21
Test Name aws linux-2core-8gb aws-windows-latest-4core-16gb azure linux-2core-8gb azure-windows-latest-4core-16gb gcp linux-2core-8gb gcp-windows-latest-4core-16gb
💚​ TestAccept 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R
Top 3 slowest tests (at least 2 minutes):
duration env testname
3:16 azure-windows-latest-4core-16gb TestAccept
3:12 aws-windows-latest-4core-16gb TestAccept
3:03 gcp-windows-latest-4core-16gb TestAccept

Co-authored-by: Isaac <no-reply@databricks.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

AIR Databricks AI Runtime CLI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants