diff --git a/acceptance/experimental/air/run-submit-deps/output.txt b/acceptance/experimental/air/run-submit-deps/output.txt index 4b2dae7b8d..5425469f42 100644 --- a/acceptance/experimental/air/run-submit-deps/output.txt +++ b/acceptance/experimental/air/run-submit-deps/output.txt @@ -5,7 +5,9 @@ Submitting experiment: deps-smoke Submitted workload with Job Run ID: 555 View job run at: [DATABRICKS_URL]/jobs/runs/555 -Tip: use --watch to stream logs until the run completes. +Tip: use --watch when submitting a run to stream logs to your terminal. +Stream logs after submission using: + databricks experimental air logs 555 === only config + command are uploaded; no requirements.yaml >>> print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt index bd917135d0..5b54d70cc4 100644 --- a/acceptance/experimental/air/run-submit/output.txt +++ b/acceptance/experimental/air/run-submit/output.txt @@ -6,7 +6,9 @@ Uploading [SNAPSHOT_TARBALL]... Submitted workload with Job Run ID: 555 View job run at: [DATABRICKS_URL]/jobs/runs/555 -Tip: use --watch to stream logs until the run completes. +Tip: use --watch when submitting a run to stream logs to your terminal. +Stream logs after submission using: + databricks experimental air logs 555 === the ai_runtime_task carries the code_source_path >>> print_requests.py //api/2.2/jobs/runs/submit diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index fbcc7b44aa..e496a1e89c 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -104,16 +104,39 @@ func runStatus(state *jobs.RunState) string { func displayRunStatus(run *jobs.Run) string { status := runStatus(run.State) - if status != string(jobs.RunLifeCycleStateRunning) || len(run.Tasks) == 0 || run.Tasks[0].State == nil { - return status + if runWaitingForCompute(run) { + return "PENDING" + } + return status +} + +func runWaitingForCompute(run *jobs.Run) bool { + if run.State == nil || run.State.ResultState != "" { + return false + } + if run.State.LifeCycleState == jobs.RunLifeCycleStatePending { + return true + } + if run.State.LifeCycleState != jobs.RunLifeCycleStateRunning || len(run.Tasks) == 0 || run.Tasks[0].State == nil { + return false } switch run.Tasks[0].State.LifeCycleState { case jobs.RunLifeCycleStatePending, jobs.RunLifeCycleStateQueued, jobs.RunLifeCycleStateWaitingForRetry, jobs.RunLifeCycleStateBlocked: - return "PENDING" + return true default: - return status + return false + } +} + +func detailedDisplayRunStatus(run *jobs.Run, statusMessage string) string { + if run.State != nil && run.State.ResultState == "" && statusMessage != "" { + return statusMessage + } + if runWaitingForCompute(run) { + return waitingForComputeStatus } + return displayRunStatus(run) } func terminationReason(run *jobs.Run) *string { diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go index 3df5c79994..4cc6dbf909 100644 --- a/experimental/air/cmd/format_test.go +++ b/experimental/air/cmd/format_test.go @@ -262,6 +262,27 @@ func TestDisplayRunStatus(t *testing.T) { })) } +func TestDetailedDisplayRunStatus(t *testing.T) { + waiting := &jobs.Run{ + State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}, + Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateQueued}}}, + } + assert.Equal(t, "Waiting for GPU capacity...", detailedDisplayRunStatus(waiting, "Waiting for GPU capacity...")) + assert.Equal(t, waitingForComputeStatus, detailedDisplayRunStatus(waiting, "")) + + started := &jobs.Run{ + State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}, + Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}}}, + } + assert.Equal(t, "RUNNING", detailedDisplayRunStatus(started, "")) + + failed := &jobs.Run{ + State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated, ResultState: jobs.RunResultStateFailed}, + Tasks: []jobs.RunTask{{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStatePending}}}, + } + assert.Equal(t, "FAILED", detailedDisplayRunStatus(failed, "stale status...")) +} + func TestTerminationReason(t *testing.T) { reason := terminationReason(&jobs.Run{ State: &jobs.RunState{ResultState: jobs.RunResultStateFailed, StateMessage: "parent reason"}, diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 3fb3288761..98dd9a656b 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -161,7 +161,8 @@ func newGetCommand() *cobra.Command { data := buildGetData(run) data.DashboardURL = dashboardURL(w.Config.Host, runID, workspaceID) - ids := mlflowIDs(ctx, w, run) + taskOutput := aiRuntimeTaskOutput(ctx, w, run) + ids := mlflowIDsFromOutput(taskOutput) if ids != nil { url := mlflowLogsURL(w.Config.Host, ids) data.MLflowURL = &url @@ -194,6 +195,11 @@ func newGetCommand() *cobra.Command { fmt.Fprintf(out, "Job Link: %s\n\n", hyperlink(ctx, out, data.DashboardURL, data.DashboardURL)) return renderEnvelope(ctx, data) } + statusMessage := "" + if taskOutput != nil { + statusMessage = normalizeStatusMessage(taskOutput.StatusMessage) + } + data.DisplayStatus = detailedDisplayRunStatus(run, statusMessage) renderRunText(ctx, out, w, run, &data, ids) return nil diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index e1afd1c0a3..63da11728a 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -105,7 +105,7 @@ func newLogsCommand() *cobra.Command { tailLines = lines } - return runLogs(ctx, cmd, logRequest{ + err = runLogs(ctx, cmd, logRequest{ runID: runID, node: node, nodeSet: cmd.Flags().Changed("node"), @@ -115,6 +115,10 @@ func newLogsCommand() *cobra.Command { downloadTo: downloadTo, jsonOutput: root.OutputType(cmd) == flags.OutputJSON, }) + if downloadTo != "" || root.OutputType(cmd) == flags.OutputJSON { + return err + } + return handleWatchResult(cmd.OutOrStdout(), cmdctx.WorkspaceClient(ctx).Config.Profile, args[0], err) } return cmd diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go index 9ea9026e01..e459779bb7 100644 --- a/experimental/air/cmd/logs_test.go +++ b/experimental/air/cmd/logs_test.go @@ -2,11 +2,13 @@ package aircmd import ( "bytes" + "context" "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" @@ -194,6 +196,37 @@ func TestLogsFallsBackToMLflow(t *testing.T) { assert.Equal(t, "line one\nline two\n", buf.String()) } +func TestLogsCommandPrintsGuidanceWhenStreamingIsInterrupted(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + 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":"RUNNING"},"tasks":[{"run_id":456}]}`)) + case strings.HasSuffix(r.URL.Path, "/logs"): + cancel() + _, _ = w.Write([]byte(`{"log_records":[]}`)) + default: + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + client := newTestWorkspaceClient(t, srv.URL) + client.Config.Profile = "team profile" + ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, nil, &buf, &buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, client) + cmd := withOutput(newLogsCommand(), flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + err := cmd.RunE(cmd, []string{"5"}) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Streaming logs interrupted.") + assert.Contains(t, buf.String(), "To check status:\ndatabricks experimental air get 5 -p 'team profile'") + assert.Contains(t, buf.String(), "To resume streaming logs:\ndatabricks experimental air logs 5 -p 'team profile'") +} + // 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 diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index aa995d40b0..ef5197c24e 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -58,6 +58,9 @@ func normalizeStatusMessage(raw string) string { if payload == "" { return "" } + if strings.EqualFold(payload, "Waiting for GPU compute capacity to become available") { + return waitingForComputeStatus + } return payload + "..." } @@ -104,11 +107,12 @@ type logRequest struct { // 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 + lifeCycleState string + firstTaskLifeCycleState string + resultState string + stateMessage string + startTimeMs int64 + endTimeMs int64 // latestAttempt is the highest attempt_number across the run's tasks. latestAttempt int } @@ -128,6 +132,25 @@ func (s logRunStatus) succeeded() bool { return s.resultState == "SUCCESS" } +func (s logRunStatus) waitingForCompute() bool { + if s.resultState != "" { + return false + } + if s.lifeCycleState == string(jobs.RunLifeCycleStatePending) { + return true + } + if s.lifeCycleState != string(jobs.RunLifeCycleStateRunning) { + return false + } + switch jobs.RunLifeCycleState(s.firstTaskLifeCycleState) { + case jobs.RunLifeCycleStatePending, jobs.RunLifeCycleStateQueued, + jobs.RunLifeCycleStateWaitingForRetry, jobs.RunLifeCycleStateBlocked: + return true + default: + return false + } +} + // 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. @@ -157,6 +180,9 @@ func projectRunStatus(run *jobs.Run) logRunStatus { s.resultState = string(run.State.ResultState) s.stateMessage = run.State.StateMessage } + if len(run.Tasks) > 0 && run.Tasks[0].State != nil { + s.firstTaskLifeCycleState = string(run.Tasks[0].State.LifeCycleState) + } for i := range run.Tasks { s.latestAttempt = max(s.latestAttempt, run.Tasks[i].AttemptNumber) } @@ -260,7 +286,7 @@ func (st *bricklensStreamer) waitingSpinnerText() string { if msg := st.serverStatusMessage(); msg != "" { return msg } - if st.status.lifeCycleState == "PENDING" { + if st.status.waitingForCompute() { return waitingForComputeStatus } return fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node) diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go index 429e5112e7..78604fe3f1 100644 --- a/experimental/air/cmd/logstream_test.go +++ b/experimental/air/cmd/logstream_test.go @@ -82,7 +82,7 @@ func TestProjectRunStatus(t *testing.T) { StateMessage: "done", }, Tasks: []jobs.RunTask{ - {AttemptNumber: 0}, + {AttemptNumber: 0, State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateQueued}}, {AttemptNumber: 2}, {AttemptNumber: 1}, }, @@ -92,6 +92,7 @@ func TestProjectRunStatus(t *testing.T) { assert.Equal(t, "TERMINATED", s.lifeCycleState) assert.Equal(t, "SUCCESS", s.resultState) assert.Equal(t, "done", s.stateMessage) + assert.Equal(t, "QUEUED", s.firstTaskLifeCycleState) assert.Equal(t, int64(1000), s.startTimeMs) assert.Equal(t, int64(2000), s.endTimeMs) assert.Equal(t, 2, s.latestAttempt) @@ -213,6 +214,7 @@ func TestNormalizeStatusMessage(t *testing.T) { }{ {"STATUS: Waiting for GPU capacity.", "Waiting for GPU capacity..."}, {"STATUS:Waiting for GPU capacity", "Waiting for GPU capacity..."}, + {"STATUS: Waiting for GPU compute capacity to become available", waitingForComputeStatus}, {"status: provisioning", "provisioning..."}, // type match is case-insensitive {"STATUS: done...", "done..."}, // trailing dots collapse to one "..." {"INFO: not a status", ""}, // other type ignored @@ -228,7 +230,7 @@ func TestNormalizeStatusMessage(t *testing.T) { func TestWaitingSpinnerText(t *testing.T) { // A server that returns the run (with a task) and a STATUS-typed status_message. - newStreamer := func(t *testing.T, statusMessage, lifeCycle string) *bricklensStreamer { + newStreamer := func(t *testing.T, statusMessage, lifeCycle, firstTaskLifeCycle string) *bricklensStreamer { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { @@ -245,21 +247,30 @@ func TestWaitingSpinnerText(t *testing.T) { ctx: t.Context(), w: newTestWorkspaceClient(t, srv.URL), req: logRequest{runID: 1, node: 0}, - status: logRunStatus{lifeCycleState: lifeCycle}, + status: logRunStatus{lifeCycleState: lifeCycle, firstTaskLifeCycleState: firstTaskLifeCycle}, } } // Server STATUS message wins. assert.Equal(t, "Waiting for GPU capacity...", - newStreamer(t, "STATUS: Waiting for GPU capacity", "PENDING").waitingSpinnerText()) + newStreamer(t, "STATUS: Waiting for GPU capacity", "PENDING", "").waitingSpinnerText()) // No status message + PENDING -> compute-capacity fallback. assert.Equal(t, waitingForComputeStatus, - newStreamer(t, "", "PENDING").waitingSpinnerText()) + newStreamer(t, "", "PENDING", "").waitingSpinnerText()) + + for _, state := range []string{"PENDING", "QUEUED", "WAITING_FOR_RETRY", "BLOCKED"} { + assert.Equal(t, waitingForComputeStatus, + newStreamer(t, "", "RUNNING", state).waitingSpinnerText(), state) + } // No status message + non-PENDING -> default "waiting for run to start". assert.Equal(t, "Waiting for run to start (node 0)...", - newStreamer(t, "", "RUNNING").waitingSpinnerText()) + newStreamer(t, "", "RUNNING", "RUNNING").waitingSpinnerText()) + assert.Equal(t, "Waiting for run to start (node 0)...", + newStreamer(t, "", "RUNNING", "").waitingSpinnerText()) + assert.Equal(t, "Waiting for run to start (node 0)...", + newStreamer(t, "", "TERMINATED", "PENDING").waitingSpinnerText()) } func TestEmitLogLineJSON(t *testing.T) { diff --git a/experimental/air/cmd/mlflow.go b/experimental/air/cmd/mlflow.go index 0eef26f80c..9316aafd05 100644 --- a/experimental/air/cmd/mlflow.go +++ b/experimental/air/cmd/mlflow.go @@ -30,11 +30,14 @@ type mlflowIdentifiers struct { // mlflowIDs fetches the MLflow IDs for a run via its latest task. Returns nil if // they can't be obtained. func mlflowIDs(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run) *mlflowIdentifiers { + return mlflowIDsFromOutput(aiRuntimeTaskOutput(ctx, w, run)) +} + +func aiRuntimeTaskOutput(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run) *jobs.AiRuntimeTaskOutput { if len(run.Tasks) == 0 { return nil } - // The MLflow output is attached to the task run, not the parent job run. - return mlflowIDsForTask(ctx, w, run.Tasks[len(run.Tasks)-1].RunId) + return aiRuntimeTaskOutputForTask(ctx, w, run.Tasks[len(run.Tasks)-1].RunId) } // mlflowIDsForTask fetches a task run's MLflow experiment and run IDs from @@ -42,6 +45,10 @@ func mlflowIDs(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run // link, so any failure (endpoint error, run not yet started, no MLflow output) // is logged and treated as "no link" rather than failing the command. func mlflowIDsForTask(ctx context.Context, w *databricks.WorkspaceClient, taskRunID int64) *mlflowIdentifiers { + return mlflowIDsFromOutput(aiRuntimeTaskOutputForTask(ctx, w, taskRunID)) +} + +func aiRuntimeTaskOutputForTask(ctx context.Context, w *databricks.WorkspaceClient, taskRunID int64) *jobs.AiRuntimeTaskOutput { if taskRunID == 0 { return nil } @@ -52,10 +59,14 @@ func mlflowIDsForTask(ctx context.Context, w *databricks.WorkspaceClient, taskRu return nil } - if o := out.AiRuntimeTaskOutput; o != nil && o.MlflowExperimentId != "" && o.MlflowRunId != "" { - return &mlflowIdentifiers{ExperimentID: o.MlflowExperimentId, RunID: o.MlflowRunId} + return out.AiRuntimeTaskOutput +} + +func mlflowIDsFromOutput(output *jobs.AiRuntimeTaskOutput) *mlflowIdentifiers { + if output == nil || output.MlflowExperimentId == "" || output.MlflowRunId == "" { + return nil } - return nil + return &mlflowIdentifiers{ExperimentID: output.MlflowExperimentId, RunID: output.MlflowRunId} } // mlflowLogsURL is the deep link to a run's node-0 logs. It is the value of the diff --git a/experimental/air/cmd/mlflow_test.go b/experimental/air/cmd/mlflow_test.go index 1501a2a59d..a3dd537629 100644 --- a/experimental/air/cmd/mlflow_test.go +++ b/experimental/air/cmd/mlflow_test.go @@ -78,6 +78,17 @@ func TestMLflowIDs(t *testing.T) { }) } +func TestAiRuntimeTaskOutput(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{"ai_runtime_task_output":{"status_message":"STATUS: Waiting for GPU capacity"}}`, &hit) + run := &jobs.Run{Tasks: []jobs.RunTask{{RunId: 99}}} + + got := aiRuntimeTaskOutput(t.Context(), newTestWorkspaceClient(t, srv.URL), run) + require.NotNil(t, got) + assert.True(t, hit) + assert.Equal(t, "STATUS: Waiting for GPU capacity", got.StatusMessage) +} + func TestMLflowIDsForTask(t *testing.T) { ctx := t.Context() diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 422d438821..6b024e011f 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -2,8 +2,11 @@ package aircmd import ( "context" + "errors" "fmt" "io" + "os" + "os/signal" "strconv" "strings" "unicode/utf8" @@ -12,6 +15,7 @@ import ( "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/shellquote" "github.com/databricks/databricks-sdk-go" "github.com/spf13/cobra" ) @@ -126,7 +130,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so if ids := resolveMLflowIDsForRun(ctx, w, runID); ids != nil { printMLflowLinks(ctx, out, w.Config.Host, ids) } - cmdio.LogString(ctx, "\nTip: use --watch to stream logs until the run completes.") + printPostSubmitGuidance(out, w.Config.Profile, runIDStr) return nil } // PENDING is the submit status, distinct from the --watch JSONL @@ -144,6 +148,9 @@ The path must be a separate argument: cobra reserves -h as a boolean, so jsonOutput: jsonOut, } + watchCtx, stop := signal.NotifyContext(ctx, os.Interrupt) + defer stop() + if !jsonOut { out := cmd.OutOrStdout() // The MLflow links stream in via the logs below, so don't poll here. @@ -152,7 +159,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so fmt.Fprintln(out) fmt.Fprintln(out, "Monitoring run and streaming logs...") printLogsDivider(ctx, out) - return runLogs(ctx, cmd, req) + return handleWatchResult(out, w.Config.Profile, runIDStr, runLogs(watchCtx, cmd, req)) } // --json: emit SUBMITTED first (so a consumer sees the run id immediately), @@ -163,7 +170,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so req.onStatusChange = func(current, previous string) { printStatusEvent(out, current, previous) } - err = runLogs(ctx, cmd, req) + err = runLogs(watchCtx, cmd, req) // Re-resolve the run for the closing envelope. STATUS events only fire on // the Bricklens path, so the terminal status must come from the run's @@ -176,6 +183,45 @@ The path must be a separate argument: cobra reserves -h as a boolean, so return cmd } +func airLogsCommand(profile, runID string) string { + args := []string{"databricks", "experimental", "air", "logs", shellquote.BashArg(runID)} + if profile != "" { + args = append(args, "-p", shellquote.BashArg(profile)) + } + return strings.Join(args, " ") +} + +func airGetCommand(profile, runID string) string { + args := []string{"databricks", "experimental", "air", "get", shellquote.BashArg(runID)} + if profile != "" { + args = append(args, "-p", shellquote.BashArg(profile)) + } + return strings.Join(args, " ") +} + +func printPostSubmitGuidance(out io.Writer, profile, runID string) { + fmt.Fprintln(out) + fmt.Fprintln(out, "Tip: use --watch when submitting a run to stream logs to your terminal.") + fmt.Fprintln(out, "Stream logs after submission using:") + fmt.Fprintln(out, " "+airLogsCommand(profile, runID)) +} + +func handleWatchResult(out io.Writer, profile, runID string, err error) error { + if !errors.Is(err, context.Canceled) { + return err + } + + fmt.Fprintln(out) + fmt.Fprintln(out, "Streaming logs interrupted.") + fmt.Fprintln(out) + fmt.Fprintln(out, "To check status:") + fmt.Fprintln(out, airGetCommand(profile, runID)) + fmt.Fprintln(out) + fmt.Fprintln(out, "To resume streaming logs:") + fmt.Fprintln(out, airLogsCommand(profile, runID)) + return root.ErrAlreadyPrinted +} + // printSubmitResult writes the green success line and Job Run link. These don't // depend on the MLflow IDs, so they print before any MLflow poll. The link is // styled (blue, underlined) and clickable, matching the `air get` view, and diff --git a/experimental/air/cmd/run_test.go b/experimental/air/cmd/run_test.go index c1efd3a8b3..a49997fb45 100644 --- a/experimental/air/cmd/run_test.go +++ b/experimental/air/cmd/run_test.go @@ -48,13 +48,19 @@ func submitServer(t *testing.T, getOutput string) *httptest.Server { } func runSubmitCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL string) error { + return runSubmitCmdWithProfile(t, out, buf, srvURL, "") +} + +func runSubmitCmdWithProfile(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL, profile string) error { t.Helper() cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) cmd := withOutput(newRunCommand(), out) require.NoError(t, cmd.Flags().Set("file", cfgPath)) ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, nil, buf, buf, "", "")) - ctx = cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srvURL)) + w := newTestWorkspaceClient(t, srvURL) + w.Config.Profile = profile + ctx = cmdctx.SetWorkspaceClient(ctx, w) cmd.SetContext(ctx) cmd.SetOut(buf) return cmd.RunE(cmd, nil) @@ -73,10 +79,33 @@ func TestRunSubmitTextOutput(t *testing.T) { assert.Contains(t, out, "Submitted workload with Job Run ID: 555") assert.Contains(t, out, "View job run at: ") assert.Contains(t, out, "/jobs/runs/555") - assert.Contains(t, out, "Tip: use --watch") + assert.Contains(t, out, "Tip: use --watch when submitting a run to stream logs to your terminal.") + assert.Contains(t, out, "Stream logs after submission using:") + assert.Contains(t, out, "databricks experimental air logs 555") assert.NotContains(t, out, "View MLflow run at:") } +func TestRunSubmitTextOutputIncludesProfileInLogsCommand(t *testing.T) { + fastMLflowPoll(t) + var buf bytes.Buffer + err := runSubmitCmdWithProfile(t, flags.OutputText, &buf, submitServer(t, `{}`).URL, "team profile") + require.NoError(t, err) + + assert.Contains(t, buf.String(), "databricks experimental air logs 555 -p 'team profile'") +} + +func TestAirLogsCommand(t *testing.T) { + assert.Equal(t, "databricks experimental air logs 123", airLogsCommand("", "123")) + assert.Equal(t, "databricks experimental air logs 123 -p profile-name", airLogsCommand("profile-name", "123")) + assert.Equal(t, "databricks experimental air logs 123 -p 'team profile'", airLogsCommand("team profile", "123")) +} + +func TestAirGetCommand(t *testing.T) { + assert.Equal(t, "databricks experimental air get 123", airGetCommand("", "123")) + assert.Equal(t, "databricks experimental air get 123 -p profile-name", airGetCommand("profile-name", "123")) + assert.Equal(t, "databricks experimental air get 123 -p 'team profile'", airGetCommand("team profile", "123")) +} + func TestRunSubmitTextOutputWithMLflowLinks(t *testing.T) { var buf bytes.Buffer srvURL := submitServer(t, `{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`).URL diff --git a/experimental/air/cmd/run_watch_test.go b/experimental/air/cmd/run_watch_test.go index a0b640fc47..6034fbdb1d 100644 --- a/experimental/air/cmd/run_watch_test.go +++ b/experimental/air/cmd/run_watch_test.go @@ -2,6 +2,9 @@ package aircmd import ( "bytes" + "context" + "errors" + "fmt" "net/http" "net/http/httptest" "strings" @@ -158,6 +161,26 @@ func TestRunWatchFailedRunExitsNonZero(t *testing.T) { assert.Contains(t, buf.String(), "step 1\nstep 2") } +func TestHandleWatchResultPrintsProfileAwareResumeCommand(t *testing.T) { + var buf bytes.Buffer + err := handleWatchResult(&buf, "team profile", "777", fmt.Errorf("stream logs: %w", context.Canceled)) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + out := buf.String() + assert.Contains(t, out, "Streaming logs interrupted.") + assert.Contains(t, out, "To check status:\ndatabricks experimental air get 777 -p 'team profile'") + assert.Contains(t, out, "To resume streaming logs:\ndatabricks experimental air logs 777 -p 'team profile'") + assert.NotContains(t, out, "The workload was not canceled") +} + +func TestHandleWatchResultPassesThroughOtherErrors(t *testing.T) { + want := errors.New("stream failed") + var buf bytes.Buffer + + assert.ErrorIs(t, handleWatchResult(&buf, "profile", "777", want), want) + assert.Empty(t, buf.String()) +} + func TestRunWatchDryRunSkipsSubmit(t *testing.T) { // --dry-run takes precedence over --watch: nothing is submitted or streamed. var got []string