From ee3a5beebbfad144345a89269b873a628ba9b052 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Thu, 27 Aug 2026 23:36:00 +0000 Subject: [PATCH] Improve AIR log retrieval reliability --- .nextchanges/cli/air-log-reliability.md | 1 + experimental/air/cmd/logdownload.go | 8 +- experimental/air/cmd/logmlflow.go | 199 +++++++++++++++++++++++- experimental/air/cmd/logmlflow_test.go | 136 ++++++++++++++++ experimental/air/cmd/logs.go | 29 ++++ experimental/air/cmd/logstream.go | 11 ++ experimental/air/cmd/logstream_test.go | 35 +++++ 7 files changed, 412 insertions(+), 7 deletions(-) create mode 100644 .nextchanges/cli/air-log-reliability.md diff --git a/.nextchanges/cli/air-log-reliability.md b/.nextchanges/cli/air-log-reliability.md new file mode 100644 index 00000000000..65863cee528 --- /dev/null +++ b/.nextchanges/cli/air-log-reliability.md @@ -0,0 +1 @@ +Improved AIR log fallback, retry-attempt selection, and Unity Catalog volume artifact handling. diff --git a/experimental/air/cmd/logdownload.go b/experimental/air/cmd/logdownload.go index 50f1c90ad3b..a0eeea33a05 100644 --- a/experimental/air/cmd/logdownload.go +++ b/experimental/air/cmd/logdownload.go @@ -80,7 +80,11 @@ func downloadLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Wri // A run with no logs is reported the same way as on the streaming path, so // the message and exit code agree between them. - ids := mlflowIDs(ctx, w, run) + attempt, taskRunID, err := resolveLogAttempt(run, req.attempt) + if err != nil { + return false, err + } + ids := mlflowIDsForTask(ctx, w, taskRunID) if ids == nil || ids.RunID == "" { emitNoLogs(out, req, status) return status.downloadOutcome(), nil @@ -96,7 +100,7 @@ func downloadLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Wri return false, fmt.Errorf("failed to create %s: %w", dir, err) } - nodeLogs, failures, err := downloadAllNodeLogs(ctx, w, ids.RunID, dir, nodes, req.attempt) + nodeLogs, failures, err := downloadAllNodeLogs(ctx, w, ids.RunID, dir, nodes, attempt) if err != nil { return false, err } diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go index debddc0d618..427c7efc824 100644 --- a/experimental/air/cmd/logmlflow.go +++ b/experimental/air/cmd/logmlflow.go @@ -3,6 +3,7 @@ package aircmd import ( "bufio" "context" + "errors" "fmt" "io" "net" @@ -12,16 +13,47 @@ import ( "regexp" "slices" "strconv" + "strings" + "sync" "time" + "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/listing" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/databricks/databricks-sdk-go/service/ml" ) +var mlflowArtifactRoots sync.Map + +func mlflowArtifactRoot(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) (string, error) { + key := strings.TrimRight(w.Config.Host, "/") + "\x00" + mlflowRunID + if cached, ok := mlflowArtifactRoots.Load(key); ok { + return cached.(string), nil + } + response, err := w.Experiments.GetRun(ctx, ml.GetRunRequest{RunId: mlflowRunID}) + if err != nil { + log.Debugf(ctx, "air logs: could not resolve MLflow artifact URI for %s: %v", mlflowRunID, err) + return "", nil + } + if response.Run == nil || response.Run.Info == nil { + return "", nil + } + root := strings.TrimRight(response.Run.Info.ArtifactUri, "/") + mlflowArtifactRoots.Store(key, root) + return root, nil +} + +func volumeArtifactRoot(root string) (string, bool) { + if strings.HasPrefix(root, "dbfs:/Volumes/") { + return strings.TrimPrefix(root, "dbfs:"), true + } + return "", false +} + // chunkFilePattern matches a log chunk file (logs-.chunk.txt); group 1 is // the chunk index. The sidecar splits stdout into 4MB chunks, index ascending. var chunkFilePattern = regexp.MustCompile(`^logs-(\d+)\.chunk\.txt$`) @@ -53,6 +85,13 @@ func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out i log.Debugf(ctx, "air logs: --minutes is not supported on the MLflow fallback path; showing the default tail") } + if !status.terminal() && !req.staticView { + return streamMLflowLogs(ctx, w, out, req, status) + } + return fetchMLflowLogTail(ctx, w, out, req, status) +} + +func fetchMLflowLogTail(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { mlflowRunID, logDir, err := resolveMLflowLogPath(ctx, w, req) if err != nil { return false, err @@ -94,19 +133,113 @@ func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out i return status.downloadOutcome(), nil } +func mlflowLogsExist(ctx context.Context, w *databricks.WorkspaceClient, req logRequest) bool { + mlflowRunID, logDir, err := resolveMLflowLogPath(ctx, w, req) + if err != nil || mlflowRunID == "" || logDir == "" { + if err != nil { + log.Debugf(ctx, "air logs: MLflow log-existence probe failed for run %d: %v", req.runID, err) + } + return false + } + chunks, err := listLogChunks(ctx, w, mlflowRunID, logDir) + if err != nil { + log.Debugf(ctx, "air logs: MLflow log-existence probe failed for run %d: %v", req.runID, err) + return false + } + return len(chunks) > 0 +} + +func streamMLflowLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + lineOffsets := make(map[string]int) + currentRunID := "" + printed := false + previousState := "" + + for { + if req.onStatusChange != nil { + current := status.displayState() + if current != previousState { + req.onStatusChange(current, previousState) + previousState = current + } + } + + mlflowRunID, logDir, err := resolveMLflowLogPath(ctx, w, req) + if err != nil { + if status.terminal() { + return false, err + } + log.Debugf(ctx, "air logs: failed to resolve MLflow log path: %v", err) + } else if mlflowRunID != "" && logDir != "" { + if currentRunID != mlflowRunID { + currentRunID = mlflowRunID + lineOffsets = make(map[string]int) + } + chunks, listErr := listLogChunks(ctx, w, mlflowRunID, logDir) + if listErr != nil { + if status.terminal() { + return false, listErr + } + log.Debugf(ctx, "air logs: failed to list MLflow log chunks: %v", listErr) + } else { + for _, chunk := range chunks { + lines, downloadErr := downloadChunkLines(ctx, w, mlflowRunID, chunk.path) + if downloadErr != nil { + if status.terminal() { + return false, downloadErr + } + log.Debugf(ctx, "air logs: failed to read MLflow log chunk %s: %v", chunk.path, downloadErr) + continue + } + offset := lineOffsets[chunk.path] + if offset > len(lines) { + offset = 0 + } + for _, line := range lines[offset:] { + emitLogLine(out, req, line) + printed = true + } + lineOffsets[chunk.path] = len(lines) + } + } + } + + if status.terminal() { + if !printed { + emitNoLogs(out, req, status) + } + return status.succeeded(), nil + } + if err := sleepOrCancel(ctx, retryCheckInterval); err != nil { + return false, err + } + refreshed, err := resolveRunStatus(ctx, w, req.runID) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) || ctx.Err() != nil { + return false, err + } + log.Debugf(ctx, "air logs: failed to refresh run status on MLflow fallback: %v", err) + continue + } + status = refreshed + } +} + // resolveMLflowLogPath returns the run's MLflow run id and per-node log directory. func resolveMLflowLogPath(ctx context.Context, w *databricks.WorkspaceClient, req logRequest) (string, string, error) { run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: req.runID}) if err != nil { return "", "", err } - ids := mlflowIDs(ctx, w, run) + attempt, taskRunID, err := resolveLogAttempt(run, req.attempt) + if err != nil { + return "", "", err + } + ids := mlflowIDsForTask(ctx, w, taskRunID) if ids == nil || ids.RunID == "" { return "", "", nil } - // -1 (latest) maps to attempt 0's directory. - attempt := max(req.attempt, 0) withAttempt, err := discoverAttemptPrefix(ctx, w, ids.RunID, attempt) if err != nil { return "", "", err @@ -224,8 +357,35 @@ func downloadChunkLines(ctx context.Context, w *databricks.WorkspaceClient, mlfl } // listArtifacts lists a run's artifacts under a path. -func listArtifacts(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, path string) ([]ml.FileInfo, error) { - it := w.Experiments.ListArtifacts(ctx, ml.ListArtifactsRequest{RunId: mlflowRunID, Path: path}) +func listArtifacts(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string) ([]ml.FileInfo, error) { + root, err := mlflowArtifactRoot(ctx, w, mlflowRunID) + if err != nil { + return nil, err + } + if volumeRoot, ok := volumeArtifactRoot(root); ok { + fc, err := filer.NewWorkspaceFilesClient(w, volumeRoot) + if err != nil { + return nil, err + } + entries, err := fc.ReadDir(ctx, artifactPath) + if err != nil { + return nil, err + } + files := make([]ml.FileInfo, 0, len(entries)) + for _, entry := range entries { + info := ml.FileInfo{Path: path.Join(artifactPath, entry.Name()), IsDir: entry.IsDir()} + if !entry.IsDir() { + stat, err := entry.Info() + if err != nil { + return nil, err + } + info.FileSize = stat.Size() + } + files = append(files, info) + } + return files, nil + } + it := w.Experiments.ListArtifacts(ctx, ml.ListArtifactsRequest{RunId: mlflowRunID, Path: artifactPath}) return listing.ToSlice(ctx, it) } @@ -247,6 +407,35 @@ type credentialsForReadResponse struct { // path. credentials-for-read returns a pre-signed URL, which we stream to disk; // that endpoint is not modeled by the SDK, so it is called via a raw client.Do. func downloadArtifact(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string) (string, error) { + root, err := mlflowArtifactRoot(ctx, w, mlflowRunID) + if err != nil { + return "", err + } + if volumeRoot, ok := volumeArtifactRoot(root); ok { + fc, err := filer.NewWorkspaceFilesClient(w, volumeRoot) + if err != nil { + return "", err + } + reader, err := fc.Read(ctx, artifactPath) + if err != nil { + return "", err + } + defer reader.Close() + tmp, err := os.CreateTemp("", "air-log-chunk-*") + if err != nil { + return "", err + } + if _, err := io.Copy(tmp, reader); err != nil { + tmp.Close() + os.Remove(tmp.Name()) + return "", err + } + if err := tmp.Close(); err != nil { + os.Remove(tmp.Name()) + return "", err + } + return tmp.Name(), nil + } apiClient, err := client.New(w.Config) if err != nil { return "", fmt.Errorf("failed to create API client: %w", err) diff --git a/experimental/air/cmd/logmlflow_test.go b/experimental/air/cmd/logmlflow_test.go index 44837999386..363f6a7e069 100644 --- a/experimental/air/cmd/logmlflow_test.go +++ b/experimental/air/cmd/logmlflow_test.go @@ -2,15 +2,46 @@ package aircmd import ( "bytes" + "fmt" "net/http" "net/http/httptest" "os" + "strings" + "sync/atomic" "testing" + "time" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +func TestResolveLogAttempt(t *testing.T) { + run := &jobs.Run{RunId: 5, Tasks: []jobs.RunTask{ + {RunId: 100, AttemptNumber: 0}, + {RunId: 101, AttemptNumber: 1}, + {RunId: 102, AttemptNumber: 2}, + }} + + attempt, taskRunID, err := resolveLogAttempt(run, -1) + require.NoError(t, err) + assert.Equal(t, 2, attempt) + assert.Equal(t, int64(102), taskRunID) + + attempt, taskRunID, err = resolveLogAttempt(run, 1) + require.NoError(t, err) + assert.Equal(t, 1, attempt) + assert.Equal(t, int64(101), taskRunID) + + _, _, err = resolveLogAttempt(run, 3) + require.ErrorContains(t, err, "available retries are 0 to 2") + _, _, err = resolveLogAttempt(run, -2) + require.ErrorContains(t, err, "must be -1 or greater") +} + func TestConstructLogPath(t *testing.T) { assert.Equal(t, "logs/node_0", constructLogPath(0, 0, false)) assert.Equal(t, "logs/node_3", constructLogPath(3, 2, false)) @@ -114,6 +145,111 @@ func TestMLflowFallbackNoLogsReflectsRunOutcome(t *testing.T) { assert.False(t, success) } +func TestStreamMLflowLogsFollowsWithoutDuplicates(t *testing.T) { + oldInterval := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = oldInterval }) + + var getRunCalls atomic.Int32 + var artifactReads atomic.Int32 + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.2/jobs/runs/get": + state := `{"life_cycle_state":"RUNNING"}` + if getRunCalls.Add(1) >= 4 { + state = `{"life_cycle_state":"TERMINATED","result_state":"SUCCESS"}` + } + _, _ = fmt.Fprintf(w, `{"run_id":5,"state":%s,"tasks":[{"run_id":101,"attempt_number":1}]}`, state) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output":{"mlflow_experiment_id":"exp","mlflow_run_id":"mlrun"}}`)) + case "/api/2.0/mlflow/runs/get": + _, _ = w.Write([]byte(`{}`)) + case "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files":[{"path":"logs/attempt_1","is_dir":true}]}`)) + } else { + _, _ = w.Write([]byte(`{"files":[{"path":"logs/attempt_1/node_0/logs-0.chunk.txt"}]}`)) + } + case "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = fmt.Fprintf(w, `{"credential_infos":[{"signed_uri":%q}]}`, base+"/artifact") + case "/artifact": + if artifactReads.Add(1) == 1 { + _, _ = w.Write([]byte("one\n")) + } else { + _, _ = w.Write([]byte("one\ntwo\n")) + } + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + + var out bytes.Buffer + success, err := streamMLflowLogs(t.Context(), newTestWorkspaceClient(t, srv.URL), &out, + logRequest{runID: 5, attempt: -1, tailLines: -1}, logRunStatus{lifeCycleState: "RUNNING", latestAttempt: 1}) + require.NoError(t, err) + assert.True(t, success) + assert.Equal(t, 1, strings.Count(out.String(), "one")) + assert.Equal(t, 1, strings.Count(out.String(), "two")) +} + +func TestVolumeArtifactListAndDownload(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + server.Handle("GET", "/api/2.0/mlflow/runs/get", func(req testserver.Request) any { + return map[string]any{"run": map[string]any{"info": map[string]any{ + "run_id": "volume-run", "artifact_uri": "dbfs:/Volumes/main/default/logs/root", + }}} + }) + testserver.AddDefaultHandlers(server) + + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + fc, err := filer.NewWorkspaceFilesClient(w, "/Volumes/main/default/logs/root") + require.NoError(t, err) + require.NoError(t, fc.Write(t.Context(), "logs/node_0/logs-0.chunk.txt", strings.NewReader("volume line\n"), filer.CreateParentDirectories)) + + chunks, err := listLogChunks(t.Context(), w, "volume-run", "logs/node_0") + require.NoError(t, err) + require.Len(t, chunks, 1) + assert.Equal(t, "logs/node_0/logs-0.chunk.txt", chunks[0].path) + + lines, err := downloadChunkLines(t.Context(), w, "volume-run", chunks[0].path) + require.NoError(t, err) + assert.Equal(t, []string{"volume line"}, lines) +} + +func TestMLflowArtifactRootCachesSuccessButNotFailure(t *testing.T) { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/2.0/mlflow/runs/get" { + _, _ = w.Write([]byte(`{}`)) + return + } + if calls.Add(1) == 1 { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code":"INTERNAL_ERROR","message":"retry"}`)) + return + } + _, _ = w.Write([]byte(`{"run":{"info":{"run_id":"cache-run","artifact_uri":"dbfs:/Volumes/main/default/v"}}}`)) + })) + t.Cleanup(srv.Close) + w := newTestWorkspaceClient(t, srv.URL) + + root, err := mlflowArtifactRoot(t.Context(), w, "cache-run") + require.NoError(t, err) + assert.Empty(t, root) + root, err = mlflowArtifactRoot(t.Context(), w, "cache-run") + require.NoError(t, err) + assert.Equal(t, "dbfs:/Volumes/main/default/v", root) + root, err = mlflowArtifactRoot(t.Context(), w, "cache-run") + require.NoError(t, err) + assert.Equal(t, "dbfs:/Volumes/main/default/v", root) + assert.Equal(t, int32(2), calls.Load()) +} + func TestDownloadArtifactSendsUnbracketedPath(t *testing.T) { var gotPath string var base string diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index 19861b9e957..e1afd1c0a33 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "io" + "slices" "strconv" "github.com/databricks/cli/cmd/root" @@ -13,6 +14,7 @@ import ( "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/spf13/cobra" ) @@ -85,6 +87,10 @@ func newLogsCommand() *cobra.Command { return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, fmt.Errorf("invalid --node %d: must not be negative", node)) } + if retry < -1 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --retry %d: must be -1 or greater", retry)) + } runID, err := strconv.ParseInt(args[0], 10, 64) if err != nil || runID <= 0 { @@ -181,6 +187,29 @@ func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error { return nil } +func resolveLogAttempt(run *jobs.Run, requested int) (int, int64, error) { + if requested < -1 { + return 0, 0, fmt.Errorf("invalid retry %d: must be -1 or greater", requested) + } + if len(run.Tasks) == 0 { + return 0, 0, nil + } + latest := latestAttemptNumber(run) + attempt := requested + if attempt < 0 { + attempt = latest + } + if attempt > latest { + return 0, 0, fmt.Errorf("invalid retry %d: available retries are 0 to %d", requested, latest) + } + for _, v := range slices.Backward(run.Tasks) { + if v.AttemptNumber == attempt { + return attempt, v.RunId, nil + } + } + return 0, 0, fmt.Errorf("run %d has no task for retry %d", run.RunId, attempt) +} + // 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) { diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index 51e7e54a768..944102d1111 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -30,6 +30,9 @@ const ( // statusMessageRefreshEveryNPolls throttles the status_message fetch so the // waiting spinner doesn't issue a get-output on every poll tick. statusMessageRefreshEveryNPolls = 5 + // bricklensEmptyMLflowProbeEveryNPolls throttles the active-run MLflow + // existence probe while Bricklens has returned no records. + bricklensEmptyMLflowProbeEveryNPolls = 10 ) // statusMessageType tags a client-facing message packed into @@ -324,6 +327,7 @@ func (st *bricklensStreamer) run() (bool, error) { // server status_message fetch to every Nth poll, and lastSpinnerText avoids // redundant spinner updates. statusRefreshCounter := 0 + emptyStreamPolls := 0 lastSpinnerText := "" for { if !firstIteration { @@ -388,6 +392,13 @@ 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.firstLogSeen { + emptyStreamPolls++ + if emptyStreamPolls%bricklensEmptyMLflowProbeEveryNPolls == 0 && mlflowLogsExist(st.ctx, st.w, st.req) { + log.Debugf(st.ctx, "air logs: MLflow artifacts exist for run %d; falling back to MLflow log stream", st.req.runID) + return false, errBricklensFeatureDisabled + } + } firstIteration = false if err := sleepOrCancel(st.ctx, retryCheckInterval); err != nil { diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go index be4a9d0c452..99ac88f94c7 100644 --- a/experimental/air/cmd/logstream_test.go +++ b/experimental/air/cmd/logstream_test.go @@ -498,6 +498,41 @@ func TestStreamBricklensTerminalWithRecordsDoesNotFallBack(t *testing.T) { assert.Contains(t, buf.String(), `"line":"hello"`) } +func TestStreamBricklensActiveEmptyProbesMLflow(t *testing.T) { + oldInterval := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = oldInterval }) + + 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":"RUNNING"},"tasks":[{"run_id":456,"attempt_number":0}]}`)) + 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/runs/get": + _, _ = w.Write([]byte(`{}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files":[{"path":"logs/node_0","is_dir":true}]}`)) + } else { + _, _ = w.Write([]byte(`{"files":[{"path":"logs/node_0/logs-0.chunk.txt"}]}`)) + } + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + + var out bytes.Buffer + _, err := streamBricklensLogs(t.Context(), newTestWorkspaceClient(t, srv.URL), &out, + logRequest{runID: 123, attempt: -1, tailLines: -1, jsonOutput: true}, + logRunStatus{lifeCycleState: "RUNNING"}) + require.ErrorIs(t, err, errBricklensFeatureDisabled) + assert.Empty(t, out.String()) +} + 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