diff --git a/acceptance/experimental/air/logs-download/out.test.toml b/acceptance/experimental/air/logs-download/out.test.toml new file mode 100644 index 0000000000..e90b6d5d1b --- /dev/null +++ b/acceptance/experimental/air/logs-download/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/logs-download/output.txt b/acceptance/experimental/air/logs-download/output.txt new file mode 100644 index 0000000000..d3458758c4 --- /dev/null +++ b/acceptance/experimental/air/logs-download/output.txt @@ -0,0 +1,16 @@ + +=== download-to reports no logs when none are available +>>> [CLI] experimental air logs 123 --download-to dl-logs +No logs available for run 123. Run terminated in state SUCCESS + +=== download-to with an out-of-range node is rejected +>>> [CLI] experimental air logs 123 --download-to dl-logs --node 5 +Error: invalid --node 5: run has 2 node(s), indexed 0 to 1 + +Exit code: 1 + +=== download-to cannot be combined with --lines +>>> [CLI] experimental air logs 123 --download-to dl-logs --lines 50 +Error: --download-to writes complete logs, so it cannot be combined with --lines or --minutes + +Exit code: 1 diff --git a/acceptance/experimental/air/logs-download/script b/acceptance/experimental/air/logs-download/script new file mode 100644 index 0000000000..027b5c8952 --- /dev/null +++ b/acceptance/experimental/air/logs-download/script @@ -0,0 +1,12 @@ +# --download-to resolves the run's node count, then downloads each node's logs. +# This run resolves no MLflow run id, so it reports no logs (the full byte +# download is covered by unit tests, since the pre-signed URL host is dynamic). + +title "download-to reports no logs when none are available" +errcode trace $CLI experimental air logs 123 --download-to dl-logs + +title "download-to with an out-of-range node is rejected" +errcode trace $CLI experimental air logs 123 --download-to dl-logs --node 5 + +title "download-to cannot be combined with --lines" +errcode trace $CLI experimental air logs 123 --download-to dl-logs --lines 50 diff --git a/acceptance/experimental/air/logs-download/test.toml b/acceptance/experimental/air/logs-download/test.toml new file mode 100644 index 0000000000..1fd40b58e2 --- /dev/null +++ b/acceptance/experimental/air/logs-download/test.toml @@ -0,0 +1,40 @@ +# The command creates this download directory; don't treat it as test output. +Ignore = ["dl-logs"] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A completed 2-node run (GPU_1xA10 x 2 = 2 nodes). +[[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, + "ai_runtime_task": { + "experiment": "dl-exp", + "deployments": [ + {"command_path": "/x/command.sh", "compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}} + ] + } + } + ] +} +''' + +# No MLflow run id resolvable, so the download reports no logs rather than +# attempting a (host-dynamic) pre-signed artifact fetch. The full download path +# is covered by unit tests. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = '{}' diff --git a/acceptance/experimental/air/logs/output.txt b/acceptance/experimental/air/logs/output.txt index 84d6202afd..9ea8deedc5 100644 --- a/acceptance/experimental/air/logs/output.txt +++ b/acceptance/experimental/air/logs/output.txt @@ -64,9 +64,3 @@ Exit code: 1 Error: invalid --node -1: must not be negative 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 index b851a8784f..167ee331f0 100644 --- a/acceptance/experimental/air/logs/script +++ b/acceptance/experimental/air/logs/script @@ -27,6 +27,3 @@ errcode trace $CLI experimental air logs notanumber title "negative node" errcode trace $CLI experimental air logs 123 --node -1 - -title "--download-to not implemented" -errcode trace $CLI experimental air logs 123 --download-to /tmp/out diff --git a/experimental/air/cmd/logdownload.go b/experimental/air/cmd/logdownload.go new file mode 100644 index 0000000000..e6764f49a6 --- /dev/null +++ b/experimental/air/cmd/logdownload.go @@ -0,0 +1,266 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "path" + "path/filepath" + "slices" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "golang.org/x/sync/errgroup" +) + +// downloadConcurrency caps how many nodes download at once, to avoid hammering +// the artifact store on a wide run. +const downloadConcurrency = 8 + +// errNodeOutOfRange marks --node naming a node the run doesn't have. It is user +// input, so the caller reports it as an invalid argument rather than a failure. +var errNodeOutOfRange = errors.New("invalid --node") + +// resolveNodeCount returns how many nodes a run used. +func resolveNodeCount(run *jobs.Run) (int, error) { + accelType, count := jobCompute(run) + if accelType == "" { + return 0, fmt.Errorf("run %d has no AI runtime compute config", run.RunId) + } + if count <= 0 { + return 0, fmt.Errorf("run %d reports %d accelerators", run.RunId, count) + } + g, err := parseGPUType(accelType) + if err != nil { + return 0, err + } + perNode, err := gpusPerNode(g) + if err != nil { + return 0, err + } + // Accelerators come in whole nodes, so a remainder means we can't map the + // count onto node indices. + if count%perNode != 0 { + return 0, fmt.Errorf("run %d reports %d %s accelerators, which is not a multiple of %d per node", run.RunId, count, accelType, perNode) + } + return count / perNode, nil +} + +// downloadLogs writes each node's logs to /logs/node_.log and +// prints a summary. An explicit --node downloads only that node; otherwise all of +// them. Logs come from MLflow artifacts, since Bricklens only streams. The +// returned bool is the run's outcome, so the exit code matches the streaming path. +func downloadLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: req.runID}) + if err != nil { + return false, err + } + + numNodes, err := resolveNodeCount(run) + if err != nil { + return false, err + } + + nodes := make([]int, 0, numNodes) + if req.nodeSet { + if req.node >= numNodes { + return false, fmt.Errorf("%w %d: run has %d node(s), indexed 0 to %d", errNodeOutOfRange, req.node, numNodes, numNodes-1) + } + nodes = append(nodes, req.node) + } else { + for n := range numNodes { + nodes = append(nodes, n) + } + } + + // A run with no logs is reported the same way as on the streaming path, so + // the message and exit code agree between them. + ids := mlflowIDs(ctx, w, run) + if ids == nil || ids.RunID == "" { + emitNoLogs(out, req, status) + return status.downloadOutcome(), nil + } + + dir, err := filepath.Abs(req.downloadTo) + if err != nil { + return false, err + } + // Created up front so a bad --download-to fails with a clear message before + // any download work happens. + if err := os.MkdirAll(dir, 0o755); err != nil { + return false, fmt.Errorf("failed to create %s: %w", dir, err) + } + + nodeLogs, failures, err := downloadAllNodeLogs(ctx, w, ids.RunID, dir, nodes, req.attempt) + if err != nil { + return false, err + } + for _, node := range sortedNodeKeys(failures) { + cmdio.LogString(ctx, fmt.Sprintf("warning: node %d: %s", node, failures[node])) + } + + if len(nodeLogs) == 0 { + // "No logs available" would be a lie when the logs exist but couldn't be + // fetched. The warnings above go to stderr, which a -o json consumer reading + // stdout never sees, so fail instead of reporting an empty run. + if len(failures) > 0 { + return false, fmt.Errorf("failed to download logs from any of %d node(s): %s", + len(nodes), failures[sortedNodeKeys(failures)[0]]) + } + emitNoLogs(out, req, status) + return status.downloadOutcome(), nil + } + + cmdio.LogString(ctx, fmt.Sprintf("Downloaded logs from %d of %d node(s) to %s", len(nodeLogs), len(nodes), dir)) + for _, node := range sortedNodeKeys(nodeLogs) { + // Flag it on the file's own line, not just in the warning above. + suffix := "" + if _, truncated := failures[node]; truncated { + suffix = " (incomplete)" + } + cmdio.LogString(ctx, fmt.Sprintf(" node %d: %s%s", node, nodeLogs[node], suffix)) + } + return status.downloadOutcome(), nil +} + +// downloadAllNodeLogs downloads the nodes' logs in parallel. It returns a +// node->path map for the nodes that had logs and a node->reason map for those +// that failed; a truncated node appears in both. The log-dir layout is run-wide, +// so it is probed once here rather than by every worker. +func downloadAllNodeLogs(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, dir string, nodes []int, attempt int) (map[int]string, map[int]string, error) { + // -1 (latest) maps to attempt 0's directory, as on the streaming path. + attemptDir := max(attempt, 0) + withAttempt, err := discoverAttemptPrefix(ctx, w, mlflowRunID, attemptDir) + if err != nil { + return nil, nil, err + } + + paths := make([]string, len(nodes)) + reasons := make([]string, len(nodes)) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(downloadConcurrency) + for i, node := range nodes { + g.Go(func() error { + path, err := downloadNodeLog(gctx, w, mlflowRunID, node, attemptDir, withAttempt, dir) + switch { + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + // Interrupting the command must not look like a node with no logs. + return err + case err != nil: + // One bad node shouldn't abort the rest. A truncated node + // returns a path as well as an error, so keep both. + reasons[i] = err.Error() + paths[i] = path + default: + paths[i] = path + } + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, nil, err + } + + nodeLogs := map[int]string{} + failures := map[int]string{} + for i, node := range nodes { + if paths[i] != "" { + nodeLogs[node] = paths[i] + } + if reasons[i] != "" { + failures[node] = reasons[i] + } + } + return nodeLogs, failures, nil +} + +// downloadNodeLog streams a node's chunks in order into dir/logs/node_.log, +// returning the path, or "" if the node logged nothing. A failed chunk is skipped +// rather than ending the walk, so a partial download returns both a path and an +// error naming the gaps. The bytes are copied verbatim: a download should +// reproduce the log exactly, so it must not round-trip through lines (which would +// rewrite line endings and cap long lines). +func downloadNodeLog(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, node, attempt int, withAttempt bool, dir string) (string, error) { + logDir := constructLogPath(node, attempt, withAttempt) + chunks, err := listLogChunks(ctx, w, mlflowRunID, logDir) + if err != nil { + return "", err + } + if len(chunks) == 0 { + // The listing can lag behind the sidecar, so fall back to chunk 0 as the + // streaming path does. + chunks = []logChunk{{index: 0, path: path.Join(logDir, chunkFileName(0))}} + } + + outPath := filepath.Join(dir, "logs", fmt.Sprintf("node_%d.log", node)) + if err := os.MkdirAll(filepath.Dir(outPath), 0o755); err != nil { + return "", err + } + f, err := os.Create(outPath) + if err != nil { + return "", err + } + defer f.Close() + + // Skip a failed chunk and keep going: the tail usually holds the failure + // signature, so losing it to an early bad chunk is worse than a gap. Cancellation + // still aborts, since every remaining chunk would fail too. + var written int64 + var missing []int + for _, chunk := range chunks { + n, err := copyArtifactTo(ctx, w, mlflowRunID, chunk.path, f) + switch { + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + os.Remove(outPath) + return "", err + case err != nil: + log.Debugf(ctx, "air logs: node %d chunk %d failed: %v", node, chunk.index, err) + missing = append(missing, chunk.index) + default: + written += n + } + } + if written == 0 { + os.Remove(outPath) + if len(missing) > 0 { + return "", fmt.Errorf("every chunk failed to download (%d total)", len(missing)) + } + return "", nil + } + if len(missing) > 0 { + return outPath, fmt.Errorf("incomplete: chunk(s) %v failed to download", missing) + } + return outPath, nil +} + +// copyArtifactTo streams one artifact's bytes into dst and returns how many were +// written. +func copyArtifactTo(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string, dst io.Writer) (int64, error) { + local, err := downloadArtifact(ctx, w, mlflowRunID, artifactPath) + if err != nil { + return 0, err + } + defer os.Remove(local) + + src, err := os.Open(local) + if err != nil { + return 0, err + } + defer src.Close() + return io.Copy(dst, src) +} + +// sortedNodeKeys returns the map's node ids in ascending order, so the summary +// prints deterministically. +func sortedNodeKeys(m map[int]string) []int { + keys := make([]int, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + slices.Sort(keys) + return keys +} diff --git a/experimental/air/cmd/logdownload_test.go b/experimental/air/cmd/logdownload_test.go new file mode 100644 index 0000000000..b7717980f4 --- /dev/null +++ b/experimental/air/cmd/logdownload_test.go @@ -0,0 +1,554 @@ +package aircmd + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "os" + "path" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// airRunWithCompute builds a run reporting the given accelerator type and count. +func airRunWithCompute(accelType string, count int) *jobs.Run { + return &jobs.Run{ + RunId: 123, + Tasks: []jobs.RunTask{{ + RunId: 456, + AiRuntimeTask: &jobs.AiRuntimeTask{ + Deployments: []jobs.DeploymentSpec{{ + Compute: jobs.ComputeSpec{ + AcceleratorType: jobs.ComputeSpecAcceleratorType(accelType), + AcceleratorCount: count, + }, + }}, + }, + }}, + } +} + +func TestResolveNodeCount(t *testing.T) { + tests := []struct { + accelType string + count int + want int + }{ + {"GPU_1xA10", 2, 2}, + {"GPU_1xH100", 4, 4}, + {"GPU_8xH100", 16, 2}, + } + for _, tt := range tests { + n, err := resolveNodeCount(airRunWithCompute(tt.accelType, tt.count)) + require.NoError(t, err) + assert.Equal(t, tt.want, n) + } + + // A run with no AI runtime compute errors. + _, err := resolveNodeCount(&jobs.Run{RunId: 1}) + require.Error(t, err) + + // A count that isn't a whole number of nodes can't be mapped to node indices, + // so it errors rather than truncating to 0. + _, err = resolveNodeCount(airRunWithCompute("GPU_8xH100", 4)) + require.Error(t, err) + assert.Contains(t, err.Error(), "not a multiple of 8") + + // A zero count is reported as such, not as a missing config. + _, err = resolveNodeCount(airRunWithCompute("GPU_1xA10", 0)) + require.Error(t, err) +} + +// downloadServer serves the MLflow artifact chain: the artifact listing, a +// pre-signed URL pointing back at itself, and the chunk bytes. +func downloadServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}, {"path": "logs/node_1", "is_dir": true}]}`)) + } else { + _, _ = w.Write([]byte(`{"files": [{"path": "` + r.URL.Query().Get("path") + `/logs-0.chunk.txt", "file_size": 12}]}`)) + } + case "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("line one\nline two\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadNodeLogWritesConcatenatedChunks(t *testing.T) { + w := newTestWorkspaceClient(t, downloadServer(t).URL) + dir := t.TempDir() + + path, err := downloadNodeLog(t.Context(), w, "run1", 0, 0, false, dir) + require.NoError(t, err) + require.NotEmpty(t, path) + + got, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "line one\nline two\n", string(got)) + assert.Equal(t, filepath.Join(dir, "logs", "node_0.log"), path) +} + +func TestDownloadAllNodeLogs(t *testing.T) { + w := newTestWorkspaceClient(t, downloadServer(t).URL) + dir := t.TempDir() + + nodeLogs, failures, err := downloadAllNodeLogs(t.Context(), w, "run1", dir, []int{0, 1}, -1) + require.NoError(t, err) + require.Empty(t, failures) + require.Len(t, nodeLogs, 2) + assert.FileExists(t, nodeLogs[0]) + assert.FileExists(t, nodeLogs[1]) +} + +// fullDownloadServer also serves the run and its output, so downloadLogs can run +// end to end against a 2-node run. +func fullDownloadServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + runGet := `{ + "run_id": 123, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"run_id": 456, "ai_runtime_task": {"deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}}]}}] + }` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}, {"path": "logs/node_1", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "` + r.URL.Query().Get("path") + `/logs-0.chunk.txt", "file_size": 6}]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("hello\n")) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadLogsAllNodes(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + dir := t.TempDir() + + success, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, attempt: -1, downloadTo: dir}, logRunStatus{resultState: "SUCCESS"}) + require.NoError(t, err) + assert.True(t, success) + assert.FileExists(t, filepath.Join(dir, "logs", "node_0.log")) + assert.FileExists(t, filepath.Join(dir, "logs", "node_1.log")) +} + +func TestDownloadLogsSingleNode(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + dir := t.TempDir() + + success, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, node: 1, nodeSet: true, attempt: -1, downloadTo: dir}, logRunStatus{resultState: "SUCCESS"}) + require.NoError(t, err) + assert.True(t, success) + assert.FileExists(t, filepath.Join(dir, "logs", "node_1.log")) + assert.NoFileExists(t, filepath.Join(dir, "logs", "node_0.log")) +} + +func TestDownloadLogsOutOfRangeNode(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + dir := t.TempDir() + + _, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, node: 5, nodeSet: true, attempt: -1, downloadTo: dir}, logRunStatus{resultState: "SUCCESS"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid --node 5") +} + +func TestDownloadLogsExplicitNodeZero(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + dir := t.TempDir() + + // An explicit --node 0 must download only node 0, unlike the default which + // downloads every node. + _, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, node: 0, nodeSet: true, attempt: -1, downloadTo: dir}, logRunStatus{resultState: "SUCCESS"}) + require.NoError(t, err) + assert.FileExists(t, filepath.Join(dir, "logs", "node_0.log")) + assert.NoFileExists(t, filepath.Join(dir, "logs", "node_1.log")) +} + +func TestDownloadLogsOutOfRangeNodeIsInvalidArgs(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + + // The sentinel lets the caller classify this as bad input rather than a + // transient, retryable failure. + _, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, node: 5, nodeSet: true, attempt: -1, downloadTo: t.TempDir()}, logRunStatus{resultState: "SUCCESS"}) + require.ErrorIs(t, err, errNodeOutOfRange) +} + +// noLogsDownloadServer serves a run with no resolvable MLflow run, so there is +// nothing to download. +func noLogsDownloadServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{ + "run_id": 123, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"run_id": 456, "ai_runtime_task": {"deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}}]}}] + }`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadLogsNoLogsMatchesStreamingPath(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, noLogsDownloadServer(t).URL) + + // A SUCCESS run with no logs reports it and still succeeds, exactly as the + // streaming path does — no error, no non-zero exit. + var buf bytes.Buffer + success, err := downloadLogs(ctx, w, &buf, logRequest{runID: 123, attempt: -1, downloadTo: t.TempDir()}, logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}) + require.NoError(t, err) + assert.True(t, success) + assert.Contains(t, buf.String(), "No logs available for run 123") +} + +// partialFailureServer fails node 1's chunk listing so one node succeeds and the +// other doesn't. +func partialFailureServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + switch p { + case "logs": + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}, {"path": "logs/node_1", "is_dir": true}]}`)) + case "logs/node_1": + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL", "message": "boom"}`)) + default: + _, _ = w.Write([]byte(`{"files": [{"path": "` + p + `/logs-0.chunk.txt"}]}`)) + } + case "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("ok\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadAllNodeLogsReportsPartialFailure(t *testing.T) { + w := newTestWorkspaceClient(t, partialFailureServer(t).URL) + dir := t.TempDir() + + // Node 1 fails, but node 0 still downloads and the failure is reported rather + // than silently dropped. + nodeLogs, failures, err := downloadAllNodeLogs(t.Context(), w, "run1", dir, []int{0, 1}, -1) + require.NoError(t, err) + require.Len(t, nodeLogs, 1) + assert.FileExists(t, nodeLogs[0]) + require.Contains(t, failures, 1) + assert.NotEmpty(t, failures[1]) +} + +func TestDownloadAllNodeLogsPropagatesCancellation(t *testing.T) { + w := newTestWorkspaceClient(t, downloadServer(t).URL) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + // A cancelled command must surface the cancellation, not look like a run with + // no logs. + _, _, err := downloadAllNodeLogs(ctx, w, "run1", t.TempDir(), []int{0, 1}, -1) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) +} + +// attemptPrefixServer serves the attempt-prefixed layout (logs/attempt_N/node_M). +func attemptPrefixServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/attempt_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "` + p + `/logs-0.chunk.txt"}]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + // Echo the path so the test can prove the attempt-prefixed dir was used. + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned?p=` + r.URL.Query().Get("path") + `"}]}`)) + case "/presigned": + _, _ = w.Write([]byte(r.URL.Query().Get("p") + "\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadAllNodeLogsAttemptPrefixedLayout(t *testing.T) { + w := newTestWorkspaceClient(t, attemptPrefixServer(t).URL) + dir := t.TempDir() + + nodeLogs, _, err := downloadAllNodeLogs(t.Context(), w, "run1", dir, []int{0}, -1) + require.NoError(t, err) + require.Len(t, nodeLogs, 1) + + body, err := os.ReadFile(nodeLogs[0]) + require.NoError(t, err) + assert.Contains(t, string(body), "logs/attempt_0/node_0") +} + +// truncatingServer lists two chunks for node 0 but fails the second one's +// credential request, so the node downloads partially. +func truncatingServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [ + {"path": "` + p + `/logs-0.chunk.txt"}, + {"path": "` + p + `/logs-1.chunk.txt"} + ]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + if strings.HasSuffix(r.URL.Query().Get("path"), "logs-1.chunk.txt") { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL", "message": "boom"}`)) + return + } + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("first chunk\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadNodeLogReportsFailedChunk(t *testing.T) { + w := newTestWorkspaceClient(t, truncatingServer(t).URL) + dir := t.TempDir() + + // Chunk 0 succeeded and chunk 1 failed: keep the bytes, but return an error + // so the gap isn't silent. + path, err := downloadNodeLog(t.Context(), w, "run1", 0, 0, false, dir) + require.Error(t, err) + assert.Contains(t, err.Error(), "chunk(s) [1]") + require.NotEmpty(t, path) + + got, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "first chunk\n", string(got)) +} + +// middleGapServer serves three chunks for node 0 and fails only the middle one, +// so the walk has to continue past a gap to reach the last chunk. +func middleGapServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [ + {"path": "` + p + `/logs-0.chunk.txt"}, + {"path": "` + p + `/logs-1.chunk.txt"}, + {"path": "` + p + `/logs-2.chunk.txt"} + ]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + p := r.URL.Query().Get("path") + if strings.HasSuffix(p, "logs-1.chunk.txt") { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL", "message": "boom"}`)) + return + } + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned?p=` + p + `"}]}`)) + case "/presigned": + _, _ = w.Write([]byte(path.Base(r.URL.Query().Get("p")) + "\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadNodeLogSkipsGapAndKeepsTail(t *testing.T) { + w := newTestWorkspaceClient(t, middleGapServer(t).URL) + dir := t.TempDir() + + // The tail usually carries the failure signature, so a bad middle chunk must + // not cost us the last one. + p, err := downloadNodeLog(t.Context(), w, "run1", 0, 0, false, dir) + require.Error(t, err) + assert.Contains(t, err.Error(), "chunk(s) [1]") + + got, err := os.ReadFile(p) + require.NoError(t, err) + assert.Equal(t, "logs-0.chunk.txt\nlogs-2.chunk.txt\n", string(got)) +} + +func TestDownloadNodeLogErrorsWhenEveryChunkFails(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "` + p + `/logs-0.chunk.txt"}]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL", "message": "boom"}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + + // Nothing downloaded because of failures, not because the node was silent: + // report it rather than returning an empty "no logs" result. + p, err := downloadNodeLog(t.Context(), newTestWorkspaceClient(t, srv.URL), "run1", 0, 0, false, t.TempDir()) + require.Error(t, err) + assert.Contains(t, err.Error(), "every chunk failed") + assert.Empty(t, p) +} + +func TestDownloadAllNodeLogsKeepsTruncatedNode(t *testing.T) { + w := newTestWorkspaceClient(t, truncatingServer(t).URL) + + // A truncated node lands in both maps, so it is still listed as downloaded + // while being reported as incomplete. + nodeLogs, failures, err := downloadAllNodeLogs(t.Context(), w, "run1", t.TempDir(), []int{0}, -1) + require.NoError(t, err) + require.Contains(t, nodeLogs, 0) + assert.FileExists(t, nodeLogs[0]) + require.Contains(t, failures, 0) + assert.Contains(t, failures[0], "chunk(s) [1]") +} + +func TestDownloadOutcomeTreatsActiveRunAsSuccess(t *testing.T) { + // An active run has no result state yet. Fetching its logs succeeded, so the + // command must exit 0 rather than report the run as failed. + for _, lc := range []string{"RUNNING", "PENDING", "QUEUED", "BLOCKED"} { + assert.True(t, logRunStatus{lifeCycleState: lc}.downloadOutcome(), lc) + } + + // A terminal run still decides the exit code by its outcome. + assert.True(t, logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}.downloadOutcome()) + assert.False(t, logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED"}.downloadOutcome()) + assert.False(t, logRunStatus{lifeCycleState: "TERMINATED", resultState: "CANCELED"}.downloadOutcome()) + assert.False(t, logRunStatus{lifeCycleState: "INTERNAL_ERROR"}.downloadOutcome()) +} + +func TestDownloadLogsActiveRunExitsZero(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + + // Downloading a still-running run's logs is not a failure. + success, err := downloadLogs(ctx, w, &bytes.Buffer{}, + logRequest{runID: 123, attempt: -1, downloadTo: t.TempDir()}, + logRunStatus{lifeCycleState: "RUNNING"}) + require.NoError(t, err) + assert.True(t, success) +} + +// allNodesFailServer serves a FAILED 2-node run whose chunk credentials all 404, +// so every node's download fails. +func allNodesFailServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{"run_id": 123, "state": {"life_cycle_state": "TERMINATED", "result_state": "FAILED"}, + "tasks": [{"run_id": 456, "ai_runtime_task": {"deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}}]}}]}`)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}, {"path": "logs/node_1", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "` + p + `/logs-0.chunk.txt"}]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code": "NOT_FOUND", "message": "gone"}`)) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadLogsAllNodesFailedIsNotReportedAsNoLogs(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, allNodesFailServer(t).URL) + var stdout bytes.Buffer + + // The logs exist but couldn't be fetched. Reporting "No logs available" would + // tell a caller the run produced nothing, so this fails instead. + _, err := downloadLogs(ctx, w, &stdout, + logRequest{runID: 123, attempt: -1, downloadTo: t.TempDir(), jsonOutput: true}, + logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to download logs from any of 2 node(s)") + assert.NotContains(t, stdout.String(), "No logs available") +} diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go index a8dd016dbb..debddc0d61 100644 --- a/experimental/air/cmd/logmlflow.go +++ b/experimental/air/cmd/logmlflow.go @@ -59,7 +59,7 @@ func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out i } if mlflowRunID == "" || logDir == "" { emitNoLogs(out, req, status) - return status.succeeded(), nil + return status.downloadOutcome(), nil } chunks, err := listLogChunks(ctx, w, mlflowRunID, logDir) @@ -73,7 +73,7 @@ func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out i target := req.tailTarget() if target <= 0 { - return status.succeeded(), nil + return status.downloadOutcome(), nil } lines, err := tailChunks(ctx, w, mlflowRunID, chunks, target) @@ -82,7 +82,7 @@ func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out i } if len(lines) == 0 { emitNoLogs(out, req, status) - return status.succeeded(), nil + return status.downloadOutcome(), nil } if len(lines) > target { @@ -91,7 +91,7 @@ func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out i for _, line := range lines { emitLogLine(out, req, line) } - return status.succeeded(), nil + return status.downloadOutcome(), nil } // resolveMLflowLogPath returns the run's MLflow run id and per-node log directory. @@ -253,10 +253,12 @@ func downloadArtifact(ctx context.Context, w *databricks.WorkspaceClient, mlflow } var resp credentialsForReadResponse + // A map query is serialized per value with %v, so a []string becomes the + // literal "[path]". The backend signs that bogus path and still returns 200, + // surfacing only as a 404 on the download. query := map[string]any{ "run_id": mlflowRunID, - // path is a repeated field, so pass a slice (serialized as path=...&path=...). - "path": []string{artifactPath}, + "path": artifactPath, } err = apiClient.Do(ctx, http.MethodGet, "/api/2.0/mlflow/artifacts/credentials-for-read", nil, nil, query, &resp) if err != nil { diff --git a/experimental/air/cmd/logmlflow_test.go b/experimental/air/cmd/logmlflow_test.go index f3de80abdb..4483799938 100644 --- a/experimental/air/cmd/logmlflow_test.go +++ b/experimental/air/cmd/logmlflow_test.go @@ -4,6 +4,7 @@ import ( "bytes" "net/http" "net/http/httptest" + "os" "testing" "github.com/stretchr/testify/assert" @@ -112,3 +113,29 @@ func TestMLflowFallbackNoLogsReflectsRunOutcome(t *testing.T) { require.NoError(t, err) assert.False(t, success) } + +func TestDownloadArtifactSendsUnbracketedPath(t *testing.T) { + var gotPath string + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/credentials-for-read": + gotPath = r.URL.Query().Get("path") + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("bytes\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + + local, err := downloadArtifact(t.Context(), newTestWorkspaceClient(t, srv.URL), "run1", "logs/node_0/logs-0.chunk.txt") + require.NoError(t, err) + t.Cleanup(func() { os.Remove(local) }) + + // The backend signs whatever path it is given and returns 200 even for a + // bracketed one, so only the download 404s. Assert on the path sent. + assert.Equal(t, "logs/node_0/logs-0.chunk.txt", gotPath) +} diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index 1bda80b572..19861b9e95 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -54,17 +54,19 @@ func newLogsCommand() *cobra.Command { 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")) - } + // --review is not yet implemented; reject rather than silently ignore. if review { return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, errors.New("--review is not implemented yet")) } + // A download always writes the full log, so a tail or time window would be + // silently dropped. + if downloadTo != "" && (cmd.Flags().Changed("lines") || minutes > 0) { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--download-to writes complete logs, so it cannot be combined with --lines or --minutes")) + } + // --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 { @@ -100,9 +102,11 @@ func newLogsCommand() *cobra.Command { return runLogs(ctx, cmd, logRequest{ runID: runID, node: node, + nodeSet: cmd.Flags().Changed("node"), attempt: retry, windowMinutes: minutes, tailLines: tailLines, + downloadTo: downloadTo, jsonOutput: root.OutputType(cmd) == flags.OutputJSON, }) } @@ -137,6 +141,22 @@ 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)) } + // --download-to writes each node's logs to disk instead of streaming. + if req.downloadTo != "" { + success, err := downloadLogs(ctx, w, cmd.OutOrStdout(), req, status) + if errors.Is(err, errNodeOutOfRange) { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, err) + } + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to download logs for run %d: %w", req.runID, err)) + } + if !success { + return root.ErrAlreadyPrinted + } + return nil + } + // 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 diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go index a9eb24b72f..9ea9026e01 100644 --- a/experimental/air/cmd/logs_test.go +++ b/experimental/air/cmd/logs_test.go @@ -7,7 +7,6 @@ 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" @@ -71,12 +70,6 @@ func TestLogsFlagValidation(t *testing.T) { 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"}, @@ -240,10 +233,10 @@ func TestLogsPastRetryOfActiveRunIsStatic(t *testing.T) { // --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. + // The run is still active, so there is no failure to report — printing the + // logs succeeded, and the command exits 0. err := runLogs(ctx, cmd, logRequest{runID: 9, node: 0, attempt: 0, tailLines: -1}) - require.ErrorIs(t, err, root.ErrAlreadyPrinted) + require.NoError(t, err) assert.Equal(t, "retry 0 log\n", buf.String()) // Exactly one runs/get: the initial status resolve in runLogs. The static diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index e5ab7ba8e8..351922d9bf 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -43,6 +43,9 @@ type logRequest struct { runID int64 // node is the node index to fetch; node 0 always exists. node int + // nodeSet distinguishes an explicit --node 0 from the default, so a download + // knows whether to fetch one node or all of them. + nodeSet bool // attempt is the retry attempt to read; -1 means latest. attempt int // windowMinutes, when > 0, restricts the fetch to the last N minutes. @@ -50,6 +53,8 @@ type logRequest struct { // tailLines caps a completed run's output to the last N lines. Negative means // --lines was unset (use the default cap); 0 prints nothing. tailLines int + // downloadTo, when set, writes logs to that directory instead of stdout. + downloadTo string // 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. @@ -88,6 +93,13 @@ func (s logRunStatus) succeeded() bool { return s.resultState == "SUCCESS" } +// downloadOutcome is the exit status for a one-shot fetch, which unlike streaming +// can run against an active run. An active run has no result state yet, and +// not-yet-finished is not a failure, so only a terminal run decides the exit code. +func (s logRunStatus) downloadOutcome() bool { + return !s.terminal() || s.succeeded() +} + // 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) { @@ -328,7 +340,7 @@ func (st *bricklensStreamer) drainStatic(toSec int64) (bool, error) { if !st.firstLogSeen { st.emitNoLogs() } - return st.status.succeeded(), nil + return st.status.downloadOutcome(), nil } // tailTarget is the number of lines a tail keeps. A negative tailLines means