AIR CLI Migration: --download-to flag for logs - #6153
Conversation
Derive a run's node count from its accelerator type and count (accelerators are allocated in whole nodes), reusing jobCompute + parseGPUType + gpusPerNode. This is the first piece of `air logs --download-to`. Co-authored-by: Isaac
downloadAllNodeLogs downloads the given nodes' logs concurrently (errgroup, bounded to downloadConcurrency) from the MLflow artifact store, probing the attempt-prefix layout once up front. downloadNodeLog concatenates a node's chunk files into dir/logs/node_<n>.log. Reuses the existing MLflow helpers (discoverAttemptPrefix, constructLogPath, listLogChunks, downloadChunkLines). Co-authored-by: Isaac
Approval status: pending
|
downloadLogs resolves the run's node count, downloads the requested node(s) in parallel, and prints a per-node summary. An explicit --node downloads just that node (rejected if out of range); omitted downloads all nodes. Replaces the "--download-to is not implemented yet" rejection in logs.go and threads the new downloadTo / nodeSet fields through logRequest. Adds the logs-download acceptance scenario (no-logs + out-of-range node; the full byte download is unit-covered since the pre-signed URL host is dynamic) and drops the stale rejection cases. Co-authored-by: Isaac
17cdb2f to
9874c6a
Compare
Correctness: - A run with no logs now reports it and exits 0, the same as the streaming path, instead of a retryable INTERNAL_ERROR and exit 1. - An out-of-range --node is classified INVALID_ARGS/PERMANENT via a sentinel, not a transient failure, and is no longer double-wrapped. - resolveNodeCount rejects an accelerator count that isn't a whole number of nodes; it previously truncated to 0 nodes and reported "no logs". - Cancellation propagates out of the parallel download instead of looking like a node with no logs. - Chunks stream byte-for-byte to the output file rather than round-tripping through lines, which rewrote line endings and capped long lines. - An empty chunk listing falls back to chunk 0, as the streaming path does. - Per-node failures are surfaced as warnings, so a truncated log set no longer looks complete. - Reject --lines/--minutes with --download-to instead of ignoring them. Tests: explicit --node 0, non-divisible and zero accelerator counts, no-logs parity with the streaming path, partial failure, cancellation, and the attempt-prefixed layout. Co-authored-by: Isaac
riddhibhagwat-db
changed the title
Integration test reportCommit: f0c9276
8 interesting tests: 4 RECOVERED, 4 SKIP
Top 6 slowest tests (at least 2 minutes):
|
The pre-signed-URL request passed the artifact path as a []string in a map query. The SDK formats map values with %v, so the path was sent as the literal "[logs/node_0/logs-0.chunk.txt]". The backend signs that path and returns 200, so every chunk download 404'd instead. This also silently broke --download-to: the per-node failure was recorded but the warnings printed after the no-logs early return, so a run with logs reported "No logs available". Report failures first. Co-authored-by: Isaac
ben-hansen-db
left a comment
ben-hansen-db
left a comment
There was a problem hiding this comment.
Looks good, just a few comments to address
| os.Remove(outPath) | ||
| return "", err | ||
| } | ||
| break |
There was a problem hiding this comment.
I realize there could be subtle bug here, where chunk X might succeed but chunk X +1 fails and then that isn't reported.
Maybe something to surface when that happens
return outPath, fmt.Errorf("truncated at chunk %d: %w", chunk.index, err)
There was a problem hiding this comment.
Also is there a retry on failed download of a chunk?
There was a problem hiding this comment.
Reimplemented this to skip a failed chunk and keep going rather than stopping at the first gap, matching the Python CLI's behavior.
The original logic was the wrong tradeoff: if chunk 1 of 20 fails, we'd lose chunks 2–20 including the tail, which is usually where the failure signature is. Skipping preserves everything else, and the gap is no longer silent: downloadNodeLog returns both the partial path and an error naming the missing chunk indices, so the summary reports the node and marks it. Cancellation still aborts immediately since every remaining chunk would fail anyway, and a node where every chunk failed returns an error rather than an empty result. Also added tests for a mid-log gap (asserts chunk 0 + chunk 2 with chunk 1 failing), the partial-download error, and the all-chunks-failed case.
| return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, | ||
| fmt.Errorf("failed to download logs for run %d: %w", req.runID, err)) | ||
| } | ||
| if !success { |
There was a problem hiding this comment.
Claude flagged this case.
Downloading a workload's logs where the run hasn't completed exits 1 rather than exits 0 b/c resultState is empty.
| // Reported before the no-logs check, so a run whose every node failed explains | ||
| // why instead of looking like a run that never logged. | ||
| for _, node := range sortedNodeKeys(failures) { | ||
| cmdio.LogString(ctx, fmt.Sprintf("warning: node %d logs could not be downloaded: %s", node, failures[node])) |
There was a problem hiding this comment.
Corner case I want to make sure we address.
2-node FAILED run, both nodes 404 or fail for some reason
Could there be a case where stdout does: "No logs available for run X. Run terminated in state FAILED"
And the above warning is missed so it seems there were no logs but they just failed to download?
this probably matters more for json case tbh where agent gets confused and thinks there's no logs
A chunk failing after earlier ones succeeded broke out of the loop and returned the path with no error, so a log cut short was listed as a normal success. The user got a short file with no indication it was incomplete. Return the path and an error naming the chunk; the caller keeps the partial log and reports it, and the summary marks the node incomplete. Co-authored-by: Isaac
6a93c62 to
2612963
Compare
Stopping at the first bad chunk cost every chunk after it, including the tail where a run's failure signature usually is. Skip the failed chunk, keep walking, and name the gaps in the returned error so the log is never silently short. Matches the Python CLI, which skips a failed chunk and concatenates the rest. Co-authored-by: Isaac
A one-shot fetch can run against an active run, which has no result state yet, so succeeded() was false and the command exited 1 even though the logs printed fine. Only a terminal run's outcome should decide the exit code. Affects --download-to and the two pre-existing one-shot paths (the MLflow fallback and a past retry's static view). Co-authored-by: Isaac
A run whose nodes all failed to download reported "No logs available", which tells the caller the run produced nothing. The per-node warnings go to stderr, so a -o json consumer reading stdout saw only the misleading ERROR event. Return an error in that case, so JSON mode renders an error envelope on stdout. Co-authored-by: Isaac
3 participants
Changes & Why
air logs --download-to DIR was a registered flag that just errored ("not implemented yet"). It exists so you can pull a failed multi-node run's logs to disk for offline inspection, instead of streaming one node at a time to your terminal. Without it, debugging a 2+ node failure meant running air logs --node N repeatedly and copy-pasting.
Design decisions and their reasons:
Tests
Unit tests:
Acceptance tests
Manual verification:
How to review this PR:
Breakdown: ~700 lines total; 437 are tests, 275 is production code
Careful reads requested on the following lines:
logmlflow.go, line 261:logdownload.go, line 46:logdownload.go:25/:70andlogs.go:147Error Classification:logdownload.golines 103-109:logdownload.go:178and:199-208: bytes are streamed verbatim via io.Copy, chunk 0 is assumed when the listing comes back empty, and a 0 byte result removes the file and reports "no logs" rather than leaving an empty artifact