[air] Show retry progress while streaming workload logs

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 {

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)

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) {

experimental/air/cmd/logstream.go

@@ -108,25 +108,85 @@ 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 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 {
+	lastReportedRetry           int
+	lastObservedAttempt         int
+	lastObservedWaitingForRetry bool
+}
+
+func newRetryTracker(status logRunStatus) *retryTracker {
+	return &retryTracker{
+		lastReportedRetry:           status.latestAttempt,
+		lastObservedAttempt:         status.latestAttempt,
+		lastObservedWaitingForRetry: status.latestTaskLifeCycleState == string(jobs.RunLifeCycleStateWaitingForRetry),
+	}
+}
+
+func (t *retryTracker) observe(status logRunStatus) (int, bool) {
+	waitingForRetry := status.latestTaskLifeCycleState == string(jobs.RunLifeCycleStateWaitingForRetry)
+	retry := 0
+	if status.latestAttempt > t.lastObservedAttempt {
+		t.lastObservedAttempt = status.latestAttempt
+		t.lastObservedWaitingForRetry = waitingForRetry
+		retry = status.latestAttempt
+	} else if status.latestAttempt == t.lastObservedAttempt {
+		if waitingForRetry && !t.lastObservedWaitingForRetry {
+			retry = status.latestAttempt + 1
+		}
+		t.lastObservedWaitingForRetry = waitingForRetry
+	}
+	if retry <= t.lastReportedRetry || retry <= 0 {
+		return 0, false
+	}
+	t.lastReportedRetry = 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 +212,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 +250,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 +459,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)

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 {

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")