diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index ee89e778d6..52e5001684 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 81f3907e4f..91dc12567a 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-mlflow-fallback/out.test.toml b/acceptance/experimental/air/logs-mlflow-fallback/out.test.toml new file mode 100644 index 0000000000..d6187dcb04 --- /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 0000000000..a511fd33bc --- /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 0000000000..ad8856a036 --- /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 0000000000..9b9bcab8e6 --- /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/acceptance/experimental/air/logs/out.test.toml b/acceptance/experimental/air/logs/out.test.toml new file mode 100644 index 0000000000..d6187dcb04 --- /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 0000000000..3ee3219faf --- /dev/null +++ b/acceptance/experimental/air/logs/output.txt @@ -0,0 +1,57 @@ + +=== 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 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 + +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 0000000000..0cf8094c18 --- /dev/null +++ b/acceptance/experimental/air/logs/script @@ -0,0 +1,26 @@ +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 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 + +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 0000000000..dfec0f46a7 --- /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 7db6ef1aec..261e0c456a 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 19dc13ffe8..ae0366a48c 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 0000000000..5e0a23f162 --- /dev/null +++ b/experimental/air/cmd/logbricklens.go @@ -0,0 +1,93 @@ +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 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. +type logRecord struct { + // 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 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 { + 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. +// 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 + // absent, so the tail fetch must send an explicit false for newest-first. + ascending bool +} + +// 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 { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + query := map[string]any{ + // Always sent: the tail path relies on an explicit false for 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 0000000000..07fd62b7ff --- /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 0000000000..15a73ee0d5 --- /dev/null +++ b/experimental/air/cmd/logmlflow.go @@ -0,0 +1,290 @@ +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 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_). 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, 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. +// +// 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") + } + + 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. + 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. +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 + } + + // -1 (latest) maps to attempt 0's directory. + attempt := max(req.attempt, 0) + withAttempt, err := discoverAttemptPrefix(ctx, w, ids.RunID, attempt) + if err != nil { + return "", "", err + } + return ids.RunID, constructLogPath(req.node, attempt, withAttempt), nil +} + +// 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 { + // 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 + } + + 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 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 { + 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 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 it: %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. +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. +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 credentials-for-read entry: a pre-signed URL 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 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 { + 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, 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) + 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 returns 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 0000000000..9fbe4f19ff --- /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.go b/experimental/air/cmd/logs.go index c34fb62a7d..4eb58dae11 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,131 @@ 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. + 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 implemented; reject rather than + // silently ignore. + 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 (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")) + } + 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 --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 them), 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)) + } + + // -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 an active run has immutable logs: render once, don't follow. + 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 { + 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; output was already + // written, so don't reprint via Cobra. + if !success { + return root.ErrAlreadyPrinted + } + return nil +} + +// 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) { + 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 0000000000..c291f76617 --- /dev/null +++ b/experimental/air/cmd/logs_test.go @@ -0,0 +1,252 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "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" + "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 is hidden. + 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()) +} + +// 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": [{"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 := 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) + + // 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.NoError(t, err) + assert.Equal(t, "line one\nline two\n", buf.String()) +} + +// 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 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()) + + // 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 new file mode 100644 index 0000000000..9f2016656b --- /dev/null +++ b/experimental/air/cmd/logstream.go @@ -0,0 +1,394 @@ +package aircmd + +import ( + "context" + "errors" + "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" +) + +const ( + // 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. + defaultCompletedRunTailLines = 10000 + // seenRecordsCap bounds the dedup set, evicting oldest-inserted entries first. + 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. +var errBricklensFeatureDisabled = errors.New("bricklens logs unavailable; falling back to mlflow") + +// 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; 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 + // 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 +} + +// logRunStatus is the subset of a run's state the log path needs, resolved once +// and reused. +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 +} + +// 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} +) + +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 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 { + return logRunStatus{}, err + } + return projectRunStatus(run), nil +} + +// 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, + 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 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. +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 `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() + } + if status.startTimeMs > 0 { + return status.startTimeMs / 1000 + } + return 0 +} + +// 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 + } + return 0 +} + +// 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, + w: w, + out: out, + req: req, + status: status, + seen: newSeenSet(seenRecordsCap), + } + return st.run() +} + +// 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 + 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) + + // 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)) + } + + 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.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) + + // 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) + } 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) + } +} + +// 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 + } + if !st.firstLogSeen { + st.emitNoLogs() + } + return st.status.succeeded(), nil +} + +// 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 + } + return defaultCompletedRunTailLines +} + +// 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 { + 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`, then reverse to print oldest -> newest. + if len(collected) > target { + collected = collected[:target] + } + for _, c := range slices.Backward(collected) { + st.emit(c.Body) + } + return nil +} + +// 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 { + resp, err := st.requestPage(pageToken, toSec, 0, true) + if err != nil { + return err + } + + for _, rec := range resp.LogRecords { + nano := rec.nano() + if nano != 0 { + // 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 + } + 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, 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{ + 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: bricklens failed %d times; falling back to mlflow", maxTransientFailures) + return nil, errBricklensFeatureDisabled + } + 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 an empty terminal run can +// report "no logs". +func (st *bricklensStreamer) emit(body string) { + st.firstLogSeen = true + emitLogLine(st.out, st.req, body) +} + +func (st *bricklensStreamer) emitNoLogs() { + emitNoLogs(st.out, st.req, st.status) +} + +// displayState is the result state, else the lifecycle state, else "UNKNOWN". +func (s logRunStatus) displayState() string { + if s.resultState != "" { + 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 0000000000..c6d4e5fe5d --- /dev/null +++ b/experimental/air/cmd/logstream_support.go @@ -0,0 +1,95 @@ +package aircmd + +import ( + "container/list" + "encoding/json" + "fmt" + "io" + "time" +) + +// 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. +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. +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)) +} + +// 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) + 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. +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) +} diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go new file mode 100644 index 0000000000..0b01ed82d6 --- /dev/null +++ b/experimental/air/cmd/logstream_test.go @@ -0,0 +1,316 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "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) { + 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 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 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") + 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 e28d7f6673..4007557ef7 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(), }