From 9f8775184c02c0f7594fe002fee350dbc275d8e6 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Fri, 17 Jul 2026 00:13:46 +0000 Subject: [PATCH 1/5] AIR CLI: port `air logs` with Bricklens-first, MLflow-fallback (phase 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement the core of `air logs`: fetch a run's training logs from the Bricklens-backed AiTrainingService endpoint, and on a backend-signaled failure fall back to the MLflow log path. The feature flag is a backend SAFE flag — the CLI never evaluates it; classifyLogError only reads the response error code (FEATURE_DISABLED / ENDPOINT_NOT_FOUND / 404), and also falls back after maxTransientFailures consecutive transient failures. The Bricklens/MLflow decision is a single consolidated try/catch in fetchLogs. Adds `--minutes N` (Bricklens-native time window) alongside `--lines N` (MLflow-native line tail); the two are mutually exclusive and both flow through the shared logRequest so either backend honors them. Streaming + status resolution live in one file (logstream.go) as a small streamer type. The MLflow fallback is a stub returning an explicit not-yet-implemented error (phase 3 fills it in); --download-to and --review are rejected explicitly rather than silently ignored. Co-authored-by: Isaac --- acceptance/experimental/air/help/output.txt | 21 + acceptance/experimental/air/help/script | 3 + .../experimental/air/logs/out.test.toml | 3 + acceptance/experimental/air/logs/output.txt | 48 ++ acceptance/experimental/air/logs/script | 20 + acceptance/experimental/air/logs/test.toml | 29 ++ .../experimental/air/unimplemented/output.txt | 6 - .../experimental/air/unimplemented/script | 3 - experimental/air/cmd/logbricklens.go | 100 +++++ experimental/air/cmd/logbricklens_test.go | 78 ++++ experimental/air/cmd/logmlflow.go | 22 + experimental/air/cmd/logs.go | 134 +++++- experimental/air/cmd/logs_test.go | 182 ++++++++ experimental/air/cmd/logstream.go | 414 ++++++++++++++++++ experimental/air/cmd/logstream_support.go | 76 ++++ experimental/air/cmd/logstream_test.go | 165 +++++++ experimental/air/cmd/stubs_test.go | 1 - 17 files changed, 1291 insertions(+), 14 deletions(-) create mode 100644 acceptance/experimental/air/logs/out.test.toml create mode 100644 acceptance/experimental/air/logs/output.txt create mode 100644 acceptance/experimental/air/logs/script create mode 100644 acceptance/experimental/air/logs/test.toml create mode 100644 experimental/air/cmd/logbricklens.go create mode 100644 experimental/air/cmd/logbricklens_test.go create mode 100644 experimental/air/cmd/logmlflow.go create mode 100644 experimental/air/cmd/logs_test.go create mode 100644 experimental/air/cmd/logstream.go create mode 100644 experimental/air/cmd/logstream_support.go create mode 100644 experimental/air/cmd/logstream_test.go diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index ee89e778d6b..52e50016843 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -47,3 +47,24 @@ Global Flags: -o, --output type output type: text or json (default text) -p, --profile string ~/.databrickscfg profile -t, --target string bundle target to use (if applicable) + +=== logs help +>>> [CLI] experimental air logs --help +Stream logs from an active run, or fetch logs from a completed run. + +Usage: + databricks experimental air logs JOB_RUN_ID [flags] + +Flags: + --download-to string Download all logs to this directory instead of printing + -h, --help help for logs + --lines int For completed runs, print the last N lines (default 10000) + --minutes int Fetch only logs from the last N minutes + --node int Fetch logs from this node + --retry int View logs from a specific retry attempt; -1 means latest (default -1) + +Global Flags: + --debug enable debug logging + -o, --output type output type: text or json (default text) + -p, --profile string ~/.databrickscfg profile + -t, --target string bundle target to use (if applicable) diff --git a/acceptance/experimental/air/help/script b/acceptance/experimental/air/help/script index 81f3907e4f5..91dc12567a6 100644 --- a/acceptance/experimental/air/help/script +++ b/acceptance/experimental/air/help/script @@ -6,3 +6,6 @@ trace $CLI experimental air --help title "list help" trace $CLI experimental air list --help + +title "logs help" +trace $CLI experimental air logs --help diff --git a/acceptance/experimental/air/logs/out.test.toml b/acceptance/experimental/air/logs/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/logs/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/logs/output.txt b/acceptance/experimental/air/logs/output.txt new file mode 100644 index 00000000000..0e1ad38ff60 --- /dev/null +++ b/acceptance/experimental/air/logs/output.txt @@ -0,0 +1,48 @@ + +=== logs (text, completed run) +>>> [CLI] experimental air logs 123 +step 1 +step 2 + +=== logs (json) +>>> [CLI] experimental air logs 123 -o json +{"type":"LOG","ts":"[TIMESTAMP]","node":0,"line":"step 1"} +{"type":"LOG","ts":"[TIMESTAMP]","node":0,"line":"step 2"} + +=== logs with --minutes +>>> [CLI] experimental air logs 123 --minutes 30 +step 1 +step 2 + +=== logs --lines and --minutes are mutually exclusive +>>> [CLI] experimental air logs 123 --lines 100 --minutes 30 +Error: cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window + +Exit code: 1 + +=== logs --lines and --minutes are mutually exclusive (json) +>>> [CLI] experimental air logs 123 --lines 100 --minutes 30 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "error": { + "code": "INVALID_ARGS", + "kind": "PERMANENT", + "message": "cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window", + "retryable": false + } +} + +Exit code: 1 + +=== invalid run id +>>> [CLI] experimental air logs notanumber +Error: invalid JOB_RUN_ID "notanumber": must be a positive integer + +Exit code: 1 + +=== --download-to not implemented +>>> [CLI] experimental air logs 123 --download-to /tmp/out +Error: --download-to is not implemented yet + +Exit code: 1 diff --git a/acceptance/experimental/air/logs/script b/acceptance/experimental/air/logs/script new file mode 100644 index 00000000000..3782221c3cf --- /dev/null +++ b/acceptance/experimental/air/logs/script @@ -0,0 +1,20 @@ +title "logs (text, completed run)" +trace $CLI experimental air logs 123 + +title "logs (json)" +trace $CLI experimental air logs 123 -o json + +title "logs with --minutes" +trace $CLI experimental air logs 123 --minutes 30 + +title "logs --lines and --minutes are mutually exclusive" +errcode trace $CLI experimental air logs 123 --lines 100 --minutes 30 + +title "logs --lines and --minutes are mutually exclusive (json)" +errcode trace $CLI experimental air logs 123 --lines 100 --minutes 30 -o json + +title "invalid run id" +errcode trace $CLI experimental air logs notanumber + +title "--download-to not implemented" +errcode trace $CLI experimental air logs 123 --download-to /tmp/out diff --git a/acceptance/experimental/air/logs/test.toml b/acceptance/experimental/air/logs/test.toml new file mode 100644 index 00000000000..dfec0f46a7c --- /dev/null +++ b/acceptance/experimental/air/logs/test.toml @@ -0,0 +1,29 @@ +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A completed run: Bricklens serves its logs via the tail drain. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"task_key": "train", "run_id": 456, "attempt_number": 0}] +} +''' + +# Bricklens log records, returned newest-first (as the tail fetch requests); +# printed oldest-first. +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows/by-run-id/123/logs" +Response.Body = ''' +{"log_records": [ + {"time_unix_nano": 1700000002000000000, "body": "step 2", "node_index": 0}, + {"time_unix_nano": 1700000001000000000, "body": "step 1", "node_index": 0} +]} +''' diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt index 7db6ef1aec2..261e0c456a9 100644 --- a/acceptance/experimental/air/unimplemented/output.txt +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -1,10 +1,4 @@ -=== logs ->>> [CLI] experimental air logs 123 -Error: `air logs` is not implemented yet - -Exit code: 1 - === register-image >>> [CLI] experimental air register-image my-image:latest Error: `air register-image` is not implemented yet diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script index 19dc13ffe85..ae0366a48c9 100644 --- a/acceptance/experimental/air/unimplemented/script +++ b/acceptance/experimental/air/unimplemented/script @@ -1,7 +1,4 @@ # Each stub must fail with "not implemented"; errcode records the exit code. -title "logs" -errcode trace $CLI experimental air logs 123 - title "register-image" errcode trace $CLI experimental air register-image my-image:latest diff --git a/experimental/air/cmd/logbricklens.go b/experimental/air/cmd/logbricklens.go new file mode 100644 index 00000000000..47ba7a54e7b --- /dev/null +++ b/experimental/air/cmd/logbricklens.go @@ -0,0 +1,100 @@ +package aircmd + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strconv" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" +) + +// bricklensLogsPathFmt is the Bricklens-backed training-workflow log endpoint, +// keyed by Jobs run id. It is called with a raw client.Do because the SDK does +// not model the AiTrainingService log surface. +const bricklensLogsPathFmt = "/api/2.0/ai-training/workflows/by-run-id/%d/logs" + +// logRecord is one log line from Bricklens. time_unix_nano stamps ordering and +// dedup; body is the raw line; node_index identifies the emitting node. +type logRecord struct { + // TimeUnixNano is nanosecond units; tolerate it arriving as a JSON number or + // string, matching the AiTrainingService index quirk in aitraining.go. + TimeUnixNano json.Number `json:"time_unix_nano"` + Body string `json:"body"` + NodeIndex int `json:"node_index"` +} + +// nano returns the record's time_unix_nano as an int64, or 0 when absent or +// unparseable (so it sorts first and never blocks dedup). +func (r logRecord) nano() int64 { + n, err := r.TimeUnixNano.Int64() + if err != nil { + return 0 + } + return n +} + +type bricklensLogsResponse struct { + LogRecords []logRecord `json:"log_records"` + NextPageToken string `json:"next_page_token"` +} + +// bricklensLogsQuery is the request-field surface of the log endpoint. The +// fields map to the query keys the Python CLI sends (ai_training_client.get_logs); +// zero-valued optionals are omitted so the endpoint applies its own defaults. +type bricklensLogsQuery struct { + // fromSeconds and toSeconds bound the query window in Unix epoch seconds. + fromSeconds int64 + toSeconds int64 + pageToken string + pageSize int + // attemptNumber selects a retry attempt (0-indexed); -1 means latest. + attemptNumber int + nodeIndex int + // ascending returns oldest-first; the endpoint defaults to ascending when + // the field is absent, so the tail fetch must send it explicitly false. + ascending bool +} + +// getBricklensLogs fetches one page of logs from the Bricklens endpoint. It +// returns the raw error so the caller can classify it (feature-gated vs +// transient vs genuine not-found) via classifyLogError. +func getBricklensLogs(ctx context.Context, w *databricks.WorkspaceClient, runID int64, q bricklensLogsQuery) (*bricklensLogsResponse, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + query := map[string]any{ + // The endpoint defaults to ascending when absent, so always send it: the + // tail path relies on an explicit false to get newest-first. + "ascending": strconv.FormatBool(q.ascending), + } + if q.fromSeconds > 0 { + query["from"] = q.fromSeconds + } + if q.toSeconds > 0 { + query["to"] = q.toSeconds + } + if q.pageToken != "" { + query["page_token"] = q.pageToken + } + if q.pageSize > 0 { + query["page_size"] = q.pageSize + } + if q.attemptNumber >= 0 { + query["ref.attempt_number"] = q.attemptNumber + } + if q.nodeIndex >= 0 { + query["filter.node_index"] = q.nodeIndex + } + + var resp bricklensLogsResponse + path := fmt.Sprintf(bricklensLogsPathFmt, runID) + if err := apiClient.Do(ctx, http.MethodGet, path, nil, nil, query, &resp); err != nil { + return nil, err + } + return &resp, nil +} diff --git a/experimental/air/cmd/logbricklens_test.go b/experimental/air/cmd/logbricklens_test.go new file mode 100644 index 00000000000..07fd62b7ffd --- /dev/null +++ b/experimental/air/cmd/logbricklens_test.go @@ -0,0 +1,78 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGetBricklensLogsQuerySerialization(t *testing.T) { + var got url.Values + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/oidc/.well-known/oauth-authorization-server" { + gotPath = r.URL.Path + got = r.URL.Query() + } + _, _ = w.Write([]byte(`{"log_records": [], "next_page_token": ""}`)) + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + _, err := getBricklensLogs(t.Context(), w, 42, bricklensLogsQuery{ + fromSeconds: 100, + toSeconds: 200, + pageToken: "tok", + pageSize: 500, + attemptNumber: 1, + nodeIndex: 3, + ascending: true, + }) + require.NoError(t, err) + + assert.Equal(t, "/api/2.0/ai-training/workflows/by-run-id/42/logs", gotPath) + assert.Equal(t, "100", got.Get("from")) + assert.Equal(t, "200", got.Get("to")) + assert.Equal(t, "tok", got.Get("page_token")) + assert.Equal(t, "500", got.Get("page_size")) + assert.Equal(t, "1", got.Get("ref.attempt_number")) + assert.Equal(t, "3", got.Get("filter.node_index")) + assert.Equal(t, "true", got.Get("ascending")) +} + +func TestGetBricklensLogsOmitsOptionals(t *testing.T) { + var got url.Values + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/oidc/.well-known/oauth-authorization-server" { + got = r.URL.Query() + } + _, _ = w.Write([]byte(`{"log_records": []}`)) + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + // attempt -1 (latest) and node 0 are the default request; from/to/page unset. + _, err := getBricklensLogs(t.Context(), w, 7, bricklensLogsQuery{attemptNumber: -1, nodeIndex: 0}) + require.NoError(t, err) + + assert.False(t, got.Has("from")) + assert.False(t, got.Has("to")) + assert.False(t, got.Has("page_token")) + assert.False(t, got.Has("page_size")) + // -1 attempt means "latest" — the field is omitted so the endpoint defaults. + assert.False(t, got.Has("ref.attempt_number")) + // node 0 is a real filter value and must be sent. + assert.Equal(t, "0", got.Get("filter.node_index")) + // ascending is always sent so the tail path can force newest-first. + assert.Equal(t, "false", got.Get("ascending")) +} + +func TestLogRecordNano(t *testing.T) { + assert.Equal(t, int64(123), logRecord{TimeUnixNano: "123"}.nano()) + assert.Equal(t, int64(0), logRecord{TimeUnixNano: ""}.nano()) + assert.Equal(t, int64(0), logRecord{TimeUnixNano: "notanumber"}.nano()) +} diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go new file mode 100644 index 00000000000..faf5e41cc03 --- /dev/null +++ b/experimental/air/cmd/logmlflow.go @@ -0,0 +1,22 @@ +package aircmd + +import ( + "context" + "fmt" + "io" + + "github.com/databricks/databricks-sdk-go" +) + +// mlflowLogFallback prints a run's logs from MLflow artifacts. It is the fallback +// used when Bricklens can't serve logs (errBricklensFeatureDisabled). Both the +// completed-run tail (--lines) and the active-run stream, plus the --minutes +// window, are honored here via req. +// +// TODO(air-logs-m3): port print_all_logs / monitor_and_stream_logs — MLflow +// artifact discovery (logs/[attempt_X/]node_Y), chunk walking (logs-N.chunk.txt), +// and credential-vending download. Until then the fallback reports that the +// legacy path is not yet available so the failure is explicit rather than silent. +func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + return false, fmt.Errorf("logs for run %d are not available via Bricklens and the MLflow fallback is not yet implemented", req.runID) +} diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index c34fb62a7df..435d58de447 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -1,7 +1,18 @@ package aircmd import ( + "context" + "errors" + "fmt" + "io" + "strconv" + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/spf13/cobra" ) @@ -9,6 +20,7 @@ func newLogsCommand() *cobra.Command { var ( node int lines int + minutes int retry int downloadTo string review bool @@ -19,18 +31,132 @@ func newLogsCommand() *cobra.Command { Args: root.ExactArgs(1), Short: "Stream or fetch logs for a run", Long: `Stream logs from an active run, or fetch logs from a completed run.`, - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("logs") - }, } cmd.Flags().IntVar(&node, "node", 0, "Fetch logs from this node") - cmd.Flags().IntVar(&lines, "lines", 10000, "For completed runs, print the last N lines") + cmd.Flags().IntVar(&lines, "lines", 0, "For completed runs, print the last N lines (default 10000)") + cmd.Flags().IntVar(&minutes, "minutes", 0, "Fetch only logs from the last N minutes") cmd.Flags().IntVar(&retry, "retry", -1, "View logs from a specific retry attempt; -1 means latest") cmd.Flags().StringVar(&downloadTo, "download-to", "", "Download all logs to this directory instead of printing") cmd.Flags().BoolVar(&review, "review", false, "Download logs from all nodes and filter for error signatures") // Hidden in the Python `air` CLI (help=argparse.SUPPRESS); keep it internal here to match. cmd.Flags().MarkHidden("review") + // In -o json mode an auth failure should be a JSON error envelope, not a bare + // error. ErrAlreadyPrinted passes through (already handled upstream). + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + err := root.MustWorkspaceClient(cmd, args) + if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { + return err + } + return authError(cmd.Context(), cmd, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + + // --download-to and --review are not yet ported. Reject them explicitly + // rather than silently ignoring, so the user knows they had no effect. + if downloadTo != "" { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--download-to is not implemented yet")) + } + if review { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--review is not implemented yet")) + } + + // --lines is a line tail (MLflow-native); --minutes is a time window + // (Bricklens-native). They answer the same "how much" question two ways, + // so reject both at once rather than silently honoring one. + if lines > 0 && minutes > 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window")) + } + if lines < 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --lines %d: must be positive", lines)) + } + if minutes < 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --minutes %d: must be positive", minutes)) + } + + runID, err := strconv.ParseInt(args[0], 10, 64) + if err != nil || runID <= 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid JOB_RUN_ID %q: must be a positive integer", args[0])) + } + + return runLogs(ctx, cmd, logRequest{ + runID: runID, + node: node, + attempt: retry, + windowMinutes: minutes, + tailLines: lines, + jsonOutput: root.OutputType(cmd) == flags.OutputJSON, + }) + } + return cmd } + +// runLogs resolves the run, validates the requested retry, and fetches logs +// Bricklens-first with an MLflow fallback. It handles error reporting; the +// backend selection lives in fetchLogs. +func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error { + w := cmdctx.WorkspaceClient(ctx) + + // Validate credentials server-side before fetching (MustWorkspaceClient only + // attaches credentials), so a bad token fails clearly here. + if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil { + return authError(ctx, cmd, err) + } + + status, err := resolveRunStatus(ctx, w, req.runID) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, + fmt.Errorf("run %d not found: check the run ID and that it is a job run ID", req.runID)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to get status for run %d: %w", req.runID, err)) + } + + // Resolve --retry against the run's attempts. -1 (default) means latest. + if req.attempt >= 0 && req.attempt > status.latestAttempt { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid retry %d: available retries are 0 to %d", req.attempt, status.latestAttempt)) + } + + out := cmd.OutOrStdout() + success, err := fetchLogs(ctx, w, out, req, status) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, + fmt.Errorf("run %d not found: check the run ID and that it is a job run ID", req.runID)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to fetch logs for run %d: %w", req.runID, err)) + } + + // A run that finished unsuccessfully exits non-zero, matching the Python CLI; + // output was already written, so don't reprint via Cobra. + if !success { + return root.ErrAlreadyPrinted + } + return nil +} + +// fetchLogs is the Bricklens-first / MLflow-fallback decision. Bricklens serves +// logs unless the backend has gated it off (a SAFE flag) or it is unavailable; +// in that case streamBricklensLogs returns errBricklensFeatureDisabled and we +// fall back to MLflow. Both paths honor the same logRequest (node, retry, +// --minutes window, --lines tail). +func fetchLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + success, err := streamBricklensLogs(ctx, w, out, req, status) + if errors.Is(err, errBricklensFeatureDisabled) { + return mlflowLogFallback(ctx, w, out, req, status) + } + return success, err +} diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go new file mode 100644 index 00000000000..10f57de4f2c --- /dev/null +++ b/experimental/air/cmd/logs_test.go @@ -0,0 +1,182 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLogsCommandShape(t *testing.T) { + cmd := newLogsCommand() + assert.Equal(t, "logs JOB_RUN_ID", cmd.Use) + assert.Empty(t, cmd.Commands(), "logs must not register subcommands") + assert.NoError(t, cmd.Args(cmd, []string{"123"})) + assert.Error(t, cmd.Args(cmd, []string{})) + assert.Error(t, cmd.Args(cmd, []string{"1", "2"})) + + // --review stays hidden to match the Python CLI. + review := cmd.Flags().Lookup("review") + require.NotNil(t, review) + assert.True(t, review.Hidden) +} + +// runLogsCmd invokes the logs command's RunE with the given flags against a mock +// (no-HTTP) workspace client. Used for input validation that fails before any +// API call. +func runLogsCmd(t *testing.T, args []string, flagsToSet map[string]string) error { + t.Helper() + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newLogsCommand(), flags.OutputText) + for k, v := range flagsToSet { + require.NoError(t, cmd.Flags().Set(k, v)) + } + cmd.SetContext(ctx) + return cmd.RunE(cmd, args) +} + +func TestLogsFlagValidation(t *testing.T) { + tests := []struct { + name string + args []string + flags map[string]string + wantMsg string + }{ + { + name: "lines and minutes are mutually exclusive", + args: []string{"5"}, + flags: map[string]string{"lines": "100", "minutes": "10"}, + wantMsg: "cannot combine --lines with --minutes", + }, + { + name: "negative lines rejected", + args: []string{"5"}, + flags: map[string]string{"lines": "-1"}, + wantMsg: "invalid --lines", + }, + { + name: "negative minutes rejected", + args: []string{"5"}, + flags: map[string]string{"minutes": "-1"}, + wantMsg: "invalid --minutes", + }, + { + name: "download-to not implemented", + args: []string{"5"}, + flags: map[string]string{"download-to": "/tmp/logs"}, + wantMsg: "--download-to is not implemented yet", + }, + { + name: "review not implemented", + args: []string{"5"}, + flags: map[string]string{"review": "true"}, + wantMsg: "--review is not implemented yet", + }, + { + name: "invalid run id", + args: []string{"abc"}, + wantMsg: "invalid JOB_RUN_ID", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := runLogsCmd(t, tt.args, tt.flags) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantMsg) + }) + } +} + +// completedRunLogsServer serves the auth probe, a terminal runs/get, and a +// single page of Bricklens logs (newest-first, as the tail fetch requests). +func completedRunLogsServer(t *testing.T) *httptest.Server { + t.Helper() + 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, + "start_time": 1000, + "end_time": 2000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"attempt_number": 0}] + }`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + // Newest-first, as drainTail requests; reversed to oldest-first on print. + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 2000000000, "body": "line two", "node_index": 0}, + {"time_unix_nano": 1000000000, "body": "line one", "node_index": 0} + ]}`)) + default: + // Me() probe and SDK config discovery. + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestLogsCompletedRunTail(t *testing.T) { + srv := completedRunLogsServer(t) + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + // Drive runLogs directly (bypassing PreRunE auth wiring) with a resolved request. + err := runLogs(ctx, cmd, logRequest{runID: 5, node: 0, attempt: -1}) + require.NoError(t, err) + + // Records print oldest-first regardless of the newest-first fetch order. + assert.Equal(t, "line one\nline two\n", buf.String()) +} + +// featureDisabledLogsServer serves a terminal runs/get and a Bricklens endpoint +// gated off by the backend SAFE flag (FEATURE_DISABLED), so the command must +// fall back to the MLflow path. +func featureDisabledLogsServer(t *testing.T) *httptest.Server { + t.Helper() + 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": [{"attempt_number": 0}] + }`)) + 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": "bricklens logs gated off"}`)) + default: + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestLogsFallsBackToMLflow(t *testing.T) { + srv := featureDisabledLogsServer(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + + // The Bricklens endpoint is gated off, so fetchLogs routes to the MLflow + // fallback. Until that path is implemented (air-logs-m3) it reports the + // explicit not-yet-available error, which proves the try/catch wiring fires. + err := runLogs(ctx, cmd, logRequest{runID: 5, node: 0, attempt: -1}) + require.Error(t, err) + assert.Contains(t, err.Error(), "MLflow fallback is not yet implemented") +} diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go new file mode 100644 index 00000000000..57acb56c9ce --- /dev/null +++ b/experimental/air/cmd/logstream.go @@ -0,0 +1,414 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "slices" + "time" + + "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/service/jobs" +) + +// Log-stream tuning constants. These mirror the Python CLI (log_streaming.py) so +// behavior stays identical after the port. +const ( + // retryCheckInterval is how long the poll loop waits between status/log polls. + retryCheckInterval = 3 * time.Second + // maxTransientFailures is how many consecutive Bricklens failures we tolerate + // before treating it as unavailable and falling back to MLflow. + maxTransientFailures = 5 + // defaultCompletedRunTailLines caps a completed run's output when neither + // --lines nor --minutes is set, so a multi-chunk log does not flood stdout. + defaultCompletedRunTailLines = 10000 + // seenRecordsCap bounds the dedup set so a large initial drain does not + // accumulate every record; oldest-inserted entries are evicted first. + seenRecordsCap = 100000 +) + +// errBricklensFeatureDisabled signals the caller to fall back to the MLflow log +// path. It is returned when Bricklens can't serve logs: the endpoint is gated +// off by a backend SAFE flag (FEATURE_DISABLED), not deployed (ENDPOINT_NOT_FOUND +// / 404), or persistently failing after maxTransientFailures retries. The flag is +// evaluated server-side; the CLI only reads the resulting error code. +var errBricklensFeatureDisabled = errors.New("bricklens logs unavailable; falling back to mlflow") + +// logRequest is the resolved, backend-agnostic description of what to fetch. It +// is shared by the Bricklens streamer and the MLflow fallback so both honor the +// same flags. windowMinutes and tailLines are mutually exclusive (validated at +// the command layer): windowMinutes drives the time window, tailLines the tail. +type logRequest struct { + runID int64 + // node is the node index to fetch; 0 by default (node 0 always exists). + node int + // attempt is the retry attempt to read; -1 means latest. + attempt int + // windowMinutes, when > 0, restricts the fetch to the last N minutes. + windowMinutes int + // tailLines, when > 0, keeps only the last N lines of a completed run. + tailLines int + jsonOutput bool +} + +// runStatus is the subset of a run's state the log path needs. It is resolved +// once from a Jobs GetRun and reused, avoiding a per-tick typed-vs-dict shuffle. +type logRunStatus struct { + lifeCycleState string + resultState string + stateMessage string + startTimeMs int64 + endTimeMs int64 + // latestAttempt is the highest attempt_number across the run's tasks. + latestAttempt int +} + +// terminalLifeCycleStates and terminalResultStates classify a finished run. A +// run is terminal when its lifecycle state is terminal, or a result state is set +// (result states only appear on terminal runs). Mirrors log_streaming.py. +var ( + terminalLifeCycleStates = map[string]bool{"TERMINATED": true, "SKIPPED": true, "INTERNAL_ERROR": true} + terminalResultStates = map[string]bool{"SUCCESS": true, "FAILED": true, "CANCELED": true} +) + +func (s logRunStatus) terminal() bool { + return terminalLifeCycleStates[s.lifeCycleState] || terminalResultStates[s.resultState] +} + +func (s logRunStatus) succeeded() bool { + return s.resultState == "SUCCESS" +} + +// resolveRunStatus fetches a run's state via the Jobs API and projects it onto +// logRunStatus. The run id being unknown surfaces as apierr.ErrResourceDoesNotExist +// so the caller can report a clean not-found. +func resolveRunStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (logRunStatus, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + return logRunStatus{}, err + } + return projectRunStatus(run), nil +} + +// projectRunStatus extracts logRunStatus from a Jobs run. Split out from +// resolveRunStatus so it can be unit-tested without an API client. +func projectRunStatus(run *jobs.Run) logRunStatus { + s := logRunStatus{ + startTimeMs: run.StartTime, + endTimeMs: run.EndTime, + } + if run.State != nil { + s.lifeCycleState = string(run.State.LifeCycleState) + s.resultState = string(run.State.ResultState) + s.stateMessage = run.State.StateMessage + } + for i := range run.Tasks { + s.latestAttempt = max(s.latestAttempt, run.Tasks[i].AttemptNumber) + } + return s +} + +// classifyLogError decides how a Bricklens failure should be handled: +// - errBricklensFeatureDisabled: fall back to MLflow (flag gated off, endpoint +// absent, or a 404 — older logs may still live in MLflow). +// - the original error: a genuine not-found, surfaced as-is. +// - nil: a transient failure the caller should retry. +// +// The backend evaluates the SAFE flag; this only reads the returned error code, +// so it never string-matches error text (repo rule). +func classifyLogError(err error) error { + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + switch apiErr.ErrorCode { + case "FEATURE_DISABLED", "ENDPOINT_NOT_FOUND": + return errBricklensFeatureDisabled + } + if apiErr.StatusCode == http.StatusNotFound { + return errBricklensFeatureDisabled + } + } + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return err + } + return nil +} + +// fromSeconds computes the Bricklens `from` bound for a request. With --minutes +// set it is now-N*60; otherwise it is the run's start second (0 when the run has +// not started, which the endpoint reads as "everything stored"). Matches +// stream_logs_via_bricklens. +func (req logRequest) fromSeconds(status logRunStatus, now time.Time) int64 { + if req.windowMinutes > 0 { + return now.Add(-time.Duration(req.windowMinutes) * time.Minute).Unix() + } + if status.startTimeMs > 0 { + return status.startTimeMs / 1000 + } + return 0 +} + +// toSeconds computes the Bricklens `to` bound. Once a run is terminal no new logs +// can appear, so we cap at the run's end second (ceil of the millisecond end time +// so records in the final partial second aren't excluded); otherwise 0 lets the +// endpoint default to now. +func (req logRequest) toSeconds(status logRunStatus) int64 { + if status.terminal() && status.endTimeMs > 0 { + return (status.endTimeMs + 999) / 1000 + } + return 0 +} + +// streamBricklensLogs fetches and prints a run's logs via Bricklens, handling +// both an already-completed run (a single bounded tail drain) and an active run +// (a poll-and-drain loop that follows until the run reaches a terminal state). +// +// It returns (success, err) where success reports whether the run finished with +// SUCCESS. A returned errBricklensFeatureDisabled means the caller should fall +// back to the MLflow path; the run genuinely not existing is returned as-is. +func streamBricklensLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + st := &bricklensStreamer{ + ctx: ctx, + w: w, + out: out, + req: req, + status: status, + seen: newSeenSet(seenRecordsCap), + } + return st.run() +} + +// bricklensStreamer carries the streaming state across the poll loop: the running +// from-second cursor, the highest emitted timestamp, and the bounded dedup set. +type bricklensStreamer struct { + ctx context.Context + w *databricks.WorkspaceClient + out io.Writer + req logRequest + status logRunStatus + + fromSec int64 + lastNano int64 + firstLogSeen bool + seen *seenSet +} + +func (st *bricklensStreamer) run() (bool, error) { + now := time.Now() + st.fromSec = st.req.fromSeconds(st.status, now) + + firstIteration := true + for { + if !firstIteration { + status, err := resolveRunStatus(st.ctx, st.w, st.req.runID) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return false, err + } + // A transient status blip should not abort a live stream; log and retry. + log.Debugf(st.ctx, "air logs: failed to refresh run status: %v", err) + time.Sleep(retryCheckInterval) + continue + } + st.status = status + } + + terminal := st.status.terminal() + toSec := st.req.toSeconds(st.status) + + // An already-completed run (terminal on the first iteration) renders as a + // tail: only the most-recent N lines, oldest-first. An active run — even + // one that terminates while we watch — streams everything and dedups so the + // final drain doesn't re-print the boundary second. A tail is line-based, so + // it only applies to the completed-run case (matching the MLflow path). + var err error + if firstIteration && terminal { + err = st.drainTail(toSec) + } else { + err = st.drainPages(toSec) + } + if err != nil { + return false, err + } + + if terminal { + if !st.firstLogSeen { + st.emitNoLogs() + } + log.Infof(st.ctx, "air logs: run %d finished in state %s", st.req.runID, st.status.displayState()) + return st.status.succeeded(), nil + } + + firstIteration = false + time.Sleep(retryCheckInterval) + } +} + +// tailTarget is the number of lines a completed-run tail should keep: --lines +// when set, else the default cap. +func (req logRequest) tailTarget() int { + if req.tailLines > 0 { + return req.tailLines + } + return defaultCompletedRunTailLines +} + +// drainTail emits the most-recent `target` records for a completed run, oldest-first. +// Bricklens returns records newest-first, so it pages until it has at least +// `target`, keeps the newest `target`, and reverses to chronological order — +// matching the MLflow --tail behavior. No dedup: a single bounded pass whose +// ordering we own client-side. +func (st *bricklensStreamer) drainTail(toSec int64) error { + target := st.req.tailTarget() + if target <= 0 { + return nil + } + + var collected []logRecord + var pageToken string + for len(collected) < target { + resp, err := st.requestPage(pageToken, toSec, target, false) + if err != nil { + return err + } + collected = append(collected, resp.LogRecords...) + pageToken = resp.NextPageToken + if pageToken == "" { + break + } + } + + // Keep the newest `target` (returned newest-first), then reverse so they print + // oldest -> newest like the MLflow tail. + if len(collected) > target { + collected = collected[:target] + } + for _, c := range slices.Backward(collected) { + st.emit(c.Body) + } + return nil +} + +// drainPages exhausts all available pages from the current from-second in +// ascending (oldest-first) order so a live run streams chronologically, +// deduping against the bounded seen-set so a re-queried boundary second is not +// re-printed. It advances fromSec to the floor-second of the newest record seen. +func (st *bricklensStreamer) drainPages(toSec int64) error { + var pageToken string + for { + resp, err := st.requestPage(pageToken, toSec, 0, true) + if err != nil { + return err + } + + for _, rec := range resp.LogRecords { + nano := rec.nano() + if nano != 0 { + // Ascending fetch, so a record older than the last emitted one is + // out of order (or a re-queried boundary record); skip it to keep + // streamed output monotonic. + if st.lastNano != 0 && nano < st.lastNano { + continue + } + if st.seen.has(nano, rec.Body) { + continue + } + } + st.emit(rec.Body) + if nano != 0 { + st.seen.add(nano, rec.Body) + st.lastNano = max(st.lastNano, nano) + } + } + + pageToken = resp.NextPageToken + if pageToken == "" { + break + } + } + + if st.lastNano != 0 { + st.fromSec = st.lastNano / 1_000_000_000 + } + return nil +} + +// requestPage fetches one page, applying the fallback classification. It retries +// transient failures up to maxTransientFailures, then falls back like a disabled +// feature (older logs may still be in MLflow). A feature-gated response or a +// genuine not-found returns immediately. +func (st *bricklensStreamer) requestPage(pageToken string, toSec int64, pageSize int, ascending bool) (*bricklensLogsResponse, error) { + q := bricklensLogsQuery{ + fromSeconds: st.fromSec, + toSeconds: toSec, + pageToken: pageToken, + pageSize: pageSize, + attemptNumber: st.req.attempt, + nodeIndex: st.req.node, + ascending: ascending, + } + + transientFailures := 0 + for { + resp, err := getBricklensLogs(st.ctx, st.w, st.req.runID, q) + if err == nil { + return resp, nil + } + + switch classified := classifyLogError(err); { + case errors.Is(classified, errBricklensFeatureDisabled): + return nil, errBricklensFeatureDisabled + case classified != nil: + return nil, classified + } + + transientFailures++ + if transientFailures >= maxTransientFailures { + log.Debugf(st.ctx, "air logs: get_logs failed %d times; falling back to mlflow", maxTransientFailures) + return nil, errBricklensFeatureDisabled + } + log.Debugf(st.ctx, "air logs: get_logs transient failure (%d/%d): %v", transientFailures, maxTransientFailures, err) + time.Sleep(retryCheckInterval) + } +} + +// emit writes one log line. In --json mode each line is a JSONL LOG event; +// otherwise it is printed raw. The first line latches firstLogSeen so a terminal +// run with no output can report "no logs". +func (st *bricklensStreamer) emit(body string) { + st.firstLogSeen = true + if st.req.jsonOutput { + printLogEvent(st.out, "LOG", st.req.node, body) + return + } + fmt.Fprintln(st.out, body) +} + +// emitNoLogs reports that a terminal run produced no logs, with its termination +// reason. In --json mode this is a JSONL ERROR so a consumer never sees an empty +// stream; otherwise a plain warning line. +func (st *bricklensStreamer) emitNoLogs() { + msg := fmt.Sprintf("No logs available for run %d. Run terminated in state %s", st.req.runID, st.status.displayState()) + if st.status.stateMessage != "" { + msg = fmt.Sprintf("%s: %s", msg, st.status.stateMessage) + } + if st.req.jsonOutput { + printLogEvent(st.out, "ERROR", st.req.node, msg) + return + } + fmt.Fprintln(st.out, msg) +} + +// displayState is the run's terminal result state, falling back to the lifecycle +// state, else "UNKNOWN". +func (s logRunStatus) displayState() string { + if s.resultState != "" { + return s.resultState + } + if s.lifeCycleState != "" { + return s.lifeCycleState + } + return "UNKNOWN" +} diff --git a/experimental/air/cmd/logstream_support.go b/experimental/air/cmd/logstream_support.go new file mode 100644 index 00000000000..54f4695184d --- /dev/null +++ b/experimental/air/cmd/logstream_support.go @@ -0,0 +1,76 @@ +package aircmd + +import ( + "container/list" + "encoding/json" + "fmt" + "io" + "time" +) + +// seenNano keys the dedup set: a (nano, body) pair. time_unix_nano is nanosecond +// units and all ranks funnel into one stream stamping from their own clocks, so +// distinct log lines can share a nano — the body disambiguates them. +type seenNano struct { + nano int64 + body string +} + +// seenSet is an insertion-ordered set bounded to a capacity, evicting the +// oldest-inserted entry first. It lets the streamer re-query a boundary second +// on the next poll without re-printing records already shown, while a large +// initial drain can't grow without bound. +type seenSet struct { + cap int + items map[seenNano]*list.Element + order *list.List +} + +func newSeenSet(capacity int) *seenSet { + return &seenSet{ + cap: capacity, + items: make(map[seenNano]*list.Element), + order: list.New(), + } +} + +func (s *seenSet) has(nano int64, body string) bool { + _, ok := s.items[seenNano{nano, body}] + return ok +} + +func (s *seenSet) add(nano int64, body string) { + key := seenNano{nano, body} + if _, ok := s.items[key]; ok { + return + } + s.items[key] = s.order.PushBack(key) + if s.order.Len() > s.cap { + oldest := s.order.Front() + s.order.Remove(oldest) + delete(s.items, oldest.Value.(seenNano)) + } +} + +// logEvent is one JSONL streaming event, matching the Python CLI's +// print_jsonl_event shape ({type, ts, node, line}). +type logEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + Node int `json:"node"` + Line string `json:"line"` +} + +// printLogEvent writes a single JSONL event line for --json streaming output. +func printLogEvent(out io.Writer, eventType string, node int, line string) { + b, err := json.Marshal(logEvent{ + Type: eventType, + TS: time.Now().UTC().Format(time.RFC3339), + Node: node, + Line: line, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go new file mode 100644 index 00000000000..fb256615674 --- /dev/null +++ b/experimental/air/cmd/logstream_test.go @@ -0,0 +1,165 @@ +package aircmd + +import ( + "errors" + "net/http" + "testing" + "time" + + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" +) + +func TestClassifyLogError(t *testing.T) { + tests := []struct { + name string + err error + want error // errBricklensFeatureDisabled, the input error, or nil + }{ + { + name: "feature disabled falls back", + err: &apierr.APIError{ErrorCode: "FEATURE_DISABLED", StatusCode: http.StatusForbidden}, + want: errBricklensFeatureDisabled, + }, + { + name: "endpoint not found falls back", + err: &apierr.APIError{ErrorCode: "ENDPOINT_NOT_FOUND", StatusCode: http.StatusNotFound}, + want: errBricklensFeatureDisabled, + }, + { + name: "bare 404 falls back", + err: &apierr.APIError{ErrorCode: "SOMETHING", StatusCode: http.StatusNotFound}, + want: errBricklensFeatureDisabled, + }, + { + name: "genuine resource-does-not-exist surfaces", + err: apierr.ErrResourceDoesNotExist, + want: apierr.ErrResourceDoesNotExist, + }, + { + name: "transient 500 is retried", + err: &apierr.APIError{ErrorCode: "INTERNAL", StatusCode: http.StatusInternalServerError}, + want: nil, + }, + { + name: "plain error is retried", + err: errors.New("connection reset"), + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := classifyLogError(tt.err) + switch tt.want { + case errBricklensFeatureDisabled: + assert.ErrorIs(t, got, errBricklensFeatureDisabled) + case nil: + assert.NoError(t, got) + default: + assert.ErrorIs(t, got, tt.want) + } + }) + } +} + +func TestProjectRunStatus(t *testing.T) { + run := &jobs.Run{ + StartTime: 1000, + EndTime: 2000, + State: &jobs.RunState{ + LifeCycleState: jobs.RunLifeCycleStateTerminated, + ResultState: jobs.RunResultStateSuccess, + StateMessage: "done", + }, + Tasks: []jobs.RunTask{ + {AttemptNumber: 0}, + {AttemptNumber: 2}, + {AttemptNumber: 1}, + }, + } + + s := projectRunStatus(run) + assert.Equal(t, "TERMINATED", s.lifeCycleState) + assert.Equal(t, "SUCCESS", s.resultState) + assert.Equal(t, "done", s.stateMessage) + assert.Equal(t, int64(1000), s.startTimeMs) + assert.Equal(t, int64(2000), s.endTimeMs) + assert.Equal(t, 2, s.latestAttempt) + assert.True(t, s.terminal()) + assert.True(t, s.succeeded()) + assert.Equal(t, "SUCCESS", s.displayState()) +} + +func TestLogRunStatusTerminal(t *testing.T) { + tests := []struct { + name string + lifeCycle string + resultState string + wantTerminal bool + }{ + {"running", "RUNNING", "", false}, + {"pending", "PENDING", "", false}, + {"terminated lifecycle", "TERMINATED", "", true}, + {"internal error lifecycle", "INTERNAL_ERROR", "", true}, + {"failed result", "TERMINATING", "FAILED", true}, + {"canceled result", "RUNNING", "CANCELED", true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := logRunStatus{lifeCycleState: tt.lifeCycle, resultState: tt.resultState} + assert.Equal(t, tt.wantTerminal, s.terminal()) + }) + } +} + +func TestLogRequestFromSeconds(t *testing.T) { + now := time.Unix(10_000, 0) + + // --minutes narrows the window to now - N*60. + req := logRequest{windowMinutes: 5} + assert.Equal(t, int64(10_000-300), req.fromSeconds(logRunStatus{startTimeMs: 1_000_000}, now)) + + // No window: from the run's start second. + req = logRequest{} + assert.Equal(t, int64(1000), req.fromSeconds(logRunStatus{startTimeMs: 1_000_000}, now)) + + // No window, run not started: everything stored (0). + assert.Equal(t, int64(0), req.fromSeconds(logRunStatus{}, now)) +} + +func TestLogRequestToSeconds(t *testing.T) { + req := logRequest{} + + // Active run: 0 lets the endpoint default to now. + assert.Equal(t, int64(0), req.toSeconds(logRunStatus{lifeCycleState: "RUNNING"})) + + // Terminal run: ceil of the end millisecond so the final partial second is kept. + terminal := logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 2001} + assert.Equal(t, int64(3), req.toSeconds(terminal)) +} + +func TestLogRequestTailTarget(t *testing.T) { + assert.Equal(t, defaultCompletedRunTailLines, logRequest{}.tailTarget()) + assert.Equal(t, 42, logRequest{tailLines: 42}.tailTarget()) +} + +func TestSeenSetEviction(t *testing.T) { + s := newSeenSet(2) + s.add(1, "a") + s.add(2, "b") + assert.True(t, s.has(1, "a")) + assert.True(t, s.has(2, "b")) + + // Adding a third evicts the oldest-inserted (1,"a"). + s.add(3, "c") + assert.False(t, s.has(1, "a")) + assert.True(t, s.has(2, "b")) + assert.True(t, s.has(3, "c")) + + // Same (nano, body) shares one entry; distinct body under the same nano does not. + s.add(3, "c") + assert.True(t, s.has(3, "c")) + assert.False(t, s.has(3, "d")) +} diff --git a/experimental/air/cmd/stubs_test.go b/experimental/air/cmd/stubs_test.go index e28d7f66730..4007557ef78 100644 --- a/experimental/air/cmd/stubs_test.go +++ b/experimental/air/cmd/stubs_test.go @@ -13,7 +13,6 @@ import ( // fails with a "not implemented" error. Drop a command here once it lands. func TestStubCommandsReturnNotImplemented(t *testing.T) { stubs := map[string]*cobra.Command{ - "logs": newLogsCommand(), "register-image": newRegisterImageCommand(), } From e6e87c030a8c646fdd6746eb06d0a15037364781 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Fri, 17 Jul 2026 00:20:59 +0000 Subject: [PATCH 2/5] AIR CLI: air logs completed-run tail + past-retry static view (phase 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the `--retry` past-retry handling: a past retry of a still-active run has immutable logs, so render them as a one-shot tail (drainStatic) instead of following the run — otherwise the poll loop would wait forever for the run, not the attempt, to finish. Mirrors handle_logs' viewing_past_retry. Exit code still follows the run's result state (matching the Python `result_state == "SUCCESS"` rule), so a past retry of a running job prints its logs then exits non-zero. Also strengthen coverage of the phase-1 tail/streaming paths: a drainPages test for boundary-second dedup and out-of-order skip, an acceptance case for --lines tailing, and one for --retry. Co-authored-by: Isaac --- acceptance/experimental/air/logs/output.txt | 9 ++++ acceptance/experimental/air/logs/script | 6 +++ experimental/air/cmd/logs.go | 6 +++ experimental/air/cmd/logs_test.go | 53 +++++++++++++++++++++ experimental/air/cmd/logstream.go | 28 ++++++++++- experimental/air/cmd/logstream_test.go | 43 +++++++++++++++++ 6 files changed, 144 insertions(+), 1 deletion(-) diff --git a/acceptance/experimental/air/logs/output.txt b/acceptance/experimental/air/logs/output.txt index 0e1ad38ff60..3ee3219faf5 100644 --- a/acceptance/experimental/air/logs/output.txt +++ b/acceptance/experimental/air/logs/output.txt @@ -14,6 +14,15 @@ step 2 step 1 step 2 +=== logs with --lines +>>> [CLI] experimental air logs 123 --lines 1 +step 2 + +=== logs from a specific retry +>>> [CLI] experimental air logs 123 --retry 0 +step 1 +step 2 + === logs --lines and --minutes are mutually exclusive >>> [CLI] experimental air logs 123 --lines 100 --minutes 30 Error: cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window diff --git a/acceptance/experimental/air/logs/script b/acceptance/experimental/air/logs/script index 3782221c3cf..0cf8094c18a 100644 --- a/acceptance/experimental/air/logs/script +++ b/acceptance/experimental/air/logs/script @@ -7,6 +7,12 @@ trace $CLI experimental air logs 123 -o json title "logs with --minutes" trace $CLI experimental air logs 123 --minutes 30 +title "logs with --lines" +trace $CLI experimental air logs 123 --lines 1 + +title "logs from a specific retry" +trace $CLI experimental air logs 123 --retry 0 + title "logs --lines and --minutes are mutually exclusive" errcode trace $CLI experimental air logs 123 --lines 100 --minutes 30 diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index 435d58de447..e311301434d 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -129,6 +129,12 @@ func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error { fmt.Errorf("invalid retry %d: available retries are 0 to %d", req.attempt, status.latestAttempt)) } + // A past retry of a still-active run has immutable logs: render them once + // rather than following the run to completion. + if req.attempt >= 0 && req.attempt < status.latestAttempt && !status.terminal() { + req.staticView = true + } + out := cmd.OutOrStdout() success, err := fetchLogs(ctx, w, out, req, status) if err != nil { diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go index 10f57de4f2c..85e62e37cca 100644 --- a/experimental/air/cmd/logs_test.go +++ b/experimental/air/cmd/logs_test.go @@ -7,6 +7,7 @@ import ( "strings" "testing" + "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" @@ -180,3 +181,55 @@ func TestLogsFallsBackToMLflow(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "MLflow fallback is not yet implemented") } + +// activeRunPastRetryServer serves a still-RUNNING run with two attempts and a +// single page of Bricklens logs. runs/get always returns RUNNING; a test that +// follows the run would poll forever, so it also asserts the static path never +// loops. +func activeRunPastRetryServer(t *testing.T, getRunHits *int) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/2.2/jobs/runs/get": + *getRunHits++ + _, _ = w.Write([]byte(`{ + "run_id": 9, + "start_time": 1700000000000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{"attempt_number": 0}, {"attempt_number": 1}] + }`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1700000001000000000, "body": "retry 0 log", "node_index": 0} + ]}`)) + default: + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestLogsPastRetryOfActiveRunIsStatic(t *testing.T) { + var getRunHits int + srv := activeRunPastRetryServer(t, &getRunHits) + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + // --retry 0 on a RUNNING run whose latest attempt is 1: the past attempt's + // logs are static, so they render once and the command returns instead of + // following the run (which would never terminate). The run itself has no + // SUCCESS result yet, so the command still exits non-zero (matching Python's + // `result_state == "SUCCESS"` exit-code rule) via ErrAlreadyPrinted — the + // logs are printed regardless. + err := runLogs(ctx, cmd, logRequest{runID: 9, node: 0, attempt: 0}) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Equal(t, "retry 0 log\n", buf.String()) + + // Exactly one runs/get: the initial status resolve in runLogs. The static + // streamer must not re-poll (that is the loop this test guards against). + assert.Equal(t, 1, getRunHits) +} diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index 57acb56c9ce..644c436d9d3 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -51,7 +51,12 @@ type logRequest struct { // windowMinutes, when > 0, restricts the fetch to the last N minutes. windowMinutes int // tailLines, when > 0, keeps only the last N lines of a completed run. - tailLines int + tailLines int + // staticView renders a one-shot tail instead of following the run. It is set + // for a past retry of a still-active run: that attempt's logs are immutable, + // so streaming them would poll forever waiting for the run (not the attempt) + // to finish. A completed run is inherently static and does not need this. + staticView bool jsonOutput bool } @@ -199,6 +204,13 @@ func (st *bricklensStreamer) run() (bool, error) { now := time.Now() st.fromSec = st.req.fromSeconds(st.status, now) + // A past retry's logs are immutable, so render them as a one-shot tail even + // when the run is still active — the attempt has ended, and following the run + // would poll forever. Mirrors handle_logs' viewing_past_retry. + if st.req.staticView { + return st.drainStatic(st.req.toSeconds(st.status)) + } + firstIteration := true for { if !firstIteration { @@ -246,6 +258,20 @@ func (st *bricklensStreamer) run() (bool, error) { } } +// drainStatic renders a single tail pass for an immutable attempt and returns +// without following the run. Success reflects the run's current result state +// (empty for an active run, so a past retry of a running job is not reported as +// a failure). +func (st *bricklensStreamer) drainStatic(toSec int64) (bool, error) { + if err := st.drainTail(toSec); err != nil { + return false, err + } + if !st.firstLogSeen { + st.emitNoLogs() + } + return st.status.succeeded(), nil +} + // tailTarget is the number of lines a completed-run tail should keep: --lines // when set, else the default cap. func (req logRequest) tailTarget() int { diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go index fb256615674..c229a4bce58 100644 --- a/experimental/air/cmd/logstream_test.go +++ b/experimental/air/cmd/logstream_test.go @@ -1,14 +1,17 @@ package aircmd import ( + "bytes" "errors" "net/http" + "net/http/httptest" "testing" "time" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestClassifyLogError(t *testing.T) { @@ -145,6 +148,46 @@ func TestLogRequestTailTarget(t *testing.T) { assert.Equal(t, 42, logRequest{tailLines: 42}.tailTarget()) } +func TestDrainPagesDedupAndOrdering(t *testing.T) { + // Two pages: page 1 has two ascending records; page 2 repeats the last record + // of page 1 (boundary re-query — must dedup) and includes an older record + // (out of order — must skip), then a genuinely newer one. + var page int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("page_token") == "" { + page = 1 + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1000, "body": "a", "node_index": 0}, + {"time_unix_nano": 2000, "body": "b", "node_index": 0} + ], "next_page_token": "p2"}`)) + return + } + page = 2 + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 2000, "body": "b", "node_index": 0}, + {"time_unix_nano": 1500, "body": "stale", "node_index": 0}, + {"time_unix_nano": 3000, "body": "c", "node_index": 0} + ]}`)) + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + st := &bricklensStreamer{ + ctx: t.Context(), + w: newTestWorkspaceClient(t, srv.URL), + out: &buf, + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + require.NoError(t, st.drainPages(0)) + require.Equal(t, 2, page) + + // "b" prints once (deduped), "stale" is skipped (older than last emitted), and + // fromSec advances to the newest record's floor-second (3000ns -> 0s here). + assert.Equal(t, "a\nb\nc\n", buf.String()) + assert.Equal(t, int64(3000), st.lastNano) +} + func TestSeenSetEviction(t *testing.T) { s := newSeenSet(2) s.add(1, "a") From 4bb6e66ae138392946c64004646ea747b763854f Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Fri, 17 Jul 2026 18:10:38 +0000 Subject: [PATCH 3/5] AIR CLI: implement the MLflow log fallback for air logs (phase 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the phase-1 stub with the real MLflow artifact path used when Bricklens is gated off. It resolves the run's MLflow run id (via the existing mlflowIDs helper), discovers the per-node log directory (logs/node_Y or the attempt-prefixed logs/attempt_X/node_Y layout), lists the logs-N.chunk.txt chunk files, and walks them newest-first — downloading each via the credential-vended credentials-for-read pre-signed URL — until it has the requested tail, then prints oldest-first. Mirrors the Python print_all_logs path. --lines drives the tail as on the Bricklens path. MLflow chunk files are not time-indexed, so --minutes can't restrict the window here; when only --minutes is set the default tail cap applies and a debug line notes it was inapplicable. Refactor the log-line and no-logs emitters into shared free functions (emitLogLine / emitNoLogs) so the Bricklens and MLflow paths produce identical raw and JSONL output instead of duplicating the logic. Co-authored-by: Isaac --- .../air/logs-mlflow-fallback/out.test.toml | 3 + .../air/logs-mlflow-fallback/output.txt | 12 + .../air/logs-mlflow-fallback/script | 9 + .../air/logs-mlflow-fallback/test.toml | 32 ++ experimental/air/cmd/logmlflow.go | 303 +++++++++++++++++- experimental/air/cmd/logmlflow_test.go | 77 +++++ experimental/air/cmd/logs_test.go | 41 ++- experimental/air/cmd/logstream.go | 26 +- experimental/air/cmd/logstream_support.go | 25 ++ 9 files changed, 488 insertions(+), 40 deletions(-) create mode 100644 acceptance/experimental/air/logs-mlflow-fallback/out.test.toml create mode 100644 acceptance/experimental/air/logs-mlflow-fallback/output.txt create mode 100644 acceptance/experimental/air/logs-mlflow-fallback/script create mode 100644 acceptance/experimental/air/logs-mlflow-fallback/test.toml create mode 100644 experimental/air/cmd/logmlflow_test.go diff --git a/acceptance/experimental/air/logs-mlflow-fallback/out.test.toml b/acceptance/experimental/air/logs-mlflow-fallback/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/logs-mlflow-fallback/output.txt b/acceptance/experimental/air/logs-mlflow-fallback/output.txt new file mode 100644 index 00000000000..a511fd33bc5 --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/output.txt @@ -0,0 +1,12 @@ + +=== logs falls back to mlflow (no logs) +>>> [CLI] experimental air logs 123 +No logs available for run 123. Run terminated in state SUCCESS + +Exit code: 1 + +=== logs falls back to mlflow (json) +>>> [CLI] experimental air logs 123 -o json +{"type":"ERROR","ts":"[TIMESTAMP]","node":0,"line":"No logs available for run 123. Run terminated in state SUCCESS"} + +Exit code: 1 diff --git a/acceptance/experimental/air/logs-mlflow-fallback/script b/acceptance/experimental/air/logs-mlflow-fallback/script new file mode 100644 index 00000000000..ad8856a036a --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/script @@ -0,0 +1,9 @@ +# Bricklens is gated off (FEATURE_DISABLED), so the command falls back to the +# MLflow log path. With no MLflow run id resolvable, the fallback reports no +# logs and exits non-zero — proving the try/catch routes to MLflow. + +title "logs falls back to mlflow (no logs)" +errcode trace $CLI experimental air logs 123 + +title "logs falls back to mlflow (json)" +errcode trace $CLI experimental air logs 123 -o json diff --git a/acceptance/experimental/air/logs-mlflow-fallback/test.toml b/acceptance/experimental/air/logs-mlflow-fallback/test.toml new file mode 100644 index 00000000000..9b9bcab8e68 --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/test.toml @@ -0,0 +1,32 @@ +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A completed run. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"task_key": "train", "run_id": 456, "attempt_number": 0}] +} +''' + +# Bricklens is gated off by the backend SAFE flag, forcing the MLflow fallback. +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows/by-run-id/123/logs" +Response.StatusCode = 403 +Response.Body = ''' +{"error_code": "FEATURE_DISABLED", "message": "training log streaming is not enabled"} +''' + +# The MLflow fallback has no run output to resolve an MLflow run id from, so it +# reports no logs rather than failing — exercising the fallback wiring end to end. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = '{}' diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go index faf5e41cc03..4a6528921d5 100644 --- a/experimental/air/cmd/logmlflow.go +++ b/experimental/air/cmd/logmlflow.go @@ -1,22 +1,309 @@ package aircmd import ( + "bufio" "context" "fmt" "io" + "net/http" + "os" + "path" + "regexp" + "slices" + "strconv" + "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" + "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" ) +// chunkFilePattern matches a sidecar log chunk file (logs-.chunk.txt); +// group 1 is the chunk index. The sidecar splits the consolidated stdout stream +// into 4MB chunks, advancing the index monotonically. +var chunkFilePattern = regexp.MustCompile(`^logs-(\d+)\.chunk\.txt$`) + +// oldFormatNodeDir matches a bare per-node log dir (logs/node_) at the top of +// the run's logs/ dir. The attempt-prefixed layout nests these under +// logs/attempt_/, so a bare logs/node_ only appears in the old layout. +var oldFormatNodeDir = regexp.MustCompile(`^logs/node_\d+$`) + // mlflowLogFallback prints a run's logs from MLflow artifacts. It is the fallback -// used when Bricklens can't serve logs (errBricklensFeatureDisabled). Both the -// completed-run tail (--lines) and the active-run stream, plus the --minutes -// window, are honored here via req. +// used when Bricklens can't serve logs (errBricklensFeatureDisabled). +// +// It resolves the run's MLflow run id, discovers the per-node log directory +// (logs/node_Y or logs/attempt_X/node_Y), lists the chunk files, and walks them +// newest-first until it has the requested tail, then prints oldest-first. This +// mirrors the Python print_all_logs path. // -// TODO(air-logs-m3): port print_all_logs / monitor_and_stream_logs — MLflow -// artifact discovery (logs/[attempt_X/]node_Y), chunk walking (logs-N.chunk.txt), -// and credential-vending download. Until then the fallback reports that the -// legacy path is not yet available so the failure is explicit rather than silent. +// The tail length is --lines when set, else the default cap. MLflow chunk files +// are not time-indexed, so --minutes cannot restrict the window here (Bricklens +// is the time-indexed backend); when only --minutes is set the default tail cap +// applies, and a debug line notes the flag was inapplicable on this path. func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { - return false, fmt.Errorf("logs for run %d are not available via Bricklens and the MLflow fallback is not yet implemented", req.runID) + if req.windowMinutes > 0 { + log.Debugf(ctx, "air logs: --minutes is not supported on the MLflow fallback path; showing the default tail") + } + + mlflowRunID, logDir, err := resolveMLflowLogPath(ctx, w, req) + if err != nil { + return false, err + } + if mlflowRunID == "" || logDir == "" { + emitNoLogs(out, req, status) + return false, nil + } + + chunks, err := listLogChunks(ctx, w, mlflowRunID, logDir) + if err != nil { + return false, err + } + if len(chunks) == 0 { + // Nothing listed yet: assume the single chunk 0, matching the Python + // legacy single-chunk fallback. + chunks = []logChunk{{index: 0, path: path.Join(logDir, chunkFileName(0))}} + } + + target := req.tailTarget() + if target <= 0 { + return status.succeeded(), nil + } + + lines, err := tailChunks(ctx, w, mlflowRunID, chunks, target) + if err != nil { + return false, err + } + if len(lines) == 0 { + emitNoLogs(out, req, status) + return false, nil + } + + if len(lines) > target { + lines = lines[len(lines)-target:] + } + for _, line := range lines { + emitLogLine(out, req, line) + } + return status.succeeded(), nil +} + +// resolveMLflowLogPath returns the run's MLflow run id and per-node log +// directory. The directory format (attempt-prefixed or not) is deployment-wide, +// so it is discovered from the artifact listing under logs/. +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) + if ids == nil || ids.RunID == "" { + return "", "", nil + } + + // attempt is -1 (latest) or an explicit retry; the log path uses attempt 0 + // when latest, matching how the sidecar names the first attempt's dir. + attempt := max(req.attempt, 0) + withAttempt, err := discoverAttemptPrefix(ctx, w, ids.RunID, attempt) + if err != nil { + return "", "", err + } + return ids.RunID, constructLogPath(req.node, attempt, withAttempt), nil +} + +// discoverAttemptPrefix probes the run's logs/ dir once to decide whether the +// deployment uses the attempt-prefixed layout (logs/attempt_X/node_Y) or the old +// one (logs/node_Y). A bare logs/node_ anywhere means old format; otherwise a +// logs/attempt_ entry means the new format. Defaults to old format when +// nothing is listed yet. +func discoverAttemptPrefix(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, attempt int) (bool, error) { + files, err := listArtifacts(ctx, w, mlflowRunID, "logs") + if err != nil { + // A listing failure here is not fatal: fall back to the old (no-attempt) + // layout, which the chunk listing will simply find empty if wrong. + log.Debugf(ctx, "air logs: could not list logs dir for format discovery: %v", err) + return false, nil + } + + attemptDir := fmt.Sprintf("logs/attempt_%d", attempt) + for _, f := range files { + if oldFormatNodeDir.MatchString(f.Path) { + return false, nil + } + if f.Path == attemptDir { + return true, nil + } + } + return false, nil +} + +// constructLogPath builds the per-node log directory for a node and attempt. +func constructLogPath(node, attempt int, withAttempt bool) string { + if withAttempt { + return fmt.Sprintf("logs/attempt_%d/node_%d", attempt, node) + } + return fmt.Sprintf("logs/node_%d", node) +} + +// chunkFileName is the artifact filename for a chunk index. +func chunkFileName(index int) string { + return fmt.Sprintf("logs-%d.chunk.txt", index) +} + +// logChunk is one listed chunk: its index and full artifact path. +type logChunk struct { + index int + path string +} + +// listLogChunks lists the chunk files under a per-node log dir, sorted ascending +// by index. Non-chunk entries are ignored. +func listLogChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, logDir string) ([]logChunk, error) { + files, err := listArtifacts(ctx, w, mlflowRunID, logDir) + if err != nil { + return nil, err + } + + var chunks []logChunk + for _, f := range files { + base := path.Base(f.Path) + m := chunkFilePattern.FindStringSubmatch(base) + if m == nil { + continue + } + idx, err := strconv.Atoi(m[1]) + if err != nil { + continue + } + chunks = append(chunks, logChunk{index: idx, path: f.Path}) + } + slices.SortFunc(chunks, func(a, b logChunk) int { return a.index - b.index }) + return chunks, nil +} + +// tailChunks walks the chunks newest-first, prepending each chunk's lines, until +// it has at least `target` lines (or runs out). This is bandwidth-optimal for the +// common case where the tail fits in the last chunk or two. If a chunk fails to +// download mid-walk, it stops rather than splice non-adjacent chunks. +func tailChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, chunks []logChunk, target int) ([]string, error) { + var accumulated []string + for _, chunk := range slices.Backward(chunks) { + lines, err := downloadChunkLines(ctx, w, mlflowRunID, chunk.path) + if err != nil { + log.Debugf(ctx, "air logs: failed to download chunk %d; showing only logs after the gap: %v", chunk.index, err) + break + } + accumulated = append(lines, accumulated...) + if len(accumulated) >= target { + break + } + } + return accumulated, nil +} + +// downloadChunkLines fetches one chunk artifact and returns its lines. The bytes +// are read via the credential-vended download so large log files are not capped +// by the proxied get-artifact endpoint. +func downloadChunkLines(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string) ([]string, error) { + f, err := downloadArtifact(ctx, w, mlflowRunID, artifactPath) + if err != nil { + return nil, err + } + defer os.Remove(f) + + file, err := os.Open(f) + if err != nil { + return nil, err + } + defer file.Close() + + var lines []string + scanner := bufio.NewScanner(file) + scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + for scanner.Scan() { + lines = append(lines, scanner.Text()) + } + return lines, scanner.Err() +} + +// listArtifacts lists a run's artifacts under a path via the typed SDK iterator. +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}) + return listing.ToSlice(ctx, it) +} + +// credentialInfo is one entry of the credentials-for-read response: a pre-signed +// URL for a run artifact plus any backend-required request headers. +type credentialInfo struct { + SignedURI string `json:"signed_uri"` + Headers []struct { + Name string `json:"name"` + Value string `json:"value"` + } `json:"headers"` +} + +type credentialsForReadResponse struct { + CredentialInfos []credentialInfo `json:"credential_infos"` +} + +// downloadArtifact downloads a single run artifact to a temp file and returns its +// path. It uses the credential-vending flow the MLflow SDK uses for +// Databricks-backed runs: credentials-for-read returns a pre-signed URL, which we +// stream to disk. The credentials-for-read 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) { + apiClient, err := client.New(w.Config) + if err != nil { + return "", fmt.Errorf("failed to create API client: %w", err) + } + + var resp credentialsForReadResponse + query := map[string]any{ + "run_id": mlflowRunID, + // path is a repeated field; the SDK serializes a slice as path=...&path=... + "path": []string{artifactPath}, + } + err = apiClient.Do(ctx, http.MethodGet, "/api/2.0/mlflow/artifacts/credentials-for-read", nil, nil, query, &resp) + if err != nil { + return "", fmt.Errorf("failed to get read credentials for %s: %w", artifactPath, err) + } + if len(resp.CredentialInfos) == 0 || resp.CredentialInfos[0].SignedURI == "" { + return "", fmt.Errorf("no download credentials returned for %s", artifactPath) + } + cred := resp.CredentialInfos[0] + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, cred.SignedURI, nil) + if err != nil { + return "", err + } + // Azure SAS / some GCS URIs require backend-supplied headers; AWS pre-signed + // URLs return none. + for _, h := range cred.Headers { + req.Header.Set(h.Name, h.Value) + } + + httpResp, err := http.DefaultClient.Do(req) + if err != nil { + return "", err + } + defer httpResp.Body.Close() + if httpResp.StatusCode >= 400 { + return "", fmt.Errorf("artifact download failed: HTTP %d", httpResp.StatusCode) + } + + tmp, err := os.CreateTemp("", "air-log-chunk-*") + if err != nil { + return "", err + } + if _, err := io.Copy(tmp, httpResp.Body); 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 } diff --git a/experimental/air/cmd/logmlflow_test.go b/experimental/air/cmd/logmlflow_test.go new file mode 100644 index 00000000000..9fbe4f19ff1 --- /dev/null +++ b/experimental/air/cmd/logmlflow_test.go @@ -0,0 +1,77 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +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)) + assert.Equal(t, "logs/attempt_2/node_3", constructLogPath(3, 2, true)) +} + +func TestChunkFileName(t *testing.T) { + assert.Equal(t, "logs-0.chunk.txt", chunkFileName(0)) + assert.Equal(t, "logs-7.chunk.txt", chunkFileName(7)) +} + +func TestChunkFilePattern(t *testing.T) { + m := chunkFilePattern.FindStringSubmatch("logs-12.chunk.txt") + require.NotNil(t, m) + assert.Equal(t, "12", m[1]) + + assert.Nil(t, chunkFilePattern.FindStringSubmatch("logs-12.chunk.txt.bak")) + assert.Nil(t, chunkFilePattern.FindStringSubmatch("node_0")) +} + +// artifactListServer serves a fixed artifacts/list response for any path. +func artifactListServer(t *testing.T, body string) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.0/mlflow/artifacts/list" { + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListLogChunksSortsAndFiltersByIndex(t *testing.T) { + // Out-of-order chunks plus a non-chunk file; result is ascending, chunk-only. + srv := artifactListServer(t, `{"files": [ + {"path": "logs/node_0/logs-2.chunk.txt"}, + {"path": "logs/node_0/other.txt"}, + {"path": "logs/node_0/logs-0.chunk.txt"}, + {"path": "logs/node_0/logs-1.chunk.txt"} + ]}`) + w := newTestWorkspaceClient(t, srv.URL) + + chunks, err := listLogChunks(t.Context(), w, "run1", "logs/node_0") + require.NoError(t, err) + require.Len(t, chunks, 3) + assert.Equal(t, 0, chunks[0].index) + assert.Equal(t, 1, chunks[1].index) + assert.Equal(t, 2, chunks[2].index) + assert.Equal(t, "logs/node_0/logs-0.chunk.txt", chunks[0].path) +} + +func TestDiscoverAttemptPrefix(t *testing.T) { + // Old format: a bare logs/node_N dir means no attempt prefix. + old := artifactListServer(t, `{"files": [{"path": "logs/node_0", "is_dir": true}]}`) + got, err := discoverAttemptPrefix(t.Context(), newTestWorkspaceClient(t, old.URL), "run1", 0) + require.NoError(t, err) + assert.False(t, got) + + // New format: a logs/attempt_N entry and no bare node dir. + newFmt := artifactListServer(t, `{"files": [{"path": "logs/attempt_0", "is_dir": true}]}`) + got, err = discoverAttemptPrefix(t.Context(), newTestWorkspaceClient(t, newFmt.URL), "run1", 0) + require.NoError(t, err) + assert.True(t, got) +} diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go index 85e62e37cca..a4426ee5dfc 100644 --- a/experimental/air/cmd/logs_test.go +++ b/experimental/air/cmd/logs_test.go @@ -144,42 +144,61 @@ func TestLogsCompletedRunTail(t *testing.T) { assert.Equal(t, "line one\nline two\n", buf.String()) } -// featureDisabledLogsServer serves a terminal runs/get and a Bricklens endpoint -// gated off by the backend SAFE flag (FEATURE_DISABLED), so the command must -// fall back to the MLflow path. -func featureDisabledLogsServer(t *testing.T) *httptest.Server { +// mlflowFallbackServer serves a terminal run whose Bricklens endpoint is gated +// off (FEATURE_DISABLED), plus the full MLflow artifact path the fallback walks: +// runs/get-output (MLflow ids), artifacts/list (logs dir + chunk file), +// credentials-for-read (pre-signed URL), and the pre-signed chunk bytes itself. +func mlflowFallbackServer(t *testing.T) *httptest.Server { t.Helper() + 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": [{"attempt_number": 0}] + "tasks": [{"run_id": 456, "attempt_number": 0}] }`)) 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": "bricklens logs gated off"}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + // The logs dir probe (format discovery) and the per-node chunk listing + // both hit this; return the old-format node dir and one chunk file. + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0/logs-0.chunk.txt", "file_size": 12}]}`)) + 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 one\nline two\n")) default: _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) } })) + base = srv.URL t.Cleanup(srv.Close) return srv } func TestLogsFallsBackToMLflow(t *testing.T) { - srv := featureDisabledLogsServer(t) + srv := mlflowFallbackServer(t) + var buf bytes.Buffer ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) cmd := withOutput(&cobra.Command{}, flags.OutputText) cmd.SetContext(ctx) + cmd.SetOut(&buf) - // The Bricklens endpoint is gated off, so fetchLogs routes to the MLflow - // fallback. Until that path is implemented (air-logs-m3) it reports the - // explicit not-yet-available error, which proves the try/catch wiring fires. + // Bricklens is gated off, so fetchLogs routes to the MLflow fallback, which + // resolves the MLflow run, lists the chunk, downloads it via the pre-signed + // URL, and prints its lines. err := runLogs(ctx, cmd, logRequest{runID: 5, node: 0, attempt: -1}) - require.Error(t, err) - assert.Contains(t, err.Error(), "MLflow fallback is not yet implemented") + require.NoError(t, err) + assert.Equal(t, "line one\nline two\n", buf.String()) } // activeRunPastRetryServer serves a still-RUNNING run with two attempts and a diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index 644c436d9d3..279e00e934d 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -3,7 +3,6 @@ package aircmd import ( "context" "errors" - "fmt" "io" "net/http" "slices" @@ -400,31 +399,16 @@ func (st *bricklensStreamer) requestPage(pageToken string, toSec int64, pageSize } } -// emit writes one log line. In --json mode each line is a JSONL LOG event; -// otherwise it is printed raw. The first line latches firstLogSeen so a terminal -// run with no output can report "no logs". +// emit writes one log line and latches firstLogSeen so a terminal run with no +// output can report "no logs". func (st *bricklensStreamer) emit(body string) { st.firstLogSeen = true - if st.req.jsonOutput { - printLogEvent(st.out, "LOG", st.req.node, body) - return - } - fmt.Fprintln(st.out, body) + emitLogLine(st.out, st.req, body) } -// emitNoLogs reports that a terminal run produced no logs, with its termination -// reason. In --json mode this is a JSONL ERROR so a consumer never sees an empty -// stream; otherwise a plain warning line. +// emitNoLogs reports that a terminal run produced no logs. func (st *bricklensStreamer) emitNoLogs() { - msg := fmt.Sprintf("No logs available for run %d. Run terminated in state %s", st.req.runID, st.status.displayState()) - if st.status.stateMessage != "" { - msg = fmt.Sprintf("%s: %s", msg, st.status.stateMessage) - } - if st.req.jsonOutput { - printLogEvent(st.out, "ERROR", st.req.node, msg) - return - } - fmt.Fprintln(st.out, msg) + emitNoLogs(st.out, st.req, st.status) } // displayState is the run's terminal result state, falling back to the lifecycle diff --git a/experimental/air/cmd/logstream_support.go b/experimental/air/cmd/logstream_support.go index 54f4695184d..f644887086d 100644 --- a/experimental/air/cmd/logstream_support.go +++ b/experimental/air/cmd/logstream_support.go @@ -74,3 +74,28 @@ func printLogEvent(out io.Writer, eventType string, node int, line string) { } fmt.Fprintln(out, string(b)) } + +// emitLogLine writes one log line, as a JSONL LOG event under --json or raw +// otherwise. Shared by the Bricklens and MLflow paths so both emit identically. +func emitLogLine(out io.Writer, req logRequest, body string) { + if req.jsonOutput { + printLogEvent(out, "LOG", req.node, body) + return + } + fmt.Fprintln(out, body) +} + +// emitNoLogs reports that a run produced no logs, with its termination reason. +// Under --json it is a JSONL ERROR (so a consumer never sees an empty stream); +// otherwise a plain line. Shared by both backends. +func emitNoLogs(out io.Writer, req logRequest, status logRunStatus) { + msg := fmt.Sprintf("No logs available for run %d. Run terminated in state %s", req.runID, status.displayState()) + if status.stateMessage != "" { + msg = fmt.Sprintf("%s: %s", msg, status.stateMessage) + } + if req.jsonOutput { + printLogEvent(out, "ERROR", req.node, msg) + return + } + fmt.Fprintln(out, msg) +} From 3f63af3e29af628f98b19e43fc4d037fca5bdbdd Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Fri, 17 Jul 2026 21:10:10 +0000 Subject: [PATCH 4/5] AIR CLI: tighten air logs comments (no logic change) Make comments concise and self-contained: drop cross-references to the Python CLI and remove restated-code noise. No behavior change. Co-authored-by: Isaac --- experimental/air/cmd/logbricklens.go | 31 ++--- experimental/air/cmd/logmlflow.go | 85 +++++------- experimental/air/cmd/logs.go | 35 ++--- experimental/air/cmd/logs_test.go | 10 +- experimental/air/cmd/logstream.go | 156 +++++++++------------- experimental/air/cmd/logstream_support.go | 18 +-- 6 files changed, 131 insertions(+), 204 deletions(-) diff --git a/experimental/air/cmd/logbricklens.go b/experimental/air/cmd/logbricklens.go index 47ba7a54e7b..5e0a23f1621 100644 --- a/experimental/air/cmd/logbricklens.go +++ b/experimental/air/cmd/logbricklens.go @@ -11,23 +11,19 @@ import ( "github.com/databricks/databricks-sdk-go/client" ) -// bricklensLogsPathFmt is the Bricklens-backed training-workflow log endpoint, -// keyed by Jobs run id. It is called with a raw client.Do because the SDK does -// not model the AiTrainingService log surface. +// bricklensLogsPathFmt is the log endpoint, keyed by Jobs run id. Called with a +// raw client.Do because the SDK does not model the AiTrainingService. const bricklensLogsPathFmt = "/api/2.0/ai-training/workflows/by-run-id/%d/logs" -// logRecord is one log line from Bricklens. time_unix_nano stamps ordering and -// dedup; body is the raw line; node_index identifies the emitting node. +// logRecord is one log line from Bricklens. type logRecord struct { - // TimeUnixNano is nanosecond units; tolerate it arriving as a JSON number or - // string, matching the AiTrainingService index quirk in aitraining.go. + // TimeUnixNano may arrive as a JSON number or string. TimeUnixNano json.Number `json:"time_unix_nano"` Body string `json:"body"` NodeIndex int `json:"node_index"` } -// nano returns the record's time_unix_nano as an int64, or 0 when absent or -// unparseable (so it sorts first and never blocks dedup). +// nano returns time_unix_nano as an int64, or 0 when absent or unparseable. func (r logRecord) nano() int64 { n, err := r.TimeUnixNano.Int64() if err != nil { @@ -41,9 +37,8 @@ type bricklensLogsResponse struct { NextPageToken string `json:"next_page_token"` } -// bricklensLogsQuery is the request-field surface of the log endpoint. The -// fields map to the query keys the Python CLI sends (ai_training_client.get_logs); -// zero-valued optionals are omitted so the endpoint applies its own defaults. +// bricklensLogsQuery is the request-field surface of the log endpoint. +// Zero-valued optionals are omitted so the endpoint applies its own defaults. type bricklensLogsQuery struct { // fromSeconds and toSeconds bound the query window in Unix epoch seconds. fromSeconds int64 @@ -53,14 +48,13 @@ type bricklensLogsQuery struct { // attemptNumber selects a retry attempt (0-indexed); -1 means latest. attemptNumber int nodeIndex int - // ascending returns oldest-first; the endpoint defaults to ascending when - // the field is absent, so the tail fetch must send it explicitly false. + // ascending returns oldest-first. The endpoint defaults to ascending when + // absent, so the tail fetch must send an explicit false for newest-first. ascending bool } -// getBricklensLogs fetches one page of logs from the Bricklens endpoint. It -// returns the raw error so the caller can classify it (feature-gated vs -// transient vs genuine not-found) via classifyLogError. +// getBricklensLogs fetches one page of logs. It returns the raw error so the +// caller can classify it via classifyLogError. func getBricklensLogs(ctx context.Context, w *databricks.WorkspaceClient, runID int64, q bricklensLogsQuery) (*bricklensLogsResponse, error) { apiClient, err := client.New(w.Config) if err != nil { @@ -68,8 +62,7 @@ func getBricklensLogs(ctx context.Context, w *databricks.WorkspaceClient, runID } query := map[string]any{ - // The endpoint defaults to ascending when absent, so always send it: the - // tail path relies on an explicit false to get newest-first. + // Always sent: the tail path relies on an explicit false for newest-first. "ascending": strconv.FormatBool(q.ascending), } if q.fromSeconds > 0 { diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go index 4a6528921d5..15a73ee0d58 100644 --- a/experimental/air/cmd/logmlflow.go +++ b/experimental/air/cmd/logmlflow.go @@ -20,28 +20,22 @@ import ( "github.com/databricks/databricks-sdk-go/service/ml" ) -// chunkFilePattern matches a sidecar log chunk file (logs-.chunk.txt); -// group 1 is the chunk index. The sidecar splits the consolidated stdout stream -// into 4MB chunks, advancing the index monotonically. +// 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$`) -// oldFormatNodeDir matches a bare per-node log dir (logs/node_) at the top of -// the run's logs/ dir. The attempt-prefixed layout nests these under -// logs/attempt_/, so a bare logs/node_ only appears in the old layout. +// oldFormatNodeDir matches a bare per-node log dir (logs/node_). The +// attempt-prefixed layout nests these under logs/attempt_/, so a bare +// logs/node_ only appears in the old layout. var oldFormatNodeDir = regexp.MustCompile(`^logs/node_\d+$`) -// mlflowLogFallback prints a run's logs from MLflow artifacts. It is the fallback -// used when Bricklens can't serve logs (errBricklensFeatureDisabled). +// mlflowLogFallback prints a run's logs from MLflow artifacts, the fallback when +// Bricklens can't serve them. It resolves the MLflow run id, discovers the +// per-node log directory, lists the chunk files, and walks them newest-first +// until it has the requested tail, then prints oldest-first. // -// It resolves the run's MLflow run id, discovers the per-node log directory -// (logs/node_Y or logs/attempt_X/node_Y), lists the chunk files, and walks them -// newest-first until it has the requested tail, then prints oldest-first. This -// mirrors the Python print_all_logs path. -// -// The tail length is --lines when set, else the default cap. MLflow chunk files -// are not time-indexed, so --minutes cannot restrict the window here (Bricklens -// is the time-indexed backend); when only --minutes is set the default tail cap -// applies, and a debug line notes the flag was inapplicable on this path. +// The tail length is --lines, else the default cap. MLflow chunks are not +// time-indexed, so --minutes cannot restrict the window here. func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { if req.windowMinutes > 0 { log.Debugf(ctx, "air logs: --minutes is not supported on the MLflow fallback path; showing the default tail") @@ -61,8 +55,7 @@ func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out i return false, err } if len(chunks) == 0 { - // Nothing listed yet: assume the single chunk 0, matching the Python - // legacy single-chunk fallback. + // Nothing listed yet: assume the single chunk 0. chunks = []logChunk{{index: 0, path: path.Join(logDir, chunkFileName(0))}} } @@ -89,9 +82,7 @@ func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out i return status.succeeded(), nil } -// resolveMLflowLogPath returns the run's MLflow run id and per-node log -// directory. The directory format (attempt-prefixed or not) is deployment-wide, -// so it is discovered from the artifact listing under logs/. +// 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 { @@ -102,8 +93,7 @@ func resolveMLflowLogPath(ctx context.Context, w *databricks.WorkspaceClient, re return "", "", nil } - // attempt is -1 (latest) or an explicit retry; the log path uses attempt 0 - // when latest, matching how the sidecar names the first attempt's dir. + // -1 (latest) maps to attempt 0's directory. attempt := max(req.attempt, 0) withAttempt, err := discoverAttemptPrefix(ctx, w, ids.RunID, attempt) if err != nil { @@ -112,16 +102,14 @@ func resolveMLflowLogPath(ctx context.Context, w *databricks.WorkspaceClient, re return ids.RunID, constructLogPath(req.node, attempt, withAttempt), nil } -// discoverAttemptPrefix probes the run's logs/ dir once to decide whether the -// deployment uses the attempt-prefixed layout (logs/attempt_X/node_Y) or the old -// one (logs/node_Y). A bare logs/node_ anywhere means old format; otherwise a -// logs/attempt_ entry means the new format. Defaults to old format when -// nothing is listed yet. +// discoverAttemptPrefix probes the logs/ dir once to decide whether the layout is +// attempt-prefixed (logs/attempt_X/node_Y) or old (logs/node_Y). A bare +// logs/node_ means old; a logs/attempt_ entry means prefixed. +// Defaults to old when nothing is listed. func discoverAttemptPrefix(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, attempt int) (bool, error) { files, err := listArtifacts(ctx, w, mlflowRunID, "logs") if err != nil { - // A listing failure here is not fatal: fall back to the old (no-attempt) - // layout, which the chunk listing will simply find empty if wrong. + // Not fatal: default to the old layout; the chunk listing finds it empty if wrong. log.Debugf(ctx, "air logs: could not list logs dir for format discovery: %v", err) return false, nil } @@ -157,8 +145,7 @@ type logChunk struct { path string } -// listLogChunks lists the chunk files under a per-node log dir, sorted ascending -// by index. Non-chunk entries are ignored. +// listLogChunks lists the chunk files under a log dir, sorted ascending by index. func listLogChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, logDir string) ([]logChunk, error) { files, err := listArtifacts(ctx, w, mlflowRunID, logDir) if err != nil { @@ -182,16 +169,15 @@ func listLogChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRun return chunks, nil } -// tailChunks walks the chunks newest-first, prepending each chunk's lines, until -// it has at least `target` lines (or runs out). This is bandwidth-optimal for the -// common case where the tail fits in the last chunk or two. If a chunk fails to -// download mid-walk, it stops rather than splice non-adjacent chunks. +// tailChunks walks chunks newest-first, prepending each chunk's lines, until it +// has `target` lines or runs out. A mid-walk download failure stops the walk +// rather than splice non-adjacent chunks. func tailChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, chunks []logChunk, target int) ([]string, error) { var accumulated []string for _, chunk := range slices.Backward(chunks) { lines, err := downloadChunkLines(ctx, w, mlflowRunID, chunk.path) if err != nil { - log.Debugf(ctx, "air logs: failed to download chunk %d; showing only logs after the gap: %v", chunk.index, err) + log.Debugf(ctx, "air logs: failed to download chunk %d; showing only logs after it: %v", chunk.index, err) break } accumulated = append(lines, accumulated...) @@ -202,9 +188,7 @@ func tailChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID return accumulated, nil } -// downloadChunkLines fetches one chunk artifact and returns its lines. The bytes -// are read via the credential-vended download so large log files are not capped -// by the proxied get-artifact endpoint. +// downloadChunkLines fetches one chunk artifact and returns its lines. func downloadChunkLines(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string) ([]string, error) { f, err := downloadArtifact(ctx, w, mlflowRunID, artifactPath) if err != nil { @@ -227,14 +211,14 @@ func downloadChunkLines(ctx context.Context, w *databricks.WorkspaceClient, mlfl return lines, scanner.Err() } -// listArtifacts lists a run's artifacts under a path via the typed SDK iterator. +// 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}) return listing.ToSlice(ctx, it) } -// credentialInfo is one entry of the credentials-for-read response: a pre-signed -// URL for a run artifact plus any backend-required request headers. +// credentialInfo is one credentials-for-read entry: a pre-signed URL plus any +// backend-required request headers. type credentialInfo struct { SignedURI string `json:"signed_uri"` Headers []struct { @@ -247,11 +231,9 @@ type credentialsForReadResponse struct { CredentialInfos []credentialInfo `json:"credential_infos"` } -// downloadArtifact downloads a single run artifact to a temp file and returns its -// path. It uses the credential-vending flow the MLflow SDK uses for -// Databricks-backed runs: credentials-for-read returns a pre-signed URL, which we -// stream to disk. The credentials-for-read endpoint is not modeled by the SDK, so -// it is called via a raw client.Do. +// downloadArtifact downloads one run artifact to a temp file and returns its +// 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) { apiClient, err := client.New(w.Config) if err != nil { @@ -261,7 +243,7 @@ func downloadArtifact(ctx context.Context, w *databricks.WorkspaceClient, mlflow var resp credentialsForReadResponse query := map[string]any{ "run_id": mlflowRunID, - // path is a repeated field; the SDK serializes a slice as path=...&path=... + // path is a repeated field, so pass a slice (serialized as path=...&path=...). "path": []string{artifactPath}, } err = apiClient.Do(ctx, http.MethodGet, "/api/2.0/mlflow/artifacts/credentials-for-read", nil, nil, query, &resp) @@ -277,8 +259,7 @@ func downloadArtifact(ctx context.Context, w *databricks.WorkspaceClient, mlflow if err != nil { return "", err } - // Azure SAS / some GCS URIs require backend-supplied headers; AWS pre-signed - // URLs return none. + // Azure SAS / some GCS URIs require backend-supplied headers; AWS returns none. for _, h := range cred.Headers { req.Header.Set(h.Name, h.Value) } diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index e311301434d..4eb58dae113 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -39,11 +39,10 @@ func newLogsCommand() *cobra.Command { cmd.Flags().IntVar(&retry, "retry", -1, "View logs from a specific retry attempt; -1 means latest") cmd.Flags().StringVar(&downloadTo, "download-to", "", "Download all logs to this directory instead of printing") cmd.Flags().BoolVar(&review, "review", false, "Download logs from all nodes and filter for error signatures") - // Hidden in the Python `air` CLI (help=argparse.SUPPRESS); keep it internal here to match. cmd.Flags().MarkHidden("review") // In -o json mode an auth failure should be a JSON error envelope, not a bare - // error. ErrAlreadyPrinted passes through (already handled upstream). + // error. ErrAlreadyPrinted passes through. cmd.PreRunE = func(cmd *cobra.Command, args []string) error { err := root.MustWorkspaceClient(cmd, args) if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { @@ -55,8 +54,8 @@ func newLogsCommand() *cobra.Command { cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() - // --download-to and --review are not yet ported. Reject them explicitly - // rather than silently ignoring, so the user knows they had no effect. + // --download-to and --review are not yet implemented; reject rather than + // silently ignore. if downloadTo != "" { return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, errors.New("--download-to is not implemented yet")) @@ -66,9 +65,8 @@ func newLogsCommand() *cobra.Command { errors.New("--review is not implemented yet")) } - // --lines is a line tail (MLflow-native); --minutes is a time window - // (Bricklens-native). They answer the same "how much" question two ways, - // so reject both at once rather than silently honoring one. + // --lines (line tail) and --minutes (time window) answer the same question + // two ways, so reject both together rather than silently honoring one. if lines > 0 && minutes > 0 { return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, errors.New("cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window")) @@ -101,14 +99,13 @@ func newLogsCommand() *cobra.Command { return cmd } -// runLogs resolves the run, validates the requested retry, and fetches logs -// Bricklens-first with an MLflow fallback. It handles error reporting; the -// backend selection lives in fetchLogs. +// runLogs resolves the run, validates --retry, and fetches logs. It handles error +// reporting; the backend selection lives in fetchLogs. func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error { w := cmdctx.WorkspaceClient(ctx) // Validate credentials server-side before fetching (MustWorkspaceClient only - // attaches credentials), so a bad token fails clearly here. + // attaches them), so a bad token fails clearly here. if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil { return authError(ctx, cmd, err) } @@ -123,14 +120,13 @@ func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error { fmt.Errorf("failed to get status for run %d: %w", req.runID, err)) } - // Resolve --retry against the run's attempts. -1 (default) means latest. + // -1 (default) means latest; reject an attempt past the newest. if req.attempt >= 0 && req.attempt > status.latestAttempt { return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, fmt.Errorf("invalid retry %d: available retries are 0 to %d", req.attempt, status.latestAttempt)) } - // A past retry of a still-active run has immutable logs: render them once - // rather than following the run to completion. + // A past retry of an active run has immutable logs: render once, don't follow. if req.attempt >= 0 && req.attempt < status.latestAttempt && !status.terminal() { req.staticView = true } @@ -146,19 +142,16 @@ func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error { fmt.Errorf("failed to fetch logs for run %d: %w", req.runID, err)) } - // A run that finished unsuccessfully exits non-zero, matching the Python CLI; - // output was already written, so don't reprint via Cobra. + // A run that finished unsuccessfully exits non-zero; output was already + // written, so don't reprint via Cobra. if !success { return root.ErrAlreadyPrinted } return nil } -// fetchLogs is the Bricklens-first / MLflow-fallback decision. Bricklens serves -// logs unless the backend has gated it off (a SAFE flag) or it is unavailable; -// in that case streamBricklensLogs returns errBricklensFeatureDisabled and we -// fall back to MLflow. Both paths honor the same logRequest (node, retry, -// --minutes window, --lines tail). +// 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) { success, err := streamBricklensLogs(ctx, w, out, req, status) if errors.Is(err, errBricklensFeatureDisabled) { diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go index a4426ee5dfc..c291f766172 100644 --- a/experimental/air/cmd/logs_test.go +++ b/experimental/air/cmd/logs_test.go @@ -25,7 +25,7 @@ func TestLogsCommandShape(t *testing.T) { assert.Error(t, cmd.Args(cmd, []string{})) assert.Error(t, cmd.Args(cmd, []string{"1", "2"})) - // --review stays hidden to match the Python CLI. + // --review is hidden. review := cmd.Flags().Lookup("review") require.NotNil(t, review) assert.True(t, review.Hidden) @@ -239,11 +239,9 @@ func TestLogsPastRetryOfActiveRunIsStatic(t *testing.T) { cmd.SetOut(&buf) // --retry 0 on a RUNNING run whose latest attempt is 1: the past attempt's - // logs are static, so they render once and the command returns instead of - // following the run (which would never terminate). The run itself has no - // SUCCESS result yet, so the command still exits non-zero (matching Python's - // `result_state == "SUCCESS"` exit-code rule) via ErrAlreadyPrinted — the - // logs are printed regardless. + // logs render once instead of following the run (which would never terminate). + // The run has no SUCCESS result yet, so it still exits non-zero via + // ErrAlreadyPrinted; the logs are printed regardless. err := runLogs(ctx, cmd, logRequest{runID: 9, node: 0, attempt: 0}) require.ErrorIs(t, err, root.ErrAlreadyPrinted) assert.Equal(t, "retry 0 log\n", buf.String()) diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index 279e00e934d..8e77ec62c74 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -14,36 +14,29 @@ import ( "github.com/databricks/databricks-sdk-go/service/jobs" ) -// Log-stream tuning constants. These mirror the Python CLI (log_streaming.py) so -// behavior stays identical after the port. const ( - // retryCheckInterval is how long the poll loop waits between status/log polls. + // retryCheckInterval is the wait between status/log polls. retryCheckInterval = 3 * time.Second - // maxTransientFailures is how many consecutive Bricklens failures we tolerate - // before treating it as unavailable and falling back to MLflow. + // maxTransientFailures is how many consecutive Bricklens failures to tolerate + // before falling back to MLflow. maxTransientFailures = 5 // defaultCompletedRunTailLines caps a completed run's output when neither - // --lines nor --minutes is set, so a multi-chunk log does not flood stdout. + // --lines nor --minutes is set. defaultCompletedRunTailLines = 10000 - // seenRecordsCap bounds the dedup set so a large initial drain does not - // accumulate every record; oldest-inserted entries are evicted first. + // seenRecordsCap bounds the dedup set, evicting oldest-inserted entries first. seenRecordsCap = 100000 ) -// errBricklensFeatureDisabled signals the caller to fall back to the MLflow log -// path. It is returned when Bricklens can't serve logs: the endpoint is gated -// off by a backend SAFE flag (FEATURE_DISABLED), not deployed (ENDPOINT_NOT_FOUND -// / 404), or persistently failing after maxTransientFailures retries. The flag is -// evaluated server-side; the CLI only reads the resulting error code. +// 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. var errBricklensFeatureDisabled = errors.New("bricklens logs unavailable; falling back to mlflow") -// logRequest is the resolved, backend-agnostic description of what to fetch. It -// is shared by the Bricklens streamer and the MLflow fallback so both honor the -// same flags. windowMinutes and tailLines are mutually exclusive (validated at -// the command layer): windowMinutes drives the time window, tailLines the tail. +// logRequest describes what to fetch, shared by both backends so they honor the +// same flags. windowMinutes and tailLines are mutually exclusive. type logRequest struct { runID int64 - // node is the node index to fetch; 0 by default (node 0 always exists). + // node is the node index to fetch; node 0 always exists. node int // attempt is the retry attempt to read; -1 means latest. attempt int @@ -51,16 +44,15 @@ type logRequest struct { windowMinutes int // tailLines, when > 0, keeps only the last N lines of a completed run. tailLines int - // staticView renders a one-shot tail instead of following the run. It is set - // for a past retry of a still-active run: that attempt's logs are immutable, - // so streaming them would poll forever waiting for the run (not the attempt) - // to finish. A completed run is inherently static and does not need this. + // staticView renders a one-shot tail instead of following the run. Set for a + // past retry of an active run: that attempt's logs are immutable, so streaming + // would poll forever waiting for the run (not the attempt) to finish. staticView bool jsonOutput bool } -// runStatus is the subset of a run's state the log path needs. It is resolved -// once from a Jobs GetRun and reused, avoiding a per-tick typed-vs-dict shuffle. +// logRunStatus is the subset of a run's state the log path needs, resolved once +// and reused. type logRunStatus struct { lifeCycleState string resultState string @@ -71,9 +63,8 @@ type logRunStatus struct { latestAttempt int } -// terminalLifeCycleStates and terminalResultStates classify a finished run. A -// run is terminal when its lifecycle state is terminal, or a result state is set -// (result states only appear on terminal runs). Mirrors log_streaming.py. +// A run is terminal when its lifecycle state is terminal, or a result state is +// set (result states only appear on terminal runs). var ( terminalLifeCycleStates = map[string]bool{"TERMINATED": true, "SKIPPED": true, "INTERNAL_ERROR": true} terminalResultStates = map[string]bool{"SUCCESS": true, "FAILED": true, "CANCELED": true} @@ -87,9 +78,8 @@ func (s logRunStatus) succeeded() bool { return s.resultState == "SUCCESS" } -// resolveRunStatus fetches a run's state via the Jobs API and projects it onto -// logRunStatus. The run id being unknown surfaces as apierr.ErrResourceDoesNotExist -// so the caller can report a clean not-found. +// resolveRunStatus fetches a run's state and projects it onto logRunStatus. An +// unknown run id surfaces as apierr.ErrResourceDoesNotExist. func resolveRunStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (logRunStatus, error) { run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) if err != nil { @@ -98,8 +88,8 @@ func resolveRunStatus(ctx context.Context, w *databricks.WorkspaceClient, runID return projectRunStatus(run), nil } -// projectRunStatus extracts logRunStatus from a Jobs run. Split out from -// resolveRunStatus so it can be unit-tested without an API client. +// projectRunStatus extracts logRunStatus from a run. Split out so it can be +// tested without an API client. func projectRunStatus(run *jobs.Run) logRunStatus { s := logRunStatus{ startTimeMs: run.StartTime, @@ -116,14 +106,11 @@ func projectRunStatus(run *jobs.Run) logRunStatus { return s } -// classifyLogError decides how a Bricklens failure should be handled: -// - errBricklensFeatureDisabled: fall back to MLflow (flag gated off, endpoint -// absent, or a 404 — older logs may still live in MLflow). +// classifyLogError maps a Bricklens failure to one of: +// - errBricklensFeatureDisabled: fall back to MLflow (gated off, endpoint +// absent, or 404). // - the original error: a genuine not-found, surfaced as-is. // - nil: a transient failure the caller should retry. -// -// The backend evaluates the SAFE flag; this only reads the returned error code, -// so it never string-matches error text (repo rule). func classifyLogError(err error) error { if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { switch apiErr.ErrorCode { @@ -140,10 +127,9 @@ func classifyLogError(err error) error { return nil } -// fromSeconds computes the Bricklens `from` bound for a request. With --minutes -// set it is now-N*60; otherwise it is the run's start second (0 when the run has -// not started, which the endpoint reads as "everything stored"). Matches -// stream_logs_via_bricklens. +// fromSeconds computes the `from` bound. With --minutes set it is now-N*60; +// otherwise the run's start second (0 before the run starts, which the endpoint +// reads as "everything stored"). func (req logRequest) fromSeconds(status logRunStatus, now time.Time) int64 { if req.windowMinutes > 0 { return now.Add(-time.Duration(req.windowMinutes) * time.Minute).Unix() @@ -154,10 +140,9 @@ func (req logRequest) fromSeconds(status logRunStatus, now time.Time) int64 { return 0 } -// toSeconds computes the Bricklens `to` bound. Once a run is terminal no new logs -// can appear, so we cap at the run's end second (ceil of the millisecond end time -// so records in the final partial second aren't excluded); otherwise 0 lets the -// endpoint default to now. +// toSeconds computes the `to` bound. A terminal run caps at its end second (ceil +// of the millisecond time, so the final partial second is kept); otherwise 0 lets +// the endpoint default to now. func (req logRequest) toSeconds(status logRunStatus) int64 { if status.terminal() && status.endTimeMs > 0 { return (status.endTimeMs + 999) / 1000 @@ -165,13 +150,10 @@ func (req logRequest) toSeconds(status logRunStatus) int64 { return 0 } -// streamBricklensLogs fetches and prints a run's logs via Bricklens, handling -// both an already-completed run (a single bounded tail drain) and an active run -// (a poll-and-drain loop that follows until the run reaches a terminal state). -// -// It returns (success, err) where success reports whether the run finished with -// SUCCESS. A returned errBricklensFeatureDisabled means the caller should fall -// back to the MLflow path; the run genuinely not existing is returned as-is. +// streamBricklensLogs fetches and prints a run's logs: a bounded tail for a +// completed run, or a poll-and-drain loop that follows an active run to +// completion. It returns whether the run finished with SUCCESS; +// errBricklensFeatureDisabled means the caller should fall back to MLflow. func streamBricklensLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { st := &bricklensStreamer{ ctx: ctx, @@ -184,8 +166,8 @@ func streamBricklensLogs(ctx context.Context, w *databricks.WorkspaceClient, out return st.run() } -// bricklensStreamer carries the streaming state across the poll loop: the running -// from-second cursor, the highest emitted timestamp, and the bounded dedup set. +// bricklensStreamer holds the poll-loop state: the from-second cursor, the +// highest emitted timestamp, and the dedup set. type bricklensStreamer struct { ctx context.Context w *databricks.WorkspaceClient @@ -203,9 +185,8 @@ func (st *bricklensStreamer) run() (bool, error) { now := time.Now() st.fromSec = st.req.fromSeconds(st.status, now) - // A past retry's logs are immutable, so render them as a one-shot tail even - // when the run is still active — the attempt has ended, and following the run - // would poll forever. Mirrors handle_logs' viewing_past_retry. + // A past retry's logs are immutable: render a one-shot tail rather than + // following the still-active run, which would poll forever. if st.req.staticView { return st.drainStatic(st.req.toSeconds(st.status)) } @@ -218,7 +199,7 @@ func (st *bricklensStreamer) run() (bool, error) { if errors.Is(err, apierr.ErrResourceDoesNotExist) { return false, err } - // A transient status blip should not abort a live stream; log and retry. + // A transient status blip should not abort a live stream. log.Debugf(st.ctx, "air logs: failed to refresh run status: %v", err) time.Sleep(retryCheckInterval) continue @@ -229,11 +210,9 @@ func (st *bricklensStreamer) run() (bool, error) { terminal := st.status.terminal() toSec := st.req.toSeconds(st.status) - // An already-completed run (terminal on the first iteration) renders as a - // tail: only the most-recent N lines, oldest-first. An active run — even - // one that terminates while we watch — streams everything and dedups so the - // final drain doesn't re-print the boundary second. A tail is line-based, so - // it only applies to the completed-run case (matching the MLflow path). + // A run already terminal on the first iteration renders as a tail (most + // recent N lines). An active run streams everything with dedup, so a run + // that terminates while we watch doesn't re-print the boundary second. var err error if firstIteration && terminal { err = st.drainTail(toSec) @@ -257,10 +236,8 @@ func (st *bricklensStreamer) run() (bool, error) { } } -// drainStatic renders a single tail pass for an immutable attempt and returns -// without following the run. Success reflects the run's current result state -// (empty for an active run, so a past retry of a running job is not reported as -// a failure). +// drainStatic renders a single tail pass without following the run. Success +// reflects the run's current result state (empty while active). func (st *bricklensStreamer) drainStatic(toSec int64) (bool, error) { if err := st.drainTail(toSec); err != nil { return false, err @@ -271,8 +248,7 @@ func (st *bricklensStreamer) drainStatic(toSec int64) (bool, error) { return st.status.succeeded(), nil } -// tailTarget is the number of lines a completed-run tail should keep: --lines -// when set, else the default cap. +// tailTarget is the number of lines a tail keeps: --lines, else the default cap. func (req logRequest) tailTarget() int { if req.tailLines > 0 { return req.tailLines @@ -280,11 +256,9 @@ func (req logRequest) tailTarget() int { return defaultCompletedRunTailLines } -// drainTail emits the most-recent `target` records for a completed run, oldest-first. -// Bricklens returns records newest-first, so it pages until it has at least -// `target`, keeps the newest `target`, and reverses to chronological order — -// matching the MLflow --tail behavior. No dedup: a single bounded pass whose -// ordering we own client-side. +// drainTail emits the most-recent `target` records oldest-first. Bricklens +// returns records newest-first, so it pages until it has `target`, keeps the +// newest `target`, and reverses to chronological order. func (st *bricklensStreamer) drainTail(toSec int64) error { target := st.req.tailTarget() if target <= 0 { @@ -305,8 +279,7 @@ func (st *bricklensStreamer) drainTail(toSec int64) error { } } - // Keep the newest `target` (returned newest-first), then reverse so they print - // oldest -> newest like the MLflow tail. + // Keep the newest `target`, then reverse to print oldest -> newest. if len(collected) > target { collected = collected[:target] } @@ -316,10 +289,9 @@ func (st *bricklensStreamer) drainTail(toSec int64) error { return nil } -// drainPages exhausts all available pages from the current from-second in -// ascending (oldest-first) order so a live run streams chronologically, -// deduping against the bounded seen-set so a re-queried boundary second is not -// re-printed. It advances fromSec to the floor-second of the newest record seen. +// drainPages exhausts all pages from the current from-second in ascending order, +// deduping against the seen-set so a re-queried boundary second is not +// re-printed, then advances fromSec to the newest record's floor-second. func (st *bricklensStreamer) drainPages(toSec int64) error { var pageToken string for { @@ -331,9 +303,8 @@ func (st *bricklensStreamer) drainPages(toSec int64) error { for _, rec := range resp.LogRecords { nano := rec.nano() if nano != 0 { - // Ascending fetch, so a record older than the last emitted one is - // out of order (or a re-queried boundary record); skip it to keep - // streamed output monotonic. + // Skip a record older than the last emitted one to keep output + // monotonic (out of order, or a re-queried boundary record). if st.lastNano != 0 && nano < st.lastNano { continue } @@ -360,9 +331,8 @@ func (st *bricklensStreamer) drainPages(toSec int64) error { return nil } -// requestPage fetches one page, applying the fallback classification. It retries -// transient failures up to maxTransientFailures, then falls back like a disabled -// feature (older logs may still be in MLflow). A feature-gated response or a +// requestPage fetches one page, retrying transient failures up to +// maxTransientFailures before falling back to MLflow. A feature-gated response or // genuine not-found returns immediately. func (st *bricklensStreamer) requestPage(pageToken string, toSec int64, pageSize int, ascending bool) (*bricklensLogsResponse, error) { q := bricklensLogsQuery{ @@ -391,28 +361,26 @@ func (st *bricklensStreamer) requestPage(pageToken string, toSec int64, pageSize transientFailures++ if transientFailures >= maxTransientFailures { - log.Debugf(st.ctx, "air logs: get_logs failed %d times; falling back to mlflow", maxTransientFailures) + log.Debugf(st.ctx, "air logs: bricklens failed %d times; falling back to mlflow", maxTransientFailures) return nil, errBricklensFeatureDisabled } - log.Debugf(st.ctx, "air logs: get_logs transient failure (%d/%d): %v", transientFailures, maxTransientFailures, err) + log.Debugf(st.ctx, "air logs: bricklens transient failure (%d/%d): %v", transientFailures, maxTransientFailures, err) time.Sleep(retryCheckInterval) } } -// emit writes one log line and latches firstLogSeen so a terminal run with no -// output can report "no logs". +// emit writes one log line and latches firstLogSeen so an empty terminal run can +// report "no logs". func (st *bricklensStreamer) emit(body string) { st.firstLogSeen = true emitLogLine(st.out, st.req, body) } -// emitNoLogs reports that a terminal run produced no logs. func (st *bricklensStreamer) emitNoLogs() { emitNoLogs(st.out, st.req, st.status) } -// displayState is the run's terminal result state, falling back to the lifecycle -// state, else "UNKNOWN". +// displayState is the result state, else the lifecycle state, else "UNKNOWN". func (s logRunStatus) displayState() string { if s.resultState != "" { return s.resultState diff --git a/experimental/air/cmd/logstream_support.go b/experimental/air/cmd/logstream_support.go index f644887086d..c6d4e5fe5d7 100644 --- a/experimental/air/cmd/logstream_support.go +++ b/experimental/air/cmd/logstream_support.go @@ -8,18 +8,15 @@ import ( "time" ) -// seenNano keys the dedup set: a (nano, body) pair. time_unix_nano is nanosecond -// units and all ranks funnel into one stream stamping from their own clocks, so -// distinct log lines can share a nano — the body disambiguates them. +// seenNano keys the dedup set. Distinct lines can share a nano (each rank stamps +// from its own clock), so the body disambiguates them. type seenNano struct { nano int64 body string } // seenSet is an insertion-ordered set bounded to a capacity, evicting the -// oldest-inserted entry first. It lets the streamer re-query a boundary second -// on the next poll without re-printing records already shown, while a large -// initial drain can't grow without bound. +// oldest-inserted entry first. type seenSet struct { cap int items map[seenNano]*list.Element @@ -52,8 +49,7 @@ func (s *seenSet) add(nano int64, body string) { } } -// logEvent is one JSONL streaming event, matching the Python CLI's -// print_jsonl_event shape ({type, ts, node, line}). +// logEvent is one JSONL streaming event. type logEvent struct { Type string `json:"type"` TS string `json:"ts"` @@ -75,8 +71,7 @@ func printLogEvent(out io.Writer, eventType string, node int, line string) { fmt.Fprintln(out, string(b)) } -// emitLogLine writes one log line, as a JSONL LOG event under --json or raw -// otherwise. Shared by the Bricklens and MLflow paths so both emit identically. +// emitLogLine writes one log line: a JSONL LOG event under --json, else raw. func emitLogLine(out io.Writer, req logRequest, body string) { if req.jsonOutput { printLogEvent(out, "LOG", req.node, body) @@ -86,8 +81,7 @@ func emitLogLine(out io.Writer, req logRequest, body string) { } // emitNoLogs reports that a run produced no logs, with its termination reason. -// Under --json it is a JSONL ERROR (so a consumer never sees an empty stream); -// otherwise a plain line. Shared by both backends. +// Under --json it is a JSONL ERROR, so a consumer never sees an empty stream. func emitNoLogs(out io.Writer, req logRequest, status logRunStatus) { msg := fmt.Sprintf("No logs available for run %d. Run terminated in state %s", req.runID, status.displayState()) if status.stateMessage != "" { From a594c75b6f1a6375069e3c687c04fe38a6a2b34f Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Fri, 17 Jul 2026 21:41:03 +0000 Subject: [PATCH 5/5] AIR CLI: close air logs unit-test coverage gaps Add direct unit tests for the JSONL emit path (emitLogLine / emitNoLogs), the requestPage transient-retry loop (retries then falls back, and retries then succeeds), and displayState. Make retryCheckInterval a var so the retry-loop tests run without the 3s production wait. Co-authored-by: Isaac --- experimental/air/cmd/logstream.go | 6 +- experimental/air/cmd/logstream_test.go | 108 +++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 2 deletions(-) diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index 8e77ec62c74..9f2016656b7 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -15,8 +15,6 @@ import ( ) const ( - // retryCheckInterval is the wait between status/log polls. - retryCheckInterval = 3 * time.Second // maxTransientFailures is how many consecutive Bricklens failures to tolerate // before falling back to MLflow. maxTransientFailures = 5 @@ -27,6 +25,10 @@ const ( seenRecordsCap = 100000 ) +// retryCheckInterval is the wait between status/log polls. A var so tests can +// shrink it. +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. diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go index c229a4bce58..0b01ed82d61 100644 --- a/experimental/air/cmd/logstream_test.go +++ b/experimental/air/cmd/logstream_test.go @@ -2,9 +2,11 @@ package aircmd import ( "bytes" + "encoding/json" "errors" "net/http" "net/http/httptest" + "strings" "testing" "time" @@ -188,6 +190,112 @@ func TestDrainPagesDedupAndOrdering(t *testing.T) { assert.Equal(t, int64(3000), st.lastNano) } +func TestDisplayState(t *testing.T) { + assert.Equal(t, "SUCCESS", logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}.displayState()) + assert.Equal(t, "RUNNING", logRunStatus{lifeCycleState: "RUNNING"}.displayState()) + assert.Equal(t, "UNKNOWN", logRunStatus{}.displayState()) +} + +func TestEmitLogLineJSON(t *testing.T) { + var buf bytes.Buffer + emitLogLine(&buf, logRequest{node: 2, jsonOutput: true}, "hello") + + var ev logEvent + require.NoError(t, json.Unmarshal(buf.Bytes(), &ev)) + assert.Equal(t, "LOG", ev.Type) + assert.Equal(t, 2, ev.Node) + assert.Equal(t, "hello", ev.Line) + assert.NotEmpty(t, ev.TS) +} + +func TestEmitLogLineText(t *testing.T) { + var buf bytes.Buffer + emitLogLine(&buf, logRequest{node: 0}, "hello") + assert.Equal(t, "hello\n", buf.String()) +} + +func TestEmitNoLogs(t *testing.T) { + status := logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED", stateMessage: "boom"} + + var text bytes.Buffer + emitNoLogs(&text, logRequest{runID: 7}, status) + assert.Equal(t, "No logs available for run 7. Run terminated in state FAILED: boom\n", text.String()) + + var jsonBuf bytes.Buffer + emitNoLogs(&jsonBuf, logRequest{runID: 7, node: 1, jsonOutput: true}, status) + var ev logEvent + require.NoError(t, json.Unmarshal(jsonBuf.Bytes(), &ev)) + assert.Equal(t, "ERROR", ev.Type) + assert.Equal(t, 1, ev.Node) + assert.Contains(t, ev.Line, "No logs available for run 7") + assert.Contains(t, ev.Line, "boom") +} + +func TestRequestPageRetriesThenFallsBack(t *testing.T) { + // Shrink the retry wait so the transient-failure loop runs fast. + orig := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = orig }) + + var calls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.URL.Path, "/logs") { + // Ignore SDK host/config probes so `calls` counts only log requests. + _, _ = w.Write([]byte(`{}`)) + return + } + calls++ + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "transient"}`)) + })) + t.Cleanup(srv.Close) + + st := &bricklensStreamer{ + ctx: t.Context(), + w: newTestWorkspaceClient(t, srv.URL), + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + _, err := st.requestPage("", 0, 0, true) + // Persistent transient failures fall back to MLflow after the retry budget. + require.ErrorIs(t, err, errBricklensFeatureDisabled) + assert.Equal(t, maxTransientFailures, calls) +} + +func TestRequestPageRetriesThenSucceeds(t *testing.T) { + orig := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = orig }) + + var calls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.URL.Path, "/logs") { + _, _ = w.Write([]byte(`{}`)) + return + } + calls++ + if calls < 3 { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "transient"}`)) + return + } + _, _ = w.Write([]byte(`{"log_records": [{"time_unix_nano": 1, "body": "ok", "node_index": 0}]}`)) + })) + t.Cleanup(srv.Close) + + st := &bricklensStreamer{ + ctx: t.Context(), + w: newTestWorkspaceClient(t, srv.URL), + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + resp, err := st.requestPage("", 0, 0, true) + require.NoError(t, err) + require.Len(t, resp.LogRecords, 1) + assert.Equal(t, "ok", resp.LogRecords[0].Body) + assert.Equal(t, 3, calls) +} + func TestSeenSetEviction(t *testing.T) { s := newSeenSet(2) s.add(1, "a")