From c1d0116ebea1862e9d28a416dbe397de29e6401b Mon Sep 17 00:00:00 2001 From: Caroline Chen <324939130+caroline-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:07:06 +0000 Subject: [PATCH 1/2] air: show retry progress while streaming logs --- experimental/air/cmd/logmlflow.go | 1 + experimental/air/cmd/logs.go | 3 + experimental/air/cmd/logs_test.go | 56 +++++++++++++++ experimental/air/cmd/logstream.go | 86 ++++++++++++++++++++--- experimental/air/cmd/logstream_support.go | 20 ++++++ experimental/air/cmd/logstream_test.go | 81 +++++++++++++++++++-- experimental/air/cmd/run.go | 18 +++-- experimental/air/cmd/run_watch_test.go | 1 + 8 files changed, 247 insertions(+), 19 deletions(-) diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go index 0e51ca8ae36..90c5c0bd5c8 100644 --- a/experimental/air/cmd/logmlflow.go +++ b/experimental/air/cmd/logmlflow.go @@ -165,6 +165,7 @@ func streamMLflowLogs(ctx context.Context, w *databricks.WorkspaceClient, out io previousState = current } } + req.reportRetry(ctx, out, status, nil) mlflowRunID, logDir, err := resolveMLflowLogPath(ctx, w, req) if err != nil { diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index b6f9b723d99..07461a70eff 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -233,6 +233,9 @@ func resolveLogAttempt(run *jobs.Run, requested int) (int, int64, error) { // fetchLogs serves logs from Bricklens, falling back to MLflow when Bricklens // returns errBricklensFeatureDisabled. func fetchLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + if req.attempt < 0 && !status.terminal() && req.retryTracker == nil { + req.retryTracker = newRetryTracker(status) + } success, err := streamBricklensLogs(ctx, w, out, req, status) if errors.Is(err, errBricklensFeatureDisabled) { return mlflowLogFallback(ctx, w, out, req, status) diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go index 73d7645299b..058ed1dd8c0 100644 --- a/experimental/air/cmd/logs_test.go +++ b/experimental/air/cmd/logs_test.go @@ -7,6 +7,7 @@ import ( "net/http/httptest" "strings" "testing" + "time" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" @@ -204,6 +205,61 @@ func TestLogsFallsBackToMLflow(t *testing.T) { assert.Equal(t, "line one\nline two\n", buf.String()) } +func TestFetchLogsRetryNotificationDeduplicatedAcrossFallback(t *testing.T) { + oldRetryInterval := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = oldRetryInterval }) + + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{"run_id":5,"state":{"life_cycle_state":"TERMINATED","result_state":"SUCCESS"},"tasks":[{"run_id":101,"attempt_number":1}]}`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code":"FEATURE_DISABLED","message":"gated off"}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output":{"mlflow_experiment_id":"exp","mlflow_run_id":"run"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files":[{"path":"logs/attempt_1","is_dir":true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files":[{"path":"logs/attempt_1/node_0/logs-0.chunk.txt"}]}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos":[{"signed_uri":"` + base + `/artifact"}]}`)) + case r.URL.Path == "/artifact": + _, _ = w.Write([]byte("retry log\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + + initial := logRunStatus{ + lifeCycleState: "RUNNING", + latestAttempt: 1, + latestTaskLifeCycleState: "WAITING_FOR_RETRY", + } + maxRetries := 3 + var out bytes.Buffer + success, err := fetchLogs(t.Context(), newTestWorkspaceClient(t, srv.URL), &out, logRequest{ + runID: 5, + attempt: -1, + tailLines: -1, + jsonOutput: true, + maxRetries: &maxRetries, + retryTracker: newRetryTracker(logRunStatus{}), + }, initial) + require.NoError(t, err) + assert.True(t, success) + assert.Equal(t, 1, strings.Count(out.String(), `"type":"RETRY"`)) + assert.Contains(t, out.String(), `"retry":1`) + assert.Contains(t, out.String(), `"max_retries":3`) + assert.Contains(t, out.String(), `"line":"retry log"`) +} + func TestLogsCommandPrintsGuidanceWhenStreamingIsInterrupted(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index cfd7b27868b..f97e85c3677 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -108,25 +108,83 @@ type logRequest struct { // boundInitialLogs limits existing output before following an active run. boundInitialLogs bool jsonOutput bool + maxRetries *int // onStatusChange, when set, is called on each lifecycle transition while // following the run (current, previous display states). Used by // `air run --watch -o json` to emit STATUS events. onStatusChange func(current, previous string) + retryTracker *retryTracker } // logRunStatus is the subset of a run's state the log path needs, resolved once // and reused. type logRunStatus struct { - lifeCycleState string - firstTaskLifeCycleState string - resultState string - stateMessage string - startTimeMs int64 - endTimeMs int64 + lifeCycleState string + latestTaskLifeCycleState string + resultState string + stateMessage string + startTimeMs int64 + endTimeMs int64 // latestAttempt is the highest attempt_number across the run's tasks. latestAttempt int } +// retryTracker reports each retry once whether WAITING_FOR_RETRY appears before +// or after Jobs exposes the next attempt_number. +type retryTracker struct { + announced int + observedAttempt int + waiting bool +} + +func newRetryTracker(status logRunStatus) *retryTracker { + return &retryTracker{ + announced: status.latestAttempt, + observedAttempt: status.latestAttempt, + waiting: status.latestTaskLifeCycleState == string(jobs.RunLifeCycleStateWaitingForRetry), + } +} + +func (t *retryTracker) observe(status logRunStatus) (int, bool) { + waiting := status.latestTaskLifeCycleState == string(jobs.RunLifeCycleStateWaitingForRetry) + retry := 0 + if status.latestAttempt > t.observedAttempt { + t.observedAttempt = status.latestAttempt + t.waiting = waiting + retry = status.latestAttempt + } else if status.latestAttempt == t.observedAttempt { + if waiting && !t.waiting { + retry = status.latestAttempt + 1 + } + t.waiting = waiting + } + if retry <= t.announced || retry <= 0 { + return 0, false + } + t.announced = retry + return retry, true +} + +func (req logRequest) reportRetry(ctx context.Context, out io.Writer, status logRunStatus, before func()) { + if req.retryTracker == nil { + return + } + retry, ok := req.retryTracker.observe(status) + if !ok { + return + } + if before != nil { + before() + } + if req.jsonOutput { + printRetryEvent(out, retry, req.maxRetries) + } else if req.maxRetries != nil { + cmdio.LogString(ctx, fmt.Sprintf("\nRetrying (%d of %d)...\n", retry, *req.maxRetries)) + } else { + cmdio.LogString(ctx, fmt.Sprintf("\nRetrying (%d)...\n", retry)) + } +} + // A run is terminal when its lifecycle state is terminal, or a result state is // set (result states only appear on terminal runs). var ( @@ -152,7 +210,7 @@ func (s logRunStatus) waitingForCompute() bool { if s.lifeCycleState != string(jobs.RunLifeCycleStateRunning) { return false } - switch jobs.RunLifeCycleState(s.firstTaskLifeCycleState) { + switch jobs.RunLifeCycleState(s.latestTaskLifeCycleState) { case jobs.RunLifeCycleStatePending, jobs.RunLifeCycleStateQueued, jobs.RunLifeCycleStateWaitingForRetry, jobs.RunLifeCycleStateBlocked: return true @@ -190,11 +248,16 @@ func projectRunStatus(run *jobs.Run) logRunStatus { s.resultState = string(run.State.ResultState) s.stateMessage = run.State.StateMessage } - if len(run.Tasks) > 0 && run.Tasks[0].State != nil { - s.firstTaskLifeCycleState = string(run.Tasks[0].State.LifeCycleState) - } for i := range run.Tasks { - s.latestAttempt = max(s.latestAttempt, run.Tasks[i].AttemptNumber) + task := &run.Tasks[i] + if i > 0 && task.AttemptNumber <= s.latestAttempt { + continue + } + s.latestAttempt = task.AttemptNumber + s.latestTaskLifeCycleState = "" + if task.State != nil { + s.latestTaskLifeCycleState = string(task.State.LifeCycleState) + } } return s } @@ -394,6 +457,7 @@ func (st *bricklensStreamer) run() (bool, error) { } st.reportStatusChange() + st.req.reportRetry(st.ctx, st.out, st.status, st.onFirstLog) terminal := st.status.terminal() toSec := st.req.toSeconds(st.status) diff --git a/experimental/air/cmd/logstream_support.go b/experimental/air/cmd/logstream_support.go index 99943e342d9..69896f71c3c 100644 --- a/experimental/air/cmd/logstream_support.go +++ b/experimental/air/cmd/logstream_support.go @@ -100,6 +100,26 @@ func printLogEvent(out io.Writer, eventType string, node int, line string) { fmt.Fprintln(out, string(b)) } +type retryEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + Retry int `json:"retry"` + MaxRetries *int `json:"max_retries,omitempty"` +} + +func printRetryEvent(out io.Writer, retry int, maxRetries *int) { + b, err := json.Marshal(retryEvent{ + Type: "RETRY", + TS: time.Now().UTC().Format(time.RFC3339), + Retry: retry, + MaxRetries: maxRetries, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + // submittedEvent is the JSONL event `air run --watch -o json` emits before the // streamed log events, so a consumer sees the run id immediately. type submittedEvent struct { diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go index 147b77f9a6e..242a0b51b38 100644 --- a/experimental/air/cmd/logstream_test.go +++ b/experimental/air/cmd/logstream_test.go @@ -12,6 +12,8 @@ import ( "testing" "time" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/service/jobs" @@ -83,7 +85,7 @@ func TestProjectRunStatus(t *testing.T) { }, Tasks: []jobs.RunTask{ {AttemptNumber: 0, State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateQueued}}, - {AttemptNumber: 2}, + {AttemptNumber: 2, State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateWaitingForRetry}}, {AttemptNumber: 1}, }, } @@ -92,7 +94,7 @@ func TestProjectRunStatus(t *testing.T) { assert.Equal(t, "TERMINATED", s.lifeCycleState) assert.Equal(t, "SUCCESS", s.resultState) assert.Equal(t, "done", s.stateMessage) - assert.Equal(t, "QUEUED", s.firstTaskLifeCycleState) + assert.Equal(t, "WAITING_FOR_RETRY", s.latestTaskLifeCycleState) assert.Equal(t, int64(1000), s.startTimeMs) assert.Equal(t, int64(2000), s.endTimeMs) assert.Equal(t, 2, s.latestAttempt) @@ -101,6 +103,63 @@ func TestProjectRunStatus(t *testing.T) { assert.Equal(t, "SUCCESS", s.displayState()) } +func TestRetryTrackerDeduplicatesWaitingAndAttemptIncrease(t *testing.T) { + tracker := newRetryTracker(logRunStatus{latestAttempt: 0}) + + retry, ok := tracker.observe(logRunStatus{ + latestAttempt: 0, + latestTaskLifeCycleState: string(jobs.RunLifeCycleStateWaitingForRetry), + }) + assert.True(t, ok) + assert.Equal(t, 1, retry) + + _, ok = tracker.observe(logRunStatus{ + latestAttempt: 1, + latestTaskLifeCycleState: string(jobs.RunLifeCycleStateWaitingForRetry), + }) + assert.False(t, ok) +} + +func TestRetryTrackerDetectsAttemptCreatedInWaitingState(t *testing.T) { + tracker := newRetryTracker(logRunStatus{latestAttempt: 0}) + waiting := logRunStatus{ + latestAttempt: 1, + latestTaskLifeCycleState: string(jobs.RunLifeCycleStateWaitingForRetry), + } + + retry, ok := tracker.observe(waiting) + assert.True(t, ok) + assert.Equal(t, 1, retry) + + _, ok = tracker.observe(waiting) + assert.False(t, ok) +} + +func TestRetryTrackerDetectsMissedWaitingState(t *testing.T) { + tracker := newRetryTracker(logRunStatus{latestAttempt: 0}) + + retry, ok := tracker.observe(logRunStatus{latestAttempt: 1}) + assert.True(t, ok) + assert.Equal(t, 1, retry) +} + +func TestReportRetryText(t *testing.T) { + var stdout, stderr bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, nil, &stdout, &stderr, "", "")) + maxRetries := 3 + req := logRequest{retryTracker: newRetryTracker(logRunStatus{latestAttempt: 0}), maxRetries: &maxRetries} + stopped := false + + req.reportRetry(ctx, &stdout, logRunStatus{ + latestAttempt: 0, + latestTaskLifeCycleState: string(jobs.RunLifeCycleStateWaitingForRetry), + }, func() { stopped = true }) + + assert.True(t, stopped) + assert.Empty(t, stdout.String()) + assert.Equal(t, "\nRetrying (1 of 3)...\n\n", stderr.String()) +} + func TestLogRunStatusTerminal(t *testing.T) { tests := []struct { name string @@ -334,7 +393,7 @@ func TestNormalizeStatusMessage(t *testing.T) { func TestWaitingSpinnerText(t *testing.T) { // A server that returns the run (with a task) and a STATUS-typed status_message. - newStreamer := func(t *testing.T, statusMessage, lifeCycle, firstTaskLifeCycle string) *bricklensStreamer { + newStreamer := func(t *testing.T, statusMessage, lifeCycle, latestTaskLifeCycle string) *bricklensStreamer { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { @@ -351,7 +410,7 @@ func TestWaitingSpinnerText(t *testing.T) { ctx: t.Context(), w: newTestWorkspaceClient(t, srv.URL), req: logRequest{runID: 1, node: 0}, - status: logRunStatus{lifeCycleState: lifeCycle, firstTaskLifeCycleState: firstTaskLifeCycle}, + status: logRunStatus{lifeCycleState: lifeCycle, latestTaskLifeCycleState: latestTaskLifeCycle}, } } @@ -389,6 +448,20 @@ func TestEmitLogLineJSON(t *testing.T) { assert.NotEmpty(t, ev.TS) } +func TestPrintRetryEvent(t *testing.T) { + var buf bytes.Buffer + maxRetries := 3 + printRetryEvent(&buf, 1, &maxRetries) + + var ev retryEvent + require.NoError(t, json.Unmarshal(buf.Bytes(), &ev)) + assert.Equal(t, "RETRY", ev.Type) + assert.Equal(t, 1, ev.Retry) + require.NotNil(t, ev.MaxRetries) + assert.Equal(t, 3, *ev.MaxRetries) + assert.NotEmpty(t, ev.TS) +} + func TestEmitLogLineText(t *testing.T) { var buf bytes.Buffer emitLogLine(&buf, logRequest{node: 0}, "hello") diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 45202a09b87..95db5366381 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -136,11 +136,14 @@ 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). + maxRetries := cfg.maxRetries() req := logRequest{ - runID: runID, - attempt: -1, - tailLines: -1, - jsonOutput: jsonOut, + runID: runID, + attempt: -1, + tailLines: -1, + jsonOutput: jsonOut, + maxRetries: &maxRetries, + retryTracker: newRetryTracker(logRunStatus{}), } watchCtx, stop := notifyInterrupt(ctx) @@ -161,6 +164,13 @@ The path must be a separate argument: cobra reserves -h as a boolean, so // Separate the submit summary from the streamed logs. fmt.Fprintln(out) fmt.Fprintln(out, monitoringMessage) + if maxRetries > 0 { + unit := "times" + if maxRetries == 1 { + unit = "time" + } + fmt.Fprintf(out, "Failed attempts will be retried up to %d %s.\n", maxRetries, unit) + } printLogsDivider(ctx, out) return handleWatchResult(out, w.Config.Profile, runIDStr, runLogs(watchCtx, cmd, req)) } diff --git a/experimental/air/cmd/run_watch_test.go b/experimental/air/cmd/run_watch_test.go index 3cf30055617..9dc97298fe8 100644 --- a/experimental/air/cmd/run_watch_test.go +++ b/experimental/air/cmd/run_watch_test.go @@ -129,6 +129,7 @@ func TestRunWatchStreamsLogs(t *testing.T) { assert.Contains(t, out, "Submitted workload with Job Run ID: 777") assert.Contains(t, out, "View job run at: ") assert.Contains(t, out, "Monitoring run and streaming logs...") + assert.Contains(t, out, "Failed attempts will be retried up to 3 times.") assert.NotContains(t, out, "from node 0") // A "Logs" divider separates the submit summary from the streamed logs. assert.Contains(t, out, "Logs") From e978741cf12282a6e57e283fbd8314d941a72b1d Mon Sep 17 00:00:00 2001 From: Caroline Chen <324939130+caroline-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 21:45:23 +0000 Subject: [PATCH 2/2] air: clarify retry tracker state --- experimental/air/cmd/logstream.go | 36 ++++++++++++++++--------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index f97e85c3677..11c03d86d47 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -129,39 +129,41 @@ type logRunStatus struct { latestAttempt int } -// retryTracker reports each retry once whether WAITING_FOR_RETRY appears before -// or after Jobs exposes the next attempt_number. +// retryTracker reports each retry once across either Jobs transition: +// +// 1. Attempt N enters WAITING_FOR_RETRY: report retry N+1. +// 2. Attempt N+1 appears: report retry N+1 only if transition 1 was missed. type retryTracker struct { - announced int - observedAttempt int - waiting bool + lastReportedRetry int + lastObservedAttempt int + lastObservedWaitingForRetry bool } func newRetryTracker(status logRunStatus) *retryTracker { return &retryTracker{ - announced: status.latestAttempt, - observedAttempt: status.latestAttempt, - waiting: status.latestTaskLifeCycleState == string(jobs.RunLifeCycleStateWaitingForRetry), + lastReportedRetry: status.latestAttempt, + lastObservedAttempt: status.latestAttempt, + lastObservedWaitingForRetry: status.latestTaskLifeCycleState == string(jobs.RunLifeCycleStateWaitingForRetry), } } func (t *retryTracker) observe(status logRunStatus) (int, bool) { - waiting := status.latestTaskLifeCycleState == string(jobs.RunLifeCycleStateWaitingForRetry) + waitingForRetry := status.latestTaskLifeCycleState == string(jobs.RunLifeCycleStateWaitingForRetry) retry := 0 - if status.latestAttempt > t.observedAttempt { - t.observedAttempt = status.latestAttempt - t.waiting = waiting + if status.latestAttempt > t.lastObservedAttempt { + t.lastObservedAttempt = status.latestAttempt + t.lastObservedWaitingForRetry = waitingForRetry retry = status.latestAttempt - } else if status.latestAttempt == t.observedAttempt { - if waiting && !t.waiting { + } else if status.latestAttempt == t.lastObservedAttempt { + if waitingForRetry && !t.lastObservedWaitingForRetry { retry = status.latestAttempt + 1 } - t.waiting = waiting + t.lastObservedWaitingForRetry = waitingForRetry } - if retry <= t.announced || retry <= 0 { + if retry <= t.lastReportedRetry || retry <= 0 { return 0, false } - t.announced = retry + t.lastReportedRetry = retry return retry, true }