From 9d8e92afc567026d85b68061845db9dd0a16843b Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Wed, 12 Aug 2026 00:28:08 +0000 Subject: [PATCH] air: fall back to MLflow when Bricklens returns no logs for a terminal run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `air logs ` printed "No logs available for run . Run terminated in state SUCCESS" and exited 0 for runs whose logs were fully retrievable: `air logs --download-to DIR` on the same run returned the complete log. This happens when Bricklens is enabled for the workspace but never ingested the run — it answers every request successfully with zero records, yet the logs are present in MLflow. The streaming (print) path only fell back to MLflow on errBricklensFeatureDisabled (gated off / not deployed / persistent failure); an empty-but-successful Bricklens response was treated as the final answer. The download path already reads from MLflow, which is why it worked. Treat "Bricklens served every request but never returned a record" the same as feature-disabled: hand off to the MLflow fallback, which owns the real no-logs report and preserves the run-derived exit code. This mirrors the Python CLI fix (databricks-eng/universe#2366012). Applies to both the terminal tail and the static (past-retry) view. Co-authored-by: Isaac --- experimental/air/cmd/logstream.go | 24 +++--- experimental/air/cmd/logstream_test.go | 104 +++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 12 deletions(-) diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index e5ab7ba8e85..e178be9d6da 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -33,8 +33,10 @@ const ( var retryCheckInterval = 3 * time.Second // errBricklensFeatureDisabled signals the caller to fall back to MLflow: Bricklens -// is gated off (FEATURE_DISABLED), not deployed (ENDPOINT_NOT_FOUND / 404), or -// persistently failing. The flag is evaluated server-side. +// is gated off (FEATURE_DISABLED), not deployed (ENDPOINT_NOT_FOUND / 404), +// persistently failing, or served every request successfully but never returned a +// record for a run whose logs may still be in MLflow. The flag is evaluated +// server-side. var errBricklensFeatureDisabled = errors.New("bricklens logs unavailable; falling back to mlflow") // logRequest describes what to fetch, shared by both backends so they honor the @@ -289,11 +291,10 @@ func (st *bricklensStreamer) run() (bool, error) { if terminal { if !st.firstLogSeen { - // Stop the spinner before the no-logs line so frames don't smear. - if st.onFirstLog != nil { - st.onFirstLog() - } - st.emitNoLogs() + // A successful but empty Bricklens stream isn't proof the run has no + // logs; they may be in MLflow (as --download-to reads). Fall back + // there, which owns the real no-logs report and the same exit code. + return false, errBricklensFeatureDisabled } log.Infof(st.ctx, "air logs: run %d finished in state %s", st.req.runID, st.status.displayState()) return st.status.succeeded(), nil @@ -326,7 +327,10 @@ func (st *bricklensStreamer) drainStatic(toSec int64) (bool, error) { return false, err } if !st.firstLogSeen { - st.emitNoLogs() + // An empty Bricklens tail doesn't mean the attempt has no logs; fall back to + // MLflow, which holds the immutable per-attempt artifacts. See the terminal + // branch in run. + return false, errBricklensFeatureDisabled } return st.status.succeeded(), nil } @@ -467,10 +471,6 @@ func (st *bricklensStreamer) emit(body string) { emitLogLine(st.out, st.req, body) } -func (st *bricklensStreamer) emitNoLogs() { - emitNoLogs(st.out, st.req, st.status) -} - // displayState is the result state, else the lifecycle state, else "UNKNOWN". func (s logRunStatus) displayState() string { if s.resultState != "" { diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go index d8a061687a6..b68c8b4d8b1 100644 --- a/experimental/air/cmd/logstream_test.go +++ b/experimental/air/cmd/logstream_test.go @@ -373,6 +373,110 @@ func TestRequestPageRetriesThenSucceeds(t *testing.T) { assert.Equal(t, 3, calls) } +// emptyLogsServer serves an empty Bricklens log response for any /logs request +// and a stub for everything else (SDK config probes, etc.). +func emptyLogsServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/logs") { + _, _ = w.Write([]byte(`{"log_records": []}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestStreamBricklensEmptyFallsBackToMLflow(t *testing.T) { + // Bricklens served every request but returned no record. That is not proof the + // run has no logs (they may be in MLflow), so the streamer must hand off via + // errBricklensFeatureDisabled and emit nothing itself, rather than reporting + // "No logs available" (the reported bug: the print path did, --download-to did not). + tests := []struct { + name string + req logRequest + status logRunStatus + }{ + { + name: "terminal run", + req: logRequest{runID: 123, node: 0, attempt: -1, tailLines: -1, jsonOutput: true}, + status: logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 1700000012000}, + }, + { + name: "static view of a past retry", + req: logRequest{runID: 123, node: 0, attempt: 0, tailLines: -1, staticView: true, jsonOutput: true}, + status: logRunStatus{lifeCycleState: "RUNNING"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + w := newTestWorkspaceClient(t, emptyLogsServer(t).URL) + _, err := streamBricklensLogs(t.Context(), w, &buf, tt.req, tt.status) + require.ErrorIs(t, err, errBricklensFeatureDisabled) + assert.Empty(t, buf.String(), "nothing should be emitted before the hand-off") + }) + } +} + +func TestStreamBricklensTerminalWithRecordsDoesNotFallBack(t *testing.T) { + // A terminal run whose Bricklens stream has records prints them and reports the + // run's outcome, without triggering the empty-result fallback. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/logs") { + _, _ = w.Write([]byte(`{"log_records": [{"time_unix_nano": 1700000001000000000, "body": "hello", "node_index": 0}]}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + w := newTestWorkspaceClient(t, srv.URL) + status := logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 1700000012000} + ok, err := streamBricklensLogs(t.Context(), w, &buf, logRequest{runID: 123, node: 0, attempt: -1, tailLines: -1, jsonOutput: true}, status) + require.NoError(t, err) + assert.True(t, ok) + assert.Contains(t, buf.String(), `"line":"hello"`) +} + +func TestFetchLogsFallsBackToMLflowWhenBricklensEmpty(t *testing.T) { + // End-to-end repro: a terminal SUCCESS run whose Bricklens stream is empty but + // whose logs are in MLflow. The print path must fall back to MLflow and print + // them, exactly as --download-to already does. + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/logs"): + _, _ = w.Write([]byte(`{"log_records": []}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{"run_id": 123, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, "tasks": [{"run_id": 456}]}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "E1", "mlflow_run_id": "R1"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0"}, {"path": "logs/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 + `/presigned"}]}`)) + case r.URL.Path == "/presigned": + _, _ = w.Write([]byte("line 1\nline 2\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + + var buf bytes.Buffer + w := newTestWorkspaceClient(t, srv.URL) + status := logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 1700000012000} + ok, err := fetchLogs(t.Context(), w, &buf, logRequest{runID: 123, node: 0, attempt: -1, tailLines: -1, jsonOutput: true}, status) + require.NoError(t, err) + assert.True(t, ok) + assert.Contains(t, buf.String(), `"line":"line 1"`) + assert.Contains(t, buf.String(), `"line":"line 2"`) +} + func TestSeenSetEviction(t *testing.T) { s := newSeenSet(2) s.add(1, "a")