diff --git a/acceptance/experimental/air/run-submit-deps/output.txt b/acceptance/experimental/air/run-submit-deps/output.txt index 98c7dc1dfa9..7a0d971b9f5 100644 --- a/acceptance/experimental/air/run-submit-deps/output.txt +++ b/acceptance/experimental/air/run-submit-deps/output.txt @@ -1,8 +1,9 @@ === submit with inline dependencies >>> [CLI] experimental air run -f run.yaml -Submitted run 555 -View at: [DATABRICKS_URL]/jobs/runs/555 +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. @@ -57,8 +58,9 @@ Tip: use --watch to stream logs until the run completes. === file-form deps: version comes from the requirements file >>> [CLI] experimental air run -f run-file.yaml -Submitted run 555 -View at: [DATABRICKS_URL]/jobs/runs/555 +Submitting experiment: deps-file-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. diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt index 2d92122c8e6..bd917135d0e 100644 --- a/acceptance/experimental/air/run-submit/output.txt +++ b/acceptance/experimental/air/run-submit/output.txt @@ -1,9 +1,10 @@ === submit with a git code_source >>> [CLI] experimental air run -f run.yaml +Submitting experiment: submit-smoke Uploading [SNAPSHOT_TARBALL]... -Submitted run 555 -View at: [DATABRICKS_URL]/jobs/runs/555 +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. diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index e5ab7ba8e85..9bfb4bbb53b 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "slices" + "strings" "time" "github.com/databricks/cli/libs/cmdio" @@ -26,8 +27,37 @@ const ( defaultCompletedRunTailLines = 10000 // seenRecordsCap bounds the dedup set, evicting oldest-inserted entries first. seenRecordsCap = 100000 + // statusMessageRefreshEveryNPolls throttles the status_message fetch so the + // waiting spinner doesn't issue a get-output on every poll tick. + statusMessageRefreshEveryNPolls = 5 ) +// statusMessageType tags a client-facing message packed into +// ai_runtime_task_output.status_message as ":"; only STATUS-typed +// messages are surfaced. +const statusMessageType = "STATUS" + +// waitingForComputeStatus is the fallback shown while a PENDING run waits for +// accelerator compute. +const waitingForComputeStatus = "Waiting for accelerator compute capacity to become available..." + +// normalizeStatusMessage returns the payload of a "STATUS:" message, +// normalized for display (trailing "." stripped, "..." suffix added), or "" for +// any other type or an empty/absent message. +func normalizeStatusMessage(raw string) string { + messageType, payload, ok := strings.Cut(raw, ":") + if !ok || !strings.EqualFold(strings.TrimSpace(messageType), statusMessageType) { + return "" + } + payload = strings.TrimSpace(payload) + payload = strings.TrimRight(payload, ".") + payload = strings.TrimSpace(payload) + if payload == "" { + return "" + } + return payload + "..." +} + // retryCheckInterval is the wait between status/log polls. A var so tests can // shrink it. var retryCheckInterval = 3 * time.Second @@ -206,6 +236,40 @@ type bricklensStreamer struct { updateSpinner func(string) } +// waitingSpinnerText returns the waiting-spinner text: the server-set STATUS +// message if present, else the compute-capacity message for a PENDING run, else +// the default "waiting for run to start". +func (st *bricklensStreamer) waitingSpinnerText() string { + if msg := st.serverStatusMessage(); msg != "" { + return msg + } + if st.status.lifeCycleState == "PENDING" { + return waitingForComputeStatus + } + return fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node) +} + +// serverStatusMessage returns the run's server-set STATUS message (normalized for +// display), or "" if unavailable. The message lives on the latest task run's +// output, re-resolved each call so a retry's new task run is picked up. +// Best-effort: any fetch failure logs at debug and returns "". +func (st *bricklensStreamer) serverStatusMessage() string { + run, err := st.w.Jobs.GetRun(st.ctx, jobs.GetRunRequest{RunId: st.req.runID}) + if err != nil || len(run.Tasks) == 0 { + return "" + } + taskRunID := run.Tasks[len(run.Tasks)-1].RunId + out, err := st.w.Jobs.GetRunOutputByRunId(st.ctx, taskRunID) + if err != nil { + log.Debugf(st.ctx, "air logs: status_message fetch failed for run %d: %v", st.req.runID, err) + return "" + } + if out.AiRuntimeTaskOutput == nil { + return "" + } + return normalizeStatusMessage(out.AiRuntimeTaskOutput.StatusMessage) +} + // reportStatusChange fires onStatusChange when the run's display state differs // from the last reported one. func (st *bricklensStreamer) reportStatusChange() { @@ -242,6 +306,11 @@ func (st *bricklensStreamer) run() (bool, error) { } firstIteration := true + // Throttled refresh of the waiting-spinner text: statusRefreshCounter gates the + // server status_message fetch to every Nth poll, and lastSpinnerText avoids + // redundant spinner updates. + statusRefreshCounter := 0 + lastSpinnerText := "" for { if !firstIteration { status, err := resolveRunStatus(st.ctx, st.w, st.req.runID) @@ -269,9 +338,17 @@ func (st *bricklensStreamer) run() (bool, error) { terminal := st.status.terminal() toSec := st.req.toSeconds(st.status) - // While waiting on a still-active run with no logs yet, refresh the spinner. + // While waiting on a still-active run with no logs yet, refresh the spinner + // with the server-set status (throttled), so a run stuck waiting for compute + // shows why rather than a generic "waiting" message. if !terminal && !st.firstLogSeen && st.updateSpinner != nil { - st.updateSpinner(fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node)) + if statusRefreshCounter%statusMessageRefreshEveryNPolls == 0 { + if desired := st.waitingSpinnerText(); desired != lastSpinnerText { + st.updateSpinner(desired) + lastSpinnerText = desired + } + } + statusRefreshCounter++ } // A run already terminal on the first iteration renders as a tail (most diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go index d8a061687a6..7568588bfeb 100644 --- a/experimental/air/cmd/logstream_test.go +++ b/experimental/air/cmd/logstream_test.go @@ -7,6 +7,7 @@ import ( "errors" "net/http" "net/http/httptest" + "strconv" "strings" "testing" "time" @@ -204,6 +205,62 @@ func TestDisplayState(t *testing.T) { assert.Equal(t, "UNKNOWN", logRunStatus{}.displayState()) } +func TestNormalizeStatusMessage(t *testing.T) { + tests := []struct { + raw string + want string + }{ + {"STATUS: Waiting for GPU capacity.", "Waiting for GPU capacity..."}, + {"STATUS:Waiting for GPU capacity", "Waiting for GPU capacity..."}, + {"status: provisioning", "provisioning..."}, // type match is case-insensitive + {"STATUS: done...", "done..."}, // trailing dots collapse to one "..." + {"INFO: not a status", ""}, // other type ignored + {"no type prefix", ""}, + {"STATUS:", ""}, // empty payload + {"STATUS: ", ""}, // whitespace-only payload + {"", ""}, + } + for _, tt := range tests { + assert.Equal(t, tt.want, normalizeStatusMessage(tt.raw), "raw=%q", tt.raw) + } +} + +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 { + 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": 1, "tasks": [{"run_id": 2}]}`)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"status_message": ` + strconv.Quote(statusMessage) + `}}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return &bricklensStreamer{ + ctx: t.Context(), + w: newTestWorkspaceClient(t, srv.URL), + req: logRequest{runID: 1, node: 0}, + status: logRunStatus{lifeCycleState: lifeCycle}, + } + } + + // Server STATUS message wins. + assert.Equal(t, "Waiting for GPU capacity...", + newStreamer(t, "STATUS: Waiting for GPU capacity", "PENDING").waitingSpinnerText()) + + // No status message + PENDING -> compute-capacity fallback. + assert.Equal(t, waitingForComputeStatus, + newStreamer(t, "", "PENDING").waitingSpinnerText()) + + // 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()) +} + func TestEmitLogLineJSON(t *testing.T) { var buf bytes.Buffer emitLogLine(&buf, logRequest{node: 2, jsonOutput: true}, "hello") diff --git a/experimental/air/cmd/mlflow.go b/experimental/air/cmd/mlflow.go index 070caefbb06..0eef26f80c4 100644 --- a/experimental/air/cmd/mlflow.go +++ b/experimental/air/cmd/mlflow.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strings" + "time" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" @@ -11,6 +12,15 @@ import ( "github.com/databricks/databricks-sdk-go/service/ml" ) +// mlflowLinkPollAttempts bounds the best-effort poll for a freshly-submitted +// run's MLflow IDs (see resolveMLflowIDsForRun), kept short so a bare `air run` +// returns promptly when the IDs aren't ready yet. +const mlflowLinkPollAttempts = 3 + +// mlflowLinkPollInterval is the delay between poll attempts. A var, not a const, +// so tests can shrink it and avoid a real sleep. +var mlflowLinkPollInterval = 500 * time.Millisecond + // mlflowIdentifiers are the experiment and run IDs MLflow assigns to a run. type mlflowIdentifiers struct { ExperimentID string @@ -62,6 +72,46 @@ func mlflowRunURL(host string, ids *mlflowIdentifiers) string { strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) } +// mlflowExperimentURL links to the MLflow experiment page. Omits the ?o= query +// for consistency with mlflowRunURL and the run-submit dashboard URL. +func mlflowExperimentURL(host string, ids *mlflowIdentifiers) string { + return fmt.Sprintf("%s/ml/experiments/%s", strings.TrimRight(host, "/"), ids.ExperimentID) +} + +// resolveMLflowIDsForRun best-effort resolves a just-submitted run's MLflow IDs, +// polling because they are assigned only once the task run starts. Returns nil +// (treated as "no link", not an error) if they don't appear within the budget or +// the context is cancelled. +func resolveMLflowIDsForRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) *mlflowIdentifiers { + // The task run id is fixed at submit time, so resolve it once; only the MLflow + // output (runs/get-output) fills in later, so that is all we poll. + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + log.Debugf(ctx, "air run: could not fetch run %d for MLflow link: %v", runID, err) + return nil + } + if len(run.Tasks) == 0 { + return nil + } + // The MLflow output is attached to the task run, not the parent job run. + taskRunID := run.Tasks[len(run.Tasks)-1].RunId + + for attempt := range mlflowLinkPollAttempts { + if ids := mlflowIDsForTask(ctx, w, taskRunID); ids != nil { + return ids + } + if attempt == mlflowLinkPollAttempts-1 { + break + } + select { + case <-ctx.Done(): + return nil + case <-time.After(mlflowLinkPollInterval): + } + } + return nil +} + // fetchMLflowRunName fetches a run's MLflow run_name via the MLflow REST API, // returning "" if it can't be obtained. Best-effort, like the rest of the MLflow // enrichment. diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index c982d06fa7b..791ec9fb4b3 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -3,7 +3,10 @@ package aircmd import ( "context" "fmt" + "io" "strconv" + "strings" + "unicode/utf8" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" @@ -98,23 +101,37 @@ The path must be a separate argument: cobra reserves -h as a boolean, so return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true}) } + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + // Announce the experiment before uploading; skipped in JSON mode to keep + // stdout a clean envelope stream. + if !jsonOut { + cmdio.LogString(ctx, "Submitting experiment: "+cfg.ExperimentName) + } + w := cmdctx.WorkspaceClient(ctx) - runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey) + runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey, !jsonOut) if err != nil { return err } runIDStr := strconv.FormatInt(runID, 10) - jsonOut := root.OutputType(cmd) == flags.OutputJSON if !watch { if !jsonOut { - cmdio.LogString(ctx, "Submitted run "+runIDStr) - cmdio.LogString(ctx, "View at: "+dashboardURL) + out := cmd.OutOrStdout() + printSubmitResult(ctx, out, runIDStr, dashboardURL) + // Append the MLflow links only if they resolve; a bare submit is not + // blocked on them since the confirmation above is already printed. + 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.") return nil } - return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + // PENDING is the submit status, distinct from the --watch JSONL + // SUBMITTED event type below. + return renderEnvelope(ctx, runResult{Status: "PENDING", RunID: runIDStr, DashboardURL: dashboardURL}) } // --watch: stream the submitted run's logs until it reaches a terminal @@ -128,15 +145,19 @@ The path must be a separate argument: cobra reserves -h as a boolean, so } if !jsonOut { - cmdio.LogString(ctx, "Submitted run "+runIDStr) - cmdio.LogString(ctx, "View at: "+dashboardURL) - cmdio.LogString(ctx, "Monitoring run and streaming logs...") + out := cmd.OutOrStdout() + // The MLflow links stream in via the logs below, so don't poll here. + printSubmitResult(ctx, out, runIDStr, dashboardURL) + // Separate the submit summary from the streamed logs. + fmt.Fprintln(out) + fmt.Fprintln(out, "Monitoring run and streaming logs...") + printLogsDivider(ctx, out) return runLogs(ctx, cmd, req) } // --json: emit SUBMITTED first (so a consumer sees the run id immediately), // STATUS events on each lifecycle transition, and a closing terminal-status - // envelope after streaming. Mirrors the Python CLI's --watch JSONL contract. + // envelope after streaming. out := cmd.OutOrStdout() printSubmittedEvent(out, runIDStr, dashboardURL) req.onStatusChange = func(current, previous string) { @@ -155,6 +176,42 @@ The path must be a separate argument: cobra reserves -h as a boolean, so return cmd } +// printSubmitResult writes the green success line and Job Run hyperlink. These +// don't depend on the MLflow IDs, so they print before any MLflow poll. Color +// and links degrade to plain text on non-rich terminals. +func printSubmitResult(ctx context.Context, out io.Writer, runIDStr, dashboardURL string) { + renderer, _ := cmdio.NewRenderer(ctx, out) + p := newPalette(renderer) + + fmt.Fprintln(out, p.green.Render("Submitted workload with Job Run ID: "+runIDStr)) + fmt.Fprintln(out, "View job run at: "+hyperlink(ctx, out, dashboardURL, dashboardURL)) +} + +// printMLflowLinks appends the MLflow run and experiment hyperlinks once their +// IDs are resolved. +func printMLflowLinks(ctx context.Context, out io.Writer, host string, ids *mlflowIdentifiers) { + runURL := mlflowRunURL(host, ids) + expURL := mlflowExperimentURL(host, ids) + fmt.Fprintln(out, "View MLflow run at: "+hyperlink(ctx, out, runURL, runURL)) + fmt.Fprintln(out, "View MLflow experiment at: "+hyperlink(ctx, out, expURL, expURL)) +} + +// logsDividerWidth is the total display width of the --watch logs divider. +const logsDividerWidth = 60 + +// printLogsDivider prints a centered "Logs" rule marking where the streamed +// --watch logs begin, separating them from the submit summary. The dim color is +// dropped on non-rich terminals; the rule characters are always printed. +func printLogsDivider(ctx context.Context, out io.Writer) { + renderer, _ := cmdio.NewRenderer(ctx, out) + p := newPalette(renderer) + + const label = " Logs " + side := max((logsDividerWidth-utf8.RuneCountInString(label))/2, 0) + rule := strings.Repeat("─", side) + label + strings.Repeat("─", side) + fmt.Fprintln(out, p.n7.Render(rule)) +} + // watchTerminalStatus resolves a watched run's final display state for the // closing --watch envelope. The run is terminal once streaming returns; if the // status can't be re-fetched, "UNKNOWN" is reported rather than guessing. diff --git a/experimental/air/cmd/run_test.go b/experimental/air/cmd/run_test.go new file mode 100644 index 00000000000..c1efd3a8b34 --- /dev/null +++ b/experimental/air/cmd/run_test.go @@ -0,0 +1,136 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fastMLflowPoll shrinks the MLflow-ID poll interval so tests that exercise the +// poll's retry/timeout path don't sleep a real half-second per attempt. +func fastMLflowPoll(t *testing.T) { + t.Helper() + orig := mlflowLinkPollInterval + mlflowLinkPollInterval = time.Millisecond + t.Cleanup(func() { mlflowLinkPollInterval = orig }) +} + +// submitServer serves a non-watch `air run` submit: runs/submit returns a run +// id, runs/get returns the run with a task, and runs/get-output returns +// getOutput (pass `{}` for "no MLflow IDs yet"). Everything else — the auth +// probe and the workspace-files upload — gets a permissive stub. +func submitServer(t *testing.T, getOutput string) *httptest.Server { + t.Helper() + runGet := `{"run_id": 555, "tasks": [{"run_id": 556, "attempt_number": 0}]}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/jobs/runs/submit"): + _, _ = w.Write([]byte(`{"run_id": 555}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(getOutput)) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func runSubmitCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL 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)) + cmd.SetContext(ctx) + cmd.SetOut(buf) + return cmd.RunE(cmd, nil) +} + +func TestRunSubmitTextOutput(t *testing.T) { + fastMLflowPoll(t) + var buf bytes.Buffer + // get-output carries no MLflow IDs, so the poll times out and the two MLflow + // links are omitted — only the Job Run link is printed. + err := runSubmitCmd(t, flags.OutputText, &buf, submitServer(t, `{}`).URL) + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, "Submitting experiment: my-run") + 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.NotContains(t, out, "View MLflow run at:") +} + +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 + require.NoError(t, runSubmitCmd(t, flags.OutputText, &buf, srvURL)) + + out := buf.String() + assert.Contains(t, out, "View MLflow run at: ") + assert.Contains(t, out, "/ml/experiments/exp1/runs/run1") + assert.Contains(t, out, "View MLflow experiment at: ") + assert.Contains(t, out, "/ml/experiments/exp1") +} + +func TestRunSubmitMLflowLinksResolveOnRetry(t *testing.T) { + fastMLflowPoll(t) + // get-output is empty on the first poll and carries the IDs on the second, + // exercising the poll's sleep-and-retry path. + var getOutputCalls int + runGet := `{"run_id": 555, "tasks": [{"run_id": 556, "attempt_number": 0}]}` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/jobs/runs/submit"): + _, _ = w.Write([]byte(`{"run_id": 555}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + getOutputCalls++ + if getOutputCalls == 1 { + _, _ = w.Write([]byte(`{}`)) + return + } + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + require.NoError(t, runSubmitCmd(t, flags.OutputText, &buf, srv.URL)) + + out := buf.String() + assert.Contains(t, out, "View MLflow run at: ") + assert.Contains(t, out, "/ml/experiments/exp1/runs/run1") + assert.GreaterOrEqual(t, getOutputCalls, 2, "should have polled get-output at least twice") +} + +func TestRunSubmitJSONStatusPending(t *testing.T) { + var buf bytes.Buffer + err := runSubmitCmd(t, flags.OutputJSON, &buf, submitServer(t, `{}`).URL) + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, `"status": "PENDING"`) + assert.Contains(t, out, `"run_id": "555"`) + // JSON stdout stays a clean envelope stream — no human-readable submit lines. + assert.NotContains(t, out, "Submitting experiment") +} diff --git a/experimental/air/cmd/run_watch_test.go b/experimental/air/cmd/run_watch_test.go index 487b0b9989a..a0b640fc479 100644 --- a/experimental/air/cmd/run_watch_test.go +++ b/experimental/air/cmd/run_watch_test.go @@ -110,8 +110,12 @@ func TestRunWatchStreamsLogs(t *testing.T) { require.NoError(t, err) out := buf.String() - assert.Contains(t, out, "Submitted run 777") + assert.Contains(t, out, "Submitted workload with Job Run ID: 777") + assert.Contains(t, out, "View job run at: ") assert.Contains(t, out, "Monitoring run and streaming logs...") + // A "Logs" divider separates the submit summary from the streamed logs. + assert.Contains(t, out, "Logs") + assert.Contains(t, out, "───") // The submitted run's logs stream through, oldest-first. assert.Contains(t, out, "step 1\nstep 2") } diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 3aa2436b827..9518a0e1a2e 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -8,6 +8,7 @@ import ( "strconv" "strings" + "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go" @@ -141,10 +142,24 @@ func submitToken(flag string, cfg *runConfig) (string, error) { return token, nil } +// withSpinner runs fn, showing an stderr spinner labeled msg when show is true. +// The spinner auto-degrades to nothing on a non-interactive terminal; show is +// false in JSON mode so the stdout envelope stream stays clean. +func withSpinner(ctx context.Context, show bool, msg string, fn func() error) error { + if !show { + return fn() + } + sp := cmdio.NewSpinner(ctx) + sp.Update(msg) + defer sp.Close() + return fn() +} + // submitWorkload runs the submit happy path: ensure the experiment directory, // upload the launch artifacts, assemble the Jobs payload, and submit it. It -// returns the new run_id and its dashboard URL. -func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { +// returns the new run_id and its dashboard URL. showProgress enables the +// stderr upload/packaging spinners (text mode only). +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (int64, string, error) { // Resolve the idempotency token first so a bad key fails before any upload, // and before the policy lookup below spends a round trip on it. token, err := submitToken(idempotencyKey, cfg) @@ -200,7 +215,9 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run if err != nil { return 0, "", err } - if err := uploadArtifacts(ctx, fc, items); err != nil { + if err := withSpinner(ctx, showProgress, "Uploading yaml configuration files…", func() error { + return uploadArtifacts(ctx, fc, items) + }); err != nil { return 0, "", err } @@ -210,7 +227,11 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run var snap snapshotResult if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { // Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh. - snap, err = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) + err = withSpinner(ctx, showProgress, "Packaging code snapshot…", func() error { + var e error + snap, e = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) + return e + }) if err != nil { return 0, "", err } diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 76ab3aeb3ac..559c3e8d188 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -225,7 +225,7 @@ func TestSubmitWorkload(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") + runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) require.NoError(t, err) assert.Equal(t, int64(777), runID) assert.Contains(t, dashboardURL, "/jobs/runs/777") @@ -268,7 +268,7 @@ func TestSubmitWorkloadHonorsOverride(t *testing.T) { cfg, err := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"compute.num_accelerators=4"}) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") + _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) require.NoError(t, err) require.Len(t, got.Tasks, 1) @@ -310,7 +310,7 @@ code_source: // The DABs upload path logs via cmdio; the real `air run` context carries it. ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -353,7 +353,7 @@ code_source: require.NoError(t, err) ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -550,7 +550,7 @@ code_source: require.NoError(t, err) ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -584,7 +584,7 @@ func TestSubmitWorkloadGuards(t *testing.T) { cfg := *base cfg.UsagePolicyName = new("nope") - _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "") + _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false) require.ErrorContains(t, err, `no usage policy named "nope"`) for _, p := range paths { assert.NotContains(t, p, "/workspace/", "no workspace write may precede policy resolution") @@ -605,7 +605,7 @@ func TestSubmitWorkloadGuards(t *testing.T) { cfg := *base cfg.Environment = &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "missing.yaml"}} - _, _, err = submitWorkload(t.Context(), tw, &cfg, cfgPath, "") + _, _, err = submitWorkload(t.Context(), tw, &cfg, cfgPath, "", false) require.ErrorContains(t, err, "failed to read requirements file") assert.False(t, uploaded, "no artifacts should be uploaded when dependency resolution fails") }) @@ -639,7 +639,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem") + _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false) require.NoError(t, err) assert.Equal(t, policyID, got.BudgetPolicyId) }) @@ -650,7 +650,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem") + _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false) require.NoError(t, err) assert.Equal(t, policyID, got.BudgetPolicyId) })