diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt index 8d52eed1dab..13aca091a96 100644 --- a/acceptance/experimental/air/run-submit/output.txt +++ b/acceptance/experimental/air/run-submit/output.txt @@ -4,6 +4,8 @@ Submitted run 555 View at: [DATABRICKS_URL]/jobs/runs/555 +Tip: use --watch to stream logs until the run completes. + === the ai_runtime_task carries the code_source_path >>> print_requests.py //api/2.2/jobs/runs/submit { diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt index e733bca091e..13c9b360ece 100644 --- a/acceptance/experimental/air/run/output.txt +++ b/acceptance/experimental/air/run/output.txt @@ -33,11 +33,9 @@ Error: compute.num_accelerators must be positive, got 0 Exit code: 1 -=== watch not yet supported +=== watch is ignored with dry-run (nothing is submitted) >>> [CLI] experimental air run -f valid.yaml --dry-run --watch -Error: --watch is not yet supported - -Exit code: 1 +Dry run: configuration for "smoke-test" is valid; not submitting. === code_source config passes validation >>> [CLI] experimental air run -f with-code-source.yaml --dry-run diff --git a/acceptance/experimental/air/run/script b/acceptance/experimental/air/run/script index 3308b5bc0fc..2a5263bceb2 100644 --- a/acceptance/experimental/air/run/script +++ b/acceptance/experimental/air/run/script @@ -13,8 +13,8 @@ errcode trace $CLI experimental air run -f valid.yaml --dry-run --override bogus title "override still runs schema validation" errcode trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0 -title "watch not yet supported" -errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch +title "watch is ignored with dry-run (nothing is submitted)" +trace $CLI experimental air run -f valid.yaml --dry-run --watch title "code_source config passes validation" trace $CLI experimental air run -f with-code-source.yaml --dry-run diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go index 0f6a0520d70..e5ab7ba8e85 100644 --- a/experimental/air/cmd/logstream.go +++ b/experimental/air/cmd/logstream.go @@ -55,6 +55,10 @@ type logRequest struct { // would poll forever waiting for the run (not the attempt) to finish. staticView bool jsonOutput bool + // onStatusChange, when set, is called on each lifecycle transition while + // following the run (current, previous display states). Used by + // `air run --watch -o json` to emit STATUS events. + onStatusChange func(current, previous string) } // logRunStatus is the subset of a run's state the log path needs, resolved once @@ -192,6 +196,8 @@ type bricklensStreamer struct { lastNano int64 firstLogSeen bool seen *seenSet + // previousState is the last display state reported to onStatusChange. + previousState string // onFirstLog, when set, is called once just before the first log line is // emitted — used to stop the "waiting for run to start" spinner before any // log byte reaches stdout. @@ -200,6 +206,20 @@ type bricklensStreamer struct { updateSpinner func(string) } +// reportStatusChange fires onStatusChange when the run's display state differs +// from the last reported one. +func (st *bricklensStreamer) reportStatusChange() { + if st.req.onStatusChange == nil { + return + } + current := st.status.displayState() + if current == st.previousState { + return + } + st.req.onStatusChange(current, st.previousState) + st.previousState = current +} + func (st *bricklensStreamer) run() (bool, error) { now := time.Now() st.fromSec = st.req.fromSeconds(st.status, now) @@ -244,6 +264,8 @@ func (st *bricklensStreamer) run() (bool, error) { st.status = status } + st.reportStatusChange() + terminal := st.status.terminal() toSec := st.req.toSeconds(st.status) diff --git a/experimental/air/cmd/logstream_support.go b/experimental/air/cmd/logstream_support.go index 1780ea790c9..ac250547e16 100644 --- a/experimental/air/cmd/logstream_support.go +++ b/experimental/air/cmd/logstream_support.go @@ -71,6 +71,78 @@ func printLogEvent(out io.Writer, eventType string, node int, line string) { fmt.Fprintln(out, string(b)) } +// submittedEvent is the JSONL event `air run --watch -o json` emits before the +// streamed log events, so a consumer sees the run id immediately. +type submittedEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + RunID string `json:"run_id"` + DashboardURL string `json:"dashboard_url"` +} + +// printSubmittedEvent writes the SUBMITTED JSONL event. +func printSubmittedEvent(out io.Writer, runID, dashboardURL string) { + b, err := json.Marshal(submittedEvent{ + Type: "SUBMITTED", + TS: time.Now().UTC().Format(time.RFC3339), + RunID: runID, + DashboardURL: dashboardURL, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// statusEvent is a JSONL event emitted on each lifecycle transition while +// following a run with --watch. +type statusEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + Status string `json:"status"` + Previous string `json:"previous_status,omitempty"` +} + +// printStatusEvent writes a STATUS JSONL event for a lifecycle transition. +func printStatusEvent(out io.Writer, current, previous string) { + b, err := json.Marshal(statusEvent{ + Type: "STATUS", + TS: time.Now().UTC().Format(time.RFC3339), + Status: current, + Previous: previous, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// terminalEvent is the closing envelope `air run --watch -o json` emits after +// streaming, carrying the run's terminal status. +type terminalEvent struct { + V int `json:"v"` + TS string `json:"ts"` + Data runResult `json:"data"` +} + +// printTerminalEvent writes the closing terminal-status envelope, matching the +// shape of renderEnvelope(runResult). +func printTerminalEvent(out io.Writer, runID, status, dashboardURL string) { + b, err := json.Marshal(terminalEvent{ + V: envelopeVersion, + TS: time.Now().UTC().Format(time.RFC3339), + Data: runResult{ + Status: status, + RunID: runID, + DashboardURL: dashboardURL, + }, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + // emitLogLine writes one log line: raw in text mode, or a JSONL LOG event under // --json. In --json mode a line matching a fatal-failure pattern also emits an // ALERT event first, giving an agent an immediate actionable signal. diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 327c4de167f..ea00368679f 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -1,7 +1,7 @@ package aircmd import ( - "errors" + "context" "fmt" "strconv" @@ -9,6 +9,7 @@ import ( "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" "github.com/spf13/cobra" ) @@ -57,11 +58,6 @@ The workload is described by a YAML config file (see --file).`, cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() - // --watch's pipeline is not ported yet; reject rather than silently ignore it. - if watch { - return errors.New("--watch is not yet supported") - } - cfg, err := loadRunConfigWithOverrides(ctx, file, overrides) if err != nil { return err @@ -82,13 +78,63 @@ The workload is described by a YAML config file (see --file).`, } runIDStr := strconv.FormatInt(runID, 10) - if root.OutputType(cmd) == flags.OutputText { + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + if !watch { + if !jsonOut { + cmdio.LogString(ctx, "Submitted run "+runIDStr) + cmdio.LogString(ctx, "View at: "+dashboardURL) + 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}) + } + + // --watch: stream the submitted run's logs until it reaches a terminal + // state, then exit with the run's outcome. This is the same pipeline as + // `air logs ` (Bricklens with MLflow fallback). + req := logRequest{ + runID: runID, + attempt: -1, + tailLines: -1, + jsonOutput: jsonOut, + } + + if !jsonOut { cmdio.LogString(ctx, "Submitted run "+runIDStr) cmdio.LogString(ctx, "View at: "+dashboardURL) - return nil + cmdio.LogString(ctx, "Monitoring run and streaming logs...") + 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. + out := cmd.OutOrStdout() + printSubmittedEvent(out, runIDStr, dashboardURL) + req.onStatusChange = func(current, previous string) { + printStatusEvent(out, current, previous) } - return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + err = runLogs(ctx, 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 + // actual state — correct whether Bricklens or the MLflow fallback served + // the logs. + printTerminalEvent(out, runIDStr, watchTerminalStatus(ctx, w, runID), dashboardURL) + return err } return cmd } + +// 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. +func watchTerminalStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) string { + status, err := resolveRunStatus(ctx, w, runID) + if err != nil { + return "UNKNOWN" + } + return status.displayState() +} diff --git a/experimental/air/cmd/run_watch_test.go b/experimental/air/cmd/run_watch_test.go new file mode 100644 index 00000000000..487b0b9989a --- /dev/null +++ b/experimental/air/cmd/run_watch_test.go @@ -0,0 +1,206 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// watchServer serves submit, the auth probe, a terminal runs/get, and a single +// page of Bricklens logs — everything `air run --watch` touches after submit. +// resultState is the terminal result the run reports (e.g. SUCCESS, FAILED). +func watchServer(t *testing.T, resultState string) *httptest.Server { + t.Helper() + runGet := `{ + "run_id": 777, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "` + resultState + `"}, + "tasks": [{"run_id": 778, "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": 777}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1700000002000000000, "body": "step 2", "node_index": 0}, + {"time_unix_nano": 1700000001000000000, "body": "step 1", "node_index": 0} + ]}`)) + default: + // Me() probe, workspace-id, SDK config discovery. + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +// watchServerMLflow serves a run whose Bricklens endpoint is gated off +// (FEATURE_DISABLED), forcing the MLflow fallback, plus the MLflow artifact +// chain (get-output, artifacts/list, credentials-for-read, the pre-signed bytes). +// STATUS events never fire on this path, so it guards the closing terminal +// envelope against relying on onStatusChange. +func watchServerMLflow(t *testing.T, resultState string) *httptest.Server { + t.Helper() + var base string + runGet := `{ + "run_id": 777, + "state": {"life_cycle_state": "TERMINATED", "result_state": "` + resultState + `"}, + "tasks": [{"run_id": 778, "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": 777}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code": "FEATURE_DISABLED", "message": "gated off"}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0/logs-0.chunk.txt", "file_size": 12}]}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case r.URL.Path == "/presigned": + _, _ = w.Write([]byte("step 1\nstep 2\n")) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func runWatchCmd(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)) + require.NoError(t, cmd.Flags().Set("watch", "true")) + + 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 TestRunWatchStreamsLogs(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputText, &buf, watchServer(t, "SUCCESS").URL) + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, "Submitted run 777") + assert.Contains(t, out, "Monitoring run and streaming logs...") + // The submitted run's logs stream through, oldest-first. + assert.Contains(t, out, "step 1\nstep 2") +} + +func TestRunWatchJSONEmitsSubmittedThenLogs(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServer(t, "SUCCESS").URL) + require.NoError(t, err) + + all := buf.String() + lines := strings.Split(strings.TrimSpace(all), "\n") + require.GreaterOrEqual(t, len(lines), 3) + // First event is SUBMITTED with the run id; then STATUS + streamed LOG events. + assert.Contains(t, lines[0], `"type":"SUBMITTED"`) + assert.Contains(t, lines[0], `"run_id":"777"`) + assert.Contains(t, all, `"type":"STATUS"`) + assert.Contains(t, all, `"type":"LOG"`) + assert.Contains(t, all, `"line":"step 1"`) + // The last line is the closing terminal-status envelope carrying SUCCESS. + assert.Contains(t, lines[len(lines)-1], `"status":"SUCCESS"`) + assert.Contains(t, lines[len(lines)-1], `"run_id":"777"`) +} + +func TestRunWatchJSONFailedRunTerminalEnvelope(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServer(t, "FAILED").URL) + // Non-zero exit is surfaced as ErrAlreadyPrinted, but the closing envelope + // still carries the terminal status. + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + assert.Contains(t, lines[len(lines)-1], `"status":"FAILED"`) +} + +func TestRunWatchFailedRunExitsNonZero(t *testing.T) { + var buf bytes.Buffer + // A run that ends FAILED streams its logs but exits non-zero, surfaced as + // ErrAlreadyPrinted (the output was already written). + err := runWatchCmd(t, flags.OutputText, &buf, watchServer(t, "FAILED").URL) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "step 1\nstep 2") +} + +func TestRunWatchDryRunSkipsSubmit(t *testing.T) { + // --dry-run takes precedence over --watch: nothing is submitted or streamed. + var got []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = append(got, r.URL.Path) + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cmd := withOutput(newRunCommand(), flags.OutputText) + require.NoError(t, cmd.Flags().Set("file", cfgPath)) + require.NoError(t, cmd.Flags().Set("watch", "true")) + require.NoError(t, cmd.Flags().Set("dry-run", "true")) + var buf bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, nil, &buf, &buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srv.URL)) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + require.NoError(t, cmd.RunE(cmd, nil)) + assert.Contains(t, buf.String(), "Dry run") + for _, p := range got { + assert.NotContains(t, p, "/jobs/runs/submit", "dry-run must not submit") + assert.NotContains(t, p, "/logs", "dry-run must not stream logs") + } +} + +func TestRunWatchJSONMLflowFallbackTerminalEnvelope(t *testing.T) { + // Regression: through the MLflow fallback the terminal status must come from + // the run's actual state, not the onStatusChange callback (which is + // Bricklens-only), so a SUCCESS run isn't mislabeled FAILED in the envelope. + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServerMLflow(t, "SUCCESS").URL) + require.NoError(t, err) + + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + assert.Contains(t, lines[0], `"type":"SUBMITTED"`) + // Logs stream through the fallback, and the closing envelope reflects the + // real terminal status. + assert.Contains(t, buf.String(), `"line":"step 1"`) + assert.Contains(t, lines[len(lines)-1], `"status":"SUCCESS"`) +} + +func TestRunWatchFlagRegistered(t *testing.T) { + cmd := newRunCommand() + f := cmd.Flags().Lookup("watch") + require.NotNil(t, f) + assert.Equal(t, "false", f.DefValue) +}