From 997f6bbc991f26114a78a475b5421d9b87ca43fa Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Wed, 17 Jun 2026 14:11:11 -0700 Subject: [PATCH 01/18] AIR CLI Integration: Scaffold experimental AIR CLI command package (#5564) ## Changes - New `cmd/experimental/air/ ` package containing an air parent command plus 7 stub subcommands: run, status, list, logs, cancel, register-image - currently, all subcommands return an air is not implemented yet error with representative flags mapped from the Python CLI. Registered under the hidden experimental group. - tools/list_embeds.py: text=True was changed to universal_newlines=True so the acceptance harness runs on Python 3.6. General tooling fix. ## Why The AI runtime CLI ships today as a separately installed Python wheel with its own auth, output, and packaging. Folding it into the main Go CLI gives users one databricks install with consistent profiles, authentication, and -o json output, and removes a parallel toolchain to maintain. Landing the package scaffold first lets the individual commands be ported in small, reviewable PRs (status is next) instead of one large drop. Every stub is wired and navigable, so the command tree and registration are reviewable now without functional code. ## Tests - Unit (cmd/experimental/air/): New() registers all six subcommands; each stub returns the not-implemented error. - Acceptance (acceptance/experimental/air/unimplemented/): runs every stub end-to-end and asserts the message + non-zero exit. test with: `go test ./cmd/experimental/air/...` `go test ./acceptance -run 'TestAccept/experimental/air'` --- Taskfile.yml | 2 +- .../experimental/air/help/out.test.toml | 3 ++ acceptance/experimental/air/help/output.txt | 29 ++++++++++++++ acceptance/experimental/air/help/script | 5 +++ acceptance/experimental/air/help/test.toml | 3 ++ .../air/unimplemented/out.test.toml | 3 ++ .../experimental/air/unimplemented/output.txt | 36 +++++++++++++++++ .../experimental/air/unimplemented/script | 19 +++++++++ .../experimental/air/unimplemented/test.toml | 3 ++ cmd/experimental/experimental.go | 2 + experimental/air/cmd/air.go | 36 +++++++++++++++++ experimental/air/cmd/air_test.go | 22 +++++++++++ experimental/air/cmd/cancel.go | 39 +++++++++++++++++++ experimental/air/cmd/get.go | 19 +++++++++ experimental/air/cmd/list.go | 31 +++++++++++++++ experimental/air/cmd/logs.go | 36 +++++++++++++++++ experimental/air/cmd/register_image.go | 33 ++++++++++++++++ experimental/air/cmd/run.go | 36 +++++++++++++++++ experimental/air/cmd/stubs_test.go | 31 +++++++++++++++ 19 files changed, 387 insertions(+), 1 deletion(-) create mode 100644 acceptance/experimental/air/help/out.test.toml create mode 100644 acceptance/experimental/air/help/output.txt create mode 100644 acceptance/experimental/air/help/script create mode 100644 acceptance/experimental/air/help/test.toml create mode 100644 acceptance/experimental/air/unimplemented/out.test.toml create mode 100644 acceptance/experimental/air/unimplemented/output.txt create mode 100644 acceptance/experimental/air/unimplemented/script create mode 100644 acceptance/experimental/air/unimplemented/test.toml create mode 100644 experimental/air/cmd/air.go create mode 100644 experimental/air/cmd/air_test.go create mode 100644 experimental/air/cmd/cancel.go create mode 100644 experimental/air/cmd/get.go create mode 100644 experimental/air/cmd/list.go create mode 100644 experimental/air/cmd/logs.go create mode 100644 experimental/air/cmd/register_image.go create mode 100644 experimental/air/cmd/run.go create mode 100644 experimental/air/cmd/stubs_test.go diff --git a/Taskfile.yml b/Taskfile.yml index 488e2284cc6..f48abdb99bb 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -4,7 +4,7 @@ vars: # Absolute path so tasks with `dir:` (lint-go-tools, lint-go-codegen) can use it. GO_TOOL: go tool -modfile={{.ROOT_DIR}}/tools/go.mod EXE_EXT: '{{if eq OS "windows"}}.exe{{end}}' - TEST_PACKAGES: ./acceptance/internal ./libs/... ./internal/... ./cmd/... ./bundle/... ./experimental/ssh/... . + TEST_PACKAGES: ./acceptance/internal ./libs/... ./internal/... ./cmd/... ./bundle/... ./experimental/air/... ./experimental/ssh/... . ACCEPTANCE_TEST_FILTER: "" # Single brace-expansion glob covering every //go:embed target in the repo, # computed by grepping `//go:embed` directives. Evaluated lazily by Task so diff --git a/acceptance/experimental/air/help/out.test.toml b/acceptance/experimental/air/help/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/help/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt new file mode 100644 index 00000000000..3a0f86e164f --- /dev/null +++ b/acceptance/experimental/air/help/output.txt @@ -0,0 +1,29 @@ + +=== help +>>> [CLI] experimental air --help +Run and manage AI runtime training workloads on Databricks serverless GPU compute. + +This command set is the Go port of the standalone Python "air" CLI. It is +experimental and may change in future versions. + +Usage: + databricks experimental air [command] + +Available Commands: + cancel Cancel one or more runs + get Show details for a run + list List recent runs + logs Stream or fetch logs for a run + register-image Mirror a Docker image into the workspace registry + run Submit a training workload from a YAML config + +Flags: + -h, --help help for air + +Global Flags: + --debug enable debug logging + -o, --output type output type: text or json (default text) + -p, --profile string ~/.databrickscfg profile + -t, --target string bundle target to use (if applicable) + +Use "databricks experimental air [command] --help" for more information about a command. diff --git a/acceptance/experimental/air/help/script b/acceptance/experimental/air/help/script new file mode 100644 index 00000000000..cd67a6fc1b1 --- /dev/null +++ b/acceptance/experimental/air/help/script @@ -0,0 +1,5 @@ +# Pin the command tree so any change to a subcommand or its short description +# shows up as a diff here. + +title "help" +trace $CLI experimental air --help diff --git a/acceptance/experimental/air/help/test.toml b/acceptance/experimental/air/help/test.toml new file mode 100644 index 00000000000..49709b578ef --- /dev/null +++ b/acceptance/experimental/air/help/test.toml @@ -0,0 +1,3 @@ +# --help prints without authenticating, so no server stubs are needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/unimplemented/out.test.toml b/acceptance/experimental/air/unimplemented/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/unimplemented/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt new file mode 100644 index 00000000000..4a07a38a378 --- /dev/null +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -0,0 +1,36 @@ + +=== run +>>> [CLI] experimental air run +Error: `air run` is not implemented yet + +Exit code: 1 + +=== get +>>> [CLI] experimental air get 123 +Error: `air get` is not implemented yet + +Exit code: 1 + +=== list +>>> [CLI] experimental air list +Error: `air list` is not implemented yet + +Exit code: 1 + +=== logs +>>> [CLI] experimental air logs 123 +Error: `air logs` is not implemented yet + +Exit code: 1 + +=== cancel +>>> [CLI] experimental air cancel 123 +Error: `air cancel` is not implemented yet + +Exit code: 1 + +=== register-image +>>> [CLI] experimental air register-image my-image:latest +Error: `air register-image` is not implemented yet + +Exit code: 1 diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script new file mode 100644 index 00000000000..2ed885c0e66 --- /dev/null +++ b/acceptance/experimental/air/unimplemented/script @@ -0,0 +1,19 @@ +# Each stub must fail with "not implemented"; errcode records the exit code. + +title "run" +errcode trace $CLI experimental air run + +title "get" +errcode trace $CLI experimental air get 123 + +title "list" +errcode trace $CLI experimental air list + +title "logs" +errcode trace $CLI experimental air logs 123 + +title "cancel" +errcode trace $CLI experimental air cancel 123 + +title "register-image" +errcode trace $CLI experimental air register-image my-image:latest diff --git a/acceptance/experimental/air/unimplemented/test.toml b/acceptance/experimental/air/unimplemented/test.toml new file mode 100644 index 00000000000..c233c30a86c --- /dev/null +++ b/acceptance/experimental/air/unimplemented/test.toml @@ -0,0 +1,3 @@ +# Stubs fail locally before any API call, so no server stubs needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] diff --git a/cmd/experimental/experimental.go b/cmd/experimental/experimental.go index 8d9827c5c94..d87c893abc5 100644 --- a/cmd/experimental/experimental.go +++ b/cmd/experimental/experimental.go @@ -1,6 +1,7 @@ package experimental import ( + aircmd "github.com/databricks/cli/experimental/air/cmd" aitoolscmd "github.com/databricks/cli/experimental/aitools/cmd" geniecmd "github.com/databricks/cli/experimental/genie/cmd" postgrescmd "github.com/databricks/cli/experimental/postgres/cmd" @@ -22,6 +23,7 @@ These commands provide early access to new features that are still under development. They may change or be removed in future versions without notice.`, } + cmd.AddCommand(aircmd.New()) cmd.AddCommand(aitoolscmd.NewAitoolsCmd()) cmd.AddCommand(geniecmd.NewGenieCmd()) cmd.AddCommand(postgrescmd.New()) diff --git a/experimental/air/cmd/air.go b/experimental/air/cmd/air.go new file mode 100644 index 00000000000..81ffb2dd346 --- /dev/null +++ b/experimental/air/cmd/air.go @@ -0,0 +1,36 @@ +package aircmd + +import ( + "fmt" + + "github.com/spf13/cobra" +) + +// New returns the root command for the experimental AI runtime CLI. +// +// Milestone 0: scaffolds the command group with every subcommand registered as a +// stub (not yet implemented), pending the port from the Python `air` CLI. +func New() *cobra.Command { + cmd := &cobra.Command{ + Use: "air", + Short: "Run and manage AI runtime training workloads", + Long: `Run and manage AI runtime training workloads on Databricks serverless GPU compute. + +This command set is the Go port of the standalone Python "air" CLI. It is +experimental and may change in future versions.`, + } + + cmd.AddCommand(newRunCommand()) + cmd.AddCommand(newGetCommand()) + cmd.AddCommand(newListCommand()) + cmd.AddCommand(newLogsCommand()) + cmd.AddCommand(newCancelCommand()) + cmd.AddCommand(newRegisterImageCommand()) + + return cmd +} + +// notImplemented returns the placeholder error used by milestone-0 stubs. +func notImplemented(name string) error { + return fmt.Errorf("`air %s` is not implemented yet", name) +} diff --git a/experimental/air/cmd/air_test.go b/experimental/air/cmd/air_test.go new file mode 100644 index 00000000000..7efac253a2b --- /dev/null +++ b/experimental/air/cmd/air_test.go @@ -0,0 +1,22 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestNewRegistersAllSubcommands asserts the `air` command wires up every +// expected subcommand, so none is accidentally dropped from New. +func TestNewRegistersAllSubcommands(t *testing.T) { + registered := make(map[string]bool) + for _, c := range New().Commands() { + registered[c.Name()] = true + } + + want := []string{"run", "get", "list", "logs", "cancel", "register-image"} + for _, name := range want { + assert.True(t, registered[name], "subcommand %q is not registered", name) + } + assert.Len(t, registered, len(want), "unexpected number of subcommands") +} diff --git a/experimental/air/cmd/cancel.go b/experimental/air/cmd/cancel.go new file mode 100644 index 00000000000..ae5514e5b04 --- /dev/null +++ b/experimental/air/cmd/cancel.go @@ -0,0 +1,39 @@ +package aircmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newCancelCommand() *cobra.Command { + var ( + all bool + yes bool + ) + + cmd := &cobra.Command{ + Use: "cancel [JOB_RUN_ID...]", + Short: "Cancel one or more runs", + Long: `Cancel one or more runs by ID, or cancel all of your active runs with --all.`, + RunE: func(cmd *cobra.Command, args []string) error { + return notImplemented("cancel") + }, + } + + cmd.Flags().BoolVar(&all, "all", false, "Cancel all of your active runs") + cmd.Flags().BoolVarP(&yes, "yes", "y", false, "Skip the confirmation prompt") + + // Require exactly one of: one or more JOB_RUN_IDs, or --all. Cobra parses flags + // before running this, so `all` reflects the user's input. + cmd.Args = func(cmd *cobra.Command, args []string) error { + switch { + case all && len(args) > 0: + return &root.InvalidArgsError{Command: cmd, Message: "cannot combine JOB_RUN_ID arguments with --all"} + case !all && len(args) == 0: + return &root.InvalidArgsError{Command: cmd, Message: "provide at least one JOB_RUN_ID, or use --all"} + } + return nil + } + + return cmd +} diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go new file mode 100644 index 00000000000..45f93df10a6 --- /dev/null +++ b/experimental/air/cmd/get.go @@ -0,0 +1,19 @@ +package aircmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newGetCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "get JOB_RUN_ID", + Args: root.ExactArgs(1), + Short: "Show details for a run", + RunE: func(cmd *cobra.Command, args []string) error { + return notImplemented("get") + }, + } + + return cmd +} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go new file mode 100644 index 00000000000..bf24cff9b23 --- /dev/null +++ b/experimental/air/cmd/list.go @@ -0,0 +1,31 @@ +package aircmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newListCommand() *cobra.Command { + var ( + limit int + active bool + allUsers bool + filters []string + ) + + cmd := &cobra.Command{ + Use: "list", + Args: root.NoArgs, + Short: "List recent runs", + RunE: func(cmd *cobra.Command, args []string) error { + return notImplemented("list") + }, + } + + cmd.Flags().IntVar(&limit, "limit", 20, "Maximum number of runs to show") + cmd.Flags().BoolVar(&active, "active", false, "Show only active runs") + cmd.Flags().BoolVar(&allUsers, "all-users", false, "Show runs from all users") + cmd.Flags().StringArrayVar(&filters, "filter", nil, "Filter runs, e.g. experiment=foo* (repeatable)") + + return cmd +} diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go new file mode 100644 index 00000000000..c34fb62a7df --- /dev/null +++ b/experimental/air/cmd/logs.go @@ -0,0 +1,36 @@ +package aircmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newLogsCommand() *cobra.Command { + var ( + node int + lines int + retry int + downloadTo string + review bool + ) + + cmd := &cobra.Command{ + Use: "logs JOB_RUN_ID", + Args: root.ExactArgs(1), + Short: "Stream or fetch logs for a run", + Long: `Stream logs from an active run, or fetch logs from a completed run.`, + RunE: func(cmd *cobra.Command, args []string) error { + return notImplemented("logs") + }, + } + + cmd.Flags().IntVar(&node, "node", 0, "Fetch logs from this node") + cmd.Flags().IntVar(&lines, "lines", 10000, "For completed runs, print the last N lines") + cmd.Flags().IntVar(&retry, "retry", -1, "View logs from a specific retry attempt; -1 means latest") + cmd.Flags().StringVar(&downloadTo, "download-to", "", "Download all logs to this directory instead of printing") + cmd.Flags().BoolVar(&review, "review", false, "Download logs from all nodes and filter for error signatures") + // Hidden in the Python `air` CLI (help=argparse.SUPPRESS); keep it internal here to match. + cmd.Flags().MarkHidden("review") + + return cmd +} diff --git a/experimental/air/cmd/register_image.go b/experimental/air/cmd/register_image.go new file mode 100644 index 00000000000..1d8b45044a7 --- /dev/null +++ b/experimental/air/cmd/register_image.go @@ -0,0 +1,33 @@ +package aircmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newRegisterImageCommand() *cobra.Command { + var ( + scope string + key string + interactiveAuth bool + tagPolicy string + timeoutMinutes int + ) + + cmd := &cobra.Command{ + Use: "register-image IMAGE_URL", + Args: root.ExactArgs(1), + Short: "Mirror a Docker image into the workspace registry", + RunE: func(cmd *cobra.Command, args []string) error { + return notImplemented("register-image") + }, + } + + cmd.Flags().StringVar(&scope, "scope", "", "Databricks secret scope holding registry credentials") + cmd.Flags().StringVar(&key, "key", "", "Databricks secret key holding registry credentials") + cmd.Flags().BoolVarP(&interactiveAuth, "interactive-authenticate", "i", false, "Prompt for registry credentials and store them as a secret") + cmd.Flags().StringVar(&tagPolicy, "tag-policy", "auto", "Image resolution policy: auto or latest") + cmd.Flags().IntVar(&timeoutMinutes, "timeout-minutes", 60, "Timeout to wait for the image to become available") + + return cmd +} diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go new file mode 100644 index 00000000000..0bc3d1fd94b --- /dev/null +++ b/experimental/air/cmd/run.go @@ -0,0 +1,36 @@ +package aircmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newRunCommand() *cobra.Command { + var ( + file string + watch bool + overrides []string + dryRun bool + idempotencyKey string + ) + + cmd := &cobra.Command{ + Use: "run", + Args: root.NoArgs, + Short: "Submit a training workload from a YAML config", + Long: `Submit a training workload to Databricks serverless GPU compute. + +The workload is described by a YAML config file (see --file).`, + RunE: func(cmd *cobra.Command, args []string) error { + return notImplemented("run") + }, + } + + cmd.Flags().StringVarP(&file, "file", "f", "", "Path to the workload YAML config") + cmd.Flags().BoolVar(&watch, "watch", false, "Stream logs until the run completes") + cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") + cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config without submitting") + cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used") + + return cmd +} diff --git a/experimental/air/cmd/stubs_test.go b/experimental/air/cmd/stubs_test.go new file mode 100644 index 00000000000..a6e24177f33 --- /dev/null +++ b/experimental/air/cmd/stubs_test.go @@ -0,0 +1,31 @@ +package aircmd + +import ( + "fmt" + "testing" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestStubCommandsReturnNotImplemented asserts each unimplemented subcommand +// fails with a "not implemented" error. Drop a command here once it lands. +func TestStubCommandsReturnNotImplemented(t *testing.T) { + stubs := map[string]*cobra.Command{ + "run": newRunCommand(), + "get": newGetCommand(), + "list": newListCommand(), + "logs": newLogsCommand(), + "cancel": newCancelCommand(), + "register-image": newRegisterImageCommand(), + } + + for name, cmd := range stubs { + t.Run(name, func(t *testing.T) { + require.NotNil(t, cmd.RunE, "command should define RunE") + err := cmd.RunE(cmd, nil) + assert.EqualError(t, err, fmt.Sprintf("`air %s` is not implemented yet", name)) + }) + } +} From b952417ee24167fa090df9b14e28590d077fbb68 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Thu, 18 Jun 2026 09:46:47 -0700 Subject: [PATCH 02/18] AIR CLI Integration: Implement the `air get` command (#5600) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Changes Implements `databricks experimental ai get RUN_ID`, the Go port of the Python `air get` command. It fetches the run via `Jobs.GetRun` and renders: - Core fields: run ID, status, submitted time, duration, retries, experiment, accelerators, creator (`User`), and the run's dashboard URL. - An MLflow deep-link, built from `jobs/runs/get-output` (the `gen_ai_compute_output` field is not modeled by the typed SDK, so it's fetched via a direct REST call). - For foreach/sweep runs, an iteration summary (counts + per iteration table) instead of the single-run view. - The run's training-config YAML, downloaded from the workspace and printed before the status (text mode only). ## Why `get` is the first real command integrated from the air cli and it sets the conventions the rest of the CLI will follow. The `{v, ts, data}` envelope mirrors the Python CLI so existing machine consumers keep working. The implementation is a faithful port of `handle_status` + the `cli_display` helpers, verified field-by-field against the Python source: - The text view shows the foreach branch (`_display_foreach_sweep_status`) and the training-config panel (`_fetch_and_display_yaml_config`); JSON output omits both, exactly matching `air get --json`. - MLflow IDs live under an unmodeled `gen_ai_compute_output` field (direct REST call), and the MLflow link / YAML fetch are best-effort (logic matches python cli) ## Tests - Unit tests cover every formatting/extraction helper, `buildGetData`, and all template branches (single-run minimal/all-fields, sweep, sweep-with-no-tasks). - Mock-backed unit tests (mirroring the Python `unittest.mock` suite) cover `buildSweepInfo`, `printConfigYAML`, `mlflowURL` (over `httptest`, since it bypasses the typed SDK), and the `RunE` invalid-id / not-found branches. - An acceptance test (`acceptance/experimental/air/get`) runs the command end-to-end against a stubbed Jobs API: text output, `-o json`, and an invalid run ID. Manual verification outputs: Successful run: Screenshot 2026-06-17 at 1 17 30 PM Screenshot 2026-06-17 at 1 16 48 PM Screenshot 2026-06-17 at 11 56 00 AM Screenshot 2026-06-17 at 2 05 21 PM Failed run: Screenshot 2026-06-17 at 1 11 31 PM Screenshot 2026-06-17 at 1 13 22 PM Screenshot 2026-06-17 at 1 15 52 PM Screenshot 2026-06-17 at 2 04 48 PM --- acceptance/experimental/air/get/out.test.toml | 3 + acceptance/experimental/air/get/output.txt | 52 ++++ acceptance/experimental/air/get/script | 11 + acceptance/experimental/air/get/test.toml | 47 +++ acceptance/experimental/air/help/output.txt | 2 +- .../experimental/air/unimplemented/output.txt | 6 - .../experimental/air/unimplemented/script | 3 - experimental/air/cmd/format.go | 265 +++++++++++++++++ experimental/air/cmd/format_test.go | 227 +++++++++++++++ experimental/air/cmd/get.go | 235 ++++++++++++++- experimental/air/cmd/get_test.go | 267 ++++++++++++++++++ experimental/air/cmd/mlflow.go | 126 +++++++++ experimental/air/cmd/mlflow_test.go | 112 ++++++++ experimental/air/cmd/output.go | 82 ++++++ experimental/air/cmd/output_test.go | 52 ++++ experimental/air/cmd/stubs_test.go | 1 - experimental/air/cmd/sweep.go | 76 +++++ experimental/air/cmd/sweep_test.go | 81 ++++++ 18 files changed, 1633 insertions(+), 15 deletions(-) create mode 100644 acceptance/experimental/air/get/out.test.toml create mode 100644 acceptance/experimental/air/get/output.txt create mode 100644 acceptance/experimental/air/get/script create mode 100644 acceptance/experimental/air/get/test.toml create mode 100644 experimental/air/cmd/format.go create mode 100644 experimental/air/cmd/format_test.go create mode 100644 experimental/air/cmd/get_test.go create mode 100644 experimental/air/cmd/mlflow.go create mode 100644 experimental/air/cmd/mlflow_test.go create mode 100644 experimental/air/cmd/output.go create mode 100644 experimental/air/cmd/output_test.go create mode 100644 experimental/air/cmd/sweep.go create mode 100644 experimental/air/cmd/sweep_test.go diff --git a/acceptance/experimental/air/get/out.test.toml b/acceptance/experimental/air/get/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/get/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/get/output.txt b/acceptance/experimental/air/get/output.txt new file mode 100644 index 00000000000..f8df1fb36bb --- /dev/null +++ b/acceptance/experimental/air/get/output.txt @@ -0,0 +1,52 @@ + +=== get (text) +>>> [CLI] experimental air get run 123 +Job Link: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] + +Run ID: 123 +Status: SUCCESS +Submitted: 2023-11-14 22:13 UTC +Retries: 0 +Duration: 12s +Experiment: my-exp +MLflow Run: my-run +User: user@example.com +Accelerators: 8x H100 + +=== get (json) +>>> [CLI] experimental air get run 123 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "run_id": "123", + "status": "SUCCESS", + "started_at": "[TIMESTAMP]", + "duration_seconds": 12, + "attempt_number": 0, + "experiment_name": "my-exp", + "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[NUMID]", + "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0" + } +} + +=== invalid run id +>>> [CLI] experimental air get run notanumber +Error: invalid JOB_RUN_ID "notanumber": must be a positive integer + +Exit code: 1 + +=== invalid run id (json) +>>> [CLI] experimental air get run notanumber -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "error": { + "code": "INVALID_ARGS", + "kind": "PERMANENT", + "message": "invalid JOB_RUN_ID \"notanumber\": must be a positive integer", + "retryable": false + } +} + +Exit code: 1 diff --git a/acceptance/experimental/air/get/script b/acceptance/experimental/air/get/script new file mode 100644 index 00000000000..b775d06a48d --- /dev/null +++ b/acceptance/experimental/air/get/script @@ -0,0 +1,11 @@ +title "get (text)" +trace $CLI experimental air get run 123 + +title "get (json)" +trace $CLI experimental air get run 123 -o json + +title "invalid run id" +errcode trace $CLI experimental air get run notanumber + +title "invalid run id (json)" +errcode trace $CLI experimental air get run notanumber -o json diff --git a/acceptance/experimental/air/get/test.toml b/acceptance/experimental/air/get/test.toml new file mode 100644 index 00000000000..6162cebdb5d --- /dev/null +++ b/acceptance/experimental/air/get/test.toml @@ -0,0 +1,47 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A single GenAI-compute run with an experiment, GPUs, and a creator. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "run_page_url": "https://my-workspace.cloud.databricks.test/jobs/runs/123", + "creator_user_name": "user@example.com", + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [ + { + "task_key": "train", + "attempt_number": 0, + "gen_ai_compute_task": { + "mlflow_experiment_name": "/Users/user@example.com/my-exp", + "compute": {"gpu_type": "GPU_8xH100", "num_gpus": 8} + } + } + ] +} +''' + +# MLflow identifiers for the deep-link (runs/get-output is not modeled by the typed SDK). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = ''' +{"gen_ai_compute_output": {"run_info": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}} +''' + +# The MLflow Run cell shows the run's name, fetched from the MLflow REST API. +[[Server]] +Pattern = "GET /api/2.0/mlflow/runs/get" +Response.Body = ''' +{"run": {"info": {"run_name": "my-run"}}} +''' diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index 3a0f86e164f..41a1d6815a8 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -11,7 +11,7 @@ Usage: Available Commands: cancel Cancel one or more runs - get Show details for a run + get Show details for a specific resource list List recent runs logs Stream or fetch logs for a run register-image Mirror a Docker image into the workspace registry diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt index 4a07a38a378..0a86360c78f 100644 --- a/acceptance/experimental/air/unimplemented/output.txt +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -5,12 +5,6 @@ Error: `air run` is not implemented yet Exit code: 1 -=== get ->>> [CLI] experimental air get 123 -Error: `air get` is not implemented yet - -Exit code: 1 - === list >>> [CLI] experimental air list Error: `air list` is not implemented yet diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script index 2ed885c0e66..e6e8d33ef9d 100644 --- a/acceptance/experimental/air/unimplemented/script +++ b/acceptance/experimental/air/unimplemented/script @@ -3,9 +3,6 @@ title "run" errcode trace $CLI experimental air run -title "get" -errcode trace $CLI experimental air get 123 - title "list" errcode trace $CLI experimental air list diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go new file mode 100644 index 00000000000..8694ed69e9f --- /dev/null +++ b/experimental/air/cmd/format.go @@ -0,0 +1,265 @@ +package aircmd + +import ( + "bytes" + "context" + "fmt" + "io" + "math" + "strings" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/service/jobs" + "go.yaml.in/yaml/v3" +) + +// na is the placeholder shown for an empty text-table cell, matching the Python CLI. +const na = "N/A" + +// orNA returns s, or "N/A" when s is empty, for text-table cells. +func orNA(s string) string { + if s == "" { + return na + } + return s +} + +// osc8Link wraps label in an OSC 8 terminal hyperlink to url. +// See https://gist.github.com/egmontkob/eb114294efbcd5adb1944c9f3cb5feda +func osc8Link(label, url string) string { + return "\x1b]8;;" + url + "\x1b\\" + label + "\x1b]8;;\x1b\\" +} + +// hyperlink renders label as a terminal hyperlink to url when out is a rich +// terminal, otherwise it returns label unchanged. This mirrors the Python CLI's +// Rich link markup, which drops the URL on non-terminals (so piped or captured +// output stays plain text). +func hyperlink(ctx context.Context, out io.Writer, label, url string) string { + if url == "" || !cmdio.SupportsColor(ctx, out) { + return label + } + return osc8Link(label, url) +} + +// reformatYAMLForDisplay re-renders a training-config YAML so multi-line strings +// (notably the `command:` field) appear as `|` block literals instead of the +// quoted "\n"-escaped single line they are stored as, which is unreadable. It +// mirrors Python's _reformat_yaml_for_display (cli_display.py); we skip the +// Rich syntax-highlighted panel and only fix the whitespace. On any parse or +// re-encode failure it returns the original content unchanged. +func reformatYAMLForDisplay(content []byte) string { + var node yaml.Node + if err := yaml.Unmarshal(content, &node); err != nil { + return string(content) + } + forceLiteralBlockStrings(&node) + + var buf bytes.Buffer + enc := yaml.NewEncoder(&buf) + enc.SetIndent(2) + if err := enc.Encode(&node); err != nil { + return string(content) + } + enc.Close() + return buf.String() +} + +// forceLiteralBlockStrings walks a YAML node tree and marks every multi-line +// string scalar for `|` block-literal rendering. The encoder automatically +// falls back to a quoted style when a value can't be represented as a block +// literal (e.g. lines with trailing whitespace), so no explicit guard is needed. +func forceLiteralBlockStrings(node *yaml.Node) { + if node.Kind == yaml.ScalarNode && node.Tag == "!!str" && strings.Contains(node.Value, "\n") { + node.Style = yaml.LiteralStyle + } + for _, child := range node.Content { + forceLiteralBlockStrings(child) + } +} + +// gpuDisplayNames maps the GPU identifiers returned by the backend to the short +// names we show to users. Unknown identifiers are shown unchanged. +var gpuDisplayNames = map[string]string{ + "h100_80gb": "H100", + "a10": "A10", + "GPU_1xA10": "A10", + "GPU_8xH100": "H100", + "GPU_1xH100": "H100", +} + +// runStatus returns the single status word to show for a run. The backend +// reports two values: a lifecycle state (e.g. PENDING, RUNNING) and, once the +// run has finished, a result state (e.g. SUCCESS, FAILED). The result state is +// the more meaningful one, so we prefer it when it is set. +func runStatus(state *jobs.RunState) string { + if state == nil { + return "UNKNOWN" + } + if state.ResultState != "" { + return string(state.ResultState) + } + if state.LifeCycleState != "" { + return string(state.LifeCycleState) + } + return "UNKNOWN" +} + +// reportedTiming returns the run's start and end times (epoch milliseconds), +// preferring the last task's window over the run-level times so a retried run +// reports its latest attempt. Mirrors Python's _reported_attempt_timing +// (cli_display.py:78-87). +func reportedTiming(run *jobs.Run) (startMillis, endMillis int64) { + startMillis, endMillis = run.StartTime, run.EndTime + if n := len(run.Tasks); n > 0 { + last := run.Tasks[n-1] + if last.StartTime > 0 { + startMillis = last.StartTime + } + if last.EndTime > 0 { + endMillis = last.EndTime + } + } + return startMillis, endMillis +} + +// startedAt returns the run's start time as a Python-isoformat string ("+00:00", +// not "Z"; microseconds only when non-zero, cli_entrypoint.py:1899), or nil if it +// hasn't started. +func startedAt(run *jobs.Run) *string { + startMillis, _ := reportedTiming(run) + if startMillis == 0 { + return nil + } + t := time.UnixMilli(startMillis).UTC() + layout := "2006-01-02T15:04:05-07:00" + if t.Nanosecond() != 0 { + layout = "2006-01-02T15:04:05.000000-07:00" + } + s := t.Format(layout) + return &s +} + +// submittedDisplay formats the run's start time for the text table as +// "2006-01-02 15:04 UTC", or "N/A" if it hasn't started. Mirrors Python's +// _format_timestamp (cli_display.py); we render in UTC for stable output rather +// than the local zone Python uses. +func submittedDisplay(run *jobs.Run) string { + startMillis, _ := reportedTiming(run) + if startMillis == 0 { + return na + } + return time.UnixMilli(startMillis).UTC().Format("2006-01-02 15:04 MST") +} + +// durationSeconds returns how long the run has taken, in whole seconds, or nil +// if it has not started. For a finished run this is the elapsed time of the +// reported attempt; for a still-running run it is the time since it started. +func durationSeconds(run *jobs.Run) *int64 { + startMillis, endMillis := reportedTiming(run) + if startMillis == 0 { + return nil + } + if endMillis == 0 { + // Still running: measure against the current time. + endMillis = time.Now().UnixMilli() + } + d := roundMillisToSeconds(endMillis - startMillis) + return &d +} + +// roundMillisToSeconds rounds milliseconds to whole seconds, half to even, to +// match Python's round() (cli_entrypoint.py:1903). +func roundMillisToSeconds(ms int64) int64 { + return int64(math.RoundToEven(float64(ms) / 1000)) +} + +// dashboardURL builds {host}/jobs/runs/{id}?o={workspace_id}, matching Python +// (cli_entrypoint.py:1911). The ?o= workspace id deep-links to the right +// workspace on multi-workspace accounts. +func dashboardURL(host string, runID, workspaceID int64) string { + return fmt.Sprintf("%s/jobs/runs/%d?o=%d", strings.TrimRight(host, "/"), runID, workspaceID) +} + +// formatDuration turns a number of seconds into a compact human string such as +// "1h 2m 3s". Trailing zero units are dropped, but a lone "0s" is kept so the +// result is never empty. +func formatDuration(totalSeconds int64) string { + hours := totalSeconds / 3600 + minutes := (totalSeconds % 3600) / 60 + seconds := totalSeconds % 60 + + var parts []string + if hours > 0 { + parts = append(parts, fmt.Sprintf("%dh", hours)) + } + if minutes > 0 { + parts = append(parts, fmt.Sprintf("%dm", minutes)) + } + if seconds > 0 || len(parts) == 0 { + parts = append(parts, fmt.Sprintf("%ds", seconds)) + } + return strings.Join(parts, " ") +} + +// latestAttemptNumber returns the retry count of the run's most recent task. +// Tasks start at attempt 0, so a value of 0 means the run has not been retried. +func latestAttemptNumber(run *jobs.Run) int { + if len(run.Tasks) == 0 { + return 0 + } + return run.Tasks[len(run.Tasks)-1].AttemptNumber +} + +// experimentName returns the MLflow experiment name for the run, or nil if there +// isn't one. Experiment names are often stored under a user's home folder (e.g. +// "/Users/me@example.com/my-experiment"); we strip that prefix so users see just +// the experiment name they chose. +func experimentName(run *jobs.Run) *string { + if len(run.Tasks) == 0 { + return nil + } + task := run.Tasks[0].GenAiComputeTask + if task == nil || task.MlflowExperimentName == "" { + return nil + } + name := stripExperimentUserPrefix(task.MlflowExperimentName) + return &name +} + +// stripExperimentUserPrefix removes a leading "/Users//" from an +// experiment name, leaving the remainder. Names without that prefix are returned +// unchanged. +func stripExperimentUserPrefix(name string) string { + if !strings.HasPrefix(name, "/Users/") { + return name + } + // Split into ["", "Users", "", ""]; keep "". + parts := strings.SplitN(name, "/", 4) + if len(parts) == 4 { + return parts[3] + } + return name +} + +// accelerators returns a short description of the GPUs the run uses, such as +// "8x H100", or an empty string if the run has no GPU compute attached. +func accelerators(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "" + } + task := run.Tasks[0].GenAiComputeTask + if task == nil || task.Compute == nil || task.Compute.NumGpus == 0 { + return "" + } + return fmt.Sprintf("%dx %s", task.Compute.NumGpus, gpuDisplayName(task.Compute.GpuType)) +} + +// gpuDisplayName returns the friendly name for a GPU identifier, falling back to +// the identifier itself when it is not one we recognize. +func gpuDisplayName(gpuType string) string { + if name, ok := gpuDisplayNames[gpuType]; ok { + return name + } + return gpuType +} diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go new file mode 100644 index 00000000000..acecd0ce901 --- /dev/null +++ b/experimental/air/cmd/format_test.go @@ -0,0 +1,227 @@ +package aircmd + +import ( + "io" + "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" +) + +func TestOrNA(t *testing.T) { + assert.Equal(t, "x", orNA("x")) + assert.Equal(t, "N/A", orNA("")) +} + +func TestSubmittedDisplay(t *testing.T) { + assert.Equal(t, "N/A", submittedDisplay(&jobs.Run{})) + // 1700000000000 ms == 2023-11-14 22:13:20 UTC. + assert.Equal(t, "2023-11-14 22:13 UTC", submittedDisplay(&jobs.Run{StartTime: 1700000000000})) +} + +func TestOSC8Link(t *testing.T) { + assert.Equal(t, "\x1b]8;;https://h.test/x\x1b\\label\x1b]8;;\x1b\\", osc8Link("label", "https://h.test/x")) +} + +func TestHyperlink(t *testing.T) { + // On a non-terminal (no color), the URL is dropped and only the label shows. + ctx := cmdio.MockDiscard(t.Context()) + assert.Equal(t, "label", hyperlink(ctx, io.Discard, "label", "https://h.test/x")) + // An empty URL is always rendered as the bare label. + assert.Equal(t, "label", hyperlink(ctx, io.Discard, "label", "")) +} + +func TestFormatDuration(t *testing.T) { + cases := []struct { + seconds int64 + want string + }{ + {0, "0s"}, + {45, "45s"}, + {60, "1m"}, + {63, "1m 3s"}, + {3600, "1h"}, + {3723, "1h 2m 3s"}, + {7260, "2h 1m"}, + } + for _, c := range cases { + assert.Equal(t, c.want, formatDuration(c.seconds)) + } +} + +func TestStripExperimentUserPrefix(t *testing.T) { + cases := []struct { + name string + want string + }{ + {"/Users/me@example.com/my-experiment", "my-experiment"}, + {"/Users/me@example.com/nested/path", "nested/path"}, + {"my-experiment", "my-experiment"}, + {"/Shared/team-experiment", "/Shared/team-experiment"}, + {"/Users/me@example.com", "/Users/me@example.com"}, + } + for _, c := range cases { + assert.Equal(t, c.want, stripExperimentUserPrefix(c.name)) + } +} + +func TestGpuDisplayName(t *testing.T) { + assert.Equal(t, "H100", gpuDisplayName("h100_80gb")) + assert.Equal(t, "A10", gpuDisplayName("GPU_1xA10")) + assert.Equal(t, "A10", gpuDisplayName("a10")) + assert.Equal(t, "H100", gpuDisplayName("GPU_8xH100")) + assert.Equal(t, "H100", gpuDisplayName("GPU_1xH100")) + // Unknown identifiers pass through unchanged. + assert.Equal(t, "b200", gpuDisplayName("b200")) + assert.Empty(t, gpuDisplayName("")) +} + +func TestRunStatusPrefersResultState(t *testing.T) { + // Result state wins once the run has finished. + assert.Equal(t, "SUCCESS", runStatus(&jobs.RunState{ + LifeCycleState: jobs.RunLifeCycleStateTerminated, + ResultState: jobs.RunResultStateSuccess, + })) + // Before completion only the lifecycle state is set. + assert.Equal(t, "RUNNING", runStatus(&jobs.RunState{ + LifeCycleState: jobs.RunLifeCycleStateRunning, + })) + // Non-nil state with neither field set, and nil state. + assert.Equal(t, "UNKNOWN", runStatus(&jobs.RunState{})) + assert.Equal(t, "UNKNOWN", runStatus(nil)) +} + +func TestReportedTiming(t *testing.T) { + // No tasks: run-level times are used. + start, end := reportedTiming(&jobs.Run{StartTime: 100, EndTime: 200}) + assert.Equal(t, int64(100), start) + assert.Equal(t, int64(200), end) + + // The last task's window is preferred over the run-level window, so a + // retried run reports its most recent attempt. + start, end = reportedTiming(&jobs.Run{ + StartTime: 100, EndTime: 200, + Tasks: []jobs.RunTask{ + {StartTime: 100, EndTime: 150}, + {StartTime: 300, EndTime: 450}, + }, + }) + assert.Equal(t, int64(300), start) + assert.Equal(t, int64(450), end) + + // A task missing a field falls back to the run-level value for that field. + start, end = reportedTiming(&jobs.Run{ + StartTime: 100, EndTime: 200, + Tasks: []jobs.RunTask{{StartTime: 300}}, + }) + assert.Equal(t, int64(300), start) + assert.Equal(t, int64(200), end) +} + +func TestStartedAt(t *testing.T) { + // Not started yet. + assert.Nil(t, startedAt(&jobs.Run{})) + // 1700000000000 ms == 2023-11-14T22:13:20+00:00 (Python isoformat, not "Z"). + got := startedAt(&jobs.Run{StartTime: 1700000000000}) + require.NotNil(t, got) + assert.Equal(t, "2023-11-14T22:13:20+00:00", *got) + // A sub-second start time carries microsecond precision. + got = startedAt(&jobs.Run{StartTime: 1700000000500}) + require.NotNil(t, got) + assert.Equal(t, "2023-11-14T22:13:20.500000+00:00", *got) + // The last attempt's start time is reported. 1700000060000 ms == 22:14:20. + got = startedAt(&jobs.Run{ + StartTime: 1700000000000, + Tasks: []jobs.RunTask{{StartTime: 1700000060000}}, + }) + require.NotNil(t, got) + assert.Equal(t, "2023-11-14T22:14:20+00:00", *got) +} + +func TestDurationSeconds(t *testing.T) { + // Not started yet. + assert.Nil(t, durationSeconds(&jobs.Run{})) + + // Finished run: reported end - start. + d := durationSeconds(&jobs.Run{StartTime: 1700000000000, EndTime: 1700000012000}) + require.NotNil(t, d) + assert.Equal(t, int64(12), *d) + + // Sub-second remainders round to the nearest second, matching Python: an + // 11,500 ms run reports 12s, not 11s. + d = durationSeconds(&jobs.Run{StartTime: 1700000000000, EndTime: 1700000011500}) + require.NotNil(t, d) + assert.Equal(t, int64(12), *d) + + // The last attempt's window drives the duration for a retried run. + d = durationSeconds(&jobs.Run{ + StartTime: 1700000000000, EndTime: 1700000012000, + Tasks: []jobs.RunTask{{StartTime: 1700000000000, EndTime: 1700000005000}}, + }) + require.NotNil(t, d) + assert.Equal(t, int64(5), *d) + + // Still running: measured against the current time, so positive. + d = durationSeconds(&jobs.Run{StartTime: 1700000000000}) + require.NotNil(t, d) + assert.Positive(t, *d) +} + +func TestDashboardURL(t *testing.T) { + // The ?o= workspace id and a trailing-slash-trimmed host, matching Python. + assert.Equal(t, "https://example.test/jobs/runs/123?o=42", dashboardURL("https://example.test/", 123, 42)) +} + +func TestLatestAttemptNumber(t *testing.T) { + assert.Equal(t, 0, latestAttemptNumber(&jobs.Run{})) + run := &jobs.Run{Tasks: []jobs.RunTask{{AttemptNumber: 0}, {AttemptNumber: 2}}} + assert.Equal(t, 2, latestAttemptNumber(run)) +} + +func TestExperimentName(t *testing.T) { + assert.Nil(t, experimentName(&jobs.Run{})) + assert.Nil(t, experimentName(&jobs.Run{Tasks: []jobs.RunTask{{}}})) + assert.Nil(t, experimentName(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{MlflowExperimentName: ""}, + }}})) + got := experimentName(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{MlflowExperimentName: "/Users/me@example.com/exp"}, + }}}) + require.NotNil(t, got) + assert.Equal(t, "exp", *got) +} + +func TestReformatYAMLForDisplay(t *testing.T) { + // A multi-line command stored as a quoted "\n"-escaped string is re-rendered + // as a `|` block literal, while key order is preserved. + in := "experiment_name: foo\ncommand: \"set -e\\npython train.py --epochs 3\\n\"\nmax_retries: 2\n" + want := "experiment_name: foo\ncommand: |\n set -e\n python train.py --epochs 3\nmax_retries: 2\n" + assert.Equal(t, want, reformatYAMLForDisplay([]byte(in))) + + // A single-line command is left as a plain scalar. + assert.Equal(t, "command: bash train.sh\n", reformatYAMLForDisplay([]byte("command: bash train.sh\n"))) + + // A multi-line value with trailing whitespace can't be a block literal, so the + // encoder falls back to a quoted style rather than emitting invalid YAML. + got := reformatYAMLForDisplay([]byte("command: \"trailing space \\nsecond\"\n")) + assert.Equal(t, "command: \"trailing space \\nsecond\"\n", got) + + // Unparseable content is returned unchanged. + assert.Equal(t, "\tnot: valid: yaml:", reformatYAMLForDisplay([]byte("\tnot: valid: yaml:"))) +} + +func TestAccelerators(t *testing.T) { + assert.Empty(t, accelerators(&jobs.Run{})) + assert.Empty(t, accelerators(&jobs.Run{Tasks: []jobs.RunTask{{}}})) + assert.Empty(t, accelerators(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{}, + }}})) + assert.Empty(t, accelerators(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{Compute: &jobs.ComputeConfig{NumGpus: 0}}, + }}})) + assert.Equal(t, "8x H100", accelerators(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{Compute: &jobs.ComputeConfig{NumGpus: 8, GpuType: "GPU_8xH100"}}, + }}})) +} diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 45f93df10a6..a302b7caf31 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -1,19 +1,246 @@ package aircmd import ( + "context" + "errors" + "fmt" + "io" + "strconv" + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/spf13/cobra" ) +// getData is the payload printed by `air get run`. The json-tagged fields form +// the machine-readable output; fields tagged `json:"-"` are shown only in the +// human-readable text view. +type getData struct { + RunID string `json:"run_id"` + Status string `json:"status"` + StartedAt *string `json:"started_at"` + DurationSeconds *int64 `json:"duration_seconds"` + AttemptNumber int `json:"attempt_number"` + ExperimentName *string `json:"experiment_name"` + DashboardURL string `json:"dashboard_url"` + MLflowURL *string `json:"mlflow_url"` + + // The fields below are pre-rendered for the text table and excluded from + // JSON (matching `air get run --json`). The table always shows every row, + // with "N/A" for missing values, in the same order as the Python CLI. The + // Run ID, Experiment, and MLflow Run cells carry terminal hyperlinks when + // stdout is a terminal, so the URLs don't appear as bare text. + RunIDDisplay string `json:"-"` + SubmittedDisplay string `json:"-"` + DurationDisplay string `json:"-"` + ExperimentDisplay string `json:"-"` + MLflowDisplay string `json:"-"` + UserDisplay string `json:"-"` + AcceleratorsDisplay string `json:"-"` + // Sweep replaces the single-run view for foreach runs. + Sweep *sweepInfo `json:"-"` +} + +// getTemplate is the text-mode layout. It reads from the JSON envelope, so +// every field is reached through ".Data". +const getTemplate = `{{- if .Data.Sweep -}} +Sweep Run ID: {{.Data.RunID}} +Status: {{.Data.Status}} +Total: {{.Data.Sweep.Total}} +Completed: {{.Data.Sweep.Completed}} +Succeeded: {{.Data.Sweep.Succeeded}} +Failed: {{.Data.Sweep.Failed}} +Active: {{.Data.Sweep.Active}} +{{- if .Data.Sweep.Tasks}} + +Sweep Tasks: +{{printf " %-24s %-14s %-12s %s" "TASK" "RUN ID" "STATUS" "EXPERIMENT"}} +{{- range .Data.Sweep.Tasks}} +{{printf " %-24s %-14s %-12s %s" .TaskKey .RunID .Status .Experiment}} +{{- end}} +{{- end}} +{{- else -}} +Run ID: {{.Data.RunIDDisplay}} +Status: {{.Data.Status}} +Submitted: {{.Data.SubmittedDisplay}} +Retries: {{.Data.AttemptNumber}} +Duration: {{.Data.DurationDisplay}} +Experiment: {{.Data.ExperimentDisplay}} +MLflow Run: {{.Data.MLflowDisplay}} +User: {{.Data.UserDisplay}} +Accelerators: {{.Data.AcceleratorsDisplay}} +{{- end}} +` + +// newGetCommand is the `get` parent group. Subcommands name the resource to +// describe, e.g. `air get run JOB_RUN_ID`, mirroring the Python CLI. func newGetCommand() *cobra.Command { cmd := &cobra.Command{ - Use: "get JOB_RUN_ID", + Use: "get", + Short: "Show details for a specific resource", + } + cmd.AddCommand(newGetRunCommand()) + return cmd +} + +func newGetRunCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "run JOB_RUN_ID", Args: root.ExactArgs(1), - Short: "Show details for a run", - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("get") + Short: "Show status, configuration, and timing details for a specific run", + Annotations: map[string]string{ + "template": getTemplate, }, } + // Match Python: a client/auth failure is a JSON error envelope in -o json mode, + // not a bare error. ErrAlreadyPrinted passes through (it was handled upstream). + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + err := root.MustWorkspaceClient(cmd, args) + if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { + return err + } + return renderError(cmd.Context(), cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + runID, err := strconv.ParseInt(args[0], 10, 64) + if err != nil || runID <= 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid JOB_RUN_ID %q: must be a positive integer", args[0])) + } + + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + // The backend returns this when the run ID is unknown to the user. + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, + fmt.Errorf("run %d not found: check the run ID and that it is a job run ID", runID)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to get status for run %d: %w", runID, err)) + } + + workspaceID, err := w.CurrentWorkspaceID(ctx) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to get workspace id for run %d: %w", runID, err)) + } + + data := buildGetData(run) + data.DashboardURL = dashboardURL(w.Config.Host, runID, workspaceID) + ids := mlflowIDs(ctx, w, run) + if ids != nil { + url := mlflowLogsURL(w.Config.Host, ids) + data.MLflowURL = &url + } + if task := findForEachTask(run); task != nil { + data.Sweep = buildSweepInfo(ctx, w, task) + } + + if root.OutputType(cmd) == flags.OutputText { + out := cmd.OutOrStdout() + addTextLinks(ctx, out, w, &data, ids) + + // Lead with the job run link (hyperlinked, falling back to the bare + // URL off a terminal), then a gap before the training config and the + // status table, mirroring the Python CLI's header. + fmt.Fprintf(out, "Job Link: %s\n\n", hyperlink(ctx, out, data.DashboardURL, data.DashboardURL)) + + // Text mode shows the training-config YAML before the status, + // mirroring `air get run`. JSON output omits it. + if path := yamlConfigPath(run); path != "" { + printConfigYAML(ctx, out, w, path) + } + } + return renderEnvelope(ctx, data) + } + return cmd } + +// buildGetData extracts the fields we display from a run. The text-table cells +// are pre-rendered here with their "N/A" fallbacks; the Run ID, Experiment, and +// MLflow Run cells are finalized later by addTextLinks once the dashboard and +// MLflow identifiers are known. +func buildGetData(run *jobs.Run) getData { + data := getData{ + RunID: strconv.FormatInt(run.RunId, 10), + Status: runStatus(run.State), + StartedAt: startedAt(run), + DurationSeconds: durationSeconds(run), + AttemptNumber: latestAttemptNumber(run), + ExperimentName: experimentName(run), + } + data.RunIDDisplay = data.RunID + data.SubmittedDisplay = submittedDisplay(run) + data.DurationDisplay = na + if data.DurationSeconds != nil { + data.DurationDisplay = formatDuration(*data.DurationSeconds) + } + data.ExperimentDisplay = na + if data.ExperimentName != nil { + data.ExperimentDisplay = *data.ExperimentName + } + data.MLflowDisplay = na + data.UserDisplay = orNA(run.CreatorUserName) + data.AcceleratorsDisplay = orNA(accelerators(run)) + return data +} + +// addTextLinks adds the terminal hyperlinks shown in text mode: the Run ID links +// to the run's dashboard page (Python embeds this on the Run ID instead of a +// separate Dashboard row), and the Experiment and MLflow Run cells link to their +// MLflow pages. On a non-terminal these degrade to plain text. +func addTextLinks(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, data *getData, ids *mlflowIdentifiers) { + data.RunIDDisplay = hyperlink(ctx, out, data.RunID, data.DashboardURL) + if ids == nil { + return + } + if data.ExperimentName != nil { + data.ExperimentDisplay = hyperlink(ctx, out, *data.ExperimentName, mlflowExperimentURL(w.Config.Host, ids)) + } + data.MLflowDisplay = hyperlink(ctx, out, mlflowRunLabel(ctx, w, ids.RunID), mlflowRunURL(w.Config.Host, ids)) +} + +// yamlConfigPath returns the run's training-config YAML path, or "" if none. +func yamlConfigPath(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "" + } + task := run.Tasks[0].GenAiComputeTask + if task == nil { + return "" + } + return task.YamlParametersFilePath +} + +// printConfigYAML downloads the run's training-config YAML and writes it to out +// (stdout), mirroring the Python `air get`. It is best-effort: a download or read +// failure is surfaced as a warning on stderr but does not fail the command. +func printConfigYAML(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, path string) { + r, err := w.Workspace.Download(ctx, path) + if err != nil { + log.Warnf(ctx, "air get: could not download training config %s: %v", path, err) + return + } + defer r.Close() + + content, err := io.ReadAll(r) + if err != nil { + log.Warnf(ctx, "air get: could not read training config %s: %v", path, err) + return + } + + fmt.Fprintln(out, "Training Configuration:") + fmt.Fprintln(out, reformatYAMLForDisplay(content)) + fmt.Fprintln(out) +} diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go new file mode 100644 index 00000000000..643c521c0e7 --- /dev/null +++ b/experimental/air/cmd/get_test.go @@ -0,0 +1,267 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "io" + "strings" + "testing" + "text/template" + + "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/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +// renderGet renders the get template against the JSON envelope, exactly as +// the command does, so the test covers the real template branches. +func renderGet(t *testing.T, data getData) string { + t.Helper() + tmpl, err := template.New("get").Parse(getTemplate) + require.NoError(t, err) + var buf bytes.Buffer + require.NoError(t, tmpl.Execute(&buf, envelope{V: envelopeVersion, Data: data})) + return buf.String() +} + +func TestGetTemplateSingleRun(t *testing.T) { + out := renderGet(t, getData{ + RunIDDisplay: "123", + Status: "RUNNING", + UserDisplay: "me@example.com", + }) + assert.Contains(t, out, "Run ID: 123") + assert.Contains(t, out, "Status: RUNNING") + assert.Contains(t, out, "User: me@example.com") + assert.NotContains(t, out, "Sweep") + // Python embeds the dashboard link on the Run ID; there is no Dashboard row. + assert.NotContains(t, out, "Dashboard") +} + +func TestGetRunInvalidID(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newGetRunCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"abc"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid JOB_RUN_ID") +} + +func TestGetRunNotFound(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return( + nil, apierr.ErrResourceDoesNotExist) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newGetRunCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"5"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "run 5 not found") +} + +func TestGetRunNotFoundJSON(t *testing.T) { + var buf bytes.Buffer + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return( + nil, apierr.ErrResourceDoesNotExist) + ctx := cmdctx.SetWorkspaceClient(t.Context(), m.WorkspaceClient) + ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) + cmd := withOutput(newGetRunCommand(), flags.OutputJSON) + cmd.SetContext(ctx) + + // In JSON mode the not-found error is a structured envelope, not a bare error. + err := cmd.RunE(cmd, []string{"5"}) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, jsonError{Code: "NOT_FOUND", Kind: "NOT_FOUND", Message: "run 5 not found: check the run ID and that it is a job run ID"}, got.Error) +} + +func TestPrintConfigYAML(t *testing.T) { + t.Run("downloads and prints to stdout", func(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + m := mocks.NewMockWorkspaceClient(t) + // The mock asserts Download is called with the resolved path. + m.GetMockWorkspaceAPI().EXPECT(). + Download(mock.Anything, "/Workspace/cfg.yaml"). + Return(io.NopCloser(strings.NewReader("epochs: 3\ncommand: \"set -e\\npython train.py\\n\"\n")), nil) + + // The config is data output and must land on stdout (the out writer), + // matching the Python `air get` behavior. + var out bytes.Buffer + printConfigYAML(ctx, &out, m.WorkspaceClient, "/Workspace/cfg.yaml") + assert.Contains(t, out.String(), "Training Configuration:") + assert.Contains(t, out.String(), "epochs: 3") + // The multi-line command is reformatted to a `|` block literal. + assert.Contains(t, out.String(), "command: |\n set -e\n python train.py") + }) + + t.Run("download failure is non-fatal and writes nothing", func(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + m := mocks.NewMockWorkspaceClient(t) + m.GetMockWorkspaceAPI().EXPECT(). + Download(mock.Anything, "/Workspace/missing.yaml"). + Return(nil, apierr.ErrResourceDoesNotExist) + + // Must not panic and must not write to stdout: a failed config fetch is + // best-effort, surfaced only as a stderr warning. + var out bytes.Buffer + printConfigYAML(ctx, &out, m.WorkspaceClient, "/Workspace/missing.yaml") + assert.Empty(t, out.String()) + }) +} + +func TestYAMLConfigPath(t *testing.T) { + // No tasks, or a task without GenAiComputeTask, yields no path. + assert.Empty(t, yamlConfigPath(&jobs.Run{})) + assert.Empty(t, yamlConfigPath(&jobs.Run{Tasks: []jobs.RunTask{{}}})) + + run := &jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{YamlParametersFilePath: "/Workspace/cfg.yaml"}, + }}} + assert.Equal(t, "/Workspace/cfg.yaml", yamlConfigPath(run)) +} + +func TestGetTemplateSweep(t *testing.T) { + out := renderGet(t, getData{ + RunID: "456", + Status: "RUNNING", + Sweep: &sweepInfo{ + Total: 4, Completed: 2, Succeeded: 1, Failed: 1, Active: 2, + Tasks: []sweepTask{ + {TaskKey: "iter_0", RunID: "789", Status: "SUCCESS", Experiment: "my-exp"}, + {TaskKey: "iter_1", RunID: "790", Status: "FAILED", Experiment: "my-exp"}, + }, + }, + }) + assert.Contains(t, out, "Sweep Run ID: 456") + assert.Contains(t, out, "Total: 4") + assert.Contains(t, out, "Sweep Tasks:") + assert.Contains(t, out, "iter_0") + assert.Contains(t, out, "iter_1") + assert.Contains(t, out, "FAILED") + assert.Contains(t, out, "my-exp") + // The single-run rows must not appear in the sweep view. + assert.NotContains(t, out, "Dashboard:") +} + +func TestGetTemplateSweepNoTasks(t *testing.T) { + // A sweep whose iterations haven't materialized yet: counts show, but the + // task table header is hidden. + out := renderGet(t, getData{ + RunID: "456", + Status: "RUNNING", + Sweep: &sweepInfo{Total: 4, Active: 4}, + }) + assert.Contains(t, out, "Sweep Run ID: 456") + assert.Contains(t, out, "Total: 4") + assert.NotContains(t, out, "Sweep Tasks:") +} + +func TestGetTemplateMinimal(t *testing.T) { + // Every row always renders; missing values show "N/A", in Python's order. + out := renderGet(t, getData{ + RunIDDisplay: "1", + Status: "PENDING", + SubmittedDisplay: na, + DurationDisplay: na, + ExperimentDisplay: na, + MLflowDisplay: na, + UserDisplay: na, + AcceleratorsDisplay: na, + }) + for _, want := range []string{ + "Run ID: 1", + "Status: PENDING", + "Submitted: N/A", + "Retries: 0", + "Duration: N/A", + "Experiment: N/A", + "MLflow Run: N/A", + "User: N/A", + "Accelerators: N/A", + } { + assert.Contains(t, out, want) + } +} + +func TestGetTemplateAllFields(t *testing.T) { + out := renderGet(t, getData{ + RunIDDisplay: "1", + Status: "SUCCESS", + SubmittedDisplay: "2023-11-14 22:13 UTC", + DurationDisplay: "12s", + AttemptNumber: 2, + ExperimentDisplay: "exp", + MLflowDisplay: "sunny-cat-42", + UserDisplay: "me@example.com", + AcceleratorsDisplay: "8x H100", + }) + for _, want := range []string{ + "Run ID: 1", + "Status: SUCCESS", + "Submitted: 2023-11-14 22:13 UTC", + "Retries: 2", + "Duration: 12s", + "Experiment: exp", + "MLflow Run: sunny-cat-42", + "User: me@example.com", + "Accelerators: 8x H100", + } { + assert.Contains(t, out, want) + } +} + +func TestBuildGetData(t *testing.T) { + run := &jobs.Run{ + RunId: 123, + CreatorUserName: "me@example.com", + StartTime: 1700000000000, + EndTime: 1700000012000, + State: &jobs.RunState{ResultState: jobs.RunResultStateSuccess}, + Tasks: []jobs.RunTask{{ + AttemptNumber: 1, + GenAiComputeTask: &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/exp", + Compute: &jobs.ComputeConfig{NumGpus: 8, GpuType: "GPU_8xH100"}, + }, + }}, + } + d := buildGetData(run) + assert.Equal(t, "123", d.RunID) + assert.Equal(t, "123", d.RunIDDisplay) + assert.Equal(t, "SUCCESS", d.Status) + assert.Equal(t, 1, d.AttemptNumber) + assert.Equal(t, "2023-11-14 22:13 UTC", d.SubmittedDisplay) + assert.Equal(t, "me@example.com", d.UserDisplay) + assert.Equal(t, "8x H100", d.AcceleratorsDisplay) + assert.Equal(t, "12s", d.DurationDisplay) + assert.Equal(t, "exp", d.ExperimentDisplay) + require.NotNil(t, d.ExperimentName) + assert.Equal(t, "exp", *d.ExperimentName) + require.NotNil(t, d.DurationSeconds) + assert.Equal(t, int64(12), *d.DurationSeconds) +} + +func TestBuildGetDataEmpty(t *testing.T) { + // A run with no tasks, creator, or timing renders every text cell as "N/A". + d := buildGetData(&jobs.Run{RunId: 7}) + assert.Equal(t, "7", d.RunIDDisplay) + assert.Equal(t, na, d.SubmittedDisplay) + assert.Equal(t, na, d.DurationDisplay) + assert.Equal(t, na, d.ExperimentDisplay) + assert.Equal(t, na, d.MLflowDisplay) + assert.Equal(t, na, d.UserDisplay) + assert.Equal(t, na, d.AcceleratorsDisplay) +} diff --git a/experimental/air/cmd/mlflow.go b/experimental/air/cmd/mlflow.go new file mode 100644 index 00000000000..7e2b5b8fd58 --- /dev/null +++ b/experimental/air/cmd/mlflow.go @@ -0,0 +1,126 @@ +package aircmd + +import ( + "context" + "fmt" + "net/http" + "strings" + + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// getRunOutputResponse is the slice of the jobs runs/get-output response we care +// about. The MLflow identifiers live under a gen_ai_compute_output field that +// the typed SDK does not model, so we call the endpoint directly and parse just +// these fields. +type getRunOutputResponse struct { + GenAiComputeOutput *struct { + RunInfo *struct { + MlflowExperimentID string `json:"mlflow_experiment_id"` + MlflowRunID string `json:"mlflow_run_id"` + } `json:"run_info"` + } `json:"gen_ai_compute_output"` +} + +// mlflowIdentifiers are the experiment and run IDs MLflow assigns to a run. +type mlflowIdentifiers struct { + ExperimentID string + RunID string +} + +// mlflowIDs fetches the run's MLflow experiment and run IDs, or nil if they +// can't be obtained. They drive a convenience link, so any failure here +// (missing task, endpoint error, run not yet started) is logged and treated as +// "no link" rather than failing the whole command. +func mlflowIDs(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run) *mlflowIdentifiers { + 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 + + apiClient, err := client.New(w.Config) + if err != nil { + log.Debugf(ctx, "air get: could not build API client for MLflow link: %v", err) + return nil + } + + // Pass run_id through the request arg (the SDK serializes it to the query + // string for GET); passing it via queryParams instead leaves a nil body that + // this endpoint rejects with "expected a map". + var out getRunOutputResponse + err = apiClient.Do(ctx, http.MethodGet, "/api/2.2/jobs/runs/get-output", + nil, nil, map[string]any{"run_id": taskRunID}, &out) + if err != nil { + log.Debugf(ctx, "air get: could not fetch run output for MLflow link: %v", err) + return nil + } + + if out.GenAiComputeOutput == nil || out.GenAiComputeOutput.RunInfo == nil { + return nil + } + info := out.GenAiComputeOutput.RunInfo + if info.MlflowExperimentID == "" || info.MlflowRunID == "" { + return nil + } + return &mlflowIdentifiers{ExperimentID: info.MlflowExperimentID, RunID: info.MlflowRunID} +} + +// mlflowLogsURL is the deep link to a run's node-0 logs. It is the value of the +// JSON `mlflow_url` field, matching the Python CLI. +func mlflowLogsURL(host string, ids *mlflowIdentifiers) string { + return fmt.Sprintf("%s/ml/experiments/%s/runs/%s/artifacts/logs/node_0", + strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) +} + +// mlflowExperimentURL links to the MLflow experiment page; mlflowRunURL links to +// the run page. These back the Experiment and MLflow Run hyperlinks in text mode. +func mlflowExperimentURL(host string, ids *mlflowIdentifiers) string { + return fmt.Sprintf("%s/ml/experiments/%s", strings.TrimRight(host, "/"), ids.ExperimentID) +} + +func mlflowRunURL(host string, ids *mlflowIdentifiers) string { + return fmt.Sprintf("%s/ml/experiments/%s/runs/%s", + strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) +} + +// mlflowRunLabel returns the MLflow run's human-readable name to use as the +// hyperlink text, falling back to "...{last 8 of run id}" when the name can't be +// fetched. Mirrors Python's _get_mlflow_run_name (cli_display.py). +func mlflowRunLabel(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) string { + if name := fetchMLflowRunName(ctx, w, mlflowRunID); name != "" { + return name + } + if len(mlflowRunID) > 8 { + return "..." + mlflowRunID[len(mlflowRunID)-8:] + } + return "..." + mlflowRunID +} + +// 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. +func fetchMLflowRunName(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) string { + apiClient, err := client.New(w.Config) + if err != nil { + log.Debugf(ctx, "air get: could not build API client for MLflow run name: %v", err) + return "" + } + var out struct { + Run struct { + Info struct { + RunName string `json:"run_name"` + } `json:"info"` + } `json:"run"` + } + err = apiClient.Do(ctx, http.MethodGet, "/api/2.0/mlflow/runs/get", + nil, nil, map[string]any{"run_id": mlflowRunID}, &out) + if err != nil { + log.Debugf(ctx, "air get: could not fetch MLflow run name: %v", err) + return "" + } + return out.Run.Info.RunName +} diff --git a/experimental/air/cmd/mlflow_test.go b/experimental/air/cmd/mlflow_test.go new file mode 100644 index 00000000000..a6a5c151074 --- /dev/null +++ b/experimental/air/cmd/mlflow_test.go @@ -0,0 +1,112 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestWorkspaceClient builds a WorkspaceClient pointed at a mock HTTP server. +// mlflowIDs calls the runs/get-output REST endpoint directly (the field it needs +// is not modeled by the typed SDK), so it must be exercised over HTTP. +func newTestWorkspaceClient(t *testing.T, host string) *databricks.WorkspaceClient { + t.Helper() + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: host, Token: "token"}) + require.NoError(t, err) + return w +} + +// runOutputServer serves the given runs/get-output body and a stub for the SDK's +// well-known config discovery request. *hit is set when get-output is called. +func runOutputServer(t *testing.T, body string, hit *bool) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.2/jobs/runs/get-output" { + *hit = true + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestMLflowIDs(t *testing.T) { + ctx := t.Context() + run := &jobs.Run{Tasks: []jobs.RunTask{{RunId: 99}}} + + t.Run("returns the identifiers on success", func(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{"gen_ai_compute_output":{"run_info":{"mlflow_experiment_id":"E1","mlflow_run_id":"R1"}}}`, &hit) + + got := mlflowIDs(ctx, newTestWorkspaceClient(t, srv.URL), run) + require.NotNil(t, got) + assert.True(t, hit, "runs/get-output should have been called") + assert.Equal(t, &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"}, got) + }) + + t.Run("nil when the run has no MLflow info", func(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{}`, &hit) + assert.Nil(t, mlflowIDs(ctx, newTestWorkspaceClient(t, srv.URL), run)) + }) + + t.Run("nil when the run has no tasks", func(t *testing.T) { + // Returns before any HTTP call, so the host is never contacted. + assert.Nil(t, mlflowIDs(ctx, newTestWorkspaceClient(t, "https://unused.invalid"), &jobs.Run{})) + }) + + t.Run("uses the latest attempt's task run", func(t *testing.T) { + // A retried run must link to the last task, not the stale first attempt. + var gotRunID string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.2/jobs/runs/get-output" { + gotRunID = r.URL.Query().Get("run_id") + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + retried := &jobs.Run{Tasks: []jobs.RunTask{{RunId: 99}, {RunId: 100}}} + mlflowIDs(ctx, newTestWorkspaceClient(t, srv.URL), retried) + assert.Equal(t, "100", gotRunID) + }) +} + +func TestMLflowURLs(t *testing.T) { + ids := &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"} + // A trailing slash on the host must not produce a double slash in the link. + assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1/artifacts/logs/node_0", mlflowLogsURL("https://h.test/", ids)) + assert.Equal(t, "https://h.test/ml/experiments/E1", mlflowExperimentURL("https://h.test", ids)) + assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1", mlflowRunURL("https://h.test", ids)) +} + +func TestMLflowRunLabel(t *testing.T) { + ctx := t.Context() + + t.Run("uses the run name when available", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.0/mlflow/runs/get" { + _, _ = w.Write([]byte(`{"run":{"info":{"run_name":"sunny-cat-42"}}}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + assert.Equal(t, "sunny-cat-42", mlflowRunLabel(ctx, newTestWorkspaceClient(t, srv.URL), "0123456789abcdef")) + }) + + t.Run("falls back to the last 8 characters of the run id", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + assert.Equal(t, "...9abcdef0", mlflowRunLabel(ctx, newTestWorkspaceClient(t, srv.URL), "0123456789abcdef0")) + }) +} diff --git a/experimental/air/cmd/output.go b/experimental/air/cmd/output.go new file mode 100644 index 00000000000..c00d870e386 --- /dev/null +++ b/experimental/air/cmd/output.go @@ -0,0 +1,82 @@ +package aircmd + +import ( + "context" + "time" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/spf13/cobra" +) + +// envelopeVersion is the envelope's format-version marker. The Python `air` CLI +// hardcodes it to 1; it lets consumers detect a future incompatible change to +// the envelope shape. +const envelopeVersion = 1 + +// envelope is the JSON shape that the AI runtime CLI prints: +// +// { "v": 1, "ts": "2024-01-15T14:30:45Z", "data": { ... } } +// +// It mirrors the envelope used by the original Python `air` CLI so existing +// consumers keep working after the port to Go. +type envelope struct { + // V is the envelope format-version marker (always 1). + V int `json:"v"` + // TS is the wall-clock time the response was produced, in RFC 3339 UTC. + // It is an absolute timestamp, not an elapsed duration. + TS string `json:"ts"` + // Data is the command-specific payload. + Data any `json:"data"` +} + +// renderEnvelope wraps data in the JSON envelope and prints it. +// Fields that should appear only in text output are tagged `json:"-"` on the payload struct. +func renderEnvelope(ctx context.Context, data any) error { + return cmdio.Render(ctx, envelope{ + V: envelopeVersion, + TS: time.Now().UTC().Format(time.RFC3339), + Data: data, + }) +} + +// jsonError is the error payload, matching the Python `air` CLI's shape (cli/json_output.py). +type jsonError struct { + Code string `json:"code"` + Kind string `json:"kind"` + Message string `json:"message"` + Retryable bool `json:"retryable"` +} + +// errorEnvelope is what a failed command prints in JSON mode: +// +// { "v": 1, "ts": "...", "error": { "code": ..., "kind": ..., "message": ..., "retryable": ... } } +type errorEnvelope struct { + V int `json:"v"` + TS string `json:"ts"` + Error jsonError `json:"error"` +} + +// renderError prints err as a JSON error envelope when output is JSON, returning +// root.ErrAlreadyPrinted so the command exits non-zero without Cobra reprinting +// it; in text mode it returns err unchanged. code/kind/retryable match the +// Python CLI's call site. +func renderError(ctx context.Context, cmd *cobra.Command, code, kind string, retryable bool, err error) error { + if root.OutputType(cmd) != flags.OutputJSON { + return err + } + if rerr := cmdio.Render(ctx, errorEnvelope{ + V: envelopeVersion, + TS: time.Now().UTC().Format(time.RFC3339), + Error: jsonError{ + Code: code, + Kind: kind, + Message: err.Error(), + Retryable: retryable, + }, + }); rerr != nil { + return rerr + } + return root.ErrAlreadyPrinted +} diff --git a/experimental/air/cmd/output_test.go b/experimental/air/cmd/output_test.go new file mode 100644 index 00000000000..3c35dc60a67 --- /dev/null +++ b/experimental/air/cmd/output_test.go @@ -0,0 +1,52 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "errors" + "testing" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRenderEnvelope(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + require.NoError(t, renderEnvelope(ctx, getData{RunID: "1", Status: "RUNNING"})) +} + +// withOutput registers the --output flag on cmd and sets it, mirroring how the +// root command wires output mode in production. Subcommand unit tests need it +// because they invoke RunE without going through the root command. +func withOutput(cmd *cobra.Command, output flags.Output) *cobra.Command { + cmd.Flags().Var(&output, "output", "") + return cmd +} + +func TestRenderErrorJSON(t *testing.T) { + var buf bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputJSON, nil, &buf, &buf, "", "")) + cmd := withOutput(&cobra.Command{}, flags.OutputJSON) + + err := renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, errors.New("run 1 not found")) + // JSON mode prints the envelope, so Cobra must stay silent but still exit non-zero. + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + // The envelope must match the Python air CLI's print_json_error shape exactly. + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, 1, got.V) + assert.NotEmpty(t, got.TS) + assert.Equal(t, jsonError{Code: "NOT_FOUND", Kind: "NOT_FOUND", Message: "run 1 not found"}, got.Error) +} + +func TestRenderErrorText(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + want := errors.New("run 1 not found") + require.Equal(t, want, renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, want)) +} diff --git a/experimental/air/cmd/stubs_test.go b/experimental/air/cmd/stubs_test.go index a6e24177f33..5e35bcdcd14 100644 --- a/experimental/air/cmd/stubs_test.go +++ b/experimental/air/cmd/stubs_test.go @@ -14,7 +14,6 @@ import ( func TestStubCommandsReturnNotImplemented(t *testing.T) { stubs := map[string]*cobra.Command{ "run": newRunCommand(), - "get": newGetCommand(), "list": newListCommand(), "logs": newLogsCommand(), "cancel": newCancelCommand(), diff --git a/experimental/air/cmd/sweep.go b/experimental/air/cmd/sweep.go new file mode 100644 index 00000000000..b346f43f1b6 --- /dev/null +++ b/experimental/air/cmd/sweep.go @@ -0,0 +1,76 @@ +package aircmd + +import ( + "context" + "strconv" + + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// sweepInfo summarizes a "foreach" run, which fans a single config out into many +// iterations (a hyperparameter sweep). It is shown only in text output. +type sweepInfo struct { + Total int + Succeeded int + Failed int + Active int + Completed int + Tasks []sweepTask +} + +// sweepTask is one iteration of a sweep. +type sweepTask struct { + TaskKey string + RunID string + Status string + Experiment string +} + +// findForEachTask returns the run's foreach task if it has one, or nil. A run is +// a sweep when one of its tasks fans out into iterations. +func findForEachTask(run *jobs.Run) *jobs.RunTask { + for i := range run.Tasks { + if run.Tasks[i].ForEachTask != nil { + return &run.Tasks[i] + } + } + return nil +} + +// buildSweepInfo gathers the iteration counts and per-iteration rows for a +// sweep. The counts come from the task we already have; the individual +// iterations require a second lookup. If that lookup fails we still return the +// counts (logging the failure) so the user sees the summary. +func buildSweepInfo(ctx context.Context, w *databricks.WorkspaceClient, task *jobs.RunTask) *sweepInfo { + info := &sweepInfo{} + if task.ForEachTask.Stats != nil && task.ForEachTask.Stats.TaskRunStats != nil { + stats := task.ForEachTask.Stats.TaskRunStats + info.Total = stats.TotalIterations + info.Succeeded = stats.SucceededIterations + info.Failed = stats.FailedIterations + info.Active = stats.ActiveIterations + info.Completed = stats.CompletedIterations + } + + // The iterations are returned as part of a run lookup on the foreach task. + iterated, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: task.RunId}) + if err != nil { + log.Debugf(ctx, "air get: could not fetch sweep iterations: %v", err) + return info + } + + for _, it := range iterated.Iterations { + row := sweepTask{ + TaskKey: it.TaskKey, + RunID: strconv.FormatInt(it.RunId, 10), + Status: runStatus(it.State), + } + if it.GenAiComputeTask != nil && it.GenAiComputeTask.MlflowExperimentName != "" { + row.Experiment = stripExperimentUserPrefix(it.GenAiComputeTask.MlflowExperimentName) + } + info.Tasks = append(info.Tasks, row) + } + return info +} diff --git a/experimental/air/cmd/sweep_test.go b/experimental/air/cmd/sweep_test.go new file mode 100644 index 00000000000..10134c0df42 --- /dev/null +++ b/experimental/air/cmd/sweep_test.go @@ -0,0 +1,81 @@ +package aircmd + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +func TestFindForEachTask(t *testing.T) { + // No tasks at all. + assert.Nil(t, findForEachTask(&jobs.Run{})) + + // A task that is not a foreach. + assert.Nil(t, findForEachTask(&jobs.Run{Tasks: []jobs.RunTask{{TaskKey: "a"}}})) + + // The foreach task is found even when it isn't first. + run := &jobs.Run{Tasks: []jobs.RunTask{ + {TaskKey: "a"}, + {TaskKey: "sweep", ForEachTask: &jobs.RunForEachTask{}}, + }} + got := findForEachTask(run) + require.NotNil(t, got) + assert.Equal(t, "sweep", got.TaskKey) +} + +func sweepTaskFixture() *jobs.RunTask { + return &jobs.RunTask{ + RunId: 99, + ForEachTask: &jobs.RunForEachTask{ + Stats: &jobs.ForEachStats{TaskRunStats: &jobs.ForEachTaskTaskRunStats{ + TotalIterations: 4, + SucceededIterations: 1, + FailedIterations: 1, + ActiveIterations: 2, + CompletedIterations: 2, + }}, + }, + } +} + +func TestBuildSweepInfo(t *testing.T) { + ctx := t.Context() + + t.Run("counts and iteration rows", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 99}).Return( + &jobs.Run{Iterations: []jobs.RunTask{{ + TaskKey: "iter_0", + RunId: 100, + State: &jobs.RunState{ResultState: jobs.RunResultStateSuccess}, + GenAiComputeTask: &jobs.GenAiComputeTask{MlflowExperimentName: "/Users/me@example.com/exp"}, + }}}, nil) + + info := buildSweepInfo(ctx, m.WorkspaceClient, sweepTaskFixture()) + assert.Equal(t, 4, info.Total) + assert.Equal(t, 2, info.Completed) + assert.Equal(t, 1, info.Succeeded) + assert.Equal(t, 1, info.Failed) + assert.Equal(t, 2, info.Active) + require.Len(t, info.Tasks, 1) + assert.Equal(t, "iter_0", info.Tasks[0].TaskKey) + assert.Equal(t, "100", info.Tasks[0].RunID) + assert.Equal(t, "SUCCESS", info.Tasks[0].Status) + assert.Equal(t, "exp", info.Tasks[0].Experiment) + }) + + t.Run("iteration lookup failure still returns counts", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 99}).Return( + nil, apierr.ErrResourceDoesNotExist) + + info := buildSweepInfo(ctx, m.WorkspaceClient, sweepTaskFixture()) + assert.Equal(t, 4, info.Total) + assert.Empty(t, info.Tasks) + }) +} From f1601b2278fc4daa9cb59a7f11daed249a32f301 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Thu, 18 Jun 2026 09:51:01 -0700 Subject: [PATCH 03/18] AIR CLI Integration: `air run` Command Pt. 1 - Add GPU accelerator type and compute config model (#5602) ## Changes Adds `experimental/air/cmd/compute.go` , which is the `gpuType` model and `compute` which is the block validation that the `air run` configuration layer depends on. Specifically: - the training service accelerator types were added (`GPU_1xA10`, `GPU_8xH100`, `GPU_1xH100`) - `parseGPUType` resolves a YAML accelerator type string - `gpusPerNode` is the per node partition count based on the type name - `computeConfig` and `validate()` are the port of the python `ComputeConfig` validators ## Why This is the first, leaf-most piece of the `air run` port for the AIR CLI and the root of the config validation layer dependencies. This piece for compute does not depend on anything else so it lands first as a small and fully unit-tested unit. Note that we also use exact case sensitive parsing since a potential typo in the user's YAML could misroute the run. Additionally, we only support `GPU_*` training service types (legacy MAPI types (eg. `h100_80gb`) are no longer supported and intentionally deprecated in this port. However, they still have their own display map for historical runs to be able to be displayed (but no new runs can use the MAPI path). Rendering them in get is unaffected since format.go keeps its own display map for historical runs. ## Tests Table-driven unit tests in compute_test.go: parseGPUType for valid types and rejected inputs (wrong casing, legacy types, unknown, empty); gpusPerNode counts plus its invalid-type error; and computeConfig.validate across valid configs and every failure mode (unknown/legacy type, non-positive count, non-multiple count, dual-pool conflict). go build, go test, and golangci-lint are clean. --- experimental/air/cmd/compute.go | 81 ++++++++++++++++++++++++++ experimental/air/cmd/compute_test.go | 86 ++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+) create mode 100644 experimental/air/cmd/compute.go create mode 100644 experimental/air/cmd/compute_test.go diff --git a/experimental/air/cmd/compute.go b/experimental/air/cmd/compute.go new file mode 100644 index 00000000000..07013c53906 --- /dev/null +++ b/experimental/air/cmd/compute.go @@ -0,0 +1,81 @@ +package aircmd + +import ( + "fmt" + "strings" +) + +// gpuType is a wire-facing accelerator type submitted to the training service. +// The number in the name is the partition count (e.g. GPU_8xH100 is 8 GPUs). +type gpuType string + +const ( + gpuType1xA10 gpuType = "GPU_1xA10" + gpuType8xH100 gpuType = "GPU_8xH100" + gpuType1xH100 gpuType = "GPU_1xH100" +) + +// gpuTypes lists every valid type. Used for validation error messages. +var gpuTypes = []gpuType{gpuType1xA10, gpuType1xH100, gpuType8xH100} + +func validGPUTypesHint() string { + names := make([]string, len(gpuTypes)) + for i, g := range gpuTypes { + names[i] = string(g) + } + return "valid types are: " + strings.Join(names, ", ") +} + +// parseGPUType resolves a YAML accelerator_type string to a gpuType. The match is +// exact: the server's lookup is case-sensitive. +func parseGPUType(value string) (gpuType, error) { + switch gpuType(value) { + case gpuType1xA10, gpuType8xH100, gpuType1xH100: + return gpuType(value), nil + } + return "", fmt.Errorf("invalid GPU type %q: %s", value, validGPUTypesHint()) +} + +// gpusPerNode returns the per-node GPU count, which is the partition count from +// the name (GPU_1xH100 -> 1, GPU_8xH100 -> 8). num_accelerators must be a +// round multiple of this since accelerators are allocated in whole nodes. +func gpusPerNode(g gpuType) (int, error) { + switch g { + case gpuType1xA10, gpuType1xH100: + return 1, nil + case gpuType8xH100: + return 8, nil + } + // Unreachable: callers resolve g through parseGPUType first, which rejects + // unknown types. Kept as a defensive guard. + return 0, fmt.Errorf("invalid GPU type %q", string(g)) +} + +// computeConfig is the `compute` block of the run YAML: which accelerators to +// use and how many. +type computeConfig struct { + NumAccelerators int `yaml:"num_accelerators"` + AcceleratorType string `yaml:"accelerator_type"` +} + +// validate checks the compute block against the backend's constraints. +func (c computeConfig) validate() error { + g, err := parseGPUType(c.AcceleratorType) + if err != nil { + return fmt.Errorf("compute.accelerator_type: %w", err) + } + + if c.NumAccelerators <= 0 { + return fmt.Errorf("compute.num_accelerators must be positive, got %d", c.NumAccelerators) + } + + perNode, err := gpusPerNode(g) + if err != nil { + return err + } + if c.NumAccelerators%perNode != 0 { + return fmt.Errorf("compute.num_accelerators for %s must be a multiple of %d, got %d", c.AcceleratorType, perNode, c.NumAccelerators) + } + + return nil +} diff --git a/experimental/air/cmd/compute_test.go b/experimental/air/cmd/compute_test.go new file mode 100644 index 00000000000..3464afbe9ea --- /dev/null +++ b/experimental/air/cmd/compute_test.go @@ -0,0 +1,86 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseGPUType(t *testing.T) { + tests := []struct { + in string + want gpuType + }{ + {"GPU_1xA10", gpuType1xA10}, + {"GPU_8xH100", gpuType8xH100}, + {"GPU_1xH100", gpuType1xH100}, + } + for _, tt := range tests { + t.Run(tt.in, func(t *testing.T) { + got, err := parseGPUType(tt.in) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestParseGPUTypeInvalid(t *testing.T) { + // Wrong casing is rejected rather than fixed up; legacy types (h100_80gb, a10) + // can no longer be submitted; unknown types are rejected. + for _, in := range []string{"gpu_1xa10", "GPU_1XA10", "GPU_2xH100", "h100_80gb", "a10", "b200", ""} { + t.Run(in, func(t *testing.T) { + _, err := parseGPUType(in) + require.Error(t, err) + assert.Contains(t, err.Error(), "valid types are") + }) + } +} + +func TestGPUsPerNode(t *testing.T) { + tests := []struct { + in gpuType + want int + }{ + {gpuType1xA10, 1}, + {gpuType1xH100, 1}, + {gpuType8xH100, 8}, + } + for _, tt := range tests { + t.Run(string(tt.in), func(t *testing.T) { + got, err := gpusPerNode(tt.in) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } + + _, err := gpusPerNode(gpuType("nonsense")) + require.Error(t, err) +} + +func TestComputeConfigValidate(t *testing.T) { + tests := []struct { + name string + cfg computeConfig + wantErr string // substring; empty means the config is valid + }{ + {"single node", computeConfig{NumAccelerators: 8, AcceleratorType: "GPU_8xH100"}, ""}, + {"multiple nodes", computeConfig{NumAccelerators: 16, AcceleratorType: "GPU_8xH100"}, ""}, + {"single-gpu partitions", computeConfig{NumAccelerators: 3, AcceleratorType: "GPU_1xH100"}, ""}, + {"unknown type", computeConfig{NumAccelerators: 8, AcceleratorType: "b200"}, "accelerator_type"}, + {"legacy type rejected", computeConfig{NumAccelerators: 8, AcceleratorType: "h100_80gb"}, "accelerator_type"}, + {"non-positive count", computeConfig{NumAccelerators: 0, AcceleratorType: "GPU_1xH100"}, "must be positive"}, + {"count not a multiple", computeConfig{NumAccelerators: 4, AcceleratorType: "GPU_8xH100"}, "multiple of 8"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.cfg.validate() + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} From e118e673e9bbeafdc93c8d89ee420ac1aa69bc3a Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Tue, 23 Jun 2026 15:58:16 -0700 Subject: [PATCH 04/18] AIR CLI Integration: render `air get run` as styled boxes (#5654) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What Replaces the plain-text view of `air get run ` with a one-shot, styled terminal renderer built on **lipgloss** (layout/styling) and **termenv** (hyperlinks + color-profile detection). It builds the full string and writes it once — no streaming, spinner, or redraw. The view is two boxes: - **Configuration** — the resolved run config YAML (inline `yaml_parameters`, the downloaded `yaml_parameters_file_path`, or a synthesized fallback), colorized line by line. - **Metadata** — Run ID, Status, Submitted, Retries, Max Retries, Duration, Experiment, MLflow Run, User, Accelerators, Environment. Run ID and MLflow Run are OSC 8 hyperlinks. ## Look & feel - Boxes share a light-purple border/title, warm Oat neutrals, and a restrained accent palette (blue for keys/links; green/amber/red reserved for the status dot). - Honors `--no-color` / `NO_COLOR` / non-TTY via `termenv.Ascii`: no escape codes, and links degrade to the bare label (the URLs remain available in `-o json` as `dashboard_url` / `mlflow_url`). ## Scope - Sweep (foreach) runs and JSON output are unchanged. - `termenv` becomes a direct dependency (annotated `// MIT` in `go.mod`, added to `NOTICE`). ## Testing - Unit tests in `render_test.go` / `mlflow_test.go` cover the box, field list, link fallback, config sourcing, and the MLflow run-name fetch. - Acceptance output regenerated (`acceptance/experimental/air/get`). - `go build ./...`, `./task lint-q` (0 issues), and the air + acceptance suites pass. This pull request and its description were written by Isaac. --------- Co-authored-by: Maggie Wang <141875985+maggiewang-db@users.noreply.github.com> --- NOTICE | 4 + acceptance/experimental/air/get/output.txt | 43 ++- acceptance/experimental/air/get/test.toml | 5 +- experimental/air/cmd/format.go | 27 ++ experimental/air/cmd/get.go | 122 ++----- experimental/air/cmd/get_test.go | 125 +------ experimental/air/cmd/mlflow.go | 38 +-- experimental/air/cmd/mlflow_test.go | 30 +- experimental/air/cmd/render.go | 379 +++++++++++++++++++++ experimental/air/cmd/render_test.go | 162 +++++++++ go.mod | 2 +- 11 files changed, 677 insertions(+), 260 deletions(-) create mode 100644 experimental/air/cmd/render.go create mode 100644 experimental/air/cmd/render_test.go diff --git a/NOTICE b/NOTICE index 922dd1fd79a..dd1e4c0c1b6 100644 --- a/NOTICE +++ b/NOTICE @@ -155,6 +155,10 @@ mattn/go-isatty - https://github.com/mattn/go-isatty Copyright (c) Yasuhiro MATSUMOTO License - https://github.com/mattn/go-isatty/blob/master/LICENSE +muesli/termenv - https://github.com/muesli/termenv +Copyright (c) 2019 Christian Muehlhaeuser +License - https://github.com/muesli/termenv/blob/master/LICENSE + sabhiram/go-gitignore - https://github.com/sabhiram/go-gitignore Copyright (c) 2015 Shaba Abhiram License - https://github.com/sabhiram/go-gitignore/blob/master/LICENSE diff --git a/acceptance/experimental/air/get/output.txt b/acceptance/experimental/air/get/output.txt index f8df1fb36bb..4bc127d79b4 100644 --- a/acceptance/experimental/air/get/output.txt +++ b/acceptance/experimental/air/get/output.txt @@ -1,17 +1,38 @@ === get (text) >>> [CLI] experimental air get run 123 -Job Link: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] - -Run ID: 123 -Status: SUCCESS -Submitted: 2023-11-14 22:13 UTC -Retries: 0 -Duration: 12s -Experiment: my-exp -MLflow Run: my-run -User: user@example.com -Accelerators: 8x H100 + +╭─ Configuration ────────────────────────────────────────────────╮ +│ │ +│ experiment_name: my-exp │ +│ compute: │ +│ accelerator_type: a10 │ +│ num_accelerators: 1 │ +│ command: | │ +│ for i in $(seq 1 10); do │ +│ echo "step $i" │ +│ done │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +╭─ Metadata ─────────────────────────────────────────────────────╮ +│ │ +│ Run ID 123 │ +│ Status ● SUCCESS │ +│ Submitted 2023-11-14 22:13 UTC │ +│ Retries 0 │ +│ Max Retries 3 │ +│ Duration 12s │ +│ Experiment my-exp │ +│ MLflow Run my-run │ +│ User user@example.com │ +│ Accelerators 1x A10 │ +│ Environment ml-runtime-gpu:1.0 │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +Run URL: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] +MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 === get (json) >>> [CLI] experimental air get run 123 -o json diff --git a/acceptance/experimental/air/get/test.toml b/acceptance/experimental/air/get/test.toml index 6162cebdb5d..456feb39dbc 100644 --- a/acceptance/experimental/air/get/test.toml +++ b/acceptance/experimental/air/get/test.toml @@ -23,9 +23,12 @@ Response.Body = ''' { "task_key": "train", "attempt_number": 0, + "max_retries": 3, "gen_ai_compute_task": { "mlflow_experiment_name": "/Users/user@example.com/my-exp", - "compute": {"gpu_type": "GPU_8xH100", "num_gpus": 8} + "compute": {"gpu_type": "GPU_1xA10", "num_gpus": 1}, + "dl_runtime_image": "ml-runtime-gpu:1.0", + "yaml_parameters": "experiment_name: my-exp\ncompute:\n accelerator_type: a10\n num_accelerators: 1\ncommand: |\n for i in $(seq 1 10); do\n echo \"step $i\"\n done\n" } } ] diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index 8694ed69e9f..57f86b80f07 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "math" + "strconv" "strings" "time" @@ -263,3 +264,29 @@ func gpuDisplayName(gpuType string) string { } return gpuType } + +// environment returns the run's runtime image (the training environment), or an +// empty string if the run has no GenAI-compute task. +func environment(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "" + } + task := run.Tasks[0].GenAiComputeTask + if task == nil { + return "" + } + return task.DlRuntimeImage +} + +// maxRetries returns the configured retry limit for the run's latest task as a +// display string: "unlimited" for the backend's -1, otherwise the count. +func maxRetries(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "0" + } + n := run.Tasks[len(run.Tasks)-1].MaxRetries + if n < 0 { + return "unlimited" + } + return strconv.Itoa(n) +} diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index a302b7caf31..98581fd5dd7 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -1,17 +1,13 @@ package aircmd import ( - "context" "errors" "fmt" - "io" "strconv" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/flags" - "github.com/databricks/cli/libs/log" - "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/spf13/cobra" @@ -30,26 +26,27 @@ type getData struct { DashboardURL string `json:"dashboard_url"` MLflowURL *string `json:"mlflow_url"` - // The fields below are pre-rendered for the text table and excluded from - // JSON (matching `air get run --json`). The table always shows every row, - // with "N/A" for missing values, in the same order as the Python CLI. The - // Run ID, Experiment, and MLflow Run cells carry terminal hyperlinks when - // stdout is a terminal, so the URLs don't appear as bare text. - RunIDDisplay string `json:"-"` + // The fields below are pre-rendered text-view cells, excluded from JSON + // (matching `air get run --json`). Each shows "N/A" when its value is + // missing. The styled single-run renderer (render.go) consumes them; the + // Run ID, Status, and MLflow Run cells it draws are styled and hyperlinked + // there rather than stored here. SubmittedDisplay string `json:"-"` DurationDisplay string `json:"-"` ExperimentDisplay string `json:"-"` - MLflowDisplay string `json:"-"` UserDisplay string `json:"-"` AcceleratorsDisplay string `json:"-"` + EnvironmentDisplay string `json:"-"` + MaxRetriesDisplay string `json:"-"` // Sweep replaces the single-run view for foreach runs. Sweep *sweepInfo `json:"-"` } -// getTemplate is the text-mode layout. It reads from the JSON envelope, so -// every field is reached through ".Data". -const getTemplate = `{{- if .Data.Sweep -}} -Sweep Run ID: {{.Data.RunID}} +// getTemplate is the text-mode layout for a sweep (foreach) run. Single runs are +// drawn by the styled renderer in render.go and never reach this template; it is +// used only when .Data.Sweep is set. It reads from the JSON envelope, so every +// field is reached through ".Data". +const getTemplate = `Sweep Run ID: {{.Data.RunID}} Status: {{.Data.Status}} Total: {{.Data.Sweep.Total}} Completed: {{.Data.Sweep.Completed}} @@ -64,17 +61,6 @@ Sweep Tasks: {{printf " %-24s %-14s %-12s %s" .TaskKey .RunID .Status .Experiment}} {{- end}} {{- end}} -{{- else -}} -Run ID: {{.Data.RunIDDisplay}} -Status: {{.Data.Status}} -Submitted: {{.Data.SubmittedDisplay}} -Retries: {{.Data.AttemptNumber}} -Duration: {{.Data.DurationDisplay}} -Experiment: {{.Data.ExperimentDisplay}} -MLflow Run: {{.Data.MLflowDisplay}} -User: {{.Data.UserDisplay}} -Accelerators: {{.Data.AcceleratorsDisplay}} -{{- end}} ` // newGetCommand is the `get` parent group. Subcommands name the resource to @@ -146,31 +132,28 @@ func newGetRunCommand() *cobra.Command { data.Sweep = buildSweepInfo(ctx, w, task) } - if root.OutputType(cmd) == flags.OutputText { - out := cmd.OutOrStdout() - addTextLinks(ctx, out, w, &data, ids) + if root.OutputType(cmd) != flags.OutputText { + return renderEnvelope(ctx, data) + } - // Lead with the job run link (hyperlinked, falling back to the bare - // URL off a terminal), then a gap before the training config and the - // status table, mirroring the Python CLI's header. + out := cmd.OutOrStdout() + if data.Sweep != nil { + // A sweep has no single status, config, or timing, so lead with the + // job run link and render the foreach summary table (getTemplate). fmt.Fprintf(out, "Job Link: %s\n\n", hyperlink(ctx, out, data.DashboardURL, data.DashboardURL)) - - // Text mode shows the training-config YAML before the status, - // mirroring `air get run`. JSON output omits it. - if path := yamlConfigPath(run); path != "" { - printConfigYAML(ctx, out, w, path) - } + return renderEnvelope(ctx, data) } - return renderEnvelope(ctx, data) + + renderRunText(ctx, out, w, run, &data, ids) + return nil } return cmd } -// buildGetData extracts the fields we display from a run. The text-table cells -// are pre-rendered here with their "N/A" fallbacks; the Run ID, Experiment, and -// MLflow Run cells are finalized later by addTextLinks once the dashboard and -// MLflow identifiers are known. +// buildGetData extracts the fields we display from a run. The text-view cells +// are pre-rendered here with their "N/A" fallbacks; the styled renderer adds the +// hyperlinks and colors once the dashboard and MLflow identifiers are known. func buildGetData(run *jobs.Run) getData { data := getData{ RunID: strconv.FormatInt(run.RunId, 10), @@ -180,7 +163,6 @@ func buildGetData(run *jobs.Run) getData { AttemptNumber: latestAttemptNumber(run), ExperimentName: experimentName(run), } - data.RunIDDisplay = data.RunID data.SubmittedDisplay = submittedDisplay(run) data.DurationDisplay = na if data.DurationSeconds != nil { @@ -190,57 +172,9 @@ func buildGetData(run *jobs.Run) getData { if data.ExperimentName != nil { data.ExperimentDisplay = *data.ExperimentName } - data.MLflowDisplay = na data.UserDisplay = orNA(run.CreatorUserName) data.AcceleratorsDisplay = orNA(accelerators(run)) + data.EnvironmentDisplay = orNA(environment(run)) + data.MaxRetriesDisplay = maxRetries(run) return data } - -// addTextLinks adds the terminal hyperlinks shown in text mode: the Run ID links -// to the run's dashboard page (Python embeds this on the Run ID instead of a -// separate Dashboard row), and the Experiment and MLflow Run cells link to their -// MLflow pages. On a non-terminal these degrade to plain text. -func addTextLinks(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, data *getData, ids *mlflowIdentifiers) { - data.RunIDDisplay = hyperlink(ctx, out, data.RunID, data.DashboardURL) - if ids == nil { - return - } - if data.ExperimentName != nil { - data.ExperimentDisplay = hyperlink(ctx, out, *data.ExperimentName, mlflowExperimentURL(w.Config.Host, ids)) - } - data.MLflowDisplay = hyperlink(ctx, out, mlflowRunLabel(ctx, w, ids.RunID), mlflowRunURL(w.Config.Host, ids)) -} - -// yamlConfigPath returns the run's training-config YAML path, or "" if none. -func yamlConfigPath(run *jobs.Run) string { - if len(run.Tasks) == 0 { - return "" - } - task := run.Tasks[0].GenAiComputeTask - if task == nil { - return "" - } - return task.YamlParametersFilePath -} - -// printConfigYAML downloads the run's training-config YAML and writes it to out -// (stdout), mirroring the Python `air get`. It is best-effort: a download or read -// failure is surfaced as a warning on stderr but does not fail the command. -func printConfigYAML(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, path string) { - r, err := w.Workspace.Download(ctx, path) - if err != nil { - log.Warnf(ctx, "air get: could not download training config %s: %v", path, err) - return - } - defer r.Close() - - content, err := io.ReadAll(r) - if err != nil { - log.Warnf(ctx, "air get: could not read training config %s: %v", path, err) - return - } - - fmt.Fprintln(out, "Training Configuration:") - fmt.Fprintln(out, reformatYAMLForDisplay(content)) - fmt.Fprintln(out) -} diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go index 643c521c0e7..7f4d7848076 100644 --- a/experimental/air/cmd/get_test.go +++ b/experimental/air/cmd/get_test.go @@ -3,8 +3,6 @@ package aircmd import ( "bytes" "encoding/json" - "io" - "strings" "testing" "text/template" @@ -20,8 +18,8 @@ import ( "github.com/stretchr/testify/require" ) -// renderGet renders the get template against the JSON envelope, exactly as -// the command does, so the test covers the real template branches. +// renderGet renders the get template against the JSON envelope, exactly as the +// command does for a sweep run, so the test covers the real template branches. func renderGet(t *testing.T, data getData) string { t.Helper() tmpl, err := template.New("get").Parse(getTemplate) @@ -31,20 +29,6 @@ func renderGet(t *testing.T, data getData) string { return buf.String() } -func TestGetTemplateSingleRun(t *testing.T) { - out := renderGet(t, getData{ - RunIDDisplay: "123", - Status: "RUNNING", - UserDisplay: "me@example.com", - }) - assert.Contains(t, out, "Run ID: 123") - assert.Contains(t, out, "Status: RUNNING") - assert.Contains(t, out, "User: me@example.com") - assert.NotContains(t, out, "Sweep") - // Python embeds the dashboard link on the Run ID; there is no Dashboard row. - assert.NotContains(t, out, "Dashboard") -} - func TestGetRunInvalidID(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) @@ -88,51 +72,6 @@ func TestGetRunNotFoundJSON(t *testing.T) { assert.Equal(t, jsonError{Code: "NOT_FOUND", Kind: "NOT_FOUND", Message: "run 5 not found: check the run ID and that it is a job run ID"}, got.Error) } -func TestPrintConfigYAML(t *testing.T) { - t.Run("downloads and prints to stdout", func(t *testing.T) { - ctx := cmdio.MockDiscard(t.Context()) - m := mocks.NewMockWorkspaceClient(t) - // The mock asserts Download is called with the resolved path. - m.GetMockWorkspaceAPI().EXPECT(). - Download(mock.Anything, "/Workspace/cfg.yaml"). - Return(io.NopCloser(strings.NewReader("epochs: 3\ncommand: \"set -e\\npython train.py\\n\"\n")), nil) - - // The config is data output and must land on stdout (the out writer), - // matching the Python `air get` behavior. - var out bytes.Buffer - printConfigYAML(ctx, &out, m.WorkspaceClient, "/Workspace/cfg.yaml") - assert.Contains(t, out.String(), "Training Configuration:") - assert.Contains(t, out.String(), "epochs: 3") - // The multi-line command is reformatted to a `|` block literal. - assert.Contains(t, out.String(), "command: |\n set -e\n python train.py") - }) - - t.Run("download failure is non-fatal and writes nothing", func(t *testing.T) { - ctx := cmdio.MockDiscard(t.Context()) - m := mocks.NewMockWorkspaceClient(t) - m.GetMockWorkspaceAPI().EXPECT(). - Download(mock.Anything, "/Workspace/missing.yaml"). - Return(nil, apierr.ErrResourceDoesNotExist) - - // Must not panic and must not write to stdout: a failed config fetch is - // best-effort, surfaced only as a stderr warning. - var out bytes.Buffer - printConfigYAML(ctx, &out, m.WorkspaceClient, "/Workspace/missing.yaml") - assert.Empty(t, out.String()) - }) -} - -func TestYAMLConfigPath(t *testing.T) { - // No tasks, or a task without GenAiComputeTask, yields no path. - assert.Empty(t, yamlConfigPath(&jobs.Run{})) - assert.Empty(t, yamlConfigPath(&jobs.Run{Tasks: []jobs.RunTask{{}}})) - - run := &jobs.Run{Tasks: []jobs.RunTask{{ - GenAiComputeTask: &jobs.GenAiComputeTask{YamlParametersFilePath: "/Workspace/cfg.yaml"}, - }}} - assert.Equal(t, "/Workspace/cfg.yaml", yamlConfigPath(run)) -} - func TestGetTemplateSweep(t *testing.T) { out := renderGet(t, getData{ RunID: "456", @@ -152,8 +91,6 @@ func TestGetTemplateSweep(t *testing.T) { assert.Contains(t, out, "iter_1") assert.Contains(t, out, "FAILED") assert.Contains(t, out, "my-exp") - // The single-run rows must not appear in the sweep view. - assert.NotContains(t, out, "Dashboard:") } func TestGetTemplateSweepNoTasks(t *testing.T) { @@ -169,60 +106,6 @@ func TestGetTemplateSweepNoTasks(t *testing.T) { assert.NotContains(t, out, "Sweep Tasks:") } -func TestGetTemplateMinimal(t *testing.T) { - // Every row always renders; missing values show "N/A", in Python's order. - out := renderGet(t, getData{ - RunIDDisplay: "1", - Status: "PENDING", - SubmittedDisplay: na, - DurationDisplay: na, - ExperimentDisplay: na, - MLflowDisplay: na, - UserDisplay: na, - AcceleratorsDisplay: na, - }) - for _, want := range []string{ - "Run ID: 1", - "Status: PENDING", - "Submitted: N/A", - "Retries: 0", - "Duration: N/A", - "Experiment: N/A", - "MLflow Run: N/A", - "User: N/A", - "Accelerators: N/A", - } { - assert.Contains(t, out, want) - } -} - -func TestGetTemplateAllFields(t *testing.T) { - out := renderGet(t, getData{ - RunIDDisplay: "1", - Status: "SUCCESS", - SubmittedDisplay: "2023-11-14 22:13 UTC", - DurationDisplay: "12s", - AttemptNumber: 2, - ExperimentDisplay: "exp", - MLflowDisplay: "sunny-cat-42", - UserDisplay: "me@example.com", - AcceleratorsDisplay: "8x H100", - }) - for _, want := range []string{ - "Run ID: 1", - "Status: SUCCESS", - "Submitted: 2023-11-14 22:13 UTC", - "Retries: 2", - "Duration: 12s", - "Experiment: exp", - "MLflow Run: sunny-cat-42", - "User: me@example.com", - "Accelerators: 8x H100", - } { - assert.Contains(t, out, want) - } -} - func TestBuildGetData(t *testing.T) { run := &jobs.Run{ RunId: 123, @@ -240,7 +123,6 @@ func TestBuildGetData(t *testing.T) { } d := buildGetData(run) assert.Equal(t, "123", d.RunID) - assert.Equal(t, "123", d.RunIDDisplay) assert.Equal(t, "SUCCESS", d.Status) assert.Equal(t, 1, d.AttemptNumber) assert.Equal(t, "2023-11-14 22:13 UTC", d.SubmittedDisplay) @@ -257,11 +139,10 @@ func TestBuildGetData(t *testing.T) { func TestBuildGetDataEmpty(t *testing.T) { // A run with no tasks, creator, or timing renders every text cell as "N/A". d := buildGetData(&jobs.Run{RunId: 7}) - assert.Equal(t, "7", d.RunIDDisplay) + assert.Equal(t, "7", d.RunID) assert.Equal(t, na, d.SubmittedDisplay) assert.Equal(t, na, d.DurationDisplay) assert.Equal(t, na, d.ExperimentDisplay) - assert.Equal(t, na, d.MLflowDisplay) assert.Equal(t, na, d.UserDisplay) assert.Equal(t, na, d.AcceleratorsDisplay) } diff --git a/experimental/air/cmd/mlflow.go b/experimental/air/cmd/mlflow.go index 7e2b5b8fd58..6eaafabfcb9 100644 --- a/experimental/air/cmd/mlflow.go +++ b/experimental/air/cmd/mlflow.go @@ -76,33 +76,16 @@ func mlflowLogsURL(host string, ids *mlflowIdentifiers) string { strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) } -// mlflowExperimentURL links to the MLflow experiment page; mlflowRunURL links to -// the run page. These back the Experiment and MLflow Run hyperlinks in text mode. -func mlflowExperimentURL(host string, ids *mlflowIdentifiers) string { - return fmt.Sprintf("%s/ml/experiments/%s", strings.TrimRight(host, "/"), ids.ExperimentID) -} - +// mlflowRunURL links to the MLflow run page; it backs the MLflow Run hyperlink +// in the single-run view. func mlflowRunURL(host string, ids *mlflowIdentifiers) string { return fmt.Sprintf("%s/ml/experiments/%s/runs/%s", strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) } -// mlflowRunLabel returns the MLflow run's human-readable name to use as the -// hyperlink text, falling back to "...{last 8 of run id}" when the name can't be -// fetched. Mirrors Python's _get_mlflow_run_name (cli_display.py). -func mlflowRunLabel(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) string { - if name := fetchMLflowRunName(ctx, w, mlflowRunID); name != "" { - return name - } - if len(mlflowRunID) > 8 { - return "..." + mlflowRunID[len(mlflowRunID)-8:] - } - return "..." + mlflowRunID -} - // 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. +// returning "" if it can't be obtained. Best-effort, like the rest of the MLflow +// enrichment. func fetchMLflowRunName(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) string { apiClient, err := client.New(w.Config) if err != nil { @@ -124,3 +107,16 @@ func fetchMLflowRunName(ctx context.Context, w *databricks.WorkspaceClient, mlfl } return out.Run.Info.RunName } + +// mlflowRunLabel is the text shown for the MLflow Run cell: the run's name, or +// "...{last 8 of run id}" when the name is unknown. Mirrors Python's +// _get_mlflow_run_name (cli_display.py). +func mlflowRunLabel(name, mlflowRunID string) string { + if name != "" { + return name + } + if len(mlflowRunID) > 8 { + return "..." + mlflowRunID[len(mlflowRunID)-8:] + } + return "..." + mlflowRunID +} diff --git a/experimental/air/cmd/mlflow_test.go b/experimental/air/cmd/mlflow_test.go index a6a5c151074..505da11cacd 100644 --- a/experimental/air/cmd/mlflow_test.go +++ b/experimental/air/cmd/mlflow_test.go @@ -83,30 +83,40 @@ func TestMLflowURLs(t *testing.T) { ids := &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"} // A trailing slash on the host must not produce a double slash in the link. assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1/artifacts/logs/node_0", mlflowLogsURL("https://h.test/", ids)) - assert.Equal(t, "https://h.test/ml/experiments/E1", mlflowExperimentURL("https://h.test", ids)) assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1", mlflowRunURL("https://h.test", ids)) } func TestMLflowRunLabel(t *testing.T) { + // Uses the run name when it is known. + assert.Equal(t, "sunny-cat-42", mlflowRunLabel("sunny-cat-42", "0123456789abcdef")) + // Falls back to the last 8 characters of a long run id. + assert.Equal(t, "...9abcdef0", mlflowRunLabel("", "0123456789abcdef0")) + // A short run id is shown in full behind the ellipsis. + assert.Equal(t, "...short", mlflowRunLabel("", "short")) +} + +func TestFetchMLflowRunName(t *testing.T) { ctx := t.Context() - t.Run("uses the run name when available", func(t *testing.T) { + mlflowServer := func(body string) *httptest.Server { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/api/2.0/mlflow/runs/get" { - _, _ = w.Write([]byte(`{"run":{"info":{"run_name":"sunny-cat-42"}}}`)) + _, _ = w.Write([]byte(body)) return } _, _ = w.Write([]byte(`{}`)) })) t.Cleanup(srv.Close) - assert.Equal(t, "sunny-cat-42", mlflowRunLabel(ctx, newTestWorkspaceClient(t, srv.URL), "0123456789abcdef")) + return srv + } + + t.Run("returns the run name", func(t *testing.T) { + srv := mlflowServer(`{"run":{"info":{"run_name":"sunny-cat-42"}}}`) + assert.Equal(t, "sunny-cat-42", fetchMLflowRunName(ctx, newTestWorkspaceClient(t, srv.URL), "run1")) }) - t.Run("falls back to the last 8 characters of the run id", func(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte(`{}`)) - })) - t.Cleanup(srv.Close) - assert.Equal(t, "...9abcdef0", mlflowRunLabel(ctx, newTestWorkspaceClient(t, srv.URL), "0123456789abcdef0")) + t.Run("empty when the run cannot be fetched", func(t *testing.T) { + srv := mlflowServer(`{}`) + assert.Empty(t, fetchMLflowRunName(ctx, newTestWorkspaceClient(t, srv.URL), "run1")) }) } diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go new file mode 100644 index 00000000000..d6c2e6a25be --- /dev/null +++ b/experimental/air/cmd/render.go @@ -0,0 +1,379 @@ +package aircmd + +import ( + "context" + "fmt" + "io" + "strconv" + "strings" + + "github.com/charmbracelet/lipgloss" + "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" + "github.com/muesli/termenv" + "go.yaml.in/yaml/v3" +) + +// Box titles, rendered into the top border of each box. +const ( + configBoxTitle = "Configuration" + metadataBoxTitle = "Metadata" +) + +// minBoxInnerWidth keeps all boxes a uniform, comfortable width; boxHPad and +// boxVPad are the horizontal and vertical padding inside each box. +const ( + minBoxInnerWidth = 60 + boxHPad = 2 + boxVPad = 1 +) + +// palette holds the lipgloss styles for the single-run view. Two layers: a +// neutral ramp for chrome and text, and an accent palette for syntax and data. +// All styles come from one renderer, so they honor its color profile (Ascii +// under --no-color / non-TTY, which strips every escape). +type palette struct { + n7 lipgloss.Style // dim: block indicator, empty progress, command-block "|" + n8 lipgloss.Style // muted: field labels + n12 lipgloss.Style // content: config and metadata values, percent + + border lipgloss.Style // box borders and titles + + blue lipgloss.Style // yaml keys and hyperlinks + green lipgloss.Style // success status and progress fill + amber lipgloss.Style // in-progress status + red lipgloss.Style // failed status +} + +func newPalette(r *lipgloss.Renderer) palette { + fg := func(hex string) lipgloss.Style { return r.NewStyle().Foreground(lipgloss.Color(hex)) } + return palette{ + n7: fg("#6E6E70"), + n8: fg("#8C8A86"), + n12: fg("#F9F7F4"), // Oat Light + border: fg("#B7A8E8"), // light purple (box borders and titles) + blue: fg("#8FB3DC"), + green: fg("#74C39A"), + amber: fg("#DCAA5C"), + red: fg("#D9756B"), + } +} + +// runView is the resolved, display-ready data the renderer draws. It is built +// from getData plus the MLflow enrichment, so the renderer itself does no API +// calls or formatting decisions. +type runView struct { + runID string + dashboardURL string + status string + submitted string + retries int + maxRetries string + duration string + experiment string + mlflowLabel string + mlflowURL string + user string + accelerators string + environment string +} + +// renderRunText writes the styled single-run view: a training-config box, a +// completed progress bar, and a field list, separated by blank lines. It is a +// one-shot renderer — it builds the full string and writes it once, with no +// streaming, spinner, or redraw. +func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, run *jobs.Run, data *getData, ids *mlflowIdentifiers) { + colorOn := cmdio.SupportsColor(ctx, out) + renderer := lipgloss.NewRenderer(out) + if !colorOn { + // Ascii emits no SGR codes; combined with the link fallback below this + // gives clean, un-escaped output under --no-color / NO_COLOR / piped stdout. + renderer.SetColorProfile(termenv.Ascii) + } + p := newPalette(renderer) + + view := runView{ + runID: data.RunID, + dashboardURL: data.DashboardURL, + status: data.Status, + submitted: data.SubmittedDisplay, + retries: data.AttemptNumber, + maxRetries: data.MaxRetriesDisplay, + duration: data.DurationDisplay, + experiment: data.ExperimentDisplay, + mlflowLabel: na, + user: data.UserDisplay, + accelerators: data.AcceleratorsDisplay, + environment: data.EnvironmentDisplay, + } + + if ids != nil { + view.mlflowLabel = mlflowRunLabel(fetchMLflowRunName(ctx, w, ids.RunID), ids.RunID) + view.mlflowURL = mlflowRunURL(w.Config.Host, ids) + } + + var sections []string + if task := genAIComputeTask(run); task != nil { + if body := colorizeConfig(p, configYAML(ctx, w, task)); body != "" { + sections = append(sections, renderBox(p, configBoxTitle, body)) + } + } + sections = append(sections, renderBox(p, metadataBoxTitle, renderFields(p, colorOn, view))) + + // A single write: a blank line before the first box and after the last, and + // one between each box. + fmt.Fprintf(out, "\n%s\n\n", strings.Join(sections, "\n\n")) + + // Bare-URL footer so the job run / MLflow links remain reachable when + // stdout is not a hyperlink-capable terminal (piped, redirected, NO_COLOR). + // In that case the OSC 8 hyperlinks on the Run ID / MLflow Run cells + // degrade to plain labels and the URLs would otherwise disappear from text + // output, breaking workflows like `air get run X > out.txt` or + // `NO_COLOR=1 air get run X` that the previous `Job Link:` line supported. + if view.dashboardURL != "" { + fmt.Fprintf(out, "Run URL: %s\n", view.dashboardURL) + } + if view.mlflowURL != "" { + fmt.Fprintf(out, "MLflow URL: %s\n", view.mlflowURL) + } +} + +// genAIComputeTask returns the run's first GenAI-compute task, or nil. +func genAIComputeTask(run *jobs.Run) *jobs.GenAiComputeTask { + if len(run.Tasks) == 0 { + return nil + } + return run.Tasks[0].GenAiComputeTask +} + +// configYAML returns the run's resolved training config as YAML for the box. The +// full config (including the command/script) lives in the run's parameters, not +// the structured task fields, so we prefer the inline parameters, then the +// parameters file, and only synthesize a minimal config as a last resort. +func configYAML(ctx context.Context, w *databricks.WorkspaceClient, task *jobs.GenAiComputeTask) string { + if task.YamlParameters != "" { + return strings.TrimRight(reformatYAMLForDisplay([]byte(task.YamlParameters)), "\n") + } + if task.YamlParametersFilePath != "" { + if raw := downloadConfig(ctx, w, task.YamlParametersFilePath); len(raw) > 0 { + return strings.TrimRight(reformatYAMLForDisplay(raw), "\n") + } + } + return synthConfigYAML(task) +} + +// downloadConfig fetches the run's training-config file, returning nil on +// failure (logged as a warning). Best-effort, like the rest of the enrichment. +func downloadConfig(ctx context.Context, w *databricks.WorkspaceClient, path string) []byte { + r, err := w.Workspace.Download(ctx, path) + if err != nil { + log.Warnf(ctx, "air get: could not download training config %s: %v", path, err) + return nil + } + defer r.Close() + content, err := io.ReadAll(r) + if err != nil { + log.Warnf(ctx, "air get: could not read training config %s: %v", path, err) + return nil + } + return content +} + +// configBox describes the synthesized config we marshal when the run exposes no +// parameters, in the order the fields are shown. +type configBox struct { + ExperimentName string `yaml:"experiment_name,omitempty"` + Compute *configCompute `yaml:"compute,omitempty"` +} + +type configCompute struct { + AcceleratorType string `yaml:"accelerator_type,omitempty"` + NumAccelerators int `yaml:"num_accelerators,omitempty"` +} + +// synthConfigYAML builds a minimal config from the structured task fields. It +// omits the command, which is only available in the run parameters. +func synthConfigYAML(task *jobs.GenAiComputeTask) string { + cfg := configBox{} + if task.MlflowExperimentName != "" { + cfg.ExperimentName = stripExperimentUserPrefix(task.MlflowExperimentName) + } + if task.Compute != nil && task.Compute.NumGpus > 0 { + cfg.Compute = &configCompute{ + AcceleratorType: task.Compute.GpuType, + NumAccelerators: task.Compute.NumGpus, + } + } + if cfg.ExperimentName == "" && cfg.Compute == nil { + return "" + } + b, err := yaml.Marshal(cfg) + if err != nil { + return "" + } + return strings.TrimRight(reformatYAMLForDisplay(b), "\n") +} + +// colorizeConfig styles a YAML config block line by line. +func colorizeConfig(p palette, body string) string { + if body == "" { + return "" + } + lines := strings.Split(body, "\n") + for i, line := range lines { + lines[i] = colorizeConfigLine(p, line) + } + return strings.Join(lines, "\n") +} + +// colorizeConfigLine colors one YAML line: keys blue, the `|` block indicator +// dim, and every value (and the command body that isn't a `key:` pair) in the +// neutral content color. +func colorizeConfigLine(p palette, line string) string { + indent := line[:len(line)-len(strings.TrimLeft(line, " "))] + trimmed := strings.TrimLeft(line, " ") + + if i := strings.IndexByte(trimmed, ':'); i > 0 && isConfigKey(trimmed[:i]) { + key := trimmed[:i] + value := strings.TrimSpace(trimmed[i+1:]) + styled := indent + p.blue.Render(key+":") + switch value { + case "": + // A mapping parent such as "compute:" has no value of its own. + case "|": + styled += " " + p.n7.Render(value) + default: + styled += " " + p.n12.Render(value) + } + return styled + } + return indent + p.n12.Render(trimmed) +} + +// isConfigKey reports whether s is a bare YAML key (lowercase, digits, and +// underscores). It guards against treating a colon inside a command body as a +// key/value separator. +func isConfigKey(s string) bool { + if s == "" { + return false + } + for _, r := range s { + if r != '_' && (r < 'a' || r > 'z') && (r < '0' || r > '9') { + return false + } + } + return true +} + +// renderBox draws a rounded-border box around body, with title rendered into the +// top border in the border color. body lines are padded to the widest one (or +// minBoxInnerWidth), with boxHPad columns and boxVPad rows of padding inside. +func renderBox(p palette, title, body string) string { + border := lipgloss.RoundedBorder() + lines := strings.Split(body, "\n") + pad := strings.Repeat(" ", boxHPad) + + titleWidth := lipgloss.Width(title) + inner := max(minBoxInnerWidth, titleWidth+2) + for _, line := range lines { + inner = max(inner, lipgloss.Width(line)) + } + + left := p.border.Render(border.Left) + right := p.border.Render(border.Right) + blank := left + strings.Repeat(" ", inner+2*boxHPad) + right + + var b strings.Builder + // Top: ╭─ ──…──╮. The dash count makes the row width match the body, + // accounting for the boxHPad columns on each side. + trailing := inner + 2*boxHPad - titleWidth - 3 + b.WriteString(p.border.Render(border.TopLeft + border.Top)) + b.WriteString(" " + p.border.Render(title) + " ") + b.WriteString(p.border.Render(strings.Repeat(border.Top, trailing) + border.TopRight)) + b.WriteByte('\n') + + for range boxVPad { + b.WriteString(blank + "\n") + } + for _, line := range lines { + fill := strings.Repeat(" ", inner-lipgloss.Width(line)) + b.WriteString(left + pad + line + fill + pad + right) + b.WriteByte('\n') + } + for range boxVPad { + b.WriteString(blank + "\n") + } + + b.WriteString(p.border.Render(border.BottomLeft + strings.Repeat(border.Bottom, inner+2*boxHPad) + border.BottomRight)) + return b.String() +} + +// renderFields draws the two-column summary: muted labels right-padded to the +// longest one, neutral values, a status-colored Status, and blue Run ID / MLflow +// Run hyperlinks. +func renderFields(p palette, colorOn bool, v runView) string { + status := statusStyle(p, v.status).Render("● " + v.status) + rows := []string{ + field(p, "Run ID", link(colorOn, p.blue, v.runID, v.dashboardURL)), + field(p, "Status", status), + field(p, "Submitted", p.n12.Render(v.submitted)), + field(p, "Retries", p.n12.Render(strconv.Itoa(v.retries))), + field(p, "Max Retries", p.n12.Render(v.maxRetries)), + field(p, "Duration", p.n12.Render(v.duration)), + field(p, "Experiment", p.n12.Render(v.experiment)), + field(p, "MLflow Run", link(colorOn, p.blue, v.mlflowLabel, v.mlflowURL)), + field(p, "User", p.n12.Render(v.user)), + field(p, "Accelerators", p.n12.Render(v.accelerators)), + field(p, "Environment", p.n12.Render(v.environment)), + } + return strings.Join(rows, "\n") +} + +// fieldLabelWidth is the width of the longest label ("Accelerators"), so values +// line up in a single column. +const fieldLabelWidth = len("Accelerators") + +func field(p palette, label, value string) string { + return p.n8.Render(label+strings.Repeat(" ", fieldLabelWidth-len(label))) + " " + value +} + +// link renders label as an OSC 8 terminal hyperlink to url in the given style +// (underlined). With color off (or no url) it is just the styled label so the +// box stays aligned; the URLs remain available in JSON output. +func link(colorOn bool, style lipgloss.Style, label, url string) string { + if !colorOn || url == "" { + return style.Render(label) + } + // Wrap the already-styled label in the hyperlink. Passing the OSC 8 escape + // through lipgloss.Render instead corrupts it: lipgloss re-styles each rune + // and splits the "\x1b]8;;" introducer, so the terminal can't parse the + // sequence and prints it literally. + return termenv.Hyperlink(url, style.Underline(true).Render(label)) +} + +// statusStyle maps a run status to its accent color: green for success, red for +// terminal failures, amber for everything still in flight. +func statusStyle(p palette, status string) lipgloss.Style { + switch { + case isSuccessStatus(status): + return p.green + case isFailedStatus(status): + return p.red + default: + return p.amber + } +} + +func isSuccessStatus(status string) bool { + return status == "SUCCESS" +} + +func isFailedStatus(status string) bool { + switch status { + case "FAILED", "TIMEDOUT", "CANCELED", "INTERNAL_ERROR", "UPSTREAM_FAILED", "UPSTREAM_CANCELED": + return true + } + return false +} diff --git a/experimental/air/cmd/render_test.go b/experimental/air/cmd/render_test.go new file mode 100644 index 00000000000..71e2732fe0f --- /dev/null +++ b/experimental/air/cmd/render_test.go @@ -0,0 +1,162 @@ +package aircmd + +import ( + "io" + "strings" + "testing" + + "github.com/charmbracelet/lipgloss" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/muesli/termenv" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" +) + +// asciiPalette returns a palette whose styles emit no escape codes, so render +// output is plain text and assertions stay readable. +func asciiPalette() palette { + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + return newPalette(r) +} + +func TestConfigYAML(t *testing.T) { + ctx := t.Context() + + t.Run("inline parameters include the command as a block literal", func(t *testing.T) { + task := &jobs.GenAiComputeTask{ + YamlParameters: "experiment_name: my-exp\ncompute:\n accelerator_type: a10\n num_accelerators: 1\ncommand: \"for i in $(seq 1 3); do echo $i; done\\n\"\n", + } + got := configYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, task) + assert.Contains(t, got, "experiment_name: my-exp") + assert.Contains(t, got, "accelerator_type: a10") + assert.Contains(t, got, "command: |") + assert.Contains(t, got, " for i in $(seq 1 3); do echo $i; done") + }) + + t.Run("downloads the parameters file when there are no inline parameters", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockWorkspaceAPI().EXPECT(). + Download(mock.Anything, "/Workspace/cfg.yaml"). + Return(io.NopCloser(strings.NewReader("experiment_name: from-file\n")), nil) + task := &jobs.GenAiComputeTask{YamlParametersFilePath: "/Workspace/cfg.yaml"} + assert.Equal(t, "experiment_name: from-file", configYAML(ctx, m.WorkspaceClient, task)) + }) + + t.Run("falls back to a synthesized config", func(t *testing.T) { + task := &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/exp", + Compute: &jobs.ComputeConfig{GpuType: "a10", NumGpus: 1}, + } + got := configYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, task) + assert.Contains(t, got, "experiment_name: exp") + assert.NotContains(t, got, "command") + }) +} + +func TestSynthConfigYAML(t *testing.T) { + task := &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/stream-latency-test", + Compute: &jobs.ComputeConfig{GpuType: "a10", NumGpus: 1}, + } + // The accelerator_type uses the raw GPU type; the command is omitted because + // it lives only in the run parameters. + want := "experiment_name: stream-latency-test\n" + + "compute:\n" + + " accelerator_type: a10\n" + + " num_accelerators: 1" + assert.Equal(t, want, synthConfigYAML(task)) + assert.Empty(t, synthConfigYAML(&jobs.GenAiComputeTask{})) +} + +func TestColorizeConfigLine(t *testing.T) { + p := asciiPalette() + // Under the Ascii profile colorization adds no escapes, so each line is + // preserved verbatim (indentation included) regardless of its role. + for _, line := range []string{ + "experiment_name: stream-latency-test", + "compute:", + " accelerator_type: a10", + " num_accelerators: 1", + "command: |", + ` for i in $(seq 1 10); do echo "step $i"; done`, + } { + assert.Equal(t, line, colorizeConfigLine(p, line)) + } +} + +func TestIsConfigKey(t *testing.T) { + assert.True(t, isConfigKey("experiment_name")) + assert.True(t, isConfigKey("num_accelerators")) + assert.False(t, isConfigKey("")) + assert.False(t, isConfigKey("for i in $(seq 1 10); do echo ")) + assert.False(t, isConfigKey("Command")) +} + +func TestRenderBox(t *testing.T) { + p := asciiPalette() + out := renderBox(p, configBoxTitle, "experiment_name: stream-latency-test\ncompute:") + lines := strings.Split(out, "\n") + + // Title sits in the top border; corners are rounded; every row is the same width. + assert.Contains(t, lines[0], "╭─ "+configBoxTitle+" ") + assert.True(t, strings.HasSuffix(lines[0], "╮")) + assert.True(t, strings.HasPrefix(lines[len(lines)-1], "╰")) + assert.Contains(t, out, "│ experiment_name: stream-latency-test") + + width := lipgloss.Width(lines[0]) + for _, l := range lines { + assert.Equal(t, width, lipgloss.Width(l)) + } +} + +func TestRenderFields(t *testing.T) { + p := asciiPalette() + out := renderFields(p, false, runView{ + runID: "836121283738861", + dashboardURL: "https://h.test/jobs/runs/836121283738861", + status: "SUCCESS", + submitted: "2026-06-03 04:17 UTC", + retries: 0, + maxRetries: "3", + duration: "1m 13s", + experiment: "stream-latency-test", + mlflowLabel: "stream-latency-test", + mlflowURL: "https://h.test/ml/experiments/E1/runs/R1", + user: "riddhi.bhagwat@databricks.com", + accelerators: "1x A10", + environment: "ml-runtime-gpu:1.0", + }) + + // Labels are padded to the longest ("Accelerators"), so values align. + assert.Contains(t, out, "Run ID ") + assert.Contains(t, out, "Accelerators 1x A10") + // Max retries and environment show alongside the other fields. + assert.Contains(t, out, "Max Retries 3") + assert.Contains(t, out, "Environment ml-runtime-gpu:1.0") + // The status carries its dot prefix. + assert.Contains(t, out, "● SUCCESS") + // Off a terminal, links render as the bare label (URLs live in JSON output). + assert.Contains(t, out, "Run ID 836121283738861") + assert.NotContains(t, out, "https://h.test") + // The field list is a tight block: no blank lines. + assert.NotContains(t, out, "\n\n") +} + +func TestLink(t *testing.T) { + p := asciiPalette() + // Color off: the bare label, no URL. + assert.Equal(t, "label", link(false, p.blue, "label", "https://h.test")) + assert.Equal(t, "label", link(false, p.blue, "label", "")) + // With color on, the label is wrapped in an OSC 8 hyperlink to the url. + assert.Contains(t, link(true, p.blue, "label", "https://h.test"), termenv.Hyperlink("https://h.test", "label")) +} + +func TestStatusStyleSelectors(t *testing.T) { + assert.True(t, isSuccessStatus("SUCCESS")) + assert.False(t, isSuccessStatus("RUNNING")) + assert.True(t, isFailedStatus("FAILED")) + assert.True(t, isFailedStatus("TIMEDOUT")) + assert.False(t, isFailedStatus("RUNNING")) +} diff --git a/go.mod b/go.mod index 9329b7ea427..090d567257b 100644 --- a/go.mod +++ b/go.mod @@ -24,6 +24,7 @@ require ( github.com/hexops/gotextdiff v1.0.3 // BSD-3-Clause github.com/jackc/pgx/v5 v5.9.2 // MIT github.com/mattn/go-isatty v0.0.22 // MIT + github.com/muesli/termenv v0.16.0 // MIT github.com/palantir/pkg/yamlpatch v1.5.0 // BSD-3-Clause github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // BSD-2-Clause github.com/quasilyte/go-ruleguard/dsl v0.3.22 // BSD-3-Clause @@ -85,7 +86,6 @@ require ( github.com/mitchellh/hashstructure/v2 v2.0.2 // indirect github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect github.com/muesli/cancelreader v0.2.2 // indirect - github.com/muesli/termenv v0.16.0 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/rivo/uniseg v0.4.7 // indirect From bd3f934eb7d176d2f506c8a6ffd8b6b680d507e9 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Tue, 23 Jun 2026 15:58:31 -0700 Subject: [PATCH 05/18] AIR CLI Integration: `air list` Functionality & UI (Interfacing with Training Service) (#5684) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Changes Add `air list` as a browsable view of the caller's recent AIR training runs. - Data source: the `AiWorkflowService.ListTrainingWorkflows` RPC (`GET /api/2.0/ai-training/workflows`), called directly via `client.Do` since the endpoint is `PUBLIC_UNDOCUMENTED` and not modeled by the SDK. The server does the AIR filtering, creator scoping, MLflow-ID resolution, and pagination, so no Jobs-API logic lives in the CLI. - Interactive table: in a terminal `air list` renders an inline, navigable table (Bubble Tea + Lip Gloss + termenv): `↑/↓` move a row, `←/→` page (20 rows/page), `Enter` opens the run's MLflow page, `q` quits. Status is colored by state and the MLflow column is a short clickable hyperlink. - Non-interactive: piped output, an explicit `--limit`, and empty results print the table once; `-o json` emits the air `{v,ts,data}` envelope unchanged. - Flags: `--limit` (default: all), `--active`, `--all-users`, and client-side `--filter` keys (`experiment`, `accelerator_type`, `num_accelerators`). Gateway timeouts (e.g. HTTP 504 on `--all-users`) return an actionable message. - Adds `cmdio.IsPagerSupported`; promotes `termenv` to a direct dependency ## Why The `ai-training` service now owns the AIR-specific run logic server-side, so `air list` should call its RPC rather than reimplementing run discovery against the Jobs API. The interactive table gives a browsable run list on par with the Python `air` CLI and `databricks jobs list-runs`. ## Tests - Unit: RPC transport, `TrainingWorkflow`→row mapping, `--filter` matching, status/accelerator/timestamp helpers, and the TUI model (navigation, paging, 20-row page cap, window scroll, quit, static render). - Acceptance: `acceptance/experimental/air/list` (text + JSON) plus `help` updates; `unimplemented` no longer covers `air list` Manual verification output: <img width="1444" height="596" alt="Screenshot 2026-06-22 at 11 52 41 AM" src="https://github.com/user-attachments/assets/2e4a5917-8562-44ed-bb1d-a1cb1398731c" /> --- acceptance/experimental/air/help/output.txt | 22 +- acceptance/experimental/air/help/script | 3 + .../experimental/air/list/out.test.toml | 3 + acceptance/experimental/air/list/output.txt | 33 +++ acceptance/experimental/air/list/script | 5 + acceptance/experimental/air/list/test.toml | 50 ++++ .../experimental/air/unimplemented/output.txt | 6 - .../experimental/air/unimplemented/script | 3 - experimental/air/cmd/aiworkflow.go | 77 ++++++ experimental/air/cmd/aiworkflow_test.go | 62 +++++ experimental/air/cmd/format.go | 65 ++++- experimental/air/cmd/format_test.go | 34 +++ experimental/air/cmd/list.go | 147 ++++++++++- experimental/air/cmd/list_filter.go | 86 +++++++ experimental/air/cmd/list_filter_test.go | 74 ++++++ experimental/air/cmd/list_format.go | 54 ++++ experimental/air/cmd/list_test.go | 180 +++++++++++++ experimental/air/cmd/list_tui.go | 185 ++++++++++++++ experimental/air/cmd/list_tui_render.go | 239 ++++++++++++++++++ experimental/air/cmd/list_tui_test.go | 175 +++++++++++++ experimental/air/cmd/stubs_test.go | 1 - libs/cmdio/io.go | 9 + 22 files changed, 1493 insertions(+), 20 deletions(-) create mode 100644 acceptance/experimental/air/list/out.test.toml create mode 100644 acceptance/experimental/air/list/output.txt create mode 100644 acceptance/experimental/air/list/script create mode 100644 acceptance/experimental/air/list/test.toml create mode 100644 experimental/air/cmd/aiworkflow.go create mode 100644 experimental/air/cmd/aiworkflow_test.go create mode 100644 experimental/air/cmd/list_filter.go create mode 100644 experimental/air/cmd/list_filter_test.go create mode 100644 experimental/air/cmd/list_format.go create mode 100644 experimental/air/cmd/list_test.go create mode 100644 experimental/air/cmd/list_tui.go create mode 100644 experimental/air/cmd/list_tui_render.go create mode 100644 experimental/air/cmd/list_tui_test.go diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index 41a1d6815a8..9348686ce41 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -12,7 +12,7 @@ Usage: Available Commands: cancel Cancel one or more runs get Show details for a specific resource - list List recent runs + list List your recent runs (active and completed) for the current profile logs Stream or fetch logs for a run register-image Mirror a Docker image into the workspace registry run Submit a training workload from a YAML config @@ -27,3 +27,23 @@ Global Flags: -t, --target string bundle target to use (if applicable) Use "databricks experimental air [command] --help" for more information about a command. + +=== list help +>>> [CLI] experimental air list --help +List your recent runs (active and completed) for the current profile + +Usage: + databricks experimental air list [flags] + +Flags: + --active Show only active runs + --all-users Show runs from all users + --filter stringArray Filter runs, e.g. experiment=foo* (repeatable) + -h, --help help for list + --limit int Maximum number of runs to show (default 20) + +Global Flags: + --debug enable debug logging + -o, --output type output type: text or json (default text) + -p, --profile string ~/.databrickscfg profile + -t, --target string bundle target to use (if applicable) diff --git a/acceptance/experimental/air/help/script b/acceptance/experimental/air/help/script index cd67a6fc1b1..81f3907e4f5 100644 --- a/acceptance/experimental/air/help/script +++ b/acceptance/experimental/air/help/script @@ -3,3 +3,6 @@ title "help" trace $CLI experimental air --help + +title "list help" +trace $CLI experimental air list --help diff --git a/acceptance/experimental/air/list/out.test.toml b/acceptance/experimental/air/list/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/list/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/list/output.txt b/acceptance/experimental/air/list/output.txt new file mode 100644 index 00000000000..9e8f93826f5 --- /dev/null +++ b/acceptance/experimental/air/list/output.txt @@ -0,0 +1,33 @@ + +=== list (text) +>>> [CLI] experimental air list + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 + [NUMID] llama-train ● FAILED [TIMESTAMP] 3m 31s - [USERNAME] 1x A10 + +=== list (json) +>>> [CLI] experimental air list -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "runs": [ + { + "run_id": "[NUMID]", + "run_name": "qwen-train", + "user": "[USERNAME]", + "status": "SUCCESS", + "started_at": "[TIMESTAMP]", + "is_sweep": false + }, + { + "run_id": "[NUMID]", + "run_name": "llama-train", + "user": "[USERNAME]", + "status": "FAILED", + "started_at": "[TIMESTAMP]", + "is_sweep": false + } + ] + } +} diff --git a/acceptance/experimental/air/list/script b/acceptance/experimental/air/list/script new file mode 100644 index 00000000000..14702b283e7 --- /dev/null +++ b/acceptance/experimental/air/list/script @@ -0,0 +1,5 @@ +title "list (text)" +trace $CLI experimental air list + +title "list (json)" +trace $CLI experimental air list -o json diff --git a/acceptance/experimental/air/list/test.toml b/acceptance/experimental/air/list/test.toml new file mode 100644 index 00000000000..f547919c3e5 --- /dev/null +++ b/acceptance/experimental/air/list/test.toml @@ -0,0 +1,50 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# ListTrainingWorkflows: the server resolves AIR runs, creator scoping, and MLflow +# IDs, so the CLI just maps the response. One completed run with resolved MLflow +# IDs (so the MLflow column is a link) and one failed run without them (so it +# falls back to "-"). +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows" +Response.Body = ''' +{ + "training_workflows": [ + { + "job_run_id": "334747067049496", + "metadata": {"creator_name": "tester@databricks.com"}, + "spec": {"compute": {"hardware_accelerator_type": "GPU_8xH100", "accelerator_count": 8}}, + "status": { + "state": "TRAINING_WORKFLOW_STATE_TERMINATED_COMPLETED", + "start_time": "2026-06-05T17:32:39.000Z", + "end_time": "2026-06-05T17:32:51.000Z", + "job": {"name": "qwen-train"}, + "mlflow": { + "experiment": "/Users/tester@databricks.com/qwen-train", + "experiment_id": "exp1", + "run_id": "run1" + } + } + }, + { + "job_run_id": "566001814929041", + "metadata": {"creator_name": "tester@databricks.com"}, + "spec": {"compute": {"hardware_accelerator_type": "GPU_1xA10", "accelerator_count": 1}}, + "status": { + "state": "TRAINING_WORKFLOW_STATE_TERMINATED_FAILED", + "start_time": "2026-06-05T18:43:24.000Z", + "end_time": "2026-06-05T18:46:55.000Z", + "job": {"name": "llama-train"}, + "mlflow": {"experiment": "/Users/tester@databricks.com/llama-train"} + } + } + ] +} +''' diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt index 0a86360c78f..66ddc34d586 100644 --- a/acceptance/experimental/air/unimplemented/output.txt +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -5,12 +5,6 @@ Error: `air run` is not implemented yet Exit code: 1 -=== list ->>> [CLI] experimental air list -Error: `air list` is not implemented yet - -Exit code: 1 - === logs >>> [CLI] experimental air logs 123 Error: `air logs` is not implemented yet diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script index e6e8d33ef9d..d00d0453689 100644 --- a/acceptance/experimental/air/unimplemented/script +++ b/acceptance/experimental/air/unimplemented/script @@ -3,9 +3,6 @@ title "run" errcode trace $CLI experimental air run -title "list" -errcode trace $CLI experimental air list - title "logs" errcode trace $CLI experimental air logs 123 diff --git a/experimental/air/cmd/aiworkflow.go b/experimental/air/cmd/aiworkflow.go new file mode 100644 index 00000000000..489a24eeb8c --- /dev/null +++ b/experimental/air/cmd/aiworkflow.go @@ -0,0 +1,77 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" +) + +// listWorkflowsPath is the ListTrainingWorkflows binding. It is +// PUBLIC_UNDOCUMENTED, so we call it directly like the endpoints in mlflow.go. +const listWorkflowsPath = "/api/2.0/ai-training/workflows" + +// trainingWorkflowsResponse is the ListTrainingWorkflows response. Only the +// fields the list table and JSON envelope consume are modeled. +type trainingWorkflowsResponse struct { + TrainingWorkflows []trainingWorkflow `json:"training_workflows"` + NextPageToken string `json:"next_page_token"` +} + +type trainingWorkflow struct { + JobRunID string `json:"job_run_id"` + + TaskRunID string `json:"task_run_id"` + Metadata struct { + CreatorName string `json:"creator_name"` + } `json:"metadata"` + Spec struct { + Compute struct { + HardwareAcceleratorType string `json:"hardware_accelerator_type"` + AcceleratorCount int `json:"accelerator_count"` + } `json:"compute"` + } `json:"spec"` + Status struct { + State string `json:"state"` + StartTime string `json:"start_time"` + EndTime string `json:"end_time"` + Job struct { + Name string `json:"name"` + } `json:"job"` + Mlflow struct { + Experiment string `json:"experiment"` + ExperimentID string `json:"experiment_id"` + RunID string `json:"run_id"` + } `json:"mlflow"` + } `json:"status"` +} + +// listTrainingWorkflows calls the ListTrainingWorkflows RPC with the given query +// parameters. +func listTrainingWorkflows(ctx context.Context, w *databricks.WorkspaceClient, query map[string]any) (*trainingWorkflowsResponse, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var resp trainingWorkflowsResponse + err = apiClient.Do(ctx, http.MethodGet, listWorkflowsPath, nil, nil, query, &resp) + if err != nil { + // A server error can arrive with an empty body, leaving %w blank, so + // surface the HTTP status to make the failure diagnosable. + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + switch apiErr.StatusCode { + case http.StatusGatewayTimeout, http.StatusServiceUnavailable, http.StatusRequestTimeout: + // Trailing %w keeps the error chain; the body is usually empty. + return nil, fmt.Errorf("timed out listing runs (HTTP %d): --all-users makes the server list every user's runs, which can exceed the gateway timeout on a busy workspace. Add --active to list only active runs, or drop --all-users to list your own.%w", apiErr.StatusCode, err) + } + return nil, fmt.Errorf("failed to list training workflows (HTTP %d): %w", apiErr.StatusCode, err) + } + return nil, fmt.Errorf("failed to list training workflows: %w", err) + } + return &resp, nil +} diff --git a/experimental/air/cmd/aiworkflow_test.go b/experimental/air/cmd/aiworkflow_test.go new file mode 100644 index 00000000000..5899d310db0 --- /dev/null +++ b/experimental/air/cmd/aiworkflow_test.go @@ -0,0 +1,62 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestListTrainingWorkflows(t *testing.T) { + var gotQuery url.Values + body := `{ + "training_workflows": [ + { + "job_run_id": "123", + "metadata": {"creator_name": "me@example.com"}, + "spec": {"compute": {"hardware_accelerator_type": "GPU_8xH100", "accelerator_count": 8}}, + "status": { + "state": "TRAINING_WORKFLOW_STATE_RUNNING", + "start_time": "2026-06-05T17:32:39.791Z", + "job": {"name": "my-run"}, + "mlflow": {"experiment": "/Users/me@example.com/exp", "experiment_id": "E1", "run_id": "R1"} + } + } + ], + "next_page_token": "tok" + }` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == listWorkflowsPath { + gotQuery = r.URL.Query() + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + resp, err := listTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), map[string]any{ + "creator_name": "me@example.com", + "active_only": true, + }) + require.NoError(t, err) + require.Len(t, resp.TrainingWorkflows, 1) + assert.Equal(t, "tok", resp.NextPageToken) + + wf := resp.TrainingWorkflows[0] + assert.Equal(t, "123", wf.JobRunID) + assert.Equal(t, "me@example.com", wf.Metadata.CreatorName) + assert.Equal(t, "GPU_8xH100", wf.Spec.Compute.HardwareAcceleratorType) + assert.Equal(t, 8, wf.Spec.Compute.AcceleratorCount) + assert.Equal(t, "TRAINING_WORKFLOW_STATE_RUNNING", wf.Status.State) + assert.Equal(t, "my-run", wf.Status.Job.Name) + assert.Equal(t, "E1", wf.Status.Mlflow.ExperimentID) + assert.Equal(t, "R1", wf.Status.Mlflow.RunID) + + // The query map is serialized into the GET query string. + assert.Equal(t, "me@example.com", gotQuery.Get("creator_name")) + assert.Equal(t, "true", gotQuery.Get("active_only")) +} diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index 57f86b80f07..8539a1794c9 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -132,13 +132,34 @@ func startedAt(run *jobs.Run) *string { if startMillis == 0 { return nil } - t := time.UnixMilli(startMillis).UTC() + s := isoFormat(time.UnixMilli(startMillis)) + return &s +} + +// isoFormat renders a time as a Python-style isoformat string in UTC ("+00:00", +// not "Z"; microseconds only when the sub-second part is non-zero), matching +// cli_entrypoint.py:1899. +func isoFormat(t time.Time) string { + t = t.UTC() layout := "2006-01-02T15:04:05-07:00" if t.Nanosecond() != 0 { layout = "2006-01-02T15:04:05.000000-07:00" } - s := t.Format(layout) - return &s + return t.Format(layout) +} + +// parseRPCTime parses an RFC3339 timestamp returned by the AiWorkflow RPC (e.g. +// "2026-06-05T18:46:55.876Z"), returning the zero time when the field is absent +// or unparseable. +func parseRPCTime(s string) time.Time { + if s == "" { + return time.Time{} + } + t, err := time.Parse(time.RFC3339, s) + if err != nil { + return time.Time{} + } + return t } // submittedDisplay formats the run's start time for the text table as @@ -250,10 +271,44 @@ func accelerators(run *jobs.Run) string { return "" } task := run.Tasks[0].GenAiComputeTask - if task == nil || task.Compute == nil || task.Compute.NumGpus == 0 { + if task == nil || task.Compute == nil { return "" } - return fmt.Sprintf("%dx %s", task.Compute.NumGpus, gpuDisplayName(task.Compute.GpuType)) + return acceleratorLabel(task.Compute.GpuType, task.Compute.NumGpus) +} + +// acceleratorLabel renders a GPU count and type as "8x H100", "" for none, or +// the count alone ("8x") when the type is unrecognized. +func acceleratorLabel(gpuType string, count int) string { + if count == 0 { + return "" + } + if name := gpuDisplayName(gpuType); name != "" { + return fmt.Sprintf("%dx %s", count, name) + } + return fmt.Sprintf("%dx", count) +} + +// trainingWorkflowStatus maps a TrainingWorkflowState enum value to the status +// word shown for a run, matching `air get run`. The server collapses Jobs +// lifecycle/result into this enum, so lossy cases (e.g. TIMEDOUT) render as FAILED. +func trainingWorkflowStatus(state string) string { + switch state { + case "TRAINING_WORKFLOW_STATE_PENDING", "TRAINING_WORKFLOW_STATE_PENDING_SENT": + return "PENDING" + case "TRAINING_WORKFLOW_STATE_RUNNING": + return "RUNNING" + case "TRAINING_WORKFLOW_STATE_TERMINATION_REQUESTED", "TRAINING_WORKFLOW_STATE_TERMINATION_SENT": + return "TERMINATING" + case "TRAINING_WORKFLOW_STATE_TERMINATED_COMPLETED": + return "SUCCESS" + case "TRAINING_WORKFLOW_STATE_TERMINATED_FAILED": + return "FAILED" + case "TRAINING_WORKFLOW_STATE_TERMINATED_STOPPED": + return "CANCELED" + default: + return "UNKNOWN" + } } // gpuDisplayName returns the friendly name for a GPU identifier, falling back to diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go index acecd0ce901..45bcbd94199 100644 --- a/experimental/air/cmd/format_test.go +++ b/experimental/air/cmd/format_test.go @@ -225,3 +225,37 @@ func TestAccelerators(t *testing.T) { GenAiComputeTask: &jobs.GenAiComputeTask{Compute: &jobs.ComputeConfig{NumGpus: 8, GpuType: "GPU_8xH100"}}, }}})) } + +func TestAcceleratorLabel(t *testing.T) { + assert.Empty(t, acceleratorLabel("GPU_8xH100", 0)) + assert.Equal(t, "8x H100", acceleratorLabel("GPU_8xH100", 8)) + assert.Equal(t, "1x A10", acceleratorLabel("GPU_1xA10", 1)) + // The RPC may report a count without a recognized type. + assert.Equal(t, "8x", acceleratorLabel("", 8)) +} + +func TestTrainingWorkflowStatus(t *testing.T) { + cases := map[string]string{ + "TRAINING_WORKFLOW_STATE_PENDING": "PENDING", + "TRAINING_WORKFLOW_STATE_PENDING_SENT": "PENDING", + "TRAINING_WORKFLOW_STATE_RUNNING": "RUNNING", + "TRAINING_WORKFLOW_STATE_TERMINATION_REQUESTED": "TERMINATING", + "TRAINING_WORKFLOW_STATE_TERMINATION_SENT": "TERMINATING", + "TRAINING_WORKFLOW_STATE_TERMINATED_COMPLETED": "SUCCESS", + "TRAINING_WORKFLOW_STATE_TERMINATED_FAILED": "FAILED", + "TRAINING_WORKFLOW_STATE_TERMINATED_STOPPED": "CANCELED", + "TRAINING_WORKFLOW_STATE_UNSPECIFIED": "UNKNOWN", + "": "UNKNOWN", + } + for state, want := range cases { + assert.Equal(t, want, trainingWorkflowStatus(state), state) + } +} + +func TestParseRPCTime(t *testing.T) { + assert.True(t, parseRPCTime("").IsZero()) + assert.True(t, parseRPCTime("not-a-time").IsZero()) + got := parseRPCTime("2026-06-05T18:46:55.876Z") + require.False(t, got.IsZero()) + assert.Equal(t, "2026-06-05T18:46:55.876000+00:00", isoFormat(got)) +} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index bf24cff9b23..91ccc6515c9 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -1,10 +1,58 @@ package aircmd import ( + "context" + "fmt" + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/spf13/cobra" ) +// maxListPages caps how many pages `air list runs` fetches (a safety bound); +// listPageBatch is the per-request page size. +const ( + maxListPages = 50 + listPageBatch = 100 +) + +// listData is the payload printed by `air list runs`. +type listData struct { + Rows []listRow `json:"runs"` +} + +// listRow is one run in the list. The json-tagged fields form the +// machine-readable output; fields tagged `json:"-"` are shown only in the +// human-readable table. +type listRow struct { + RunID string `json:"run_id"` + RunName string `json:"run_name"` + User string `json:"user"` + Status string `json:"status"` + StartedAt *string `json:"started_at"` + IsSweep bool `json:"is_sweep"` + + // Experiment, Duration, MLflowURL and Accelerators are table-only columns, + // omitted from JSON to match `air list runs --json`. + Experiment string `json:"-"` + Duration string `json:"-"` + MLflowURL string `json:"-"` + Accelerators string `json:"-"` +} + +// listQuery holds the resolved inputs to listAirRuns. +type listQuery struct { + activeOnly bool + allUsers bool + userFilter string + filters listFilters + limit int +} + func newListCommand() *cobra.Command { var ( limit int @@ -16,16 +64,107 @@ func newListCommand() *cobra.Command { cmd := &cobra.Command{ Use: "list", Args: root.NoArgs, - Short: "List recent runs", - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("list") - }, + Short: "List your recent runs (active and completed) for the current profile", } + cmd.PreRunE = root.MustWorkspaceClient + cmd.Flags().IntVar(&limit, "limit", 20, "Maximum number of runs to show") cmd.Flags().BoolVar(&active, "active", false, "Show only active runs") cmd.Flags().BoolVar(&allUsers, "all-users", false, "Show runs from all users") cmd.Flags().StringArrayVar(&filters, "filter", nil, "Filter runs, e.g. experiment=foo* (repeatable)") + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + if limit <= 0 { + return fmt.Errorf("invalid --limit %d: must be a positive integer", limit) + } + + f, err := parseListFilters(filters) + if err != nil { + return err + } + + // An explicit user= filter wins; otherwise default to the current user + // unless --all-users is set. The resolved name is sent to the server as + // creator_name, which scopes the listing. + userFilter := f.User + if userFilter == "" && !allUsers { + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return fmt.Errorf("failed to resolve current user: %w", err) + } + userFilter = me.UserName + } + + rows, err := listAirRuns(ctx, w, listQuery{ + activeOnly: active, + allUsers: allUsers, + userFilter: userFilter, + filters: f, + limit: limit, + }) + if err != nil { + return err + } + + // JSON keeps the air envelope; text renders the table (interactive and + // navigable in a terminal, printed once when piped). + if root.OutputType(cmd) == flags.OutputJSON { + return renderEnvelope(ctx, listData{Rows: rows}) + } + return renderListText(cmd, rows) + } + return cmd } + +// listAirRuns pages through ListTrainingWorkflows, applies the client-side +// --filter keys, and accumulates matches up to q.limit (0 = all). It stops at the +// limit, when the server runs out of pages, or at maxListPages. +func listAirRuns(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) ([]listRow, error) { + // Fetch in bounded batches and follow the page token, so "all" doesn't ask + // the server for an unbounded page. With client-side filters we fetch full + // batches, since most rows may be dropped before reaching the limit. + pageSize := q.limit + if pageSize <= 0 || pageSize > listPageBatch || q.filters.hasClientFilters() { + pageSize = listPageBatch + } + query := map[string]any{ + "active_only": q.activeOnly, + "include_all_users": q.allUsers, + "page_size": pageSize, + } + if q.userFilter != "" { + query["creator_name"] = q.userFilter + } + + var rows []listRow + for range maxListPages { + resp, err := listTrainingWorkflows(ctx, w, query) + if err != nil { + return nil, err + } + + for i := range resp.TrainingWorkflows { + wf := &resp.TrainingWorkflows[i] + if !q.filters.matches(wf) { + continue + } + rows = append(rows, buildListRow(wf, w.Config.Host)) + if q.limit > 0 && len(rows) >= q.limit { + return rows, nil + } + } + + if resp.NextPageToken == "" { + return rows, nil + } + query["page_token"] = resp.NextPageToken + } + + log.Warnf(ctx, "air list: stopped after %d pages; results may be incomplete", maxListPages) + return rows, nil +} diff --git a/experimental/air/cmd/list_filter.go b/experimental/air/cmd/list_filter.go new file mode 100644 index 00000000000..44ffa114e70 --- /dev/null +++ b/experimental/air/cmd/list_filter.go @@ -0,0 +1,86 @@ +package aircmd + +import ( + "fmt" + "path" + "strconv" + "strings" +) + +// supportedFilterKeys are the keys accepted by `air list --filter KEY=VALUE`. +var supportedFilterKeys = []string{"accelerator_type", "experiment", "num_accelerators", "user"} + +// listFilters holds the parsed `--filter` values for `air list`. +type listFilters struct { + // User is an exact creator-email match + User string + // Experiment is a case-insensitive glob + Experiment string + // AcceleratorType is a case-insensitive substring matched against the + // display GPU name (e.g. "H100"). + AcceleratorType string + // NumAccelerators is an exact match against the GPU count. + NumAccelerators *int +} + +func parseListFilters(raw []string) (listFilters, error) { + var f listFilters + for _, item := range raw { + key, value, ok := strings.Cut(item, "=") + if !ok || key == "" { + return listFilters{}, fmt.Errorf("invalid --filter %q: expected KEY=VALUE", item) + } + switch key { + case "user": + f.User = value + case "experiment": + f.Experiment = value + case "accelerator_type": + f.AcceleratorType = value + case "num_accelerators": + n, err := strconv.Atoi(value) + if err != nil || n <= 0 { + return listFilters{}, fmt.Errorf("invalid --filter num_accelerators=%q: must be a positive integer", value) + } + f.NumAccelerators = &n + default: + return listFilters{}, fmt.Errorf("unsupported --filter key %q: supported keys are %s", key, strings.Join(supportedFilterKeys, ", ")) + } + } + return f, nil +} + +// hasClientFilters reports whether any client-side filter (those applied in +// matches) is set. The user filter is excluded — the server handles it. +func (f listFilters) hasClientFilters() bool { + return f.Experiment != "" || f.AcceleratorType != "" || f.NumAccelerators != nil +} + +// matches reports whether a workflow satisfies the experiment, accelerator-type +// and accelerator-count filters. These have no ListTrainingWorkflows equivalent, +// so they are applied client-side to the response. The user filter is handled by +// the server (via creator_name), so it is not re-checked here. +func (f listFilters) matches(w *trainingWorkflow) bool { + if f.Experiment != "" { + name := stripExperimentUserPrefix(w.Status.Mlflow.Experiment) + matched, err := path.Match(strings.ToLower(f.Experiment), strings.ToLower(name)) + if err != nil || !matched { + return false + } + } + + if f.AcceleratorType != "" || f.NumAccelerators != nil { + compute := w.Spec.Compute + if f.AcceleratorType != "" { + display := strings.ToLower(gpuDisplayName(compute.HardwareAcceleratorType)) + if !strings.Contains(display, strings.ToLower(f.AcceleratorType)) { + return false + } + } + if f.NumAccelerators != nil && compute.AcceleratorCount != *f.NumAccelerators { + return false + } + } + + return true +} diff --git a/experimental/air/cmd/list_filter_test.go b/experimental/air/cmd/list_filter_test.go new file mode 100644 index 00000000000..a475e6576bd --- /dev/null +++ b/experimental/air/cmd/list_filter_test.go @@ -0,0 +1,74 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseListFilters(t *testing.T) { + t.Run("valid", func(t *testing.T) { + f, err := parseListFilters([]string{ + "user=me@example.com", + "experiment=qwen*", + "accelerator_type=H100", + "num_accelerators=8", + }) + require.NoError(t, err) + assert.Equal(t, "me@example.com", f.User) + assert.Equal(t, "qwen*", f.Experiment) + assert.Equal(t, "H100", f.AcceleratorType) + require.NotNil(t, f.NumAccelerators) + assert.Equal(t, 8, *f.NumAccelerators) + }) + + t.Run("unknown key", func(t *testing.T) { + _, err := parseListFilters([]string{"region=us"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported --filter key") + }) + + t.Run("malformed pair", func(t *testing.T) { + _, err := parseListFilters([]string{"experiment"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "expected KEY=VALUE") + }) + + t.Run("bad num_accelerators", func(t *testing.T) { + _, err := parseListFilters([]string{"num_accelerators=lots"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "num_accelerators") + }) + + t.Run("non-positive num_accelerators", func(t *testing.T) { + _, err := parseListFilters([]string{"num_accelerators=0"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "positive integer") + }) +} + +func TestListFiltersMatches(t *testing.T) { + wf := testWorkflow(1, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/qwen-eval") + + cases := []struct { + name string + f listFilters + want bool + }{ + {"no filters", listFilters{}, true}, + {"experiment prefix glob", listFilters{Experiment: "qwen*"}, true}, + {"experiment suffix glob", listFilters{Experiment: "*-eval"}, true}, + {"experiment case-insensitive", listFilters{Experiment: "QWEN*"}, true}, + {"experiment no match", listFilters{Experiment: "llama*"}, false}, + {"accelerator type substring", listFilters{AcceleratorType: "h100"}, true}, + {"accelerator type no match", listFilters{AcceleratorType: "a10"}, false}, + {"num accelerators match", listFilters{NumAccelerators: new(8)}, true}, + {"num accelerators no match", listFilters{NumAccelerators: new(4)}, false}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + assert.Equal(t, c.want, c.f.matches(&wf)) + }) + } +} diff --git a/experimental/air/cmd/list_format.go b/experimental/air/cmd/list_format.go new file mode 100644 index 00000000000..398a2523467 --- /dev/null +++ b/experimental/air/cmd/list_format.go @@ -0,0 +1,54 @@ +package aircmd + +import ( + "time" +) + +// buildListRow maps a TrainingWorkflow to a table/JSON row. Optional text +// columns fall back to "-"; host builds the MLflow deep link. +func buildListRow(w *trainingWorkflow, host string) listRow { + experiment := "-" + if e := stripExperimentUserPrefix(w.Status.Mlflow.Experiment); e != "" { + experiment = e + } + + var startedAt *string + duration := "-" + if start := parseRPCTime(w.Status.StartTime); !start.IsZero() { + s := isoFormat(start) + startedAt = &s + + // A still-running run has no end_time, so measure against the current time. + endMillis := time.Now().UnixMilli() + if end := parseRPCTime(w.Status.EndTime); !end.IsZero() { + endMillis = end.UnixMilli() + } + duration = formatDuration(roundMillisToSeconds(endMillis - start.UnixMilli())) + } + + mlflowURL := "-" + if w.Status.Mlflow.ExperimentID != "" && w.Status.Mlflow.RunID != "" { + mlflowURL = mlflowLogsURL(host, &mlflowIdentifiers{ + ExperimentID: w.Status.Mlflow.ExperimentID, + RunID: w.Status.Mlflow.RunID, + }) + } + + accel := "-" + if a := acceleratorLabel(w.Spec.Compute.HardwareAcceleratorType, w.Spec.Compute.AcceleratorCount); a != "" { + accel = a + } + + return listRow{ + RunID: w.JobRunID, + RunName: w.Status.Job.Name, + User: w.Metadata.CreatorName, + Status: trainingWorkflowStatus(w.Status.State), + StartedAt: startedAt, + IsSweep: w.TaskRunID != "", + Experiment: experiment, + Duration: duration, + MLflowURL: mlflowURL, + Accelerators: accel, + } +} diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go new file mode 100644 index 00000000000..989c05fc6da --- /dev/null +++ b/experimental/air/cmd/list_test.go @@ -0,0 +1,180 @@ +package aircmd + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// testWorkflow builds a single-task AIR workflow for a given user and compute, +// as ListTrainingWorkflows would return it. +func testWorkflow(id int64, user, gpuType string, count int, experiment string) trainingWorkflow { + var w trainingWorkflow + w.JobRunID = strconv.FormatInt(id, 10) + w.Metadata.CreatorName = user + w.Spec.Compute.HardwareAcceleratorType = gpuType + w.Spec.Compute.AcceleratorCount = count + w.Status.State = "TRAINING_WORKFLOW_STATE_RUNNING" + w.Status.Job.Name = "run-" + strconv.FormatInt(id, 10) + w.Status.Mlflow.Experiment = experiment + return w +} + +// workflowsBody marshals a ListTrainingWorkflows response page. +func workflowsBody(t *testing.T, nextToken string, wfs ...trainingWorkflow) string { + t.Helper() + b, err := json.Marshal(trainingWorkflowsResponse{TrainingWorkflows: wfs, NextPageToken: nextToken}) + require.NoError(t, err) + return string(b) +} + +// workflowsServer serves one response body per call to the workflows endpoint, +// repeating the last body once exhausted, and a stub for any other request (the +// SDK's well-known config discovery). +func workflowsServer(t *testing.T, bodies ...string) *httptest.Server { + t.Helper() + call := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == listWorkflowsPath { + body := bodies[min(call, len(bodies)-1)] + call++ + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListAirRunsAppliesClientFilters(t *testing.T) { + qwen := testWorkflow(1, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/qwen-train") + llama := testWorkflow(2, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/llama-train") + srv := workflowsServer(t, workflowsBody(t, "", qwen, llama)) + + rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + limit: 10, + filters: listFilters{Experiment: "qwen*"}, + }) + require.NoError(t, err) + require.Len(t, rows, 1) + assert.Equal(t, "1", rows[0].RunID) +} + +func TestListAirRunsLimitTruncates(t *testing.T) { + wfs := []trainingWorkflow{ + testWorkflow(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), + testWorkflow(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), + testWorkflow(3, "me@example.com", "GPU_1xH100", 1, "exp-c"), + } + srv := workflowsServer(t, workflowsBody(t, "", wfs...)) + + rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{limit: 2}) + require.NoError(t, err) + require.Len(t, rows, 2) + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "2", rows[1].RunID) +} + +func TestListAirRunsNoLimitFetchesAll(t *testing.T) { + // limit 0 means "all": follow the page token to the end with no early stop. + page1 := workflowsBody(t, "tok", testWorkflow(1, "me@example.com", "GPU_1xH100", 1, "exp-a")) + page2 := workflowsBody(t, "", testWorkflow(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) + srv := workflowsServer(t, page1, page2) + + rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{limit: 0}) + require.NoError(t, err) + require.Len(t, rows, 2) +} + +func TestListAirRunsPaginates(t *testing.T) { + // First page returns one run and a continuation token; the loop must follow it + // and stop when the token is empty. + page1 := workflowsBody(t, "tok", testWorkflow(1, "me@example.com", "GPU_1xH100", 1, "exp-a")) + page2 := workflowsBody(t, "", testWorkflow(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) + srv := workflowsServer(t, page1, page2) + + rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{limit: 10}) + require.NoError(t, err) + require.Len(t, rows, 2) + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "2", rows[1].RunID) +} + +func TestBuildListRow(t *testing.T) { + var w trainingWorkflow + w.JobRunID = "123" + w.Metadata.CreatorName = "me@example.com" + w.Spec.Compute.HardwareAcceleratorType = "GPU_8xH100" + w.Spec.Compute.AcceleratorCount = 8 + w.Status.State = "TRAINING_WORKFLOW_STATE_TERMINATED_COMPLETED" + w.Status.StartTime = "2026-06-05T17:32:39.000Z" + w.Status.EndTime = "2026-06-05T17:32:51.000Z" + w.Status.Job.Name = "my-run" + w.Status.Mlflow.Experiment = "/Users/me@example.com/exp" + w.Status.Mlflow.ExperimentID = "E1" + w.Status.Mlflow.RunID = "R1" + + row := buildListRow(&w, "https://h.test") + assert.Equal(t, "123", row.RunID) + assert.Equal(t, "my-run", row.RunName) + assert.Equal(t, "me@example.com", row.User) + assert.Equal(t, "SUCCESS", row.Status) + assert.Equal(t, "exp", row.Experiment) + assert.Equal(t, "12s", row.Duration) + assert.Equal(t, "8x H100", row.Accelerators) + assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1/artifacts/logs/node_0", row.MLflowURL) + assert.False(t, row.IsSweep) + require.NotNil(t, row.StartedAt) + assert.Equal(t, "2026-06-05T17:32:39+00:00", *row.StartedAt) +} + +func TestBuildListRowDashFallbacks(t *testing.T) { + // A workflow with no experiment, compute, MLflow IDs, or start time falls back + // to dashes for the optional columns and UNKNOWN for the unset state. + row := buildListRow(&trainingWorkflow{}, "https://h.test") + assert.Equal(t, "-", row.Experiment) + assert.Equal(t, "-", row.Duration) + assert.Equal(t, "-", row.Accelerators) + assert.Equal(t, "-", row.MLflowURL) + assert.Equal(t, "UNKNOWN", row.Status) + assert.Nil(t, row.StartedAt) +} + +func TestBuildListRowSweep(t *testing.T) { + // task_run_id is set only for sweeps, so its presence marks the row. + w := trainingWorkflow{TaskRunID: "456"} + assert.True(t, buildListRow(&w, "https://h.test").IsSweep) +} + +func TestListInvalidLimit(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := newListCommand() + cmd.SetContext(ctx) + require.NoError(t, cmd.Flags().Set("limit", "0")) + + err := cmd.RunE(cmd, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid --limit") +} + +func TestListInvalidFilter(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := newListCommand() + cmd.SetContext(ctx) + require.NoError(t, cmd.Flags().Set("filter", "bogus=1")) + + err := cmd.RunE(cmd, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported --filter key") +} diff --git a/experimental/air/cmd/list_tui.go b/experimental/air/cmd/list_tui.go new file mode 100644 index 00000000000..369205f5992 --- /dev/null +++ b/experimental/air/cmd/list_tui.go @@ -0,0 +1,185 @@ +package aircmd + +import ( + "fmt" + "io" + "strings" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" + "github.com/databricks/cli/libs/cmdio" + "github.com/muesli/termenv" + "github.com/pkg/browser" + "github.com/spf13/cobra" +) + +// renderListText renders the table for text output: an inline navigable table +// in a terminal, otherwise printed once. JSON is handled by the caller. +func renderListText(cmd *cobra.Command, rows []listRow) error { + ctx := cmd.Context() + out := cmd.OutOrStdout() + + color := cmdio.SupportsColor(ctx, out) + r := lipgloss.NewRenderer(out) + if !color { + r.SetColorProfile(termenv.Ascii) + } + + // Navigate only with a full color TTY, at least one row, and no explicit + // --limit (which means "just print these N"). Everything else — piped, + // NO_COLOR, --limit, empty — prints once. + interactive := len(rows) > 0 && + color && + cmdio.IsPagerSupported(ctx) && + !cmd.Flags().Changed("limit") + + if interactive { + _, err := tea.NewProgram( + newListModel(r, rows, color), + tea.WithContext(ctx), + tea.WithInput(cmd.InOrStdin()), + tea.WithOutput(out), + ).Run() + return err + } + + _, err := io.WriteString(out, staticListTable(r, rows, color)) + return err +} + +// staticListTable renders the whole table once, with no selection — used when +// piped or non-interactive. +func staticListTable(r *lipgloss.Renderer, rows []listRow, links bool) string { + if len(rows) == 0 { + return "No runs found.\n" + } + styles := newListStyles(r) + cols := computeListCols(rows) + var b strings.Builder + b.WriteString(styles.renderHeader(cols)) + b.WriteByte('\n') + for _, row := range rows { + b.WriteString(styles.renderRow(cols, row, false, links)) + b.WriteByte('\n') + } + return b.String() +} + +// listModel is the inline, navigable runs table. +type listModel struct { + rows []listRow + styles listStyles + cols listCols + links bool + + cursor int + offset int // index of the first visible row + height int // terminal height, for windowing +} + +func newListModel(r *lipgloss.Renderer, rows []listRow, links bool) listModel { + return listModel{ + rows: rows, + styles: newListStyles(r), + cols: computeListCols(rows), + links: links, + } +} + +func (m listModel) Init() tea.Cmd { return nil } + +// listPageRows is the most rows shown per page. +const listPageRows = 20 + +// visibleCount is how many rows a page shows: at most listPageRows, and never +// more than fits below the header and hint. +func (m listModel) visibleCount() int { + n := min(listPageRows, len(m.rows)) + if m.height > 0 { + n = min(n, m.height-3) + } + return max(1, n) +} + +func (m listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.height = msg.Height + m.offset = m.clampedOffset() + return m, nil + + case tea.KeyMsg: + switch msg.String() { + case "q", "ctrl+c", "esc": + return m, tea.Quit + case "up", "k": + if m.cursor > 0 { + m.cursor-- + } + case "down", "j": + if m.cursor < len(m.rows)-1 { + m.cursor++ + } + case "right": + m.cursor = min(m.cursor+m.visibleCount(), len(m.rows)-1) + case "left": + m.cursor = max(m.cursor-m.visibleCount(), 0) + case "home", "g": + m.cursor = 0 + case "end", "G": + m.cursor = len(m.rows) - 1 + case "enter": + // Open the selected run's MLflow page in the browser. + if len(m.rows) > 0 { + if url := m.rows[m.cursor].MLflowURL; url != "" && url != "-" { + return m, openURL(url) + } + } + } + m.offset = m.clampedOffset() + } + return m, nil +} + +// clampedOffset returns the scroll offset that keeps the cursor visible. +func (m listModel) clampedOffset() int { + visible := m.visibleCount() + offset := min(m.offset, m.cursor) + if m.cursor >= offset+visible { + offset = m.cursor - visible + 1 + } + return max(offset, 0) +} + +func (m listModel) View() string { + if len(m.rows) == 0 { + return m.styles.r.NewStyle().Foreground(colN9).Render("No runs found.") + "\n" + } + + visible := m.visibleCount() + lines := []string{m.styles.renderHeader(m.cols)} + for i := m.offset; i < m.offset+visible && i < len(m.rows); i++ { + lines = append(lines, m.styles.renderRow(m.cols, m.rows[i], i == m.cursor, m.links)) + } + lines = append(lines, m.renderHint()) + return strings.Join(lines, "\n") + "\n" +} + +// renderHint is the faint one-line key legend, with a scroll position when the +// list is windowed. +func (m listModel) renderHint() string { + faint := m.styles.r.NewStyle().Foreground(colN7) + hint := "↑/↓ navigate · ←/→ page · ↵ mlflow · q quit" + if m.visibleCount() < len(m.rows) { + hint += fmt.Sprintf(" · row %d/%d", m.cursor+1, len(m.rows)) + } + return faint.Render(hint) +} + +// openURL opens a URL in the user's default browser, best-effort. +func openURL(url string) tea.Cmd { + return func() tea.Msg { + _ = browser.OpenURL(url) + return nil + } +} diff --git a/experimental/air/cmd/list_tui_render.go b/experimental/air/cmd/list_tui_render.go new file mode 100644 index 00000000000..6a81deceb5c --- /dev/null +++ b/experimental/air/cmd/list_tui_render.go @@ -0,0 +1,239 @@ +package aircmd + +import ( + "strings" + + "github.com/charmbracelet/lipgloss" + "github.com/muesli/termenv" +) + +// DuBois dark-first palette. Neutrals (n7..n12) carry structure; the state +// colors are reserved for the Status column. +const ( + colN7 = lipgloss.Color("#69696B") // gutter / dim + colN9 = lipgloss.Color("#ABABAE") // secondary values + colN11 = lipgloss.Color("#E0E0E3") // body + colN12 = lipgloss.Color("#F1F1F4") // selected text + colOverlay = lipgloss.Color("#1F1F23") // selected-row fill + colRunID = lipgloss.Color("#B7A8E8") // run id + colGreen = lipgloss.Color("#4CD964") // success + colAmber = lipgloss.Color("#E8B84A") // running / pending + colRed = lipgloss.Color("#EB6B6B") // failed + colBlue = lipgloss.Color("#6CA8F0") // MLflow link +) + +const mlflowColWidth = 18 + +// listStyles renders the runs table. The renderer carries the color profile, so +// styles render plain under --no-color / non-tty. +type listStyles struct { + r *lipgloss.Renderer +} + +func newListStyles(r *lipgloss.Renderer) listStyles { + return listStyles{r: r} +} + +// listCols holds the computed width of each variable-width column. MLflow is +// fixed (a short link) and the gutter is one cell. +type listCols struct { + runID, experiment, status, started, duration, user, accel int +} + +// columnCap bounds the widest free-text columns so one long value can't dominate +// the row. +const columnCap = 36 + +func computeListCols(rows []listRow) listCols { + c := listCols{ + runID: len("Run ID"), + experiment: len("Experiment"), + status: len("Status"), + started: len("Started"), + duration: len("Duration"), + user: len("User"), + accel: len("Accelerators"), + } + for _, r := range rows { + c.runID = max(c.runID, lipgloss.Width(r.RunID)) + c.experiment = min(columnCap, max(c.experiment, lipgloss.Width(r.Experiment))) + c.status = max(c.status, lipgloss.Width("● "+r.Status)) + c.started = max(c.started, lipgloss.Width(startedDisplay(r))) + c.duration = max(c.duration, lipgloss.Width(r.Duration)) + c.user = min(columnCap, max(c.user, lipgloss.Width(r.User))) + c.accel = max(c.accel, lipgloss.Width(r.Accelerators)) + } + return c +} + +// renderHeader renders the muted column-title row. +func (s listStyles) renderHeader(cols listCols) string { + h := func(text string, width int, right bool) string { + return s.r.NewStyle().Foreground(colN9).Render(pad(text, width, right)) + } + cells := []string{ + pad(" ", 1, false), + h("Run ID", cols.runID, false), + h("Experiment", cols.experiment, false), + h("Status", cols.status, false), + h("Started", cols.started, false), + h("Duration", cols.duration, true), + h("MLflow", mlflowColWidth, false), + h("User", cols.user, false), + h("Accelerators", cols.accel, false), + } + // Trim trailing pad on the final column so rows carry no trailing whitespace. + return strings.TrimRight(strings.Join(cells, " "), " ") +} + +// renderRow renders one run. The selected row uses a subtle overlay fill + n12 +// text (not full inversion, not per-state color). +func (s listStyles) renderRow(cols listCols, r listRow, selected, links bool) string { + base := s.r.NewStyle() + if selected { + base = base.Background(colOverlay) + } + fg := func(c lipgloss.Color) lipgloss.Color { + if selected { + return colN12 + } + return c + } + + gutter := " " + if selected { + gutter = "▸" + } + + cells := []string{ + s.cell(base, gutter, 1, fg(colN7), false, false, ""), + s.cell(base, r.RunID, cols.runID, fg(colRunID), false, false, ""), + s.cell(base, r.Experiment, cols.experiment, fg(colN11), false, false, ""), + s.cell(base, "● "+r.Status, cols.status, fg(statusColor(r.Status)), false, false, ""), + s.cell(base, startedDisplay(r), cols.started, fg(colN9), false, false, ""), + s.cell(base, r.Duration, cols.duration, fg(colN9), true, false, ""), + s.mlflowCell(base, r, selected, links), + s.cell(base, r.User, cols.user, fg(colN9), false, false, ""), + s.cell(base, r.Accelerators, cols.accel, fg(colN9), false, false, ""), + } + // Trim the final column's trailing pad so rows carry no trailing whitespace. + return strings.TrimRight(strings.Join(cells, base.Render(" ")), " ") +} + +// cell renders one padded, colored cell. The text is truncated to width, then +// padded with (background-only) spaces so columns align even when the text is +// styled or hyperlinked. +func (s listStyles) cell(base lipgloss.Style, text string, width int, fg lipgloss.Color, right, underline bool, link string) string { + text = truncate(text, width) + style := base.Foreground(fg) + if underline { + style = style.Underline(true) + } + rendered := style.Render(text) + if link != "" { + rendered = termenv.Hyperlink(link, rendered) + } + gap := max(width-lipgloss.Width(text), 0) + padStr := base.Render(strings.Repeat(" ", gap)) + if right { + return padStr + rendered + } + return rendered + padStr +} + +// mlflowCell renders the fixed-width MLflow column: a short, blue, underlined +// OSC 8 hyperlink (when links are enabled), or "-" when the run has no link. +func (s listStyles) mlflowCell(base lipgloss.Style, r listRow, selected, links bool) string { + if r.MLflowURL == "" || r.MLflowURL == "-" { + fg := colN9 + if selected { + fg = colN12 + } + return s.cell(base, "-", mlflowColWidth, fg, false, false, "") + } + fg := colBlue + if selected { + fg = colN12 + } + link := "" + if links { + link = r.MLflowURL + } + return s.cell(base, mlflowDisplay(r.MLflowURL), mlflowColWidth, fg, false, true, link) +} + +// statusColor maps an air run status word to its data color. +func statusColor(status string) lipgloss.Color { + switch status { + case "SUCCESS": + return colGreen + case "RUNNING", "PENDING", "TERMINATING": + return colAmber + case "FAILED": + return colRed + default: // CANCELED / UNKNOWN + return colN7 + } +} + +// startedDisplay trims the row's ISO start timestamp to second precision +// ("2006-01-02T15:04:05"), or "-" when the run hasn't started. +func startedDisplay(r listRow) string { + if r.StartedAt == nil { + return "-" + } + s := *r.StartedAt + if len(s) >= 19 { + return s[:19] + } + return s +} + +// mlflowDisplay shortens an MLflow run URL to a "…/runs/<id-prefix>" label; the +// OSC 8 target keeps the full URL. +func mlflowDisplay(url string) string { + id := mlflowRunID(url) + if id == "" { + return truncate(url, mlflowColWidth) + } + if len(id) > 8 { + id = id[:8] + "…" + } + return "…/runs/" + id +} + +// mlflowRunID extracts the run-id path segment from an MLflow URL. +func mlflowRunID(url string) string { + _, after, ok := strings.Cut(url, "/runs/") + if !ok { + return "" + } + id, _, _ := strings.Cut(after, "/") + return id +} + +// pad pads (or truncates) s to a visible width of n, right-aligned when right is +// set. It measures visible width, so it is safe on styled strings. +func pad(s string, n int, right bool) string { + s = truncate(s, n) + gap := max(n-lipgloss.Width(s), 0) + if gap == 0 { + return s + } + fill := strings.Repeat(" ", gap) + if right { + return fill + s + } + return s + fill +} + +// truncate shortens s to a visible width of n, appending "…" on overflow. +func truncate(s string, n int) string { + if lipgloss.Width(s) <= n { + return s + } + if n <= 1 { + return "…" + } + return string([]rune(s)[:n-1]) + "…" +} diff --git a/experimental/air/cmd/list_tui_test.go b/experimental/air/cmd/list_tui_test.go new file mode 100644 index 00000000000..b5b7ca2cbc8 --- /dev/null +++ b/experimental/air/cmd/list_tui_test.go @@ -0,0 +1,175 @@ +package aircmd + +import ( + "io" + "strconv" + "testing" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" + "github.com/muesli/termenv" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// testListRows is a small fixture covering each status color, a present and an +// absent MLflow link, and a still-running (no end) row. +func testListRows() []listRow { + return []listRow{ + {RunID: "1", Experiment: "qwen-train", User: "me@example.com", Status: "SUCCESS", StartedAt: new("2026-06-05T17:32:39.000000+00:00"), Duration: "1m 14s", MLflowURL: "https://h/ml/experiments/E/runs/04c41514fbb0/artifacts/logs/node_0", Accelerators: "8x H100"}, + {RunID: "2", Experiment: "llama-train", User: "me@example.com", Status: "RUNNING", StartedAt: new("2026-06-05T18:43:24.000000+00:00"), Duration: "3m 32s", MLflowURL: "-", Accelerators: "1x A10"}, + {RunID: "3", Experiment: "mixtral", User: "me@example.com", Status: "FAILED", StartedAt: nil, Duration: "-", MLflowURL: "-", Accelerators: "-"}, + } +} + +func testListModel() listModel { + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + return newListModel(r, testListRows(), false) +} + +func key(t *testing.T, m listModel, s string) listModel { + t.Helper() + next, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(s)}) + return next.(listModel) +} + +func TestListModelNavigation(t *testing.T) { + m := testListModel() + require.Equal(t, 0, m.cursor) + + m = key(t, m, "j") + assert.Equal(t, 1, m.cursor) + m = key(t, key(t, m, "k"), "k") // clamp at top + assert.Equal(t, 0, m.cursor) + + for range len(m.rows) + 2 { // clamp at bottom + m = key(t, m, "j") + } + assert.Equal(t, len(m.rows)-1, m.cursor) +} + +func TestListModelWindowScrolls(t *testing.T) { + m := testListModel() + // Height 5 leaves room for ~2 rows (header + hint reserved). + next, _ := m.Update(tea.WindowSizeMsg{Width: 200, Height: 5}) + m = next.(listModel) + require.Equal(t, 2, m.visibleCount()) + require.Equal(t, 0, m.offset) + + m = key(t, key(t, m, "j"), "j") // move to row index 2, past the window + assert.Equal(t, 2, m.cursor) + assert.Equal(t, 1, m.offset, "window scrolled to keep the cursor visible") +} + +func TestListModelPageCap(t *testing.T) { + rows := make([]listRow, 50) + for i := range rows { + rows[i] = listRow{RunID: strconv.Itoa(i)} + } + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + m := newListModel(r, rows, false) + + // A tall terminal still shows at most listPageRows per page. + next, _ := m.Update(tea.WindowSizeMsg{Width: 80, Height: 100}) + assert.Equal(t, listPageRows, next.(listModel).visibleCount()) +} + +func TestListModelPaging(t *testing.T) { + rows := make([]listRow, 10) + for i := range rows { + rows[i] = listRow{RunID: strconv.Itoa(i)} + } + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + m := newListModel(r, rows, false) + + // Height 7 leaves a 4-row window (header + hint reserved). + next, _ := m.Update(tea.WindowSizeMsg{Width: 80, Height: 7}) + m = next.(listModel) + require.Equal(t, 4, m.visibleCount()) + + page := func(k tea.KeyType) { + n, _ := m.Update(tea.KeyMsg{Type: k}) + m = n.(listModel) + } + page(tea.KeyRight) + assert.Equal(t, 4, m.cursor) + page(tea.KeyEnd) + assert.Equal(t, 9, m.cursor) + page(tea.KeyLeft) + assert.Equal(t, 5, m.cursor) + page(tea.KeyHome) + assert.Equal(t, 0, m.cursor) +} + +func TestListModelQuit(t *testing.T) { + m := testListModel() + _, cmd := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("q")}) + require.NotNil(t, cmd) + assert.Equal(t, tea.QuitMsg{}, cmd()) +} + +func TestListModelView(t *testing.T) { + next, _ := testListModel().Update(tea.WindowSizeMsg{Width: 200, Height: 24}) + out := next.(listModel).View() + + assert.NotContains(t, out, "\x1b", "Ascii profile + no links should produce no escapes") + for _, want := range []string{ + "Run ID", "Experiment", "Status", "Started", "Duration", "MLflow", "User", "Accelerators", + "qwen-train", "● SUCCESS", "● RUNNING", "● FAILED", + "…/runs/04c41514…", // shortened MLflow link + "2026-06-05T17:32:39", // started trimmed to seconds + "▸", // selection gutter on the first row + "↑/↓ navigate", // hint line + } { + assert.Contains(t, out, want) + } +} + +func TestStaticListTable(t *testing.T) { + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + out := staticListTable(r, testListRows(), false) + + assert.NotContains(t, out, "\x1b") + assert.NotContains(t, out, "▸", "static table has no selection") + for _, want := range []string{"Run ID", "1", "qwen-train", "…/runs/04c41514…", "Accelerators"} { + assert.Contains(t, out, want) + } + + assert.Equal(t, "No runs found.\n", staticListTable(r, nil, false)) +} + +func TestStatusColor(t *testing.T) { + assert.Equal(t, colGreen, statusColor("SUCCESS")) + assert.Equal(t, colAmber, statusColor("RUNNING")) + assert.Equal(t, colAmber, statusColor("PENDING")) + assert.Equal(t, colRed, statusColor("FAILED")) + assert.Equal(t, colN7, statusColor("CANCELED")) + assert.Equal(t, colN7, statusColor("UNKNOWN")) +} + +func TestStartedDisplay(t *testing.T) { + assert.Equal(t, "-", startedDisplay(listRow{})) + assert.Equal(t, "2026-06-05T17:32:39", startedDisplay(listRow{StartedAt: new("2026-06-05T17:32:39.000000+00:00")})) +} + +func TestMLflowDisplay(t *testing.T) { + assert.Equal(t, "…/runs/04c41514…", mlflowDisplay("https://h/ml/experiments/E/runs/04c41514fbb0/artifacts/logs/node_0")) + assert.Equal(t, "…/runs/run1", mlflowDisplay("https://h/ml/experiments/E/runs/run1/artifacts/logs/node_0")) + assert.LessOrEqual(t, lipgloss.Width(mlflowDisplay("https://h/no-runs/here")), mlflowColWidth) +} + +func TestMLflowRunID(t *testing.T) { + assert.Equal(t, "abc123", mlflowRunID("https://h/ml/experiments/1/runs/abc123/artifacts")) + assert.Empty(t, mlflowRunID("https://h/no-runs-here")) +} + +func TestPadAndTruncate(t *testing.T) { + assert.Equal(t, "ab ", pad("ab", 5, false)) + assert.Equal(t, " ab", pad("ab", 5, true)) + assert.Equal(t, "abcd…", truncate("abcdefgh", 5)) + assert.Equal(t, "abc", truncate("abc", 5)) +} diff --git a/experimental/air/cmd/stubs_test.go b/experimental/air/cmd/stubs_test.go index 5e35bcdcd14..b9f5c330f00 100644 --- a/experimental/air/cmd/stubs_test.go +++ b/experimental/air/cmd/stubs_test.go @@ -14,7 +14,6 @@ import ( func TestStubCommandsReturnNotImplemented(t *testing.T) { stubs := map[string]*cobra.Command{ "run": newRunCommand(), - "list": newListCommand(), "logs": newLogsCommand(), "cancel": newCancelCommand(), "register-image": newRegisterImageCommand(), diff --git a/libs/cmdio/io.go b/libs/cmdio/io.go index e57c90974b4..2b9b395ce13 100644 --- a/libs/cmdio/io.go +++ b/libs/cmdio/io.go @@ -83,6 +83,15 @@ func IsPromptSupported(ctx context.Context) bool { return c.capabilities.SupportsPrompt() } +// IsPagerSupported reports whether stdin, stdout, and stderr are all interactive +// terminals. This is the requirement for a full-screen or navigable output +// program: unlike IsPromptSupported it also checks stdout, so it returns false +// when stdout is piped or redirected. +func IsPagerSupported(ctx context.Context) bool { + c := fromContext(ctx) + return c.capabilities.SupportsPager() +} + // SupportsColor returns true if the given writer supports colored output. // This checks both TTY status and environment variables (NO_COLOR, TERM=dumb). func SupportsColor(ctx context.Context, w io.Writer) bool { From ca7c0f3ae1e86ae32f1514ce6fe036e4f7bf235a Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Wed, 24 Jun 2026 16:09:40 -0700 Subject: [PATCH 06/18] AIR CLI Integration: collapse `air get run` back to `air get JOB_RUN_ID` (#5685) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why We decided to cut the `get run` sub-resource. The run-status command is now just `air get <id>` — flat, with no `run` subcommand. ## Changes - Removed the `get` parent group and its `run` subcommand; `newGetCommand` is the run-status command itself (`Use: "get JOB_RUN_ID"`, `ExactArgs(1)`). - No change to output behavior — the styled config box, `JOB_RUN_ID` naming, `Job Link` header, status table, and sweep view are all unchanged. - Regenerated the `experimental/air/get` and `experimental/air/help` acceptance outputs; updated doc comments and tests that referenced `air get run`. ## Tests - Added `TestGetCommandShape`: asserts `Use == "get JOB_RUN_ID"`, no registered subcommands, and exactly one arg required. - Updated the existing `get` unit tests (invalid id, not-found text/JSON, templates, `buildGetData`) to the new entry point. - `experimental/air/{get,help}` acceptance regenerated; full air unit + acceptance suites pass. This pull request and its description were written by Isaac. --- acceptance/experimental/air/get/output.txt | 8 ++++---- acceptance/experimental/air/get/script | 8 ++++---- acceptance/experimental/air/help/output.txt | 2 +- experimental/air/cmd/get.go | 19 +++++-------------- experimental/air/cmd/get_test.go | 19 ++++++++++++++++--- experimental/air/cmd/render.go | 4 ++-- 6 files changed, 32 insertions(+), 28 deletions(-) diff --git a/acceptance/experimental/air/get/output.txt b/acceptance/experimental/air/get/output.txt index 4bc127d79b4..6e51d7debf9 100644 --- a/acceptance/experimental/air/get/output.txt +++ b/acceptance/experimental/air/get/output.txt @@ -1,6 +1,6 @@ === get (text) ->>> [CLI] experimental air get run 123 +>>> [CLI] experimental air get 123 ╭─ Configuration ────────────────────────────────────────────────╮ │ │ @@ -35,7 +35,7 @@ Run URL: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 === get (json) ->>> [CLI] experimental air get run 123 -o json +>>> [CLI] experimental air get 123 -o json { "v": 1, "ts": "[TIMESTAMP]", @@ -52,13 +52,13 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 } === invalid run id ->>> [CLI] experimental air get run notanumber +>>> [CLI] experimental air get notanumber Error: invalid JOB_RUN_ID "notanumber": must be a positive integer Exit code: 1 === invalid run id (json) ->>> [CLI] experimental air get run notanumber -o json +>>> [CLI] experimental air get notanumber -o json { "v": 1, "ts": "[TIMESTAMP]", diff --git a/acceptance/experimental/air/get/script b/acceptance/experimental/air/get/script index b775d06a48d..ee66b4aff04 100644 --- a/acceptance/experimental/air/get/script +++ b/acceptance/experimental/air/get/script @@ -1,11 +1,11 @@ title "get (text)" -trace $CLI experimental air get run 123 +trace $CLI experimental air get 123 title "get (json)" -trace $CLI experimental air get run 123 -o json +trace $CLI experimental air get 123 -o json title "invalid run id" -errcode trace $CLI experimental air get run notanumber +errcode trace $CLI experimental air get notanumber title "invalid run id (json)" -errcode trace $CLI experimental air get run notanumber -o json +errcode trace $CLI experimental air get notanumber -o json diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index 9348686ce41..8eac074581d 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -11,7 +11,7 @@ Usage: Available Commands: cancel Cancel one or more runs - get Show details for a specific resource + get Show status, configuration, and timing details for a specific run list List your recent runs (active and completed) for the current profile logs Stream or fetch logs for a run register-image Mirror a Docker image into the workspace registry diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 98581fd5dd7..0da1f082e12 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -13,7 +13,7 @@ import ( "github.com/spf13/cobra" ) -// getData is the payload printed by `air get run`. The json-tagged fields form +// getData is the payload printed by `air get`. The json-tagged fields form // the machine-readable output; fields tagged `json:"-"` are shown only in the // human-readable text view. type getData struct { @@ -27,7 +27,7 @@ type getData struct { MLflowURL *string `json:"mlflow_url"` // The fields below are pre-rendered text-view cells, excluded from JSON - // (matching `air get run --json`). Each shows "N/A" when its value is + // (matching `air get --json`). Each shows "N/A" when its value is // missing. The styled single-run renderer (render.go) consumes them; the // Run ID, Status, and MLflow Run cells it draws are styled and hyperlinked // there rather than stored here. @@ -63,20 +63,11 @@ Sweep Tasks: {{- end}} ` -// newGetCommand is the `get` parent group. Subcommands name the resource to -// describe, e.g. `air get run JOB_RUN_ID`, mirroring the Python CLI. +// newGetCommand returns the `air get JOB_RUN_ID` command, which shows status, +// configuration, and timing details for a specific run. func newGetCommand() *cobra.Command { cmd := &cobra.Command{ - Use: "get", - Short: "Show details for a specific resource", - } - cmd.AddCommand(newGetRunCommand()) - return cmd -} - -func newGetRunCommand() *cobra.Command { - cmd := &cobra.Command{ - Use: "run JOB_RUN_ID", + Use: "get JOB_RUN_ID", Args: root.ExactArgs(1), Short: "Show status, configuration, and timing details for a specific run", Annotations: map[string]string{ diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go index 7f4d7848076..8ff864e387c 100644 --- a/experimental/air/cmd/get_test.go +++ b/experimental/air/cmd/get_test.go @@ -29,10 +29,23 @@ func renderGet(t *testing.T, data getData) string { return buf.String() } +// TestGetCommandShape locks in that `get` takes the run id directly as +// `air get JOB_RUN_ID` and has no `run` subcommand (it was collapsed back into +// `get`). The acceptance test exercises the happy path end to end. +func TestGetCommandShape(t *testing.T) { + cmd := newGetCommand() + assert.Equal(t, "get JOB_RUN_ID", cmd.Use) + assert.Empty(t, cmd.Commands(), "get must not register subcommands") + // ExactArgs(1): exactly one run id is required. + assert.NoError(t, cmd.Args(cmd, []string{"123"})) + assert.Error(t, cmd.Args(cmd, []string{})) + assert.Error(t, cmd.Args(cmd, []string{"1", "2"})) +} + func TestGetRunInvalidID(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) - cmd := withOutput(newGetRunCommand(), flags.OutputText) + cmd := withOutput(newGetCommand(), flags.OutputText) cmd.SetContext(ctx) err := cmd.RunE(cmd, []string{"abc"}) @@ -45,7 +58,7 @@ func TestGetRunNotFound(t *testing.T) { m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return( nil, apierr.ErrResourceDoesNotExist) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) - cmd := withOutput(newGetRunCommand(), flags.OutputText) + cmd := withOutput(newGetCommand(), flags.OutputText) cmd.SetContext(ctx) err := cmd.RunE(cmd, []string{"5"}) @@ -60,7 +73,7 @@ func TestGetRunNotFoundJSON(t *testing.T) { nil, apierr.ErrResourceDoesNotExist) ctx := cmdctx.SetWorkspaceClient(t.Context(), m.WorkspaceClient) ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) - cmd := withOutput(newGetRunCommand(), flags.OutputJSON) + cmd := withOutput(newGetCommand(), flags.OutputJSON) cmd.SetContext(ctx) // In JSON mode the not-found error is a structured envelope, not a bare error. diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go index d6c2e6a25be..ff7c0298671 100644 --- a/experimental/air/cmd/render.go +++ b/experimental/air/cmd/render.go @@ -130,8 +130,8 @@ func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceCl // stdout is not a hyperlink-capable terminal (piped, redirected, NO_COLOR). // In that case the OSC 8 hyperlinks on the Run ID / MLflow Run cells // degrade to plain labels and the URLs would otherwise disappear from text - // output, breaking workflows like `air get run X > out.txt` or - // `NO_COLOR=1 air get run X` that the previous `Job Link:` line supported. + // output, breaking workflows like `air get X > out.txt` or + // `NO_COLOR=1 air get X` that the previous `Job Link:` line supported. if view.dashboardURL != "" { fmt.Fprintf(out, "Run URL: %s\n", view.dashboardURL) } From 60adcaa04a4ac10284b608ad04beba16e4824268 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Tue, 30 Jun 2026 14:35:53 -0700 Subject: [PATCH 07/18] AIR CLI Integration: Adding support for air run configuration (#5657) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Changes Ports the air run YAML config schema and its structural validation from the Python CLI (cli/sdk/config.py) to Go, under experimental/air/cmd/. - Schema (runconfig.go): the top-level runConfig plus the nested environment (with docker_image), code_source/snapshot/git, and permission blocks. Reuses the compute model from the parent branch. Includes custom YAML unmarshalers for the three polymorphic fields that don't map to a single Go type: environment.dependencies (string path or inline list), environment.version (string or int), and git.remote (bool or remote-name string). - Loader (runconfig_load.go): loadRunConfig decodes a YAML file with KnownFields(true) — mirroring pydantic's extra="forbid" so unknown keys are rejected — then runs the validation pass. - Validation: every structural rule from the Python schema — required fields, the experiment_name/mlflow_run_name task-key regex and length caps, secret-ref scope/key format, the environment docker-image/dependencies/version exclusivity rules, git branch-xor-commit and remote-requires-branch rules, code_source snapshot requirements, and include_paths relative/no-traversal checks. Two deliberate divergences from the Python schema, both following from the training-service-only port: - The compute.node_pool_id / compute.pool_name fields were already dropped on the parent branch. - The top-level priority field is dropped here: it's a node-pool queue-ordering knob (it requires a pool in Python) with no meaning for serverless workloads. ## Why "Structural" validation (types, required fields, format/cross-field rules) needs no workspace access, so it's a self-contained, fully unit-testable unit that's worth landing on its own ahead of the launch logic. Splitting it out keeps the upcoming handle_run PR focused on orchestration rather than mixing in ~900 lines of schema. The extra="forbid" / KnownFields behavior is load-bearing: it's what turns a typo'd or stale config key into an actionable error instead of a silently-ignored field, so it's preserved faithfully. This is stacked on air-integration-m2-1 (the compute model). ## Tests New unit tests in runconfig_test.go (62 subtests, table-driven), covering: - Loading a minimal config and a full-featured config (all blocks populated). - Each polymorphic union decoding both of its forms (dependencies string vs list, git.remote bool vs string, default-unset). - Unknown-field rejection at top level and nested — including explicit cases asserting the dropped priority field and the not-yet-ported _bases_ key surface as errors. - Every validation rule's failure mode, plus file-level errors (missing file, empty file). go test ./experimental/air/... passes; ./task lint-q reports 0 issues. --- experimental/air/cmd/runconfig.go | 466 +++++++++++++++++++++++++ experimental/air/cmd/runconfig_load.go | 52 +++ experimental/air/cmd/runconfig_test.go | 419 ++++++++++++++++++++++ 3 files changed, 937 insertions(+) create mode 100644 experimental/air/cmd/runconfig.go create mode 100644 experimental/air/cmd/runconfig_load.go create mode 100644 experimental/air/cmd/runconfig_test.go diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go new file mode 100644 index 00000000000..5a0b6f27e1e --- /dev/null +++ b/experimental/air/cmd/runconfig.go @@ -0,0 +1,466 @@ +package aircmd + +import ( + "errors" + "fmt" + "maps" + "regexp" + "slices" + "strings" + + "go.yaml.in/yaml/v3" +) + +// This file ports the run YAML schema and its structural validation from the +// Python CLI's sdk/config.py. "Structural" means types, required fields, and +// format/cross-field rules that need no workspace access. Online checks (e.g. +// GPU availability) and git/filesystem checks run at launch time and are +// intentionally not ported here. +// +// Divergences from the Python schema: compute.node_pool_id / compute.pool_name +// (see compute.go) and the top-level `priority` field are dropped because AIR +// does not support node-pool placement. priority is a pool-queue-ordering knob, +// so it goes with the pool fields. + +// REGEX_TASK_KEY_CHARS: ASCII alphanumeric, hyphen, underscore only (no periods). +// Explicit ASCII class, not \w: \w matches Unicode letters that the ASCII-only +// Jobs API task_key rejects. +var taskKeyRe = regexp.MustCompile(`^[A-Za-z0-9_-]+$`) + +// gitRefRe guards branch/remote names against command injection. Only safe ref +// characters are allowed. +var gitRefRe = regexp.MustCompile(`^[\w./-]+$`) + +// runConfig is the top-level run YAML schema: experiment_name + compute / +// environment / code_source plus the command and run options. +type runConfig struct { + ExperimentName string `yaml:"experiment_name"` + Compute *computeConfig `yaml:"compute"` + Environment *environmentConfig `yaml:"environment"` + Command *string `yaml:"command"` + EnvVariables map[string]string `yaml:"env_variables"` + Secrets map[string]string `yaml:"secrets"` + CodeSource *codeSourceConfig `yaml:"code_source"` + // MaxRetries defaults to 3 when unset; default-filling is a normalization + // concern handled at launch, so a nil pointer is left as-is here. + MaxRetries *int `yaml:"max_retries"` + TimeoutMinutes *int `yaml:"timeout_minutes"` + IdempotencyToken *string `yaml:"idempotency_token"` + Parameters map[string]any `yaml:"parameters"` + MLflowRunName *string `yaml:"mlflow_run_name"` + MLflowExperimentDirectory *string `yaml:"mlflow_experiment_directory"` + Permissions []permission `yaml:"permissions"` + UsagePolicyName *string `yaml:"usage_policy_name"` + UsagePolicyID *string `yaml:"usage_policy_id"` +} + +// validate runs structural validation over the whole config, returning the first +// failure. Fields are checked in declaration order to keep error output stable. +func (c *runConfig) validate() error { + if err := validateExperimentName(c.ExperimentName); err != nil { + return err + } + + if c.Compute == nil { + return errors.New("compute: section is required") + } + if err := c.Compute.validate(); err != nil { + return err + } + + if c.Environment != nil { + if err := c.Environment.validate(); err != nil { + return err + } + } + + // command is optional in the type system but required in practice, matching + // the Python validate_script_fields model validator. + if c.Command == nil { + return errors.New("command is required") + } + if err := validateCommand(*c.Command); err != nil { + return err + } + + if err := validateSecretRefs(c.Secrets); err != nil { + return err + } + + // A name can't be both a plain env var and a secret: the precedence would be + // ambiguous and could leak the secret. Sorted for a stable error. + for _, name := range slices.Sorted(maps.Keys(c.EnvVariables)) { + if _, ok := c.Secrets[name]; ok { + return fmt.Errorf("%q is set in both env_variables and secrets; remove it from one", name) + } + } + + if c.CodeSource != nil { + if err := c.CodeSource.validate(); err != nil { + return err + } + } + + if c.MaxRetries != nil && *c.MaxRetries < 0 { + return fmt.Errorf("max_retries must be >= 0, got %d", *c.MaxRetries) + } + + if c.TimeoutMinutes != nil && *c.TimeoutMinutes < 1 { + return fmt.Errorf("timeout_minutes must be >= 1, got %d", *c.TimeoutMinutes) + } + + if c.IdempotencyToken != nil { + v := strings.TrimSpace(*c.IdempotencyToken) + if v == "" { + return errors.New("idempotency_token cannot be empty") + } + if len(v) > 64 { + return errors.New("idempotency_token must be 64 characters or less") + } + } + + if c.MLflowRunName != nil { + v := strings.TrimSpace(*c.MLflowRunName) + if v == "" { + return errors.New("mlflow_run_name cannot be empty") + } + if len(v) > 100 { + return fmt.Errorf("mlflow_run_name must be 100 characters or less (got %d)", len(v)) + } + if !taskKeyRe.MatchString(v) { + return fmt.Errorf("invalid mlflow_run_name %q: only alphanumeric characters, hyphens, and underscores are allowed", v) + } + } + + if c.MLflowExperimentDirectory != nil { + v := strings.TrimSpace(*c.MLflowExperimentDirectory) + if v == "" { + return errors.New("mlflow_experiment_directory cannot be empty") + } + // MLflow experiments live under the workspace tree. + if !strings.HasPrefix(v, "/Workspace") { + return fmt.Errorf("mlflow_experiment_directory must start with '/Workspace', got: %s", v) + } + } + + for i := range c.Permissions { + if err := c.Permissions[i].validate(); err != nil { + return err + } + } + + // A usage policy is given by name or id, never both; the name resolves to an + // id at launch. + if c.UsagePolicyName != nil && c.UsagePolicyID != nil { + return errors.New("usage_policy_name and usage_policy_id are mutually exclusive; set only one") + } + if c.UsagePolicyName != nil { + v := strings.TrimSpace(*c.UsagePolicyName) + if v == "" { + return errors.New("usage_policy_name must not be empty") + } + // 127 matches the server-side max_length on the policy name filter. + if len(v) > 127 { + return fmt.Errorf("usage_policy_name must be at most 127 characters, got %d", len(v)) + } + } + if c.UsagePolicyID != nil && strings.TrimSpace(*c.UsagePolicyID) == "" { + return errors.New("usage_policy_id must not be empty") + } + + return nil +} + +// validateExperimentName enforces the Databricks Jobs API task_key constraints: +// the experiment_name becomes a task key, which caps at 100 characters and allows +// only alphanumerics, hyphens, and underscores. +func validateExperimentName(v string) error { + if v == "" { + return errors.New("experiment_name cannot be empty") + } + if len(v) > 100 { + return fmt.Errorf("experiment_name must be 100 characters or less (got %d); this is the Jobs API task_key length limit", len(v)) + } + if !taskKeyRe.MatchString(v) { + return fmt.Errorf("invalid experiment_name %q: only alphanumeric characters, hyphens (-), and underscores (_) are allowed", v) + } + return nil +} + +// validateCommand enforces command is non-empty and within the line-count cap. +func validateCommand(v string) error { + if strings.TrimSpace(v) == "" { + return errors.New("command cannot be empty") + } + lineCount := strings.Count(v, "\n") + 1 + if lineCount > 1000 { + return fmt.Errorf("command is too long (%d lines); maximum is 1000 lines — move complex logic into a script in your code_source", lineCount) + } + return nil +} + +// validateSecretRefs checks that secret references use the "scope/key" format. +func validateSecretRefs(secrets map[string]string) error { + for varName, ref := range secrets { + parts := strings.Split(ref, "/") + if len(parts) != 2 { + return fmt.Errorf("invalid secret reference %q for variable %q: expected format 'scope/key' (e.g., my_scope/hf_token)", ref, varName) + } + if parts[0] == "" || parts[1] == "" { + return fmt.Errorf("invalid secret reference %q for variable %q: scope and key cannot be empty", ref, varName) + } + } + return nil +} + +// environmentConfig is the `environment` block: dependencies and/or a custom +// docker image. +type environmentConfig struct { + Dependencies dependencies `yaml:"dependencies"` + Version stringOrInt `yaml:"version"` + DockerImage *dockerImageConfig `yaml:"docker_image"` +} + +func (e *environmentConfig) validate() error { + // docker_image is exclusive with dependencies/version: the image already pins + // the full runtime. + if e.DockerImage != nil { + var conflicting []string + if e.Dependencies.set { + conflicting = append(conflicting, "dependencies") + } + if e.Version.set { + conflicting = append(conflicting, "version") + } + if len(conflicting) > 0 { + return fmt.Errorf("when 'docker_image' is specified under 'environment', these fields are not allowed: %s", strings.Join(conflicting, ", ")) + } + return e.DockerImage.validate() + } + + // version pins the client image version, which is only meaningful for an + // inline (list) dependency set — a requirements.yaml file carries its own. + if e.Version.set { + if e.Dependencies.set && !e.Dependencies.isList { + return errors.New("'environment.version' is only valid with inline dependencies (a list); when 'dependencies' points to a requirements.yaml file, set the version inside that file") + } + if !e.Dependencies.set { + return errors.New("'environment.version' requires inline 'dependencies' (a list of packages)") + } + } + + return nil +} + +// dependencies is environment.dependencies, which is polymorphic: a string is a +// path to a requirements.yaml file; a list is an inline package list. +type dependencies struct { + set bool + isList bool + path string + list []string +} + +func (d *dependencies) UnmarshalYAML(node *yaml.Node) error { + switch node.Kind { + case yaml.ScalarNode: + d.set, d.isList = true, false + return node.Decode(&d.path) + case yaml.SequenceNode: + d.set, d.isList = true, true + return node.Decode(&d.list) + default: + return errors.New("environment.dependencies must be a string path or a list of packages") + } +} + +// stringOrInt holds a scalar that may be a string or an integer in YAML +// (environment.version). The raw text is kept; integer-format validation is a +// launch-time concern. +type stringOrInt struct { + set bool + raw string +} + +func (s *stringOrInt) UnmarshalYAML(node *yaml.Node) error { + if node.Kind != yaml.ScalarNode { + return errors.New("environment.version must be a string or integer") + } + s.set = true + s.raw = node.Value + return nil +} + +// dockerImageConfig is environment.docker_image. +type dockerImageConfig struct { + URL string `yaml:"url"` +} + +func (d *dockerImageConfig) validate() error { + if strings.TrimSpace(d.URL) == "" { + return errors.New("docker_image.url cannot be empty") + } + return nil +} + +// codeSourceConfig is the `code_source` block. Only the "snapshot" type exists. +type codeSourceConfig struct { + Type string `yaml:"type"` + Snapshot *snapshotSourceConfig `yaml:"snapshot"` +} + +func (c *codeSourceConfig) validate() error { + if c.Type != "snapshot" { + return fmt.Errorf("code_source.type must be 'snapshot', got %q", c.Type) + } + if c.Snapshot == nil { + return errors.New("code_source.type='snapshot' requires a snapshot configuration") + } + return c.Snapshot.validate() +} + +// snapshotSourceConfig describes a local directory to tar and upload. +type snapshotSourceConfig struct { + RootPath string `yaml:"root_path"` + RemoteVolume *string `yaml:"remote_volume"` + Git *gitRef `yaml:"git"` + IncludePaths []string `yaml:"include_paths"` +} + +func (s *snapshotSourceConfig) validate() error { + if strings.TrimSpace(s.RootPath) == "" { + return errors.New("code_source.snapshot.root_path cannot be empty") + } + + if s.RemoteVolume != nil && !strings.HasPrefix(*s.RemoteVolume, "/Volumes/") { + return errors.New("code_source.snapshot.remote_volume must start with '/Volumes/'") + } + + // A non-nil but empty include_paths is an explicit mistake (omit it instead). + if s.IncludePaths != nil && len(s.IncludePaths) == 0 { + return errors.New("code_source.snapshot.include_paths cannot be an empty list; either omit it or provide paths") + } + for _, p := range s.IncludePaths { + p = strings.TrimSpace(p) + if p == "" { + return errors.New("code_source.snapshot.include_paths entry cannot be empty") + } + if strings.HasPrefix(p, "/") { + return fmt.Errorf("code_source.snapshot.include_paths must be relative paths, got: %s", p) + } + // No parent traversal: snapshots must stay within root_path. + if slices.Contains(strings.Split(p, "/"), "..") { + return fmt.Errorf("code_source.snapshot.include_paths cannot contain '..' traversal, got: %s", p) + } + } + + if s.Git != nil { + return s.Git.validate() + } + return nil +} + +// gitRef pins a snapshot to a specific git ref. branch and commit are mutually +// exclusive; remote is only meaningful with branch. +type gitRef struct { + Branch *string `yaml:"branch"` + Commit *string `yaml:"commit"` + Remote gitRemote `yaml:"remote"` +} + +func (g *gitRef) validate() error { + if g.Branch != nil && !gitRefRe.MatchString(*g.Branch) { + return fmt.Errorf("invalid git.branch format %q: only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", *g.Branch) + } + if g.Remote.isString { + if g.Remote.name == "" { + return errors.New("git.remote string cannot be empty; use 'true' to auto-detect") + } + if !gitRefRe.MatchString(g.Remote.name) { + return fmt.Errorf("invalid git.remote name %q: only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", g.Remote.name) + } + } + + if g.Branch == nil && g.Commit == nil { + return errors.New("git: must specify either 'branch' or 'commit'") + } + if g.Branch != nil && g.Commit != nil { + return errors.New("git: 'branch' and 'commit' are mutually exclusive — specify only one") + } + if g.Remote.truthy() && g.Branch == nil { + return errors.New("git.remote requires git.branch (only valid with branch refs)") + } + return nil +} + +// gitRemote is git.remote: false (default, use local HEAD), true (auto-detect the +// remote), or a remote name string. +type gitRemote struct { + set bool + isString bool + name string + enabled bool +} + +func (r *gitRemote) UnmarshalYAML(node *yaml.Node) error { + if node.Kind != yaml.ScalarNode { + return errors.New("git.remote must be a boolean or a remote name string") + } + r.set = true + if node.Tag == "!!bool" { + return node.Decode(&r.enabled) + } + r.isString = true + r.name = node.Value + return nil +} + +// truthy reports whether remote requests a remote fetch (mirrors Python's +// truthiness of the bool|str union). +func (r *gitRemote) truthy() bool { + if r.isString { + return r.name != "" + } + return r.enabled +} + +// permission is a DABs-compatible permission grant: exactly one principal plus a +// level. +type permission struct { + UserName *string `yaml:"user_name"` + GroupName *string `yaml:"group_name"` + ServicePrincipalName *string `yaml:"service_principal_name"` + // Level is a databricks PermissionLevel (e.g. CAN_VIEW, CAN_MANAGE). Enum + // membership is validated server-side; here we only require it to be set. + Level string `yaml:"level"` +} + +func (p *permission) validate() error { + principals := map[string]*string{ + "user_name": p.UserName, + "group_name": p.GroupName, + "service_principal_name": p.ServicePrincipalName, + } + var set []string + for name, val := range principals { + if val != nil { + set = append(set, name) + } + } + switch len(set) { + case 0: + return errors.New("permissions: one of 'user_name', 'group_name', or 'service_principal_name' must be specified") + case 1: + name := set[0] + if strings.TrimSpace(*principals[name]) == "" { + return fmt.Errorf("permissions: '%s' cannot be empty", name) + } + default: + return errors.New("permissions: only one of 'user_name', 'group_name', or 'service_principal_name' can be specified") + } + + if strings.TrimSpace(p.Level) == "" { + return errors.New("permissions: 'level' is required") + } + return nil +} diff --git a/experimental/air/cmd/runconfig_load.go b/experimental/air/cmd/runconfig_load.go new file mode 100644 index 00000000000..81b07d3ca50 --- /dev/null +++ b/experimental/air/cmd/runconfig_load.go @@ -0,0 +1,52 @@ +package aircmd + +import ( + "errors" + "fmt" + "io" + "os" + + "go.yaml.in/yaml/v3" +) + +// decodeRunConfig reads and decodes the run YAML into the schema. Unknown keys +// are rejected (KnownFields), mirroring the Python schema's extra="forbid". +// +// The `_bases_` composition feature and CLI `--override` handling are not yet +// ported; a config using `_bases_` is currently rejected as an unknown field. +func decodeRunConfig(path string) (*runConfig, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + dec := yaml.NewDecoder(f) + dec.KnownFields(true) + + var cfg runConfig + if err := dec.Decode(&cfg); err != nil { + if errors.Is(err, io.EOF) { + return nil, fmt.Errorf("config %s is empty", path) + } + return nil, fmt.Errorf("invalid config %s: %w", path, err) + } + return &cfg, nil +} + +// validateRunConfig runs structural validation over a decoded config. +func validateRunConfig(cfg *runConfig) error { + return cfg.validate() +} + +// loadRunConfig decodes and structurally validates a run YAML config file. +func loadRunConfig(path string) (*runConfig, error) { + cfg, err := decodeRunConfig(path) + if err != nil { + return nil, err + } + if err := validateRunConfig(cfg); err != nil { + return nil, err + } + return cfg, nil +} diff --git a/experimental/air/cmd/runconfig_test.go b/experimental/air/cmd/runconfig_test.go new file mode 100644 index 00000000000..1dd5ce1ea39 --- /dev/null +++ b/experimental/air/cmd/runconfig_test.go @@ -0,0 +1,419 @@ +package aircmd + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeConfig writes content to a temp YAML file and returns its path. +func writeConfig(t *testing.T, content string) string { + t.Helper() + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) + return path +} + +// minimalConfig is the smallest valid config: the three required pieces. +const minimalConfig = ` +experiment_name: my-run +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +` + +func TestLoadRunConfig_Minimal(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, minimalConfig)) + require.NoError(t, err) + assert.Equal(t, "my-run", cfg.ExperimentName) + require.NotNil(t, cfg.Command) + assert.Equal(t, "python train.py", *cfg.Command) + require.NotNil(t, cfg.Compute) + assert.Equal(t, "GPU_1xH100", cfg.Compute.AcceleratorType) + assert.Equal(t, 1, cfg.Compute.NumAccelerators) +} + +func TestLoadRunConfig_FullFeatured(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, ` +experiment_name: full_run +command: | + python train.py + echo done +compute: + accelerator_type: GPU_8xH100 + num_accelerators: 16 +environment: + dependencies: + - torch==2.3.0 + - numpy + version: 5 +env_variables: + FOO: bar +secrets: + HF_TOKEN: my_scope/hf_token +code_source: + type: snapshot + snapshot: + root_path: project_root/src + remote_volume: /Volumes/main/default/code + git: + branch: main + remote: origin + include_paths: + - src + - configs/train.yaml +max_retries: 5 +timeout_minutes: 120 +idempotency_token: abc-123 +mlflow_run_name: full_run_v2 +mlflow_experiment_directory: /Workspace/Users/me/exp +usage_policy_name: my-policy +permissions: + - group_name: users + level: CAN_VIEW + - user_name: alice@example.com + level: CAN_MANAGE +`)) + require.NoError(t, err) + assert.Equal(t, gpuType8xH100, gpuType(cfg.Compute.AcceleratorType)) + require.NotNil(t, cfg.Environment) + assert.True(t, cfg.Environment.Dependencies.isList) + assert.Equal(t, []string{"torch==2.3.0", "numpy"}, cfg.Environment.Dependencies.list) + assert.True(t, cfg.Environment.Version.set) + assert.Equal(t, "5", cfg.Environment.Version.raw) + require.NotNil(t, cfg.CodeSource) + require.NotNil(t, cfg.CodeSource.Snapshot) + require.NotNil(t, cfg.CodeSource.Snapshot.Git) + require.NotNil(t, cfg.CodeSource.Snapshot.Git.Branch) + assert.Equal(t, "main", *cfg.CodeSource.Snapshot.Git.Branch) + assert.True(t, cfg.CodeSource.Snapshot.Git.Remote.isString) + assert.Equal(t, "origin", cfg.CodeSource.Snapshot.Git.Remote.name) + assert.Len(t, cfg.Permissions, 2) +} + +// TestLoadRunConfig_PolymorphicFields exercises the str|list, str|int, and +// bool|str unions decoded by custom UnmarshalYAML. +func TestLoadRunConfig_PolymorphicFields(t *testing.T) { + t.Run("dependencies as string path", func(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, minimalConfig+` +environment: + dependencies: requirements.yaml +`)) + require.NoError(t, err) + assert.True(t, cfg.Environment.Dependencies.set) + assert.False(t, cfg.Environment.Dependencies.isList) + assert.Equal(t, "requirements.yaml", cfg.Environment.Dependencies.path) + }) + + t.Run("git remote as bool true", func(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, minimalConfig+` +code_source: + type: snapshot + snapshot: + root_path: . + git: + branch: main + remote: true +`)) + require.NoError(t, err) + r := cfg.CodeSource.Snapshot.Git.Remote + assert.False(t, r.isString) + assert.True(t, r.enabled) + assert.True(t, r.truthy()) + }) + + t.Run("git remote defaults to false when unset", func(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, minimalConfig+` +code_source: + type: snapshot + snapshot: + root_path: . + git: + commit: deadbeef +`)) + require.NoError(t, err) + assert.False(t, cfg.CodeSource.Snapshot.Git.Remote.truthy()) + }) +} + +func TestLoadRunConfig_UnknownFieldRejected(t *testing.T) { + tests := []struct { + name string + extra string + errFrag string + }{ + {"top-level typo", "extra_field: nope\n", "extra_field"}, + // priority was intentionally dropped from the schema (pool-only concept). + {"dropped priority field", "priority: 100\n", "priority"}, + // _bases_ composition is not yet ported, so it surfaces as unknown. + {"unported _bases_", "_bases_: [base.yaml]\n", "_bases_"}, + {"nested typo", "environment:\n bogus: 1\n", "bogus"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, minimalConfig+tt.extra)) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestLoadRunConfig_Errors(t *testing.T) { + tests := []struct { + name string + yaml string + errFrag string + }{ + { + "missing experiment_name", + "command: x\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\n", + "experiment_name cannot be empty", + }, + { + "experiment_name bad chars", + "experiment_name: my.run\ncommand: x\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\n", + "invalid experiment_name", + }, + { + "missing compute", + "experiment_name: r\ncommand: x\n", + "compute: section is required", + }, + { + "missing command", + "experiment_name: r\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\n", + "command is required", + }, + { + "bad gpu type", + "experiment_name: r\ncommand: x\ncompute:\n accelerator_type: a100\n num_accelerators: 1\n", + "invalid GPU type", + }, + { + "num_accelerators not a multiple", + "experiment_name: r\ncommand: x\ncompute:\n accelerator_type: GPU_8xH100\n num_accelerators: 3\n", + "must be a multiple of 8", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, tt.yaml)) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +// TestRunConfigValidate_FieldRules unit-tests validation rules directly, away +// from YAML decoding, to keep each rule's failure mode explicit. +func TestRunConfigValidate_FieldRules(t *testing.T) { + str := func(s string) *string { return &s } + intp := func(i int) *int { return &i } + base := func() *runConfig { + return &runConfig{ + ExperimentName: "r", + Command: str("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, + } + } + + tests := []struct { + name string + mutate func(c *runConfig) + errFrag string + }{ + {"ok baseline", func(c *runConfig) {}, ""}, + {"empty command", func(c *runConfig) { c.Command = str(" ") }, "command cannot be empty"}, + {"negative max_retries", func(c *runConfig) { c.MaxRetries = intp(-1) }, "max_retries must be >= 0"}, + {"zero timeout", func(c *runConfig) { c.TimeoutMinutes = intp(0) }, "timeout_minutes must be >= 1"}, + {"empty idempotency", func(c *runConfig) { c.IdempotencyToken = str(" ") }, "idempotency_token cannot be empty"}, + {"long idempotency", func(c *runConfig) { c.IdempotencyToken = str(string(make([]byte, 65))) }, "64 characters or less"}, + {"bad mlflow_run_name", func(c *runConfig) { c.MLflowRunName = str("bad name") }, "invalid mlflow_run_name"}, + {"bad experiment dir", func(c *runConfig) { c.MLflowExperimentDirectory = str("/Users/me") }, "must start with '/Workspace'"}, + {"empty usage policy", func(c *runConfig) { c.UsagePolicyName = str(" ") }, "usage_policy_name must not be empty"}, + {"bad secret ref", func(c *runConfig) { c.Secrets = map[string]string{"T": "noslash"} }, "expected format 'scope/key'"}, + {"empty secret scope", func(c *runConfig) { c.Secrets = map[string]string{"T": "/key"} }, "scope and key cannot be empty"}, + {"env var and secret collide", func(c *runConfig) { + c.EnvVariables = map[string]string{"TOK": "v"} + c.Secrets = map[string]string{"TOK": "scope/key"} + }, `"TOK" is set in both env_variables and secrets`}, + {"long mlflow_run_name", func(c *runConfig) { c.MLflowRunName = str(strings.Repeat("a", 101)) }, "100 characters or less"}, + {"usage policy name and id", func(c *runConfig) { + c.UsagePolicyName = str("p") + c.UsagePolicyID = str("id") + }, "mutually exclusive"}, + {"empty usage_policy_id", func(c *runConfig) { c.UsagePolicyID = str(" ") }, "usage_policy_id must not be empty"}, + {"usage_policy_id alone is ok", func(c *runConfig) { c.UsagePolicyID = str("policy-uuid") }, ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := base() + tt.mutate(c) + err := c.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestEnvironmentConfigValidate(t *testing.T) { + tests := []struct { + name string + env environmentConfig + errFrag string + }{ + { + "docker image alone ok", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag"}}, + "", + }, + { + "docker image with deps conflicts", + environmentConfig{ + DockerImage: &dockerImageConfig{URL: "org/repo:tag"}, + Dependencies: dependencies{set: true, isList: true, list: []string{"torch"}}, + }, + "not allowed: dependencies", + }, + { + "empty docker url", + environmentConfig{DockerImage: &dockerImageConfig{URL: " "}}, + "docker_image.url cannot be empty", + }, + { + "version with file deps", + environmentConfig{ + Version: stringOrInt{set: true, raw: "5"}, + Dependencies: dependencies{set: true, isList: false, path: "req.yaml"}, + }, + "only valid with inline dependencies", + }, + { + "version without deps", + environmentConfig{Version: stringOrInt{set: true, raw: "5"}}, + "requires inline 'dependencies'", + }, + { + "version with inline deps ok", + environmentConfig{ + Version: stringOrInt{set: true, raw: "5"}, + Dependencies: dependencies{set: true, isList: true, list: []string{"torch"}}, + }, + "", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.env.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestGitRefValidate(t *testing.T) { + str := func(s string) *string { return &s } + tests := []struct { + name string + ref gitRef + errFrag string + }{ + {"branch only ok", gitRef{Branch: str("main")}, ""}, + {"commit only ok", gitRef{Commit: str("abc123")}, ""}, + {"branch with remote ok", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: true}}, ""}, + {"neither branch nor commit", gitRef{}, "must specify either 'branch' or 'commit'"}, + {"both branch and commit", gitRef{Branch: str("main"), Commit: str("abc")}, "mutually exclusive"}, + {"remote without branch", gitRef{Commit: str("abc"), Remote: gitRemote{set: true, isString: true, name: "origin"}}, "requires git.branch"}, + {"bad branch chars", gitRef{Branch: str("bad branch")}, "invalid git.branch"}, + {"empty remote string", gitRef{Branch: str("main"), Remote: gitRemote{set: true, isString: true, name: ""}}, "cannot be empty"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.ref.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestSnapshotSourceConfigValidate(t *testing.T) { + tests := []struct { + name string + snap snapshotSourceConfig + errFrag string + }{ + {"ok", snapshotSourceConfig{RootPath: "src"}, ""}, + {"empty root_path", snapshotSourceConfig{RootPath: " "}, "root_path cannot be empty"}, + {"bad volume", snapshotSourceConfig{RootPath: "src", RemoteVolume: new("/mnt/x")}, "must start with '/Volumes/'"}, + {"empty include list", snapshotSourceConfig{RootPath: "src", IncludePaths: []string{}}, "cannot be an empty list"}, + {"absolute include", snapshotSourceConfig{RootPath: "src", IncludePaths: []string{"/etc"}}, "must be relative"}, + {"traversal include", snapshotSourceConfig{RootPath: "src", IncludePaths: []string{"../x"}}, "'..' traversal"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.snap.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestPermissionValidate(t *testing.T) { + str := func(s string) *string { return &s } + tests := []struct { + name string + perm permission + errFrag string + }{ + {"ok user", permission{UserName: str("alice@example.com"), Level: "CAN_VIEW"}, ""}, + {"no principal", permission{Level: "CAN_VIEW"}, "must be specified"}, + {"two principals", permission{UserName: str("a"), GroupName: str("g"), Level: "CAN_VIEW"}, "only one of"}, + {"empty principal", permission{UserName: str(" "), Level: "CAN_VIEW"}, "cannot be empty"}, + {"missing level", permission{GroupName: str("users")}, "'level' is required"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.perm.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestLoadRunConfig_FileErrors(t *testing.T) { + t.Run("missing file", func(t *testing.T) { + _, err := loadRunConfig(filepath.Join(t.TempDir(), "nope.yaml")) + assert.Error(t, err) + }) + t.Run("empty file", func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, "")) + require.Error(t, err) + assert.Contains(t, err.Error(), "is empty") + }) +} From fc0ba3e1a73151510f12dfc8d62a6b21a7e31066 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Tue, 30 Jun 2026 14:46:44 -0700 Subject: [PATCH 08/18] AIR CLI Integration: `air run` end to end command (#5710) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Changes Implements the `air run` happy path on top of the config schema (#5657), submitting a one-time training run through the Jobs API. Five commits, one per phase: 1. run config launch accessors: flatten the validated config into launch values (timeout seconds, retry default, requirements file-vs-inline, runtime version). 2. wire run command (load, validate, dry-run): air run -f <config> loads + structurally validates the YAML; `--dry-run` validates offline (no workspace/auth) and returns; `--override/--watch` are rejected for now with clear errors (ported in future PR). 3. pre-submit resolution: resolve current user / workspace home / a unique cli_launch dir, and ensure a custom `experiment_directory` exists. 4. upload launch artifacts: write training_config.yaml (1 MB cap), command.sh, requirements.yaml (file or synthesized from inline deps), `env_vars.json` / `secret_env_vars.json`, and hyperparameters.yaml into the launch dir via a workspace filer. 5. assemble + submit: build the native `ai_runtime_task` payload and `POST /api/2.2/jobs/runs/submit` directly, then print the run id + dashboard URL (or a JSON envelope). Submission uses the **native `ai_runtime_task`** task (BYOT task type) and it talks only to the Jobs API (which internally routes to training service endpoint) and has no genai-mapi forwarding (the MAPI path is deprecated). It isn't modeled by the typed SDK in go, so the payload is a custom struct posted to the raw endpoint. The proto is lean: env vars and secrets ship as co-located `env_vars.json` / `secret_env_vars.json` files rather than inline, and `requirements.yaml` / `hyperparameters.yaml` are derived server-side from the command directory. **Deferred, with explicit "not yet supported" errors (no silent drops):** `code_source` snapshot packaging, `--watch` log streaming, and `usage_policy_name`. `environment.docker_image` is accepted by the schema as scaffolding but not conveyed in the payload (the native path has no docker field). `node_pool_id` / `pool_name` / `priority` remain dropped (new AIR CLI does not support pool placement). ## Why `air run` is the core of the migration for AIR CLI. Splitting it into per-phase commits keeps each reviewable in isolation, and stacking on the schema PR keeps that PR focused. Regarding some specific decisions: - We maintain the native ai_runtime_task (and not the genai_compute_task interfacing with mapi) as a hand built struct posted to the raw endpoint. This is so that we can interface with jobs directly (and jobs.SubmitTask only knows gen_ai_compute_task and this typed struct also omits the env-vars/secrets/requirements fields that are needed for the run) and make sure we also stay off the deprecated genai-mapi forwarding path. - `--dry-run` is decoupled from auth. It validates the config locally and returns before any workspace call, so config validation works fully offline (matching the Python CLI). Only actual submission requires an authenticated workspace client. ## Tests - Unit tests for every phase: launch accessors, pre-submit resolution (incl. ensureExperimentDirectory create/exists/not-a-directory), artifact assembly + upload, payload assembly, and submitWorkload end-to-end against a fake workspace. - New acceptance/experimental/air/run test covering --dry-run (text + JSON), the --override/--watch guards, an invalid config, and missing --file. - Updated the unimplemented acceptance test (removed run, now implemented). `go test ./experimental/air/...`, `go test ./acceptance -run TestAccept/experimental/air`, and `./task lint-q` all pass. **Manual verification tests (all pass):** - Dry run (offline, no auth) > - command only > - full run config > - json output - actual run submission > - throws error when profile is not set > - submission loop: submitted, can see the run in `air list` and `air get` and mlflow environment was created > - same run id gets ouputted when run submitted with the SAME idempotency key > - new run gets created when run submitted with SAME config but DIFFERENT idempotency key - `--watch` and `--override` return an informative error message (since they are not supported yet, but are valid flags) - usage_policy_name set in config throws error: usage_policy_name is not yet supported - code_source set in config throws error: code_source is not yet supported - missing --file throws informative error: required flag(s) "file" not set - invalid config (e.g. experiment_name: bad.name, or num_accelerators not a multiple of the per-node count) throws field-specific validation error **How to test locally for manual verification:** Checkout & build: ```bash git fetch origin git checkout air-integration-m2-3 # this PR (stacked on air-integration-m2-2) ./task build ``` Sample configs: ```bash cat > /tmp/min.yaml <<'YAML' experiment_name: air-cuj command: python train.py compute: {accelerator_type: GPU_1xH100, num_accelerators: 1} YAML ``` ```bash cat > /tmp/full.yaml <<'YAML' experiment_name: full-run command: | pip install -r requirements.txt python train.py compute: {accelerator_type: GPU_8xH100, num_accelerators: 16} environment: {dependencies: [torch==2.3.0], version: 5} env_variables: {WANDB_PROJECT: demo} secrets: {HF_TOKEN: my_scope/hf_token} parameters: {lr: 0.001, epochs: 3} mlflow_run_name: full-run-v2 max_retries: 2 timeout_minutes: 120 YAML ``` Automated tests ```bash go test ./experimental/air/... # unit (incl. submitWorkload vs a fake workspace) go test ./acceptance -run TestAccept/experimental/air # acceptance (run + unimplemented) ./task lint-q # lint changed files ``` Dry run: ```bash ./cli experimental air run -f /tmp/min.yaml --dry-run # note that this command will, in the final version, be databricks experimental air run ./cli experimental air run -f /tmp/full.yaml --dry-run ./cli experimental air run -f /tmp/min.yaml --dry-run -o json ``` Actual run submission: ```bash PROFILE=<your-dev-profile> # no auth configured → fails fast (exit 1) env -u DATABRICKS_HOST -u DATABRICKS_TOKEN ./cli experimental air run -f /tmp/min.yaml #> Error: ... (cannot configure default credentials / auth) # submit → prints run_id + dashboard URL ./cli experimental air run -f /tmp/min.yaml -p $PROFILE -o json #> { "data": { "status":"SUBMITTED", "run_id":"<id>", "dashboard_url":"<host>/jobs/runs/<id>" } } # verify in the workspace: open dashboard_url (run exists), and the MLflow experiment was created. ./cli experimental air get <run_id> -p $PROFILE # run state ./cli experimental air list -p $PROFILE # run appears in the list # idempotency — SAME key returns the SAME run_id (no new run) ./cli experimental air run -f /tmp/min.yaml -p $PROFILE --idempotency-key demo-key-1 -o json # run_id = X ./cli experimental air run -f /tmp/min.yaml -p $PROFILE --idempotency-key demo-key-1 -o json # run_id = X (same) # idempotency — DIFFERENT key creates a NEW run ./cli experimental air run -f /tmp/min.yaml -p $PROFILE --idempotency-key demo-key-2 -o json # run_id = Y (new) ``` Unsupported flags (asserting that error is thrown): ```bash ./cli experimental air run -f /tmp/min.yaml --dry-run --watch #> Error: --watch is not yet supported ./cli experimental air run -f /tmp/min.yaml --dry-run --override compute.num_accelerators=8 #> Error: --override is not yet supported # usage_policy_name (needs a workspace to reach the submit guard) printf 'experiment_name: t\ncommand: x\ncompute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\nusage_policy_name: my-policy\n' > /tmp/policy.yaml ./cli experimental air run -f /tmp/policy.yaml -p $PROFILE #> Error: usage_policy_name is not yet supported # code_source printf 'experiment_name: t\ncommand: x\ncompute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\ncode_source: {type: snapshot, snapshot: {root_path: .}}\n' > /tmp/code.yaml air run -f /tmp/code.yaml -p $PROFILE #> Error: code_source is not yet supported ``` Validation errors for field-specific message (exit 1, offline): ```bash # missing --file air run --dry-run #> Error: required flag(s) "file" not set # invalid experiment_name + num_accelerators not a multiple of the per-node count printf 'experiment_name: bad.name\ncommand: x\ncompute: {accelerator_type: GPU_8xH100, num_accelerators: 3}\n' > /tmp/bad.yaml air run -f /tmp/bad.yaml --dry-run #> Error: invalid experiment_name "bad.name": only alphanumeric characters, hyphens (-), and underscores (_) are allowed # (and, once the name is fixed: compute.num_accelerators for GPU_8xH100 must be a multiple of 8, got 3) ``` --- acceptance/experimental/air/run/invalid.yaml | 5 + acceptance/experimental/air/run/out.test.toml | 3 + acceptance/experimental/air/run/output.txt | 39 +++ acceptance/experimental/air/run/script | 17 ++ acceptance/experimental/air/run/test.toml | 4 + acceptance/experimental/air/run/valid.yaml | 5 + .../experimental/air/unimplemented/output.txt | 6 - .../experimental/air/unimplemented/script | 3 - experimental/air/cmd/run.go | 68 ++++- experimental/air/cmd/runconfig_launch.go | 65 +++++ experimental/air/cmd/runconfig_launch_test.go | 80 ++++++ experimental/air/cmd/runlaunch.go | 73 +++++ experimental/air/cmd/runlaunch_test.go | 65 +++++ experimental/air/cmd/runsubmit.go | 250 ++++++++++++++++++ experimental/air/cmd/runsubmit_test.go | 177 +++++++++++++ experimental/air/cmd/runupload.go | 170 ++++++++++++ experimental/air/cmd/runupload_test.go | 155 +++++++++++ experimental/air/cmd/stubs_test.go | 1 - 18 files changed, 1173 insertions(+), 13 deletions(-) create mode 100644 acceptance/experimental/air/run/invalid.yaml create mode 100644 acceptance/experimental/air/run/out.test.toml create mode 100644 acceptance/experimental/air/run/output.txt create mode 100644 acceptance/experimental/air/run/script create mode 100644 acceptance/experimental/air/run/test.toml create mode 100644 acceptance/experimental/air/run/valid.yaml create mode 100644 experimental/air/cmd/runconfig_launch.go create mode 100644 experimental/air/cmd/runconfig_launch_test.go create mode 100644 experimental/air/cmd/runlaunch.go create mode 100644 experimental/air/cmd/runlaunch_test.go create mode 100644 experimental/air/cmd/runsubmit.go create mode 100644 experimental/air/cmd/runsubmit_test.go create mode 100644 experimental/air/cmd/runupload.go create mode 100644 experimental/air/cmd/runupload_test.go diff --git a/acceptance/experimental/air/run/invalid.yaml b/acceptance/experimental/air/run/invalid.yaml new file mode 100644 index 00000000000..c011fc81b37 --- /dev/null +++ b/acceptance/experimental/air/run/invalid.yaml @@ -0,0 +1,5 @@ +experiment_name: bad.name +command: x +compute: + accelerator_type: GPU_8xH100 + num_accelerators: 3 diff --git a/acceptance/experimental/air/run/out.test.toml b/acceptance/experimental/air/run/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/run/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt new file mode 100644 index 00000000000..180886290ba --- /dev/null +++ b/acceptance/experimental/air/run/output.txt @@ -0,0 +1,39 @@ + +=== dry-run (text) +>>> [CLI] experimental air run -f valid.yaml --dry-run +Dry run: configuration for "smoke-test" is valid; not submitting. + +=== dry-run (json) +>>> [CLI] experimental air run -f valid.yaml --dry-run -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "status": "DRY_RUN_OK", + "dry_run": true + } +} + +=== override not yet supported +>>> [CLI] experimental air run -f valid.yaml --dry-run --override a=b +Error: --override is not yet supported + +Exit code: 1 + +=== watch not yet supported +>>> [CLI] experimental air run -f valid.yaml --dry-run --watch +Error: --watch is not yet supported + +Exit code: 1 + +=== invalid config is rejected +>>> [CLI] experimental air run -f invalid.yaml --dry-run +Error: invalid experiment_name "bad.name": only alphanumeric characters, hyphens (-), and underscores (_) are allowed + +Exit code: 1 + +=== missing --file +>>> [CLI] experimental air run --dry-run +Error: required flag(s) "file" not set + +Exit code: 1 diff --git a/acceptance/experimental/air/run/script b/acceptance/experimental/air/run/script new file mode 100644 index 00000000000..806bd321e6d --- /dev/null +++ b/acceptance/experimental/air/run/script @@ -0,0 +1,17 @@ +title "dry-run (text)" +trace $CLI experimental air run -f valid.yaml --dry-run + +title "dry-run (json)" +trace $CLI experimental air run -f valid.yaml --dry-run -o json + +title "override not yet supported" +errcode trace $CLI experimental air run -f valid.yaml --dry-run --override a=b + +title "watch not yet supported" +errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch + +title "invalid config is rejected" +errcode trace $CLI experimental air run -f invalid.yaml --dry-run + +title "missing --file" +errcode trace $CLI experimental air run --dry-run diff --git a/acceptance/experimental/air/run/test.toml b/acceptance/experimental/air/run/test.toml new file mode 100644 index 00000000000..2f971c3ed21 --- /dev/null +++ b/acceptance/experimental/air/run/test.toml @@ -0,0 +1,4 @@ +# `air run --dry-run` validates the config locally and makes no workspace calls, +# so no engine matrix or server stubs are needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/run/valid.yaml b/acceptance/experimental/air/run/valid.yaml new file mode 100644 index 00000000000..b82a321b051 --- /dev/null +++ b/acceptance/experimental/air/run/valid.yaml @@ -0,0 +1,5 @@ +experiment_name: smoke-test +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt index 66ddc34d586..21c3c891af8 100644 --- a/acceptance/experimental/air/unimplemented/output.txt +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -1,10 +1,4 @@ -=== run ->>> [CLI] experimental air run -Error: `air run` is not implemented yet - -Exit code: 1 - === logs >>> [CLI] experimental air logs 123 Error: `air logs` is not implemented yet diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script index d00d0453689..4c53586b16a 100644 --- a/acceptance/experimental/air/unimplemented/script +++ b/acceptance/experimental/air/unimplemented/script @@ -1,8 +1,5 @@ # Each stub must fail with "not implemented"; errcode records the exit code. -title "run" -errcode trace $CLI experimental air run - title "logs" errcode trace $CLI experimental air logs 123 diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 0bc3d1fd94b..bd32810e9bc 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -1,10 +1,25 @@ package aircmd import ( + "errors" + "fmt" + "strconv" + "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/spf13/cobra" ) +// runResult is the JSON payload for `air run`. +type runResult struct { + Status string `json:"status"` + DryRun bool `json:"dry_run,omitempty"` + RunID string `json:"run_id,omitempty"` + DashboardURL string `json:"dashboard_url,omitempty"` +} + func newRunCommand() *cobra.Command { var ( file string @@ -21,9 +36,6 @@ func newRunCommand() *cobra.Command { Long: `Submit a training workload to Databricks serverless GPU compute. The workload is described by a YAML config file (see --file).`, - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("run") - }, } cmd.Flags().StringVarP(&file, "file", "f", "", "Path to the workload YAML config") @@ -31,6 +43,56 @@ The workload is described by a YAML config file (see --file).`, cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config without submitting") cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used") + _ = cmd.MarkFlagRequired("file") + + // --dry-run only validates the config locally, so it needs no workspace. + // Submission requires an authenticated client. + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + if dryRun { + return nil + } + return root.MustWorkspaceClient(cmd, args) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + + // These flags' pipelines are not ported yet; reject rather than silently + // ignore them. + if len(overrides) > 0 { + return errors.New("--override is not yet supported") + } + if watch { + return errors.New("--watch is not yet supported") + } + + cfg, err := loadRunConfig(file) + if err != nil { + return err + } + + if dryRun { + if root.OutputType(cmd) == flags.OutputText { + cmdio.LogString(ctx, fmt.Sprintf("Dry run: configuration for %q is valid; not submitting.", cfg.ExperimentName)) + return nil + } + return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true}) + } + + w := cmdctx.WorkspaceClient(ctx) + runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey) + if err != nil { + return err + } + + runIDStr := strconv.FormatInt(runID, 10) + if root.OutputType(cmd) == flags.OutputText { + cmdio.LogString(ctx, "Submitted run "+runIDStr) + cmdio.LogString(ctx, "View at: "+dashboardURL) + return nil + } + return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + } return cmd } diff --git a/experimental/air/cmd/runconfig_launch.go b/experimental/air/cmd/runconfig_launch.go new file mode 100644 index 00000000000..1408b600736 --- /dev/null +++ b/experimental/air/cmd/runconfig_launch.go @@ -0,0 +1,65 @@ +package aircmd + +// This file flattens the validated runConfig schema into the derived values the +// launch path consumes, replacing the Python CLI's _convert_to_run_config step. +// There is no separate internal config type: handle_run reads runConfig directly, +// using these accessors for the values that need computing rather than a plain +// field read. + +const defaultMaxRetries = 3 + +// timeoutSeconds converts timeout_minutes to seconds. Zero means the user set no +// timeout and the backend default applies. +func (c *runConfig) timeoutSeconds() int { + if c.TimeoutMinutes == nil { + return 0 + } + return *c.TimeoutMinutes * 60 +} + +// maxRetries returns the retry count, applying the schema default when unset. +func (c *runConfig) maxRetries() int { + if c.MaxRetries == nil { + return defaultMaxRetries + } + return *c.MaxRetries +} + +// dockerImageURL returns the custom docker image URL, or "" when none is set. +// +// TODO: not wired into submission yet — the native ai_runtime_task carries no +// docker field, and full support needs image registration (pending the DCS work). +func (c *runConfig) dockerImageURL() string { + if c.Environment != nil && c.Environment.DockerImage != nil { + return c.Environment.DockerImage.URL + } + return "" +} + +// requirementsFile returns the path to a requirements file when +// environment.dependencies is a string, and whether it was set. +func (c *runConfig) requirementsFile() (string, bool) { + if c.Environment == nil || !c.Environment.Dependencies.set || c.Environment.Dependencies.isList { + return "", false + } + return c.Environment.Dependencies.path, true +} + +// inlineDependencies returns the inline package list when +// environment.dependencies is a list, and whether it was set. +func (c *runConfig) inlineDependencies() ([]string, bool) { + if c.Environment == nil || !c.Environment.Dependencies.set || !c.Environment.Dependencies.isList { + return nil, false + } + return c.Environment.Dependencies.list, true +} + +// runtimeVersion returns the client image version from environment.version when +// set. For a requirements-file dependency set, the version lives in that file and +// is resolved at launch, not here. +func (c *runConfig) runtimeVersion() (string, bool) { + if c.Environment == nil || !c.Environment.Version.set { + return "", false + } + return c.Environment.Version.raw, true +} diff --git a/experimental/air/cmd/runconfig_launch_test.go b/experimental/air/cmd/runconfig_launch_test.go new file mode 100644 index 00000000000..289db91c7de --- /dev/null +++ b/experimental/air/cmd/runconfig_launch_test.go @@ -0,0 +1,80 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestRunConfigTimeoutSeconds(t *testing.T) { + c := &runConfig{} + assert.Equal(t, 0, c.timeoutSeconds()) + + c.TimeoutMinutes = new(2) + assert.Equal(t, 120, c.timeoutSeconds()) +} + +func TestRunConfigMaxRetries(t *testing.T) { + c := &runConfig{} + assert.Equal(t, defaultMaxRetries, c.maxRetries()) + + c.MaxRetries = new(0) + assert.Equal(t, 0, c.maxRetries()) + + c.MaxRetries = new(7) + assert.Equal(t, 7, c.maxRetries()) +} + +func TestRunConfigDockerImageURL(t *testing.T) { + c := &runConfig{} + assert.Empty(t, c.dockerImageURL()) + + c.Environment = &environmentConfig{} + assert.Empty(t, c.dockerImageURL()) + + c.Environment.DockerImage = &dockerImageConfig{URL: "org/repo:tag"} + assert.Equal(t, "org/repo:tag", c.dockerImageURL()) +} + +func TestRunConfigDependencies(t *testing.T) { + t.Run("unset", func(t *testing.T) { + c := &runConfig{} + _, ok := c.requirementsFile() + assert.False(t, ok) + _, ok = c.inlineDependencies() + assert.False(t, ok) + }) + + t.Run("file path", func(t *testing.T) { + c := &runConfig{Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: false, path: "req.yaml"}, + }} + path, ok := c.requirementsFile() + assert.True(t, ok) + assert.Equal(t, "req.yaml", path) + _, ok = c.inlineDependencies() + assert.False(t, ok) + }) + + t.Run("inline list", func(t *testing.T) { + c := &runConfig{Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: true, list: []string{"torch", "numpy"}}, + }} + list, ok := c.inlineDependencies() + assert.True(t, ok) + assert.Equal(t, []string{"torch", "numpy"}, list) + _, ok = c.requirementsFile() + assert.False(t, ok) + }) +} + +func TestRunConfigRuntimeVersion(t *testing.T) { + c := &runConfig{} + _, ok := c.runtimeVersion() + assert.False(t, ok) + + c.Environment = &environmentConfig{Version: stringOrInt{set: true, raw: "5"}} + v, ok := c.runtimeVersion() + assert.True(t, ok) + assert.Equal(t, "5", v) +} diff --git a/experimental/air/cmd/runlaunch.go b/experimental/air/cmd/runlaunch.go new file mode 100644 index 00000000000..b2a7215e66a --- /dev/null +++ b/experimental/air/cmd/runlaunch.go @@ -0,0 +1,73 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "path" + "strings" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/workspace" + "github.com/google/uuid" +) + +// userWorkspaceDirEnv overrides the per-user workspace directory; mirrors the +// Python CLI's DATABRICKS_INTERNAL_USER_WORKSPACE_DIR escape hatch. +const userWorkspaceDirEnv = "DATABRICKS_INTERNAL_USER_WORKSPACE_DIR" + +// currentUserEmail returns the authenticated user's email (works for any domain). +func currentUserEmail(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return "", fmt.Errorf("failed to resolve current user: %w", err) + } + return me.UserName, nil +} + +// userWorkspaceDir returns the user's workspace home, honoring the env override. +func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + if override := env.Get(ctx, userWorkspaceDirEnv); override != "" { + return override, nil + } + email, err := currentUserEmail(ctx, w) + if err != nil { + return "", err + } + return "/Workspace/Users/" + email, nil +} + +// cliLaunchDir returns a unique workspace directory for a run's launch artifacts: +// <base>/.air/cli_launch/<experiment>/<run>_<uuid>. run defaults to experiment. +func cliLaunchDir(base, experiment, run string) string { + if run == "" { + run = experiment + } + unique := strings.ReplaceAll(uuid.NewString(), "-", "")[:16] + return path.Join(base, ".air", "cli_launch", experiment, run+"_"+unique) +} + +// ensureExperimentDirectory creates experimentDir if it is missing, matching the +// CLI's convention for its other artifact directories. Without this, a missing +// parent surfaces only as a server-side INTERNAL_ERROR after the run is wasted. +// An empty dir means the default (/Users/<user>/...), which always exists. +func ensureExperimentDirectory(ctx context.Context, w *databricks.WorkspaceClient, experimentDir string) error { + if experimentDir == "" { + return nil + } + + info, err := w.Workspace.GetStatusByPath(ctx, experimentDir) + if errors.Is(err, apierr.ErrNotFound) { + return w.Workspace.MkdirsByPath(ctx, experimentDir) + } + if err != nil { + return fmt.Errorf("failed to check experiment_directory %q: %w", experimentDir, err) + } + if info.ObjectType != workspace.ObjectTypeDirectory { + return fmt.Errorf("experiment_directory %q is not a directory (object_type=%s)", experimentDir, info.ObjectType) + } + return nil +} diff --git a/experimental/air/cmd/runlaunch_test.go b/experimental/air/cmd/runlaunch_test.go new file mode 100644 index 00000000000..af6f0f70d31 --- /dev/null +++ b/experimental/air/cmd/runlaunch_test.go @@ -0,0 +1,65 @@ +package aircmd + +import ( + "strings" + "testing" + + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCliLaunchDir(t *testing.T) { + dir := cliLaunchDir("/Workspace/Users/me@example.com", "my-exp", "") + assert.True(t, strings.HasPrefix(dir, "/Workspace/Users/me@example.com/.air/cli_launch/my-exp/my-exp_"), dir) + // run name overrides the leaf; the unique suffix keeps successive dirs distinct. + withRun := cliLaunchDir("/base", "exp", "run1") + assert.True(t, strings.HasPrefix(withRun, "/base/.air/cli_launch/exp/run1_"), withRun) + assert.NotEqual(t, dir, cliLaunchDir("/Workspace/Users/me@example.com", "my-exp", "")) +} + +func newFakeWorkspaceClient(t *testing.T) *databricks.WorkspaceClient { + server := testserver.New(t) + t.Cleanup(server.Close) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + return w +} + +func TestUserWorkspaceDir(t *testing.T) { + w := newFakeWorkspaceClient(t) + dir, err := userWorkspaceDir(t.Context(), w) + require.NoError(t, err) + assert.True(t, strings.HasPrefix(dir, "/Workspace/Users/"), dir) + + // The env override wins without an API call. + t.Setenv(userWorkspaceDirEnv, "/Workspace/custom") + dir, err = userWorkspaceDir(t.Context(), w) + require.NoError(t, err) + assert.Equal(t, "/Workspace/custom", dir) +} + +func TestEnsureExperimentDirectory(t *testing.T) { + ctx := t.Context() + w := newFakeWorkspaceClient(t) + + // Empty means default (always exists) — no API call, no error. + require.NoError(t, ensureExperimentDirectory(ctx, w, "")) + + // A missing path is created. + require.NoError(t, ensureExperimentDirectory(ctx, w, "/Workspace/Users/me/exp")) + + // An existing directory is accepted as-is. + require.NoError(t, w.Workspace.MkdirsByPath(ctx, "/Workspace/Users/me/existing")) + require.NoError(t, ensureExperimentDirectory(ctx, w, "/Workspace/Users/me/existing")) + + // A path that exists but is a file is rejected. + fc, err := filer.NewWorkspaceFilesClient(w, "/Workspace/Users/me") + require.NoError(t, err) + require.NoError(t, fc.Write(ctx, "afile", strings.NewReader("x"))) + err = ensureExperimentDirectory(ctx, w, "/Workspace/Users/me/afile") + require.ErrorContains(t, err, "is not a directory") +} diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go new file mode 100644 index 00000000000..b32ebd622c6 --- /dev/null +++ b/experimental/air/cmd/runsubmit.go @@ -0,0 +1,250 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "path" + "strconv" + "strings" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" + "github.com/google/uuid" +) + +// jobsRunsSubmitPath is the Jobs one-time-run endpoint. air builds the full +// payload and POSTs it here directly — the native ai_runtime_task is not modeled +// by the typed SDK, and we want no genai-mapi forwarding. +const jobsRunsSubmitPath = "/api/2.2/jobs/runs/submit" + +// dlRuntimeImageEnv overrides the default deep-learning runtime image. +const dlRuntimeImageEnv = "DATABRICKS_DL_RUNTIME_IMAGE" + +const defaultDlRuntimeImage = "CLIENT-GPU-4" + +// aiRuntimeEnvironmentKey ties the task to the serverless environment that +// carries the runtime channel. +const aiRuntimeEnvironmentKey = "default" + +// aiRuntimeCompute is a deployment's accelerator request. +type aiRuntimeCompute struct { + AcceleratorType string `json:"accelerator_type"` + AcceleratorCount int `json:"accelerator_count"` +} + +// aiRuntimeDeployment is one worker deployment of the run. +type aiRuntimeDeployment struct { + CommandPath string `json:"command_path"` + Compute aiRuntimeCompute `json:"compute"` +} + +// aiRuntimeTask is the native AI Runtime task. It routes straight to the training +// service — no genai-mapi forwarding. The proto is lean: env vars, secrets, +// requirements, and hyperparameters are staged as workspace files co-located with +// command.sh (see runupload.go), not carried inline. +type aiRuntimeTask struct { + Experiment string `json:"experiment"` + Deployments []aiRuntimeDeployment `json:"deployments"` + MlflowRun string `json:"mlflow_run,omitempty"` + MlflowExperimentDirectory string `json:"mlflow_experiment_directory,omitempty"` +} + +// environmentSpec carries the bare runtime channel ("4", "5", ...). +type environmentSpec struct { + EnvironmentVersion string `json:"environment_version"` +} + +// jobEnvironment is the serverless environment a task references for its runtime. +type jobEnvironment struct { + EnvironmentKey string `json:"environment_key"` + Spec environmentSpec `json:"spec"` +} + +// submitTask is the single task air submits: a native ai_runtime_task. +type submitTask struct { + TaskKey string `json:"task_key"` + RunIf string `json:"run_if"` + AiRuntimeTask aiRuntimeTask `json:"ai_runtime_task"` + EnvironmentKey string `json:"environment_key"` + MaxRetries int `json:"max_retries"` + RetryOnTimeout bool `json:"retry_on_timeout,omitempty"` +} + +// jobsSubmitRun is the Jobs runs/submit payload. +type jobsSubmitRun struct { + RunName string `json:"run_name"` + TimeoutSeconds int `json:"timeout_seconds,omitempty"` + Tasks []submitTask `json:"tasks"` + Environments []jobEnvironment `json:"environments"` + BudgetPolicyID string `json:"budget_policy_id,omitempty"` + IdempotencyToken string `json:"idempotency_token,omitempty"` +} + +// dlRuntimeImage resolves the bare runtime channel (config version, else env, +// else default), always stripping the CLIENT-GPU- prefix. +func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { + img := runtimeVersion + if img == "" { + img = env.Get(ctx, dlRuntimeImageEnv) + } + if img == "" { + img = defaultDlRuntimeImage + } + return strings.TrimPrefix(img, "CLIENT-GPU-") +} + +// buildSubmitPayload assembles the runs/submit payload. commandPath is the +// workspace path of the uploaded command.sh; dlImage is the runtime channel. +func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string) jobsSubmitRun { + task := aiRuntimeTask{ + Experiment: cfg.ExperimentName, + Deployments: []aiRuntimeDeployment{{ + CommandPath: commandPath, + Compute: aiRuntimeCompute{ + AcceleratorType: cfg.Compute.AcceleratorType, + AcceleratorCount: cfg.Compute.NumAccelerators, + }, + }}, + } + if cfg.MLflowRunName != nil { + task.MlflowRun = *cfg.MLflowRunName + } + if cfg.MLflowExperimentDirectory != nil { + task.MlflowExperimentDirectory = *cfg.MLflowExperimentDirectory + } + + st := submitTask{ + TaskKey: cfg.ExperimentName, + RunIf: "ALL_SUCCESS", + AiRuntimeTask: task, + EnvironmentKey: aiRuntimeEnvironmentKey, + MaxRetries: cfg.maxRetries(), + } + // max_retries 0 (no retries) is sent explicitly; retry_on_timeout only + // applies when retries are allowed. + st.RetryOnTimeout = st.MaxRetries > 0 + + return jobsSubmitRun{ + RunName: cfg.ExperimentName, + TimeoutSeconds: cfg.timeoutSeconds(), + Tasks: []submitTask{st}, + Environments: []jobEnvironment{{ + EnvironmentKey: aiRuntimeEnvironmentKey, + Spec: environmentSpec{EnvironmentVersion: dlImage}, + }}, + } +} + +// jobsSubmitClient submits one-time runs through the Jobs API. +type jobsSubmitClient struct { + c *client.DatabricksClient +} + +func newJobsSubmitClient(w *databricks.WorkspaceClient) (*jobsSubmitClient, error) { + c, err := client.New(w.Config) + if err != nil { + return nil, err + } + return &jobsSubmitClient{c: c}, nil +} + +type submitRunResponse struct { + RunID int64 `json:"run_id,omitempty"` +} + +// submit POSTs the payload to runs/submit and returns the new run_id. +func (j *jobsSubmitClient) submit(ctx context.Context, payload jobsSubmitRun) (int64, error) { + var resp submitRunResponse + if err := j.c.Do(ctx, http.MethodPost, jobsRunsSubmitPath, auth.WorkspaceIDHeaders(j.c.Config), nil, payload, &resp); err != nil { + return 0, err + } + return resp.RunID, nil +} + +// submitToken resolves the idempotency token: the --idempotency-key flag wins, +// then the config's token, else a generated one. Over-long tokens error rather +// than truncate, since truncation could make two distinct tokens collide. +func submitToken(flag string, cfg *runConfig) (string, error) { + token := flag + if token == "" && cfg.IdempotencyToken != nil { + token = *cfg.IdempotencyToken + } + if token == "" { + token = uuid.NewString() + } + if len(token) > 64 { + return "", fmt.Errorf("idempotency token must be 64 characters or less, got %d", len(token)) + } + return token, nil +} + +// 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) { + // Resolving usage_policy_name to a budget policy id and packaging a + // code_source snapshot are not ported yet; reject rather than silently drop. + if cfg.UsagePolicyName != nil { + return 0, "", errors.New("usage_policy_name is not yet supported") + } + if cfg.CodeSource != nil { + return 0, "", errors.New("code_source is not yet supported") + } + + // Resolve the idempotency token first so a bad key fails before any upload. + token, err := submitToken(idempotencyKey, cfg) + if err != nil { + return 0, "", err + } + + experimentDir := "" + if cfg.MLflowExperimentDirectory != nil { + experimentDir = *cfg.MLflowExperimentDirectory + } + if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil { + return 0, "", err + } + + base, err := userWorkspaceDir(ctx, w) + if err != nil { + return 0, "", err + } + runName := "" + if cfg.MLflowRunName != nil { + runName = *cfg.MLflowRunName + } + funcDir := cliLaunchDir(base, cfg.ExperimentName, runName) + + fc, err := filer.NewWorkspaceFilesClient(w, funcDir) + if err != nil { + return 0, "", err + } + items, err := buildArtifacts(cfg, configPath) + if err != nil { + return 0, "", err + } + if err := uploadArtifacts(ctx, fc, items); err != nil { + return 0, "", err + } + + runtimeVersion, _ := cfg.runtimeVersion() + payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion)) + payload.IdempotencyToken = token + + jc, err := newJobsSubmitClient(w) + if err != nil { + return 0, "", err + } + runID, err := jc.submit(ctx, payload) + if err != nil { + return 0, "", err + } + + dashboardURL := strings.TrimRight(w.Config.Host, "/") + "/jobs/runs/" + strconv.FormatInt(runID, 10) + return runID, dashboardURL, nil +} diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go new file mode 100644 index 00000000000..bfd92dcd58e --- /dev/null +++ b/experimental/air/cmd/runsubmit_test.go @@ -0,0 +1,177 @@ +package aircmd + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDlRuntimeImage(t *testing.T) { + ctx := t.Context() + // A config runtime version wins and is used bare. + assert.Equal(t, "5", dlRuntimeImage(ctx, "5")) + // The CLIENT-GPU- prefix is always stripped, even from the config version. + assert.Equal(t, "5", dlRuntimeImage(ctx, "CLIENT-GPU-5")) + // Default, with the prefix stripped. + assert.Equal(t, "4", dlRuntimeImage(ctx, "")) + // Env override, prefix stripped. + t.Setenv(dlRuntimeImageEnv, "CLIENT-GPU-7") + assert.Equal(t, "7", dlRuntimeImage(ctx, "")) +} + +func TestBuildSubmitPayload(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_8xH100", NumAccelerators: 16}, + MaxRetries: new(2), + TimeoutMinutes: new(30), + MLflowRunName: new("run-v2"), + MLflowExperimentDirectory: new("/Workspace/Users/me/exp"), + } + + p := buildSubmitPayload(cfg, "/d/command.sh", "5") + + assert.Equal(t, "exp", p.RunName) + assert.Equal(t, 1800, p.TimeoutSeconds) + require.Len(t, p.Environments, 1) + assert.Equal(t, aiRuntimeEnvironmentKey, p.Environments[0].EnvironmentKey) + assert.Equal(t, "5", p.Environments[0].Spec.EnvironmentVersion) + + require.Len(t, p.Tasks, 1) + task := p.Tasks[0] + assert.Equal(t, "exp", task.TaskKey) + assert.Equal(t, "ALL_SUCCESS", task.RunIf) + assert.Equal(t, aiRuntimeEnvironmentKey, task.EnvironmentKey) + assert.Equal(t, 2, task.MaxRetries) + assert.True(t, task.RetryOnTimeout) + + at := task.AiRuntimeTask + assert.Equal(t, "exp", at.Experiment) + assert.Equal(t, "run-v2", at.MlflowRun) + assert.Equal(t, "/Workspace/Users/me/exp", at.MlflowExperimentDirectory) + require.Len(t, at.Deployments, 1) + assert.Equal(t, "/d/command.sh", at.Deployments[0].CommandPath) + assert.Equal(t, aiRuntimeCompute{AcceleratorType: "GPU_8xH100", AcceleratorCount: 16}, at.Deployments[0].Compute) +} + +func TestBuildSubmitPayload_NoRetries(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, + MaxRetries: new(0), + } + + task := buildSubmitPayload(cfg, "/d/command.sh", "4").Tasks[0] + assert.Equal(t, 0, task.MaxRetries) + assert.False(t, task.RetryOnTimeout) + + // max_retries: 0 must be sent, not omitted, so the server honors "no retries". + b, err := json.Marshal(task) + require.NoError(t, err) + assert.Contains(t, string(b), `"max_retries":0`) +} + +func TestSubmitToken(t *testing.T) { + cfg := &runConfig{IdempotencyToken: new("from-config")} + + tok, err := submitToken("from-flag", cfg) // flag wins + require.NoError(t, err) + assert.Equal(t, "from-flag", tok) + + tok, err = submitToken("", cfg) // then config + require.NoError(t, err) + assert.Equal(t, "from-config", tok) + + tok, err = submitToken("", &runConfig{}) // else generated + require.NoError(t, err) + assert.NotEmpty(t, tok) + + // An over-long token errors instead of being truncated. + _, err = submitToken(strings.Repeat("a", 65), cfg) + require.ErrorContains(t, err, "64 characters or less") +} + +func TestJobsSubmitClient(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var got jobsSubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return submitRunResponse{RunID: 999} + }) + + w := &databricks.WorkspaceClient{Config: &config.Config{Host: server.URL, Token: "token"}} + jc, err := newJobsSubmitClient(w) + require.NoError(t, err) + + runID, err := jc.submit(t.Context(), jobsSubmitRun{RunName: "exp", Tasks: []submitTask{{TaskKey: "exp"}}}) + require.NoError(t, err) + assert.Equal(t, int64(999), runID) + assert.Equal(t, "exp", got.RunName) +} + +func TestSubmitWorkload(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + testserver.AddDefaultHandlers(server) + + var got jobsSubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return submitRunResponse{RunID: 777} + }) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cfg, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") + require.NoError(t, err) + assert.Equal(t, int64(777), runID) + assert.Contains(t, dashboardURL, "/jobs/runs/777") + + // The submitted payload is a native ai_runtime_task pointing at the uploaded + // command.sh under the run's launch directory. + assert.Equal(t, "my-run", got.RunName) + assert.Equal(t, "idem-key", got.IdempotencyToken) + require.Len(t, got.Environments, 1) + require.Len(t, got.Tasks, 1) + at := got.Tasks[0].AiRuntimeTask + require.Len(t, at.Deployments, 1) + d := at.Deployments[0] + assert.True(t, strings.HasSuffix(d.CommandPath, "/"+commandScriptName), d.CommandPath) + assert.Contains(t, d.CommandPath, "/.air/cli_launch/") + assert.Equal(t, aiRuntimeCompute{AcceleratorType: "GPU_1xH100", AcceleratorCount: 1}, d.Compute) +} + +func TestSubmitWorkloadGuards(t *testing.T) { + w := newFakeWorkspaceClient(t) + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + base, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + t.Run("usage_policy_name rejected", func(t *testing.T) { + cfg := *base + cfg.UsagePolicyName = new("p") + _, _, err := submitWorkload(t.Context(), w, &cfg, cfgPath, "") + require.ErrorContains(t, err, "usage_policy_name is not yet supported") + }) + + t.Run("code_source rejected", func(t *testing.T) { + cfg := *base + cfg.CodeSource = &codeSourceConfig{Type: "snapshot"} + _, _, err := submitWorkload(t.Context(), w, &cfg, cfgPath, "") + require.ErrorContains(t, err, "code_source is not yet supported") + }) +} diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go new file mode 100644 index 00000000000..fb9ca00b987 --- /dev/null +++ b/experimental/air/cmd/runupload.go @@ -0,0 +1,170 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "maps" + "os" + "path/filepath" + "slices" + "strings" + + "github.com/databricks/cli/libs/filer" + "go.yaml.in/yaml/v3" +) + +// Launch artifact basenames, uploaded into the run's cli_launch directory. The +// server-side launcher derives requirements.yaml / hyperparameters.yaml from the +// same directory, so these names are part of the contract. +const ( + trainingConfigName = "training_config.yaml" + commandScriptName = "command.sh" + requirementsName = "requirements.yaml" + hyperparametersName = "hyperparameters.yaml" + envVarsName = "env_vars.json" + secretEnvVarsName = "secret_env_vars.json" +) + +// maxConfigYAMLBytes caps training_config.yaml. It is referenced by the Jobs +// payload and rendered on the run page, so an oversized parameters/command block +// is rejected here; full parameters still ship in hyperparameters.yaml. +const maxConfigYAMLBytes = 1024 * 1024 + +// uploadItem is a single artifact to write into the launch directory. +type uploadItem struct { + name string + data []byte +} + +// fileWriter is the subset of filer.Filer the upload path needs; a narrow +// interface keeps buildArtifacts/upload testable without a live workspace. +type fileWriter interface { + Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error +} + +// requirementsDoc mirrors the on-disk requirements.yaml format so the worker +// parses synthesized inline dependencies identically to a user-provided file. +type requirementsDoc struct { + Version string `yaml:"version,omitempty"` + Dependencies []string `yaml:"dependencies"` +} + +// buildArtifacts assembles the files to upload for a run: the merged config, the +// inline command as a script, requirements (from a file or synthesized from +// inline dependencies), and hyperparameters. configPath is the local YAML path. +func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { + // TODO(DABs): with no _bases_/overrides ported yet, the merged config is the + // file as-is; once those land, upload the re-serialized merged YAML instead. + configData, err := os.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("failed to read config %s: %w", configPath, err) + } + if len(configData) > maxConfigYAMLBytes { + return nil, fmt.Errorf("config YAML is %.2f MB, over the %d MB limit; reduce 'parameters' or 'command'", + float64(len(configData))/(1024*1024), maxConfigYAMLBytes/(1024*1024)) + } + + items := []uploadItem{ + {trainingConfigName, configData}, + {commandScriptName, []byte(*cfg.Command)}, + } + + switch reqPath, ok := cfg.requirementsFile(); { + case ok: + // Resolve a relative requirements path against the config's directory. + if !filepath.IsAbs(reqPath) { + reqPath = filepath.Join(filepath.Dir(configPath), reqPath) + } + data, err := os.ReadFile(reqPath) + if err != nil { + return nil, fmt.Errorf("failed to read requirements file %s: %w", reqPath, err) + } + items = append(items, uploadItem{requirementsName, data}) + default: + if deps, ok := cfg.inlineDependencies(); ok { + version, _ := cfg.runtimeVersion() + data, err := yaml.Marshal(requirementsDoc{Version: version, Dependencies: deps}) + if err != nil { + return nil, fmt.Errorf("failed to synthesize requirements.yaml: %w", err) + } + items = append(items, uploadItem{requirementsName, data}) + } + } + + if len(cfg.Parameters) > 0 { + data, err := yaml.Marshal(cfg.Parameters) + if err != nil { + return nil, fmt.Errorf("failed to serialize parameters: %w", err) + } + items = append(items, uploadItem{hyperparametersName, data}) + } + + // The ai_runtime_task proto carries no inline env vars or secrets; stage them + // as JSON files co-located with command.sh for the server-side launcher. + if len(cfg.EnvVariables) > 0 { + data, err := json.Marshal(envVarEntries(cfg.EnvVariables)) + if err != nil { + return nil, fmt.Errorf("failed to serialize env_variables: %w", err) + } + items = append(items, uploadItem{envVarsName, data}) + } + if len(cfg.Secrets) > 0 { + data, err := json.Marshal(secretEnvVarEntries(cfg.Secrets)) + if err != nil { + return nil, fmt.Errorf("failed to serialize secrets: %w", err) + } + items = append(items, uploadItem{secretEnvVarsName, data}) + } + + return items, nil +} + +// envVarEntry is one entry in env_vars.json. +type envVarEntry struct { + Name string `json:"name"` + Value string `json:"value"` +} + +// secretEnvVarEntry is one entry in secret_env_vars.json. The YAML side is +// {ENV_VAR: "scope/key"}; the launcher wants the split form. +type secretEnvVarEntry struct { + Name string `json:"name"` + SecretScope string `json:"secret_scope"` + SecretKey string `json:"secret_key"` +} + +// envVarEntries renders env_variables sorted by name for deterministic output. +func envVarEntries(vars map[string]string) []envVarEntry { + out := make([]envVarEntry, 0, len(vars)) + for _, name := range slices.Sorted(maps.Keys(vars)) { + out = append(out, envVarEntry{Name: name, Value: vars[name]}) + } + return out +} + +// secretEnvVarEntries renders secrets sorted by name for deterministic output. +func secretEnvVarEntries(secrets map[string]string) []secretEnvVarEntry { + out := make([]secretEnvVarEntry, 0, len(secrets)) + for _, name := range slices.Sorted(maps.Keys(secrets)) { + scope, key, _ := strings.Cut(secrets[name], "/") + out = append(out, secretEnvVarEntry{Name: name, SecretScope: scope, SecretKey: key}) + } + return out +} + +// uploadArtifacts writes each artifact into the launch directory, overwriting and +// creating parents as needed. +// +// TODO(DABs): this client-side upload could move onto libs/sync / a bundle deploy +// so the CLI reuses DABs' file-staging machinery instead of writing files itself. +func uploadArtifacts(ctx context.Context, w fileWriter, items []uploadItem) error { + for _, it := range items { + if err := w.Write(ctx, it.name, bytes.NewReader(it.data), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return fmt.Errorf("failed to upload %s: %w", it.name, err) + } + } + return nil +} diff --git a/experimental/air/cmd/runupload_test.go b/experimental/air/cmd/runupload_test.go new file mode 100644 index 00000000000..0c87524735d --- /dev/null +++ b/experimental/air/cmd/runupload_test.go @@ -0,0 +1,155 @@ +package aircmd + +import ( + "context" + "errors" + "io" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/filer" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeWriter records artifact writes in place of a workspace filer. +type fakeWriter struct { + written map[string]string +} + +func (f *fakeWriter) Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error { + if f.written == nil { + f.written = map[string]string{} + } + data, err := io.ReadAll(reader) + if err != nil { + return err + } + f.written[name] = string(data) + return nil +} + +func writeConfigFile(t *testing.T, name, content string) string { + t.Helper() + path := filepath.Join(t.TempDir(), name) + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) + return path +} + +func itemNames(items []uploadItem) []string { + names := make([]string, len(items)) + for i, it := range items { + names[i] = it.name + } + return names +} + +func TestBuildArtifacts_CommandAndConfig(t *testing.T) { + path := writeConfigFile(t, "run.yaml", minimalConfig) + cfg := &runConfig{Command: new("python train.py")} + + items, err := buildArtifacts(cfg, path) + require.NoError(t, err) + assert.Equal(t, []string{trainingConfigName, commandScriptName}, itemNames(items)) + assert.Equal(t, minimalConfig, string(items[0].data)) + assert.Equal(t, "python train.py", string(items[1].data)) +} + +func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { + path := writeConfigFile(t, "run.yaml", "x: y\n") + cfg := &runConfig{ + Command: new("echo hi"), + Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: true, list: []string{"torch", "numpy"}}, + Version: stringOrInt{set: true, raw: "5"}, + }, + Parameters: map[string]any{"lr": 0.1}, + } + + items, err := buildArtifacts(cfg, path) + require.NoError(t, err) + assert.Equal(t, []string{trainingConfigName, commandScriptName, requirementsName, hyperparametersName}, itemNames(items)) + + var reqIdx int + for i, it := range items { + if it.name == requirementsName { + reqIdx = i + } + } + req := string(items[reqIdx].data) + assert.Contains(t, req, "version: \"5\"") + assert.Contains(t, req, "- torch") +} + +func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { + path := writeConfigFile(t, "run.yaml", "x: y\n") + cfg := &runConfig{ + Command: new("echo hi"), + EnvVariables: map[string]string{"WANDB": "demo"}, + Secrets: map[string]string{"HF_TOKEN": "myscope/hf"}, + } + + items, err := buildArtifacts(cfg, path) + require.NoError(t, err) + assert.Subset(t, itemNames(items), []string{envVarsName, secretEnvVarsName}) + + byName := map[string][]byte{} + for _, it := range items { + byName[it.name] = it.data + } + assert.JSONEq(t, `[{"name":"WANDB","value":"demo"}]`, string(byName[envVarsName])) + assert.JSONEq(t, `[{"name":"HF_TOKEN","secret_scope":"myscope","secret_key":"hf"}]`, string(byName[secretEnvVarsName])) +} + +func TestBuildArtifacts_RequirementsFile(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "run.yaml"), []byte("x: y\n"), 0o600)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "reqs.yaml"), []byte("version: 4\n"), 0o600)) + cfg := &runConfig{ + Command: new("echo hi"), + Environment: &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "reqs.yaml"}}, + } + + items, err := buildArtifacts(cfg, filepath.Join(dir, "run.yaml")) + require.NoError(t, err) + assert.Contains(t, itemNames(items), requirementsName) +} + +func TestBuildArtifacts_OversizeConfigRejected(t *testing.T) { + path := writeConfigFile(t, "run.yaml", strings.Repeat("a", maxConfigYAMLBytes+1)) + _, err := buildArtifacts(&runConfig{Command: new("x")}, path) + require.Error(t, err) + assert.Contains(t, err.Error(), "over the 1 MB limit") +} + +func TestUploadArtifacts(t *testing.T) { + w := &fakeWriter{} + items := []uploadItem{{trainingConfigName, []byte("cfg")}, {commandScriptName, []byte("cmd")}} + require.NoError(t, uploadArtifacts(t.Context(), w, items)) + assert.Equal(t, "cfg", w.written[trainingConfigName]) + assert.Equal(t, "cmd", w.written[commandScriptName]) +} + +// errWriter fails every Write, exercising the upload error path. +type errWriter struct{} + +func (errWriter) Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error { + return errors.New("boom") +} + +func TestUploadArtifacts_WriteError(t *testing.T) { + err := uploadArtifacts(t.Context(), errWriter{}, []uploadItem{{trainingConfigName, []byte("x")}}) + require.ErrorContains(t, err, "failed to upload "+trainingConfigName) +} + +func TestBuildArtifacts_MissingRequirementsFile(t *testing.T) { + cfgPath := writeConfigFile(t, "run.yaml", "x: y\n") + cfg := &runConfig{ + Command: new("echo hi"), + Environment: &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "nope.yaml"}}, + } + _, err := buildArtifacts(cfg, cfgPath) + require.ErrorContains(t, err, "failed to read requirements file") +} diff --git a/experimental/air/cmd/stubs_test.go b/experimental/air/cmd/stubs_test.go index b9f5c330f00..4607d7d9eac 100644 --- a/experimental/air/cmd/stubs_test.go +++ b/experimental/air/cmd/stubs_test.go @@ -13,7 +13,6 @@ import ( // fails with a "not implemented" error. Drop a command here once it lands. func TestStubCommandsReturnNotImplemented(t *testing.T) { stubs := map[string]*cobra.Command{ - "run": newRunCommand(), "logs": newLogsCommand(), "cancel": newCancelCommand(), "register-image": newRegisterImageCommand(), From 8b3f0d5838b2ba12d9ebac31f8175da3bac96a1e Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Wed, 1 Jul 2026 12:15:38 -0700 Subject: [PATCH 09/18] experimental/air: `air list` command (Jobs API, interactive UI) (#5793) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Changes Adds `air list` — a browsable view of the caller's recent AIR training runs. - **Data source:** reads Jobs `runs/list` directly (`expand_tasks`) and filters to AIR runs. The current AIR task type is `ai_runtime_task`, which the typed Jobs SDK doesn't model, so the response is parsed raw for experiment, accelerators, status, timing, and user. - **Interactive table:** in a terminal `air list` renders an inline, navigable table (Bubble Tea + Lip Gloss + termenv) — `↑/↓` move a row, `←/→` page (20/screen), `Enter` opens the run's MLflow page, `q` quits. Status is colored by state; the MLflow column is a short clickable hyperlink. - **Non-interactive:** piped output, an explicit `--limit`, and empty results print the table once; `-o json` emits the air `{v,ts,data}` envelope. `--limit` defaults to 20. - **Flags:** `--active`, `--all-users`, and client-side `--filter` keys (`experiment`, `accelerator_type`, `num_accelerators`). - **MLflow links** are resolved per run from `runs/get-output` (`ai_runtime_task_output`, with the legacy `gen_ai_compute_output` as fallback) in text mode. This also fixes `air get`'s MLflow column for `ai_runtime_task` runs. Also adds `cmdio.IsPagerSupported` (stdin+stdout+stderr TTY) and promotes `termenv` to a direct dependency (NOTICE updated). ## Why The `air` CLI needs a run list on par with the Python `air` CLI and `databricks jobs list-runs`. It talks to the Jobs API directly rather than the AiWorkflowService, whose list RPC returns only run identifiers. ## Tests - Unit: raw `runs/list` parse (incl. `ai_runtime_task`), row mapping, `--filter` matching, the TUI model (navigation, paging, 20-row cap, window scroll, quit), MLflow-id resolution (ai_runtime + legacy), and status/accelerator helpers. - Acceptance: `acceptance/experimental/air/list` (text + JSON, non-AIR filtered out, MLflow link resolved) plus `help`/`get` fixture updates. - Manual: verified against an e2-dogfood workspace (text, JSON, filters, `--all-users`, MLflow links). This pull request and its description were written by Isaac. --- acceptance/experimental/air/get/test.toml | 1 + acceptance/experimental/air/list/output.txt | 13 +- acceptance/experimental/air/list/test.toml | 59 +++---- experimental/air/cmd/aiworkflow.go | 77 --------- experimental/air/cmd/aiworkflow_test.go | 62 ------- experimental/air/cmd/format.go | 50 ++---- experimental/air/cmd/format_test.go | 28 +--- experimental/air/cmd/joblist.go | 172 +++++++++++++++++++ experimental/air/cmd/list.go | 141 +++++++++++----- experimental/air/cmd/list_filter.go | 24 +-- experimental/air/cmd/list_filter_test.go | 4 +- experimental/air/cmd/list_format.go | 46 +++--- experimental/air/cmd/list_test.go | 174 ++++++++++++-------- experimental/air/cmd/mlflow.go | 44 +++-- experimental/air/cmd/mlflow_test.go | 26 +++ 15 files changed, 505 insertions(+), 416 deletions(-) delete mode 100644 experimental/air/cmd/aiworkflow.go delete mode 100644 experimental/air/cmd/aiworkflow_test.go create mode 100644 experimental/air/cmd/joblist.go diff --git a/acceptance/experimental/air/get/test.toml b/acceptance/experimental/air/get/test.toml index 456feb39dbc..606d34ec970 100644 --- a/acceptance/experimental/air/get/test.toml +++ b/acceptance/experimental/air/get/test.toml @@ -22,6 +22,7 @@ Response.Body = ''' "tasks": [ { "task_key": "train", + "run_id": 456, "attempt_number": 0, "max_retries": 3, "gen_ai_compute_task": { diff --git a/acceptance/experimental/air/list/output.txt b/acceptance/experimental/air/list/output.txt index 9e8f93826f5..3f84d281ad4 100644 --- a/acceptance/experimental/air/list/output.txt +++ b/acceptance/experimental/air/list/output.txt @@ -1,9 +1,8 @@ === list (text) >>> [CLI] experimental air list - Run ID Experiment Status Started Duration MLflow User Accelerators - [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 - [NUMID] llama-train ● FAILED [TIMESTAMP] 3m 31s - [USERNAME] 1x A10 + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 === list (json) >>> [CLI] experimental air list -o json @@ -19,14 +18,6 @@ "status": "SUCCESS", "started_at": "[TIMESTAMP]", "is_sweep": false - }, - { - "run_id": "[NUMID]", - "run_name": "llama-train", - "user": "[USERNAME]", - "status": "FAILED", - "started_at": "[TIMESTAMP]", - "is_sweep": false } ] } diff --git a/acceptance/experimental/air/list/test.toml b/acceptance/experimental/air/list/test.toml index f547919c3e5..23dd0f6ba5b 100644 --- a/acceptance/experimental/air/list/test.toml +++ b/acceptance/experimental/air/list/test.toml @@ -8,43 +8,46 @@ DATABRICKS_BUNDLE_ENGINE = [] Pattern = "HEAD /" Response.Body = '' -# ListTrainingWorkflows: the server resolves AIR runs, creator scoping, and MLflow -# IDs, so the CLI just maps the response. One completed run with resolved MLflow -# IDs (so the MLflow column is a link) and one failed run without them (so it -# falls back to "-"). +# `air list` reads Jobs runs/list directly (expand_tasks), filtering to AIR runs. +# The current runs carry an ai_runtime_task (not modeled by the typed SDK), which +# is why the CLI parses the response raw. Both runs belong to the default user +# (tester@databricks.com, from the built-in scim/v2/Me handler). The second run +# has no AI task and must be filtered out. [[Server]] -Pattern = "GET /api/2.0/ai-training/workflows" +Pattern = "GET /api/2.2/jobs/runs/list" Response.Body = ''' { - "training_workflows": [ + "runs": [ { - "job_run_id": "334747067049496", - "metadata": {"creator_name": "tester@databricks.com"}, - "spec": {"compute": {"hardware_accelerator_type": "GPU_8xH100", "accelerator_count": 8}}, - "status": { - "state": "TRAINING_WORKFLOW_STATE_TERMINATED_COMPLETED", - "start_time": "2026-06-05T17:32:39.000Z", - "end_time": "2026-06-05T17:32:51.000Z", - "job": {"name": "qwen-train"}, - "mlflow": { + "run_id": 334747067049496, + "run_name": "qwen-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717608759000, + "end_time": 1717608771000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{ + "run_id": 334747067049497, + "ai_runtime_task": { "experiment": "/Users/tester@databricks.com/qwen-train", - "experiment_id": "exp1", - "run_id": "run1" + "deployments": [{"compute": {"accelerator_type": "GPU_8xH100", "accelerator_count": 8}}] } - } + }] }, { - "job_run_id": "566001814929041", - "metadata": {"creator_name": "tester@databricks.com"}, - "spec": {"compute": {"hardware_accelerator_type": "GPU_1xA10", "accelerator_count": 1}}, - "status": { - "state": "TRAINING_WORKFLOW_STATE_TERMINATED_FAILED", - "start_time": "2026-06-05T18:43:24.000Z", - "end_time": "2026-06-05T18:46:55.000Z", - "job": {"name": "llama-train"}, - "mlflow": {"experiment": "/Users/tester@databricks.com/llama-train"} - } + "run_id": 999000999000, + "run_name": "not-an-air-run", + "creator_user_name": "tester@databricks.com", + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{"notebook_task": {"notebook_path": "/x"}}] } ] } ''' + +# MLflow IDs for the deep link, fetched per AIR run (text mode). The current +# task type resolves them under ai_runtime_task_output. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = ''' +{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}} +''' diff --git a/experimental/air/cmd/aiworkflow.go b/experimental/air/cmd/aiworkflow.go deleted file mode 100644 index 489a24eeb8c..00000000000 --- a/experimental/air/cmd/aiworkflow.go +++ /dev/null @@ -1,77 +0,0 @@ -package aircmd - -import ( - "context" - "errors" - "fmt" - "net/http" - - "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/apierr" - "github.com/databricks/databricks-sdk-go/client" -) - -// listWorkflowsPath is the ListTrainingWorkflows binding. It is -// PUBLIC_UNDOCUMENTED, so we call it directly like the endpoints in mlflow.go. -const listWorkflowsPath = "/api/2.0/ai-training/workflows" - -// trainingWorkflowsResponse is the ListTrainingWorkflows response. Only the -// fields the list table and JSON envelope consume are modeled. -type trainingWorkflowsResponse struct { - TrainingWorkflows []trainingWorkflow `json:"training_workflows"` - NextPageToken string `json:"next_page_token"` -} - -type trainingWorkflow struct { - JobRunID string `json:"job_run_id"` - - TaskRunID string `json:"task_run_id"` - Metadata struct { - CreatorName string `json:"creator_name"` - } `json:"metadata"` - Spec struct { - Compute struct { - HardwareAcceleratorType string `json:"hardware_accelerator_type"` - AcceleratorCount int `json:"accelerator_count"` - } `json:"compute"` - } `json:"spec"` - Status struct { - State string `json:"state"` - StartTime string `json:"start_time"` - EndTime string `json:"end_time"` - Job struct { - Name string `json:"name"` - } `json:"job"` - Mlflow struct { - Experiment string `json:"experiment"` - ExperimentID string `json:"experiment_id"` - RunID string `json:"run_id"` - } `json:"mlflow"` - } `json:"status"` -} - -// listTrainingWorkflows calls the ListTrainingWorkflows RPC with the given query -// parameters. -func listTrainingWorkflows(ctx context.Context, w *databricks.WorkspaceClient, query map[string]any) (*trainingWorkflowsResponse, error) { - apiClient, err := client.New(w.Config) - if err != nil { - return nil, fmt.Errorf("failed to create API client: %w", err) - } - - var resp trainingWorkflowsResponse - err = apiClient.Do(ctx, http.MethodGet, listWorkflowsPath, nil, nil, query, &resp) - if err != nil { - // A server error can arrive with an empty body, leaving %w blank, so - // surface the HTTP status to make the failure diagnosable. - if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { - switch apiErr.StatusCode { - case http.StatusGatewayTimeout, http.StatusServiceUnavailable, http.StatusRequestTimeout: - // Trailing %w keeps the error chain; the body is usually empty. - return nil, fmt.Errorf("timed out listing runs (HTTP %d): --all-users makes the server list every user's runs, which can exceed the gateway timeout on a busy workspace. Add --active to list only active runs, or drop --all-users to list your own.%w", apiErr.StatusCode, err) - } - return nil, fmt.Errorf("failed to list training workflows (HTTP %d): %w", apiErr.StatusCode, err) - } - return nil, fmt.Errorf("failed to list training workflows: %w", err) - } - return &resp, nil -} diff --git a/experimental/air/cmd/aiworkflow_test.go b/experimental/air/cmd/aiworkflow_test.go deleted file mode 100644 index 5899d310db0..00000000000 --- a/experimental/air/cmd/aiworkflow_test.go +++ /dev/null @@ -1,62 +0,0 @@ -package aircmd - -import ( - "net/http" - "net/http/httptest" - "net/url" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestListTrainingWorkflows(t *testing.T) { - var gotQuery url.Values - body := `{ - "training_workflows": [ - { - "job_run_id": "123", - "metadata": {"creator_name": "me@example.com"}, - "spec": {"compute": {"hardware_accelerator_type": "GPU_8xH100", "accelerator_count": 8}}, - "status": { - "state": "TRAINING_WORKFLOW_STATE_RUNNING", - "start_time": "2026-06-05T17:32:39.791Z", - "job": {"name": "my-run"}, - "mlflow": {"experiment": "/Users/me@example.com/exp", "experiment_id": "E1", "run_id": "R1"} - } - } - ], - "next_page_token": "tok" - }` - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == listWorkflowsPath { - gotQuery = r.URL.Query() - _, _ = w.Write([]byte(body)) - return - } - _, _ = w.Write([]byte(`{}`)) - })) - t.Cleanup(srv.Close) - - resp, err := listTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), map[string]any{ - "creator_name": "me@example.com", - "active_only": true, - }) - require.NoError(t, err) - require.Len(t, resp.TrainingWorkflows, 1) - assert.Equal(t, "tok", resp.NextPageToken) - - wf := resp.TrainingWorkflows[0] - assert.Equal(t, "123", wf.JobRunID) - assert.Equal(t, "me@example.com", wf.Metadata.CreatorName) - assert.Equal(t, "GPU_8xH100", wf.Spec.Compute.HardwareAcceleratorType) - assert.Equal(t, 8, wf.Spec.Compute.AcceleratorCount) - assert.Equal(t, "TRAINING_WORKFLOW_STATE_RUNNING", wf.Status.State) - assert.Equal(t, "my-run", wf.Status.Job.Name) - assert.Equal(t, "E1", wf.Status.Mlflow.ExperimentID) - assert.Equal(t, "R1", wf.Status.Mlflow.RunID) - - // The query map is serialized into the GET query string. - assert.Equal(t, "me@example.com", gotQuery.Get("creator_name")) - assert.Equal(t, "true", gotQuery.Get("active_only")) -} diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index 8539a1794c9..c32184517ba 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -97,11 +97,17 @@ func runStatus(state *jobs.RunState) string { if state == nil { return "UNKNOWN" } - if state.ResultState != "" { - return string(state.ResultState) + return statusWord(string(state.LifeCycleState), string(state.ResultState)) +} + +// statusWord picks the status word to show from a run's lifecycle and result +// states: the result state is the more meaningful one, so it wins when set. +func statusWord(lifeCycle, result string) string { + if result != "" { + return result } - if state.LifeCycleState != "" { - return string(state.LifeCycleState) + if lifeCycle != "" { + return lifeCycle } return "UNKNOWN" } @@ -148,20 +154,6 @@ func isoFormat(t time.Time) string { return t.Format(layout) } -// parseRPCTime parses an RFC3339 timestamp returned by the AiWorkflow RPC (e.g. -// "2026-06-05T18:46:55.876Z"), returning the zero time when the field is absent -// or unparseable. -func parseRPCTime(s string) time.Time { - if s == "" { - return time.Time{} - } - t, err := time.Parse(time.RFC3339, s) - if err != nil { - return time.Time{} - } - return t -} - // submittedDisplay formats the run's start time for the text table as // "2006-01-02 15:04 UTC", or "N/A" if it hasn't started. Mirrors Python's // _format_timestamp (cli_display.py); we render in UTC for stable output rather @@ -289,28 +281,6 @@ func acceleratorLabel(gpuType string, count int) string { return fmt.Sprintf("%dx", count) } -// trainingWorkflowStatus maps a TrainingWorkflowState enum value to the status -// word shown for a run, matching `air get run`. The server collapses Jobs -// lifecycle/result into this enum, so lossy cases (e.g. TIMEDOUT) render as FAILED. -func trainingWorkflowStatus(state string) string { - switch state { - case "TRAINING_WORKFLOW_STATE_PENDING", "TRAINING_WORKFLOW_STATE_PENDING_SENT": - return "PENDING" - case "TRAINING_WORKFLOW_STATE_RUNNING": - return "RUNNING" - case "TRAINING_WORKFLOW_STATE_TERMINATION_REQUESTED", "TRAINING_WORKFLOW_STATE_TERMINATION_SENT": - return "TERMINATING" - case "TRAINING_WORKFLOW_STATE_TERMINATED_COMPLETED": - return "SUCCESS" - case "TRAINING_WORKFLOW_STATE_TERMINATED_FAILED": - return "FAILED" - case "TRAINING_WORKFLOW_STATE_TERMINATED_STOPPED": - return "CANCELED" - default: - return "UNKNOWN" - } -} - // gpuDisplayName returns the friendly name for a GPU identifier, falling back to // the identifier itself when it is not one we recognize. func gpuDisplayName(gpuType string) string { diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go index 45bcbd94199..62a6d7ac580 100644 --- a/experimental/air/cmd/format_test.go +++ b/experimental/air/cmd/format_test.go @@ -234,28 +234,8 @@ func TestAcceleratorLabel(t *testing.T) { assert.Equal(t, "8x", acceleratorLabel("", 8)) } -func TestTrainingWorkflowStatus(t *testing.T) { - cases := map[string]string{ - "TRAINING_WORKFLOW_STATE_PENDING": "PENDING", - "TRAINING_WORKFLOW_STATE_PENDING_SENT": "PENDING", - "TRAINING_WORKFLOW_STATE_RUNNING": "RUNNING", - "TRAINING_WORKFLOW_STATE_TERMINATION_REQUESTED": "TERMINATING", - "TRAINING_WORKFLOW_STATE_TERMINATION_SENT": "TERMINATING", - "TRAINING_WORKFLOW_STATE_TERMINATED_COMPLETED": "SUCCESS", - "TRAINING_WORKFLOW_STATE_TERMINATED_FAILED": "FAILED", - "TRAINING_WORKFLOW_STATE_TERMINATED_STOPPED": "CANCELED", - "TRAINING_WORKFLOW_STATE_UNSPECIFIED": "UNKNOWN", - "": "UNKNOWN", - } - for state, want := range cases { - assert.Equal(t, want, trainingWorkflowStatus(state), state) - } -} - -func TestParseRPCTime(t *testing.T) { - assert.True(t, parseRPCTime("").IsZero()) - assert.True(t, parseRPCTime("not-a-time").IsZero()) - got := parseRPCTime("2026-06-05T18:46:55.876Z") - require.False(t, got.IsZero()) - assert.Equal(t, "2026-06-05T18:46:55.876000+00:00", isoFormat(got)) +func TestStatusWord(t *testing.T) { + assert.Equal(t, "SUCCESS", statusWord("TERMINATED", "SUCCESS")) // result wins + assert.Equal(t, "RUNNING", statusWord("RUNNING", "")) // falls back to lifecycle + assert.Equal(t, "UNKNOWN", statusWord("", "")) } diff --git a/experimental/air/cmd/joblist.go b/experimental/air/cmd/joblist.go new file mode 100644 index 00000000000..8a50921264e --- /dev/null +++ b/experimental/air/cmd/joblist.go @@ -0,0 +1,172 @@ +package aircmd + +import ( + "context" + "fmt" + "net/http" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" +) + +// jobsRunsListPath is the Jobs runs/list endpoint. We call it directly (rather +// than via the typed SDK) because the SDK's RunTask omits ai_runtime_task, the +// task type the AI runtime now submits. +const jobsRunsListPath = "/api/2.2/jobs/runs/list" + +type jobsRunsListResponse struct { + Runs []jobRun `json:"runs"` + NextPageToken string `json:"next_page_token"` +} + +type jobRun struct { + RunID int64 `json:"run_id"` + RunName string `json:"run_name"` + CreatorUserName string `json:"creator_user_name"` + StartTime int64 `json:"start_time"` + EndTime int64 `json:"end_time"` + State jobState `json:"state"` + Tasks []jobTask `json:"tasks"` +} + +type jobState struct { + LifeCycleState string `json:"life_cycle_state"` + ResultState string `json:"result_state"` +} + +type jobTask struct { + RunID int64 `json:"run_id"` + StartTime int64 `json:"start_time"` + EndTime int64 `json:"end_time"` + AiRuntimeTask *jobAiRuntimeTask `json:"ai_runtime_task"` + GenAiComputeTask *genAiComputeTask `json:"gen_ai_compute_task"` + ForEachTask *forEachTask `json:"for_each_task"` +} + +type forEachTask struct { + Task jobTask `json:"task"` +} + +// jobAiRuntimeTask is the current AI runtime task shape; deployments[0].compute +// carries the accelerator info. +type jobAiRuntimeTask struct { + Experiment string `json:"experiment"` + Deployments []aiRuntimeDeploy `json:"deployments"` +} + +type aiRuntimeDeploy struct { + Compute airCompute `json:"compute"` +} + +type airCompute struct { + AcceleratorType string `json:"accelerator_type"` + AcceleratorCount int `json:"accelerator_count"` +} + +// genAiComputeTask is the legacy task shape, still recognized for older runs. +type genAiComputeTask struct { + TrainingScriptPath string `json:"training_script_path"` + MlflowExperimentName string `json:"mlflow_experiment_name"` + Compute *genAiCompute `json:"compute"` +} + +type genAiCompute struct { + GpuType string `json:"gpu_type"` + NumGpus int `json:"num_gpus"` +} + +// airTask returns the run's AI runtime / legacy GenAI task, unwrapping a foreach +// sweep when present. +func (r *jobRun) airTask() (*jobAiRuntimeTask, *genAiComputeTask) { + if len(r.Tasks) == 0 { + return nil, nil + } + t := r.Tasks[0] + if t.AiRuntimeTask != nil || t.GenAiComputeTask != nil { + return t.AiRuntimeTask, t.GenAiComputeTask + } + if t.ForEachTask != nil { + return t.ForEachTask.Task.AiRuntimeTask, t.ForEachTask.Task.GenAiComputeTask + } + return nil, nil +} + +// isAirRun reports whether a run is an AI runtime workload: an ai_runtime_task, +// or a legacy gen_ai_compute_task with a training script. +func isAirRun(r *jobRun) bool { + ai, gen := r.airTask() + return ai != nil || (gen != nil && gen.TrainingScriptPath != "") +} + +// isSweep reports whether the run's first task fans out into iterations. +func isSweep(r *jobRun) bool { + return len(r.Tasks) > 0 && r.Tasks[0].ForEachTask != nil +} + +// taskRunID returns the run id of the AIR task, used to fetch its MLflow output. +func taskRunID(r *jobRun) int64 { + if len(r.Tasks) == 0 { + return 0 + } + t := r.Tasks[0] + if t.ForEachTask != nil { + return t.ForEachTask.Task.RunID + } + return t.RunID +} + +// jobExperiment returns the run's MLflow experiment name (user-folder prefix +// stripped), or "" when there is none. +func jobExperiment(r *jobRun) string { + ai, gen := r.airTask() + switch { + case ai != nil && ai.Experiment != "": + return stripExperimentUserPrefix(ai.Experiment) + case gen != nil && gen.MlflowExperimentName != "": + return stripExperimentUserPrefix(gen.MlflowExperimentName) + } + return "" +} + +// jobCompute returns the run's accelerator type and count, or ("", 0) when it +// has none. +func jobCompute(r *jobRun) (string, int) { + ai, gen := r.airTask() + switch { + case ai != nil && len(ai.Deployments) > 0: + c := ai.Deployments[0].Compute + return c.AcceleratorType, c.AcceleratorCount + case gen != nil && gen.Compute != nil: + return gen.Compute.GpuType, gen.Compute.NumGpus + } + return "", 0 +} + +// jobTiming returns the run's start and end times (epoch ms), preferring the +// first task's window so a run reports its task attempt rather than the wrapper. +func jobTiming(r *jobRun) (startMillis, endMillis int64) { + startMillis, endMillis = r.StartTime, r.EndTime + if len(r.Tasks) > 0 { + if t := r.Tasks[0]; t.StartTime > 0 { + startMillis = t.StartTime + endMillis = t.EndTime + } + } + return startMillis, endMillis +} + +// fetchJobRunsPage fetches one page of Jobs runs/list. query carries the request +// params (and page_token across calls). +func fetchJobRunsPage(ctx context.Context, w *databricks.WorkspaceClient, query map[string]any) (*jobsRunsListResponse, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var resp jobsRunsListResponse + err = apiClient.Do(ctx, http.MethodGet, jobsRunsListPath, nil, nil, query, &resp) + if err != nil { + return nil, fmt.Errorf("failed to list runs: %w", err) + } + return &resp, nil +} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 91ccc6515c9..360fc97e678 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -11,16 +11,22 @@ import ( "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/spf13/cobra" + "golang.org/x/sync/errgroup" ) -// maxListPages caps how many pages `air list runs` fetches (a safety bound); -// listPageBatch is the per-request page size. +// maxListScan bounds how many runs `air list` inspects while looking for AIR runs +// that match the filters. runs/list returns runs of every kind, so this caps the +// work on a workspace with a large run history. +const maxListScan = 2000 + +// jobsPageLimit is the per-request page size for runs/list; enrichConcurrency +// bounds the parallel MLflow lookups. const ( - maxListPages = 50 - listPageBatch = 100 + jobsPageLimit = 25 + enrichConcurrency = 8 ) -// listData is the payload printed by `air list runs`. +// listData is the payload printed by `air list`. type listData struct { Rows []listRow `json:"runs"` } @@ -37,20 +43,28 @@ type listRow struct { IsSweep bool `json:"is_sweep"` // Experiment, Duration, MLflowURL and Accelerators are table-only columns, - // omitted from JSON to match `air list runs --json`. + // omitted from JSON to match `air list --json`. Experiment string `json:"-"` Duration string `json:"-"` MLflowURL string `json:"-"` Accelerators string `json:"-"` } +// listedRun pairs a row with its task run id, so the MLflow link can be fetched +// after the run has been filtered in. +type listedRun struct { + row listRow + taskRunID int64 +} + // listQuery holds the resolved inputs to listAirRuns. type listQuery struct { - activeOnly bool - allUsers bool - userFilter string - filters listFilters - limit int + activeOnly bool + allUsers bool + userFilter string + filters listFilters + limit int + fetchMLflow bool } func newListCommand() *cobra.Command { @@ -88,8 +102,8 @@ func newListCommand() *cobra.Command { } // An explicit user= filter wins; otherwise default to the current user - // unless --all-users is set. The resolved name is sent to the server as - // creator_name, which scopes the listing. + // unless --all-users is set. runs/list has no creator param, so the + // creator is matched while scanning. userFilter := f.User if userFilter == "" && !allUsers { me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) @@ -100,11 +114,12 @@ func newListCommand() *cobra.Command { } rows, err := listAirRuns(ctx, w, listQuery{ - activeOnly: active, - allUsers: allUsers, - userFilter: userFilter, - filters: f, - limit: limit, + activeOnly: active, + allUsers: allUsers, + userFilter: userFilter, + filters: f, + limit: limit, + fetchMLflow: root.OutputType(cmd) == flags.OutputText, }) if err != nil { return err @@ -121,50 +136,86 @@ func newListCommand() *cobra.Command { return cmd } -// listAirRuns pages through ListTrainingWorkflows, applies the client-side -// --filter keys, and accumulates matches up to q.limit (0 = all). It stops at the -// limit, when the server runs out of pages, or at maxListPages. +// listAirRuns pages through Jobs runs/list, keeps the AIR runs that match the +// user and filters, and stops once it has enough matches or has scanned +// maxListScan runs. Detail comes straight from runs/list (expand_tasks), so no +// per-run calls are needed. func listAirRuns(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) ([]listRow, error) { - // Fetch in bounded batches and follow the page token, so "all" doesn't ask - // the server for an unbounded page. With client-side filters we fetch full - // batches, since most rows may be dropped before reaching the limit. - pageSize := q.limit - if pageSize <= 0 || pageSize > listPageBatch || q.filters.hasClientFilters() { - pageSize = listPageBatch - } query := map[string]any{ - "active_only": q.activeOnly, - "include_all_users": q.allUsers, - "page_size": pageSize, + "run_type": "SUBMIT_RUN", + "expand_tasks": true, + "limit": jobsPageLimit, } - if q.userFilter != "" { - query["creator_name"] = q.userFilter + if q.activeOnly { + query["active_only"] = true } - var rows []listRow - for range maxListPages { - resp, err := listTrainingWorkflows(ctx, w, query) + var entries []listedRun + scanned := 0 + done := false + for !done && scanned < maxListScan { + resp, err := fetchJobRunsPage(ctx, w, query) if err != nil { return nil, err } - for i := range resp.TrainingWorkflows { - wf := &resp.TrainingWorkflows[i] - if !q.filters.matches(wf) { + for i := range resp.Runs { + run := &resp.Runs[i] + scanned++ + + if !isAirRun(run) { continue } - rows = append(rows, buildListRow(wf, w.Config.Host)) - if q.limit > 0 && len(rows) >= q.limit { - return rows, nil + if q.userFilter != "" && run.CreatorUserName != q.userFilter { + continue + } + if !q.filters.matches(run) { + continue + } + + entries = append(entries, listedRun{row: buildListRow(run), taskRunID: taskRunID(run)}) + if len(entries) >= q.limit { + done = true + break } } - if resp.NextPageToken == "" { - return rows, nil + if done || resp.NextPageToken == "" { + break } query["page_token"] = resp.NextPageToken } - log.Warnf(ctx, "air list: stopped after %d pages; results may be incomplete", maxListPages) + if !done && scanned >= maxListScan { + log.Warnf(ctx, "air list: stopped after scanning %d runs; results may be incomplete", maxListScan) + } + + // MLflow links appear only in the text table, so the per-run get-output + // lookups are skipped for JSON output (which omits the column anyway). + if q.fetchMLflow { + setMLflowLinks(ctx, w, entries) + } + + rows := make([]listRow, len(entries)) + for i, e := range entries { + rows[i] = e.row + } return rows, nil } + +// setMLflowLinks fills in each row's MLflow link in parallel, best-effort: a row +// whose IDs can't be resolved keeps its "-" placeholder. +func setMLflowLinks(ctx context.Context, w *databricks.WorkspaceClient, entries []listedRun) { + var g errgroup.Group + g.SetLimit(enrichConcurrency) + for i := range entries { + g.Go(func() error { + if ids := mlflowIDsForTask(ctx, w, entries[i].taskRunID); ids != nil { + entries[i].row.MLflowURL = mlflowLogsURL(w.Config.Host, ids) + } + return nil + }) + } + // mlflowIDsForTask never returns an error (it logs and yields nil), so Wait can't fail. + _ = g.Wait() +} diff --git a/experimental/air/cmd/list_filter.go b/experimental/air/cmd/list_filter.go index 44ffa114e70..6ac47c74b0e 100644 --- a/experimental/air/cmd/list_filter.go +++ b/experimental/air/cmd/list_filter.go @@ -50,34 +50,26 @@ func parseListFilters(raw []string) (listFilters, error) { return f, nil } -// hasClientFilters reports whether any client-side filter (those applied in -// matches) is set. The user filter is excluded — the server handles it. -func (f listFilters) hasClientFilters() bool { - return f.Experiment != "" || f.AcceleratorType != "" || f.NumAccelerators != nil -} - -// matches reports whether a workflow satisfies the experiment, accelerator-type -// and accelerator-count filters. These have no ListTrainingWorkflows equivalent, -// so they are applied client-side to the response. The user filter is handled by -// the server (via creator_name), so it is not re-checked here. -func (f listFilters) matches(w *trainingWorkflow) bool { +// matches reports whether a run satisfies the experiment, accelerator-type and +// accelerator-count filters. The user filter is applied separately while +// scanning, since it maps onto the run's creator rather than its task. +func (f listFilters) matches(run *jobRun) bool { if f.Experiment != "" { - name := stripExperimentUserPrefix(w.Status.Mlflow.Experiment) - matched, err := path.Match(strings.ToLower(f.Experiment), strings.ToLower(name)) + matched, err := path.Match(strings.ToLower(f.Experiment), strings.ToLower(jobExperiment(run))) if err != nil || !matched { return false } } if f.AcceleratorType != "" || f.NumAccelerators != nil { - compute := w.Spec.Compute + gpuType, count := jobCompute(run) if f.AcceleratorType != "" { - display := strings.ToLower(gpuDisplayName(compute.HardwareAcceleratorType)) + display := strings.ToLower(gpuDisplayName(gpuType)) if !strings.Contains(display, strings.ToLower(f.AcceleratorType)) { return false } } - if f.NumAccelerators != nil && compute.AcceleratorCount != *f.NumAccelerators { + if f.NumAccelerators != nil && count != *f.NumAccelerators { return false } } diff --git a/experimental/air/cmd/list_filter_test.go b/experimental/air/cmd/list_filter_test.go index a475e6576bd..f8351d33d8d 100644 --- a/experimental/air/cmd/list_filter_test.go +++ b/experimental/air/cmd/list_filter_test.go @@ -49,7 +49,7 @@ func TestParseListFilters(t *testing.T) { } func TestListFiltersMatches(t *testing.T) { - wf := testWorkflow(1, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/qwen-eval") + run := airJobRun(1, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/qwen-eval") cases := []struct { name string @@ -68,7 +68,7 @@ func TestListFiltersMatches(t *testing.T) { } for _, c := range cases { t.Run(c.name, func(t *testing.T) { - assert.Equal(t, c.want, c.f.matches(&wf)) + assert.Equal(t, c.want, c.f.matches(&run)) }) } } diff --git a/experimental/air/cmd/list_format.go b/experimental/air/cmd/list_format.go index 398a2523467..c728ca31945 100644 --- a/experimental/air/cmd/list_format.go +++ b/experimental/air/cmd/list_format.go @@ -1,54 +1,46 @@ package aircmd import ( + "strconv" "time" ) -// buildListRow maps a TrainingWorkflow to a table/JSON row. Optional text -// columns fall back to "-"; host builds the MLflow deep link. -func buildListRow(w *trainingWorkflow, host string) listRow { +// buildListRow extracts the columns shown for one run. Optional text columns fall +// back to "-" so the table stays aligned. MLflow links aren't carried by +// runs/list, so the column shows "-". +func buildListRow(run *jobRun) listRow { experiment := "-" - if e := stripExperimentUserPrefix(w.Status.Mlflow.Experiment); e != "" { + if e := jobExperiment(run); e != "" { experiment = e } var startedAt *string duration := "-" - if start := parseRPCTime(w.Status.StartTime); !start.IsZero() { - s := isoFormat(start) + if start, end := jobTiming(run); start > 0 { + s := isoFormat(time.UnixMilli(start)) startedAt = &s - - // A still-running run has no end_time, so measure against the current time. - endMillis := time.Now().UnixMilli() - if end := parseRPCTime(w.Status.EndTime); !end.IsZero() { - endMillis = end.UnixMilli() + if end == 0 { + // Still running: measure against the current time. + end = time.Now().UnixMilli() } - duration = formatDuration(roundMillisToSeconds(endMillis - start.UnixMilli())) - } - - mlflowURL := "-" - if w.Status.Mlflow.ExperimentID != "" && w.Status.Mlflow.RunID != "" { - mlflowURL = mlflowLogsURL(host, &mlflowIdentifiers{ - ExperimentID: w.Status.Mlflow.ExperimentID, - RunID: w.Status.Mlflow.RunID, - }) + duration = formatDuration(roundMillisToSeconds(end - start)) } accel := "-" - if a := acceleratorLabel(w.Spec.Compute.HardwareAcceleratorType, w.Spec.Compute.AcceleratorCount); a != "" { + if a := acceleratorLabel(jobCompute(run)); a != "" { accel = a } return listRow{ - RunID: w.JobRunID, - RunName: w.Status.Job.Name, - User: w.Metadata.CreatorName, - Status: trainingWorkflowStatus(w.Status.State), + RunID: strconv.FormatInt(run.RunID, 10), + RunName: run.RunName, + User: run.CreatorUserName, + Status: statusWord(run.State.LifeCycleState, run.State.ResultState), StartedAt: startedAt, - IsSweep: w.TaskRunID != "", + IsSweep: isSweep(run), Experiment: experiment, Duration: duration, - MLflowURL: mlflowURL, + MLflowURL: "-", Accelerators: accel, } } diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go index 989c05fc6da..a8858937658 100644 --- a/experimental/air/cmd/list_test.go +++ b/experimental/air/cmd/list_test.go @@ -14,36 +14,38 @@ import ( "github.com/stretchr/testify/require" ) -// testWorkflow builds a single-task AIR workflow for a given user and compute, -// as ListTrainingWorkflows would return it. -func testWorkflow(id int64, user, gpuType string, count int, experiment string) trainingWorkflow { - var w trainingWorkflow - w.JobRunID = strconv.FormatInt(id, 10) - w.Metadata.CreatorName = user - w.Spec.Compute.HardwareAcceleratorType = gpuType - w.Spec.Compute.AcceleratorCount = count - w.Status.State = "TRAINING_WORKFLOW_STATE_RUNNING" - w.Status.Job.Name = "run-" + strconv.FormatInt(id, 10) - w.Status.Mlflow.Experiment = experiment - return w +// airJobRun builds a single-task AIR run (ai_runtime_task), as runs/list returns +// it with expand_tasks. +func airJobRun(id int64, user, accelType string, count int, experiment string) jobRun { + return jobRun{ + RunID: id, + RunName: "run-" + strconv.FormatInt(id, 10), + CreatorUserName: user, + State: jobState{LifeCycleState: "RUNNING"}, + Tasks: []jobTask{{AiRuntimeTask: &jobAiRuntimeTask{ + Experiment: experiment, + Deployments: []aiRuntimeDeploy{{ + Compute: airCompute{AcceleratorType: accelType, AcceleratorCount: count}, + }}, + }}}, + } } -// workflowsBody marshals a ListTrainingWorkflows response page. -func workflowsBody(t *testing.T, nextToken string, wfs ...trainingWorkflow) string { +// runsListBody marshals one runs/list response page. +func runsListBody(t *testing.T, nextToken string, runs ...jobRun) string { t.Helper() - b, err := json.Marshal(trainingWorkflowsResponse{TrainingWorkflows: wfs, NextPageToken: nextToken}) + b, err := json.Marshal(jobsRunsListResponse{Runs: runs, NextPageToken: nextToken}) require.NoError(t, err) return string(b) } -// workflowsServer serves one response body per call to the workflows endpoint, -// repeating the last body once exhausted, and a stub for any other request (the -// SDK's well-known config discovery). -func workflowsServer(t *testing.T, bodies ...string) *httptest.Server { +// runsServer serves one runs/list response body per call, repeating the last +// once exhausted, and a stub for any other request (the SDK config probe). +func runsServer(t *testing.T, bodies ...string) *httptest.Server { t.Helper() call := 0 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == listWorkflowsPath { + if r.URL.Path == jobsRunsListPath { body := bodies[min(call, len(bodies)-1)] call++ _, _ = w.Write([]byte(body)) @@ -55,10 +57,31 @@ func workflowsServer(t *testing.T, bodies ...string) *httptest.Server { return srv } -func TestListAirRunsAppliesClientFilters(t *testing.T) { - qwen := testWorkflow(1, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/qwen-train") - llama := testWorkflow(2, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/llama-train") - srv := workflowsServer(t, workflowsBody(t, "", qwen, llama)) +func TestListAirRunsFiltersUserAndType(t *testing.T) { + runs := []jobRun{ + airJobRun(1, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/exp-a"), + {RunID: 2, CreatorUserName: "me@example.com", Tasks: []jobTask{{}}}, // not an AIR run + airJobRun(3, "other@example.com", "GPU_1xA10", 1, "/Users/other/exp-b"), // wrong user + airJobRun(5, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/exp-c"), + } + srv := runsServer(t, runsListBody(t, "", runs...)) + + rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", + limit: 10, + }) + require.NoError(t, err) + require.Len(t, rows, 2) + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "5", rows[1].RunID) +} + +func TestListAirRunsExperimentFilter(t *testing.T) { + runs := []jobRun{ + airJobRun(1, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/qwen-train"), + airJobRun(2, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/llama-train"), + } + srv := runsServer(t, runsListBody(t, "", runs...)) rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ limit: 10, @@ -70,12 +93,12 @@ func TestListAirRunsAppliesClientFilters(t *testing.T) { } func TestListAirRunsLimitTruncates(t *testing.T) { - wfs := []trainingWorkflow{ - testWorkflow(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), - testWorkflow(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), - testWorkflow(3, "me@example.com", "GPU_1xH100", 1, "exp-c"), + runs := []jobRun{ + airJobRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), + airJobRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), + airJobRun(3, "me@example.com", "GPU_1xH100", 1, "exp-c"), } - srv := workflowsServer(t, workflowsBody(t, "", wfs...)) + srv := runsServer(t, runsListBody(t, "", runs...)) rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{limit: 2}) require.NoError(t, err) @@ -84,23 +107,10 @@ func TestListAirRunsLimitTruncates(t *testing.T) { assert.Equal(t, "2", rows[1].RunID) } -func TestListAirRunsNoLimitFetchesAll(t *testing.T) { - // limit 0 means "all": follow the page token to the end with no early stop. - page1 := workflowsBody(t, "tok", testWorkflow(1, "me@example.com", "GPU_1xH100", 1, "exp-a")) - page2 := workflowsBody(t, "", testWorkflow(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) - srv := workflowsServer(t, page1, page2) - - rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{limit: 0}) - require.NoError(t, err) - require.Len(t, rows, 2) -} - func TestListAirRunsPaginates(t *testing.T) { - // First page returns one run and a continuation token; the loop must follow it - // and stop when the token is empty. - page1 := workflowsBody(t, "tok", testWorkflow(1, "me@example.com", "GPU_1xH100", 1, "exp-a")) - page2 := workflowsBody(t, "", testWorkflow(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) - srv := workflowsServer(t, page1, page2) + page1 := runsListBody(t, "tok", airJobRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a")) + page2 := runsListBody(t, "", airJobRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) + srv := runsServer(t, page1, page2) rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{limit: 10}) require.NoError(t, err) @@ -109,38 +119,64 @@ func TestListAirRunsPaginates(t *testing.T) { assert.Equal(t, "2", rows[1].RunID) } +// TestFetchJobRunsParsesAiRuntimeTask pins the raw parse against the real +// runs/get shape, since the typed SDK omits ai_runtime_task. +func TestFetchJobRunsParsesAiRuntimeTask(t *testing.T) { + body := `{"runs":[{ + "run_id": 842552489592352, + "run_name": "my-first-air-run", + "creator_user_name": "me@example.com", + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{ + "ai_runtime_task": { + "experiment": "my-first-air-run", + "deployments": [{"compute": {"accelerator_count": 1, "accelerator_type": "GPU_1xA10"}}] + } + }] + }]}` + srv := runsServer(t, body) + + resp, err := fetchJobRunsPage(t.Context(), newTestWorkspaceClient(t, srv.URL), map[string]any{}) + require.NoError(t, err) + require.Len(t, resp.Runs, 1) + run := &resp.Runs[0] + assert.True(t, isAirRun(run)) + assert.Equal(t, "my-first-air-run", jobExperiment(run)) + gpu, count := jobCompute(run) + assert.Equal(t, "GPU_1xA10", gpu) + assert.Equal(t, 1, count) + + row := buildListRow(run) + assert.Equal(t, "842552489592352", row.RunID) + assert.Equal(t, "SUCCESS", row.Status) + assert.Equal(t, "my-first-air-run", row.Experiment) + assert.Equal(t, "1x A10", row.Accelerators) + assert.Equal(t, "12s", row.Duration) +} + func TestBuildListRow(t *testing.T) { - var w trainingWorkflow - w.JobRunID = "123" - w.Metadata.CreatorName = "me@example.com" - w.Spec.Compute.HardwareAcceleratorType = "GPU_8xH100" - w.Spec.Compute.AcceleratorCount = 8 - w.Status.State = "TRAINING_WORKFLOW_STATE_TERMINATED_COMPLETED" - w.Status.StartTime = "2026-06-05T17:32:39.000Z" - w.Status.EndTime = "2026-06-05T17:32:51.000Z" - w.Status.Job.Name = "my-run" - w.Status.Mlflow.Experiment = "/Users/me@example.com/exp" - w.Status.Mlflow.ExperimentID = "E1" - w.Status.Mlflow.RunID = "R1" - - row := buildListRow(&w, "https://h.test") + run := airJobRun(123, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/exp") + run.StartTime = 1700000000000 + run.EndTime = 1700000012000 + run.State = jobState{ResultState: "SUCCESS"} + + row := buildListRow(&run) assert.Equal(t, "123", row.RunID) - assert.Equal(t, "my-run", row.RunName) assert.Equal(t, "me@example.com", row.User) assert.Equal(t, "SUCCESS", row.Status) assert.Equal(t, "exp", row.Experiment) assert.Equal(t, "12s", row.Duration) assert.Equal(t, "8x H100", row.Accelerators) - assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1/artifacts/logs/node_0", row.MLflowURL) + assert.Equal(t, "-", row.MLflowURL) assert.False(t, row.IsSweep) require.NotNil(t, row.StartedAt) - assert.Equal(t, "2026-06-05T17:32:39+00:00", *row.StartedAt) } func TestBuildListRowDashFallbacks(t *testing.T) { - // A workflow with no experiment, compute, MLflow IDs, or start time falls back - // to dashes for the optional columns and UNKNOWN for the unset state. - row := buildListRow(&trainingWorkflow{}, "https://h.test") + // A run with no task, compute, or start time falls back to dashes and UNKNOWN. + row := buildListRow(&jobRun{RunID: 7}) assert.Equal(t, "-", row.Experiment) assert.Equal(t, "-", row.Duration) assert.Equal(t, "-", row.Accelerators) @@ -150,9 +186,11 @@ func TestBuildListRowDashFallbacks(t *testing.T) { } func TestBuildListRowSweep(t *testing.T) { - // task_run_id is set only for sweeps, so its presence marks the row. - w := trainingWorkflow{TaskRunID: "456"} - assert.True(t, buildListRow(&w, "https://h.test").IsSweep) + run := jobRun{RunID: 9, Tasks: []jobTask{{ + ForEachTask: &forEachTask{Task: jobTask{AiRuntimeTask: &jobAiRuntimeTask{Experiment: "sweep"}}}, + }}} + assert.True(t, buildListRow(&run).IsSweep) + assert.Equal(t, "sweep", buildListRow(&run).Experiment) } func TestListInvalidLimit(t *testing.T) { diff --git a/experimental/air/cmd/mlflow.go b/experimental/air/cmd/mlflow.go index 6eaafabfcb9..10609e97d70 100644 --- a/experimental/air/cmd/mlflow.go +++ b/experimental/air/cmd/mlflow.go @@ -13,10 +13,14 @@ import ( ) // getRunOutputResponse is the slice of the jobs runs/get-output response we care -// about. The MLflow identifiers live under a gen_ai_compute_output field that -// the typed SDK does not model, so we call the endpoint directly and parse just -// these fields. +// about. The MLflow identifiers live under ai_runtime_task_output (current) or +// gen_ai_compute_output.run_info (legacy), neither modeled by the typed SDK, so +// we call the endpoint directly and parse just these fields. type getRunOutputResponse struct { + AiRuntimeTaskOutput *struct { + MlflowExperimentID string `json:"mlflow_experiment_id"` + MlflowRunID string `json:"mlflow_run_id"` + } `json:"ai_runtime_task_output"` GenAiComputeOutput *struct { RunInfo *struct { MlflowExperimentID string `json:"mlflow_experiment_id"` @@ -31,20 +35,28 @@ type mlflowIdentifiers struct { RunID string } -// mlflowIDs fetches the run's MLflow experiment and run IDs, or nil if they -// can't be obtained. They drive a convenience link, so any failure here -// (missing task, endpoint error, run not yet started) is logged and treated as -// "no link" rather than failing the whole command. +// 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 { 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 + return mlflowIDsForTask(ctx, w, run.Tasks[len(run.Tasks)-1].RunId) +} + +// mlflowIDsForTask fetches a task run's MLflow experiment and run IDs from +// runs/get-output, or nil if they can't be obtained. They drive a convenience +// 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 { + if taskRunID == 0 { + return nil + } apiClient, err := client.New(w.Config) if err != nil { - log.Debugf(ctx, "air get: could not build API client for MLflow link: %v", err) + log.Debugf(ctx, "air: could not build API client for MLflow link: %v", err) return nil } @@ -55,18 +67,18 @@ func mlflowIDs(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run err = apiClient.Do(ctx, http.MethodGet, "/api/2.2/jobs/runs/get-output", nil, nil, map[string]any{"run_id": taskRunID}, &out) if err != nil { - log.Debugf(ctx, "air get: could not fetch run output for MLflow link: %v", err) + log.Debugf(ctx, "air: could not fetch run output for MLflow link: %v", err) return nil } - if out.GenAiComputeOutput == nil || out.GenAiComputeOutput.RunInfo == nil { - return nil + if o := out.AiRuntimeTaskOutput; o != nil && o.MlflowExperimentID != "" && o.MlflowRunID != "" { + return &mlflowIdentifiers{ExperimentID: o.MlflowExperimentID, RunID: o.MlflowRunID} } - info := out.GenAiComputeOutput.RunInfo - if info.MlflowExperimentID == "" || info.MlflowRunID == "" { - return nil + if o := out.GenAiComputeOutput; o != nil && o.RunInfo != nil && + o.RunInfo.MlflowExperimentID != "" && o.RunInfo.MlflowRunID != "" { + return &mlflowIdentifiers{ExperimentID: o.RunInfo.MlflowExperimentID, RunID: o.RunInfo.MlflowRunID} } - return &mlflowIdentifiers{ExperimentID: info.MlflowExperimentID, RunID: info.MlflowRunID} + return nil } // 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 505da11cacd..06529d5273f 100644 --- a/experimental/air/cmd/mlflow_test.go +++ b/experimental/air/cmd/mlflow_test.go @@ -79,6 +79,32 @@ func TestMLflowIDs(t *testing.T) { }) } +func TestMLflowIDsForTask(t *testing.T) { + ctx := t.Context() + + t.Run("parses ai_runtime_task_output", func(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{"ai_runtime_task_output":{"mlflow_experiment_id":"E1","mlflow_run_id":"R1"}}`, &hit) + got := mlflowIDsForTask(ctx, newTestWorkspaceClient(t, srv.URL), 99) + require.NotNil(t, got) + assert.True(t, hit) + assert.Equal(t, &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"}, got) + }) + + t.Run("parses legacy gen_ai_compute_output", func(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{"gen_ai_compute_output":{"run_info":{"mlflow_experiment_id":"E2","mlflow_run_id":"R2"}}}`, &hit) + got := mlflowIDsForTask(ctx, newTestWorkspaceClient(t, srv.URL), 99) + require.NotNil(t, got) + assert.Equal(t, &mlflowIdentifiers{ExperimentID: "E2", RunID: "R2"}, got) + }) + + t.Run("nil when no task run id", func(t *testing.T) { + // Returns before any HTTP call, so the host is never contacted. + assert.Nil(t, mlflowIDsForTask(ctx, newTestWorkspaceClient(t, "https://unused.invalid"), 0)) + }) +} + func TestMLflowURLs(t *testing.T) { ids := &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"} // A trailing slash on the host must not produce a double slash in the link. From f9f70e29add95689f494e2f0339a3f411512d5c1 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Mon, 6 Jul 2026 10:07:26 -0700 Subject: [PATCH 10/18] experimental/air: lazily page older runs in `air list` (#5811) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The interactive `air list` table only ever held the newest `--limit` (20) runs, so its ←/→ paging just scrolled that fixed window — older runs (e.g. anything before a recent date) were never fetched from the server and could never be reached. ## Fix Replace the one-shot `listAirRuns` with a stateful `runFetcher` that pages Jobs `runs/list` on demand: - It buffers a page's leftover runs and keeps its page-token cursor, so successive `next()` calls resume where the last stopped. - The interactive table (`list_tui.go`) holds the fetcher and calls `next(listPageRows)` in the background as the cursor nears the end of the loaded rows (`maybeFetch` → `fetchCmd` → `moreRowsMsg`). New rows and their MLflow links are appended and column widths recomputed; only one fetch runs at a time. - The hint line shows `row N/M` plus `(loading…)` / `(load failed)`. - `maxListScan` (2000) remains the safety ceiling on total runs scanned. One-shot output paths (JSON, piped, and explicit `--limit`) are unchanged: they call `fetcher.next(limit)` once, so acceptance output is identical. ## Tests - Updated the existing `list`/model tests to the new `runFetcher`/`newListModel` signatures. - Added `TestRunFetcherResumesAcrossCalls`: a `next()` that stops mid-page buffers the rest, hands it back on the next call, then reports exhaustion — without re-fetching. - `air list` / `air get` acceptance tests unchanged and green; build + vet clean. This pull request and its description were written by Isaac. --- experimental/air/cmd/list.go | 139 +++++++++++++++++--------- experimental/air/cmd/list_format.go | 5 +- experimental/air/cmd/list_test.go | 43 ++++++-- experimental/air/cmd/list_tui.go | 120 +++++++++++++++++----- experimental/air/cmd/list_tui_test.go | 52 +++++++++- experimental/air/cmd/mlflow.go | 4 +- 6 files changed, 271 insertions(+), 92 deletions(-) diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 360fc97e678..8c716f5556a 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -57,13 +57,11 @@ type listedRun struct { taskRunID int64 } -// listQuery holds the resolved inputs to listAirRuns. +// listQuery holds the resolved inputs to a runFetcher. type listQuery struct { activeOnly bool - allUsers bool userFilter string filters listFilters - limit int fetchMLflow bool } @@ -113,34 +111,47 @@ func newListCommand() *cobra.Command { userFilter = me.UserName } - rows, err := listAirRuns(ctx, w, listQuery{ + fetcher := newRunFetcher(ctx, w, listQuery{ activeOnly: active, - allUsers: allUsers, userFilter: userFilter, filters: f, - limit: limit, fetchMLflow: root.OutputType(cmd) == flags.OutputText, }) - if err != nil { - return err - } - // JSON keeps the air envelope; text renders the table (interactive and - // navigable in a terminal, printed once when piped). + // JSON prints the newest `limit` runs once. Text renders the table: + // navigable in a terminal (paging in older runs on demand), printed once + // when piped. if root.OutputType(cmd) == flags.OutputJSON { + rows, err := fetcher.next(limit) + if err != nil { + return err + } + warnIfTruncated(ctx, fetcher) return renderEnvelope(ctx, listData{Rows: rows}) } - return renderListText(cmd, rows) + return renderListText(cmd, fetcher, limit) } return cmd } -// listAirRuns pages through Jobs runs/list, keeps the AIR runs that match the -// user and filters, and stops once it has enough matches or has scanned -// maxListScan runs. Detail comes straight from runs/list (expand_tasks), so no -// per-run calls are needed. -func listAirRuns(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) ([]listRow, error) { +// runFetcher pages Jobs runs/list on demand, yielding AIR runs that match the +// user and filters. It buffers a page's leftover runs so successive next() calls +// resume where the last stopped — driving both one-shot output and lazy paging. +type runFetcher struct { + ctx context.Context + w *databricks.WorkspaceClient + query map[string]any + userFilter string + filters listFilters + fetchMLflow bool + + pending []jobRun // runs from the last page not yet inspected + scanned int + exhausted bool +} + +func newRunFetcher(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *runFetcher { query := map[string]any{ "run_type": "SUBMIT_RUN", "expand_tasks": true, @@ -149,51 +160,55 @@ func listAirRuns(ctx context.Context, w *databricks.WorkspaceClient, q listQuery if q.activeOnly { query["active_only"] = true } + return &runFetcher{ + ctx: ctx, + w: w, + query: query, + userFilter: q.userFilter, + filters: q.filters, + fetchMLflow: q.fetchMLflow, + } +} +// next returns up to want more matching rows, paging runs/list (and buffering the +// leftover runs of a page) until it has enough, the server has no more pages, or +// it has scanned maxListScan runs. MLflow links are filled in for text output. +func (f *runFetcher) next(want int) ([]listRow, error) { var entries []listedRun - scanned := 0 - done := false - for !done && scanned < maxListScan { - resp, err := fetchJobRunsPage(ctx, w, query) - if err != nil { - return nil, err - } - - for i := range resp.Runs { - run := &resp.Runs[i] - scanned++ - - if !isAirRun(run) { - continue - } - if q.userFilter != "" && run.CreatorUserName != q.userFilter { - continue - } - if !q.filters.matches(run) { - continue - } - - entries = append(entries, listedRun{row: buildListRow(run), taskRunID: taskRunID(run)}) - if len(entries) >= q.limit { - done = true + for len(entries) < want { + if len(f.pending) == 0 { + if f.exhausted || f.scanned >= maxListScan { break } + if err := f.fetchPage(); err != nil { + return nil, err + } + continue } - if done || resp.NextPageToken == "" { - break + run := &f.pending[0] + f.pending = f.pending[1:] + f.scanned++ + + if !isAirRun(run) { + continue } - query["page_token"] = resp.NextPageToken + if f.userFilter != "" && run.CreatorUserName != f.userFilter { + continue + } + if !f.filters.matches(run) { + continue + } + entries = append(entries, listedRun{row: buildListRow(run), taskRunID: taskRunID(run)}) } - - if !done && scanned >= maxListScan { - log.Warnf(ctx, "air list: stopped after scanning %d runs; results may be incomplete", maxListScan) + if f.scanned >= maxListScan { + f.exhausted = true } // MLflow links appear only in the text table, so the per-run get-output // lookups are skipped for JSON output (which omits the column anyway). - if q.fetchMLflow { - setMLflowLinks(ctx, w, entries) + if f.fetchMLflow { + setMLflowLinks(f.ctx, f.w, entries) } rows := make([]listRow, len(entries)) @@ -203,6 +218,30 @@ func listAirRuns(ctx context.Context, w *databricks.WorkspaceClient, q listQuery return rows, nil } +// fetchPage loads the next runs/list page into the pending buffer, marking the +// fetcher exhausted once the server reports no further pages. +func (f *runFetcher) fetchPage() error { + resp, err := fetchJobRunsPage(f.ctx, f.w, f.query) + if err != nil { + return err + } + f.pending = resp.Runs + if resp.NextPageToken == "" { + f.exhausted = true + } else { + f.query["page_token"] = resp.NextPageToken + } + return nil +} + +// warnIfTruncated logs when the scan hit maxListScan, so one-shot output signals +// its results may be incomplete. +func warnIfTruncated(ctx context.Context, f *runFetcher) { + if f.scanned >= maxListScan { + log.Warnf(ctx, "air list: stopped after scanning %d runs; results may be incomplete", maxListScan) + } +} + // setMLflowLinks fills in each row's MLflow link in parallel, best-effort: a row // whose IDs can't be resolved keeps its "-" placeholder. func setMLflowLinks(ctx context.Context, w *databricks.WorkspaceClient, entries []listedRun) { diff --git a/experimental/air/cmd/list_format.go b/experimental/air/cmd/list_format.go index c728ca31945..5dcb21bfc62 100644 --- a/experimental/air/cmd/list_format.go +++ b/experimental/air/cmd/list_format.go @@ -5,9 +5,8 @@ import ( "time" ) -// buildListRow extracts the columns shown for one run. Optional text columns fall -// back to "-" so the table stays aligned. MLflow links aren't carried by -// runs/list, so the column shows "-". +// buildListRow extracts the columns shown for one run. Optional cells fall back +// to "-"; MLflowURL starts as "-" and setMLflowLinks fills it in for text output. func buildListRow(run *jobRun) listRow { experiment := "-" if e := jobExperiment(run); e != "" { diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go index a8858937658..11185528a2c 100644 --- a/experimental/air/cmd/list_test.go +++ b/experimental/air/cmd/list_test.go @@ -66,10 +66,9 @@ func TestListAirRunsFiltersUserAndType(t *testing.T) { } srv := runsServer(t, runsListBody(t, "", runs...)) - rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ userFilter: "me@example.com", - limit: 10, - }) + }).next(10) require.NoError(t, err) require.Len(t, rows, 2) assert.Equal(t, "1", rows[0].RunID) @@ -83,10 +82,9 @@ func TestListAirRunsExperimentFilter(t *testing.T) { } srv := runsServer(t, runsListBody(t, "", runs...)) - rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - limit: 10, + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ filters: listFilters{Experiment: "qwen*"}, - }) + }).next(10) require.NoError(t, err) require.Len(t, rows, 1) assert.Equal(t, "1", rows[0].RunID) @@ -100,7 +98,7 @@ func TestListAirRunsLimitTruncates(t *testing.T) { } srv := runsServer(t, runsListBody(t, "", runs...)) - rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{limit: 2}) + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{}).next(2) require.NoError(t, err) require.Len(t, rows, 2) assert.Equal(t, "1", rows[0].RunID) @@ -112,13 +110,42 @@ func TestListAirRunsPaginates(t *testing.T) { page2 := runsListBody(t, "", airJobRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) srv := runsServer(t, page1, page2) - rows, err := listAirRuns(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{limit: 10}) + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{}).next(10) require.NoError(t, err) require.Len(t, rows, 2) assert.Equal(t, "1", rows[0].RunID) assert.Equal(t, "2", rows[1].RunID) } +// TestRunFetcherResumesAcrossCalls covers the lazy paging the interactive table +// relies on: a next() that stops mid-page must buffer the rest and hand it back on +// the following call, then report exhaustion — without re-fetching. +func TestRunFetcherResumesAcrossCalls(t *testing.T) { + runs := []jobRun{ + airJobRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), + airJobRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), + airJobRun(3, "me@example.com", "GPU_1xH100", 1, "exp-c"), + } + srv := runsServer(t, runsListBody(t, "", runs...)) + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{}) + + first, err := f.next(2) + require.NoError(t, err) + require.Len(t, first, 2) + assert.Equal(t, "1", first[0].RunID) + assert.Equal(t, "2", first[1].RunID) + + second, err := f.next(2) + require.NoError(t, err) + require.Len(t, second, 1) // only the buffered leftover remains + assert.Equal(t, "3", second[0].RunID) + assert.True(t, f.exhausted) + + third, err := f.next(2) + require.NoError(t, err) + assert.Empty(t, third) +} + // TestFetchJobRunsParsesAiRuntimeTask pins the raw parse against the real // runs/get shape, since the typed SDK omits ai_runtime_task. func TestFetchJobRunsParsesAiRuntimeTask(t *testing.T) { diff --git a/experimental/air/cmd/list_tui.go b/experimental/air/cmd/list_tui.go index 369205f5992..88e405c867a 100644 --- a/experimental/air/cmd/list_tui.go +++ b/experimental/air/cmd/list_tui.go @@ -13,9 +13,10 @@ import ( "github.com/spf13/cobra" ) -// renderListText renders the table for text output: an inline navigable table -// in a terminal, otherwise printed once. JSON is handled by the caller. -func renderListText(cmd *cobra.Command, rows []listRow) error { +// renderListText renders the table for text output: an inline navigable table in +// a terminal (paging in older runs on demand), otherwise printed once. JSON is +// handled by the caller. +func renderListText(cmd *cobra.Command, f *runFetcher, limit int) error { ctx := cmd.Context() out := cmd.OutOrStdout() @@ -25,17 +26,24 @@ func renderListText(cmd *cobra.Command, rows []listRow) error { r.SetColorProfile(termenv.Ascii) } - // Navigate only with a full color TTY, at least one row, and no explicit - // --limit (which means "just print these N"). Everything else — piped, - // NO_COLOR, --limit, empty — prints once. - interactive := len(rows) > 0 && - color && + // Navigate only with a full color TTY and no explicit --limit (which means + // "just print these N"). Everything else — piped, NO_COLOR, --limit — prints + // once. + interactive := color && cmdio.IsPagerSupported(ctx) && !cmd.Flags().Changed("limit") if interactive { - _, err := tea.NewProgram( - newListModel(r, rows, color), + first, err := f.next(listPageRows) + if err != nil { + return err + } + if len(first) == 0 { + _, err := io.WriteString(out, "No runs found.\n") + return err + } + _, err = tea.NewProgram( + newListModel(r, f, first, color), tea.WithContext(ctx), tea.WithInput(cmd.InOrStdin()), tea.WithOutput(out), @@ -43,7 +51,12 @@ func renderListText(cmd *cobra.Command, rows []listRow) error { return err } - _, err := io.WriteString(out, staticListTable(r, rows, color)) + rows, err := f.next(limit) + if err != nil { + return err + } + warnIfTruncated(ctx, f) + _, err = io.WriteString(out, staticListTable(r, rows, color)) return err } @@ -65,29 +78,64 @@ func staticListTable(r *lipgloss.Renderer, rows []listRow, links bool) string { return b.String() } -// listModel is the inline, navigable runs table. +// listModel is the inline, navigable runs table. It lazily pages older runs from +// the fetcher as the cursor nears the end of the loaded rows. fetcher is nil for +// a fixed, non-paging table (e.g. in tests). type listModel struct { - rows []listRow - styles listStyles - cols listCols - links bool + rows []listRow + styles listStyles + cols listCols + links bool + fetcher *runFetcher + loading bool + loadErr error cursor int offset int // index of the first visible row height int // terminal height, for windowing } -func newListModel(r *lipgloss.Renderer, rows []listRow, links bool) listModel { +func newListModel(r *lipgloss.Renderer, f *runFetcher, rows []listRow, links bool) listModel { return listModel{ - rows: rows, - styles: newListStyles(r), - cols: computeListCols(rows), - links: links, + rows: rows, + styles: newListStyles(r), + cols: computeListCols(rows), + links: links, + fetcher: f, } } func (m listModel) Init() tea.Cmd { return nil } +// moreRowsMsg carries a lazily fetched batch of rows, or the error that ended paging. +type moreRowsMsg struct { + rows []listRow + err error +} + +// fetchCmd pulls the next batch of rows in the background; sets loading so only +// one runs at a time, and returns the updated model with the fetch command. +func (m listModel) fetchCmd() (listModel, tea.Cmd) { + m.loading = true + f := m.fetcher + return m, func() tea.Msg { + rows, err := f.next(listPageRows) + return moreRowsMsg{rows: rows, err: err} + } +} + +// maybeFetch starts a fetch when the cursor nears the end of the loaded rows and +// more runs may still exist. +func (m listModel) maybeFetch() (listModel, tea.Cmd) { + if m.fetcher == nil || m.loading || m.loadErr != nil || m.fetcher.exhausted { + return m, nil + } + if m.cursor < len(m.rows)-m.visibleCount() { + return m, nil + } + return m.fetchCmd() +} + // listPageRows is the most rows shown per page. const listPageRows = 20 @@ -106,6 +154,22 @@ func (m listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case tea.WindowSizeMsg: m.height = msg.Height m.offset = m.clampedOffset() + return m.maybeFetch() + + case moreRowsMsg: + m.loading = false + if msg.err != nil { + m.loadErr = msg.err + return m, nil + } + m.rows = append(m.rows, msg.rows...) + m.cols = computeListCols(m.rows) + m.offset = m.clampedOffset() + // A page with no matches but more to scan: keep paging so the cursor isn't + // stuck at the end of the loaded rows. + if len(msg.rows) == 0 && !m.fetcher.exhausted { + return m.fetchCmd() + } return m, nil case tea.KeyMsg: @@ -137,6 +201,7 @@ func (m listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } } m.offset = m.clampedOffset() + return m.maybeFetch() } return m, nil } @@ -165,13 +230,16 @@ func (m listModel) View() string { return strings.Join(lines, "\n") + "\n" } -// renderHint is the faint one-line key legend, with a scroll position when the -// list is windowed. +// renderHint is the faint one-line key legend, with the cursor position and the +// paging state (loading / load failed). func (m listModel) renderHint() string { faint := m.styles.r.NewStyle().Foreground(colN7) - hint := "↑/↓ navigate · ←/→ page · ↵ mlflow · q quit" - if m.visibleCount() < len(m.rows) { - hint += fmt.Sprintf(" · row %d/%d", m.cursor+1, len(m.rows)) + hint := fmt.Sprintf("↑/↓ navigate · ←/→ page · ↵ mlflow · q quit · row %d/%d", m.cursor+1, len(m.rows)) + switch { + case m.loadErr != nil: + hint += " (load failed)" + case m.loading: + hint += " (loading…)" } return faint.Render(hint) } diff --git a/experimental/air/cmd/list_tui_test.go b/experimental/air/cmd/list_tui_test.go index b5b7ca2cbc8..f1cff1e59f0 100644 --- a/experimental/air/cmd/list_tui_test.go +++ b/experimental/air/cmd/list_tui_test.go @@ -25,7 +25,7 @@ func testListRows() []listRow { func testListModel() listModel { r := lipgloss.NewRenderer(io.Discard) r.SetColorProfile(termenv.Ascii) - return newListModel(r, testListRows(), false) + return newListModel(r, nil, testListRows(), false) } func key(t *testing.T, m listModel, s string) listModel { @@ -69,7 +69,7 @@ func TestListModelPageCap(t *testing.T) { } r := lipgloss.NewRenderer(io.Discard) r.SetColorProfile(termenv.Ascii) - m := newListModel(r, rows, false) + m := newListModel(r, nil, rows, false) // A tall terminal still shows at most listPageRows per page. next, _ := m.Update(tea.WindowSizeMsg{Width: 80, Height: 100}) @@ -83,7 +83,7 @@ func TestListModelPaging(t *testing.T) { } r := lipgloss.NewRenderer(io.Discard) r.SetColorProfile(termenv.Ascii) - m := newListModel(r, rows, false) + m := newListModel(r, nil, rows, false) // Height 7 leaves a 4-row window (header + hint reserved). next, _ := m.Update(tea.WindowSizeMsg{Width: 80, Height: 7}) @@ -104,6 +104,52 @@ func TestListModelPaging(t *testing.T) { assert.Equal(t, 0, m.cursor) } +func TestListModelMoreRows(t *testing.T) { + m := testListModel() + m.loading = true + before := len(m.rows) + + next, cmd := m.Update(moreRowsMsg{rows: []listRow{{RunID: "4"}, {RunID: "5"}}}) + m = next.(listModel) + + assert.False(t, m.loading, "loading cleared after a batch arrives") + assert.NoError(t, m.loadErr) + require.Len(t, m.rows, before+2) + assert.Equal(t, "5", m.rows[len(m.rows)-1].RunID, "new rows appended") + assert.Nil(t, cmd) +} + +func TestListModelMoreRowsError(t *testing.T) { + m := testListModel() + m.loading = true + before := len(m.rows) + + next, cmd := m.Update(moreRowsMsg{err: io.ErrUnexpectedEOF}) + m = next.(listModel) + + assert.False(t, m.loading) + assert.ErrorIs(t, m.loadErr, io.ErrUnexpectedEOF) + assert.Len(t, m.rows, before, "rows unchanged on error") + assert.Nil(t, cmd) +} + +func TestListModelMoreRowsEmptyKeepsPaging(t *testing.T) { + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + + // An empty page while more runs remain re-fetches; once exhausted it stops. + m := newListModel(r, &runFetcher{}, testListRows(), false) + m.loading = true + next, cmd := m.Update(moreRowsMsg{}) + m = next.(listModel) + assert.NotNil(t, cmd, "empty page with more to scan keeps paging") + + m.fetcher.exhausted = true + m.loading = true + _, cmd = m.Update(moreRowsMsg{}) + assert.Nil(t, cmd, "empty page stops once the fetcher is exhausted") +} + func TestListModelQuit(t *testing.T) { m := testListModel() _, cmd := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("q")}) diff --git a/experimental/air/cmd/mlflow.go b/experimental/air/cmd/mlflow.go index 10609e97d70..02a7956d7bd 100644 --- a/experimental/air/cmd/mlflow.go +++ b/experimental/air/cmd/mlflow.go @@ -12,8 +12,8 @@ import ( "github.com/databricks/databricks-sdk-go/service/jobs" ) -// getRunOutputResponse is the slice of the jobs runs/get-output response we care -// about. The MLflow identifiers live under ai_runtime_task_output (current) or +// getRunOutputResponse is the part of the jobs runs/get-output response we use. +// The MLflow identifiers live under ai_runtime_task_output (current) or // gen_ai_compute_output.run_info (legacy), neither modeled by the typed SDK, so // we call the endpoint directly and parse just these fields. type getRunOutputResponse struct { From 18fed34bb8755b8303b459ed3bb3e83c7edc37d4 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Mon, 6 Jul 2026 15:17:02 -0700 Subject: [PATCH 11/18] AIR CLI Integration: implement the `air cancel` command (#5812) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Ports the Python `air` CLI's `handle_cancel` to the Go `air cancel` command (stacked on `air-integration-m1-2`). - Cancel one or more runs by ID, or all of the current user's active runs with `--all`. - `--all` resolves the current user, lists their active runs via the Jobs-API run fetcher, shows a preview table, and prompts for confirmation unless `-y` is set (fetches up to `maxListScan` so every active run is cancelled). - Cancellation goes through the typed SDK `w.Jobs.CancelRun`. - Not-found is detected via `errors.Is(apierr.ErrResourceDoesNotExist)` and, for the cancel endpoint's `400 INVALID_PARAMETER_VALUE` shape, a typed `errors.As` check — restoring the Python CLI's friendly "Run X not found" guidance. - Text and JSON (air envelope) output; exits non-zero on any failure. ## Tests - Unit (`cancel_test.go`): arg validation, by-ID success/not-found (both error shapes), partial-failure JSON exit code, `--all` (no runs / confirm / abort / read error), current-user + list errors, and the preview table. - Acceptance (`acceptance/experimental/air/cancel`): by-ID text+JSON, multiple IDs, and `--all -y`. - Drops the now-implemented `cancel` case from the stub unit test and the `unimplemented` acceptance test. This pull request and its description were written by Isaac. --- .../experimental/air/cancel/out.test.toml | 3 + acceptance/experimental/air/cancel/output.txt | 29 ++ acceptance/experimental/air/cancel/script | 11 + acceptance/experimental/air/cancel/test.toml | 53 ++++ .../experimental/air/unimplemented/output.txt | 6 - .../experimental/air/unimplemented/script | 3 - experimental/air/cmd/cancel.go | 189 +++++++++++- experimental/air/cmd/cancel_test.go | 289 ++++++++++++++++++ experimental/air/cmd/stubs_test.go | 1 - 9 files changed, 571 insertions(+), 13 deletions(-) create mode 100644 acceptance/experimental/air/cancel/out.test.toml create mode 100644 acceptance/experimental/air/cancel/output.txt create mode 100644 acceptance/experimental/air/cancel/script create mode 100644 acceptance/experimental/air/cancel/test.toml create mode 100644 experimental/air/cmd/cancel_test.go diff --git a/acceptance/experimental/air/cancel/out.test.toml b/acceptance/experimental/air/cancel/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/cancel/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/cancel/output.txt b/acceptance/experimental/air/cancel/output.txt new file mode 100644 index 00000000000..9fd8a055f13 --- /dev/null +++ b/acceptance/experimental/air/cancel/output.txt @@ -0,0 +1,29 @@ + +=== cancel by id (text) +>>> [CLI] experimental air cancel 123 +Successfully requested cancellation for run 123 + +=== cancel by id (json) +>>> [CLI] experimental air cancel 123 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "cancelled": [ + "123" + ] + } +} + +=== cancel multiple ids +>>> [CLI] experimental air cancel 123 456 +Successfully requested cancellation for run 123 +Successfully requested cancellation for run 456 +Successfully requested cancellation for 2 run(s). + +=== cancel --all +>>> [CLI] experimental air cancel --all -y +Searching active runs for [USERNAME] in [DATABRICKS_URL]... +Successfully requested cancellation for run [NUMID] +Successfully requested cancellation for run [NUMID] +Successfully requested cancellation for 2 run(s). diff --git a/acceptance/experimental/air/cancel/script b/acceptance/experimental/air/cancel/script new file mode 100644 index 00000000000..ce04a8977fd --- /dev/null +++ b/acceptance/experimental/air/cancel/script @@ -0,0 +1,11 @@ +title "cancel by id (text)" +trace $CLI experimental air cancel 123 + +title "cancel by id (json)" +trace $CLI experimental air cancel 123 -o json + +title "cancel multiple ids" +trace $CLI experimental air cancel 123 456 + +title "cancel --all" +trace $CLI experimental air cancel --all -y diff --git a/acceptance/experimental/air/cancel/test.toml b/acceptance/experimental/air/cancel/test.toml new file mode 100644 index 00000000000..c73594501d8 --- /dev/null +++ b/acceptance/experimental/air/cancel/test.toml @@ -0,0 +1,53 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# CancelRun accepts the request and returns an empty body. +[[Server]] +Pattern = "POST /api/2.2/jobs/runs/cancel" +Response.Body = '{}' + +# Jobs runs/list backs `cancel --all`: two active AIR runs for the current user +# (tester@databricks.com, from the built-in scim/v2/Me handler). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/list" +Response.Body = ''' +{ + "runs": [ + { + "run_id": 334747067049496, + "run_name": "qwen-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717608759000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{ + "run_id": 334747067049497, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/qwen-train", + "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}] + } + }] + }, + { + "run_id": 566001814929041, + "run_name": "llama-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717612404000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{ + "run_id": 566001814929042, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/llama-train", + "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}] + } + }] + } + ] +} +''' diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt index 21c3c891af8..7db6ef1aec2 100644 --- a/acceptance/experimental/air/unimplemented/output.txt +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -5,12 +5,6 @@ Error: `air logs` is not implemented yet Exit code: 1 -=== cancel ->>> [CLI] experimental air cancel 123 -Error: `air cancel` is not implemented yet - -Exit code: 1 - === register-image >>> [CLI] experimental air register-image my-image:latest Error: `air register-image` is not implemented yet diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script index 4c53586b16a..19dc13ffe85 100644 --- a/acceptance/experimental/air/unimplemented/script +++ b/acceptance/experimental/air/unimplemented/script @@ -3,8 +3,5 @@ title "logs" errcode trace $CLI experimental air logs 123 -title "cancel" -errcode trace $CLI experimental air cancel 123 - title "register-image" errcode trace $CLI experimental air register-image my-image:latest diff --git a/experimental/air/cmd/cancel.go b/experimental/air/cmd/cancel.go index ae5514e5b04..519a7a82068 100644 --- a/experimental/air/cmd/cancel.go +++ b/experimental/air/cmd/cancel.go @@ -1,10 +1,40 @@ package aircmd import ( + "context" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "text/tabwriter" + "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/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/spf13/cobra" ) +// cancelData is the JSON payload printed by `air cancel`. `all` is set only for +// --all, `workspace` only when --all finds no active runs, and `failed` only +// when a run could not be cancelled. +type cancelData struct { + Cancelled []string `json:"cancelled"` + All bool `json:"all,omitempty"` + Workspace string `json:"workspace,omitempty"` + Failed []cancelFailure `json:"failed,omitempty"` +} + +type cancelFailure struct { + RunID string `json:"run_id"` + Error string `json:"error"` +} + func newCancelCommand() *cobra.Command { var ( all bool @@ -15,9 +45,6 @@ func newCancelCommand() *cobra.Command { Use: "cancel [JOB_RUN_ID...]", Short: "Cancel one or more runs", Long: `Cancel one or more runs by ID, or cancel all of your active runs with --all.`, - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("cancel") - }, } cmd.Flags().BoolVar(&all, "all", false, "Cancel all of your active runs") @@ -35,5 +62,161 @@ func newCancelCommand() *cobra.Command { return nil } + // In -o json mode an auth failure should be a JSON error envelope, not a bare + // error. ErrAlreadyPrinted passes through (already handled upstream). + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + err := root.MustWorkspaceClient(cmd, args) + if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { + return err + } + return renderError(cmd.Context(), cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + runIDs := args + data := cancelData{Cancelled: []string{}} + + if all { + data.All = true + + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to resolve current user: %w", err)) + } + host := strings.TrimRight(w.Config.Host, "/") + + if !jsonOut { + cmdio.LogString(ctx, fmt.Sprintf("Searching active runs for %s in %s...", me.UserName, host)) + } + + // Fetch every active run (up to the scan bound) so --all cancels all + // of them, not just the first page. + fetcher := newRunFetcher(ctx, w, listQuery{activeOnly: true, userFilter: me.UserName}) + rows, err := fetcher.next(maxListScan) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to list active runs: %w", err)) + } + + runIDs = make([]string, 0, len(rows)) + for i := range rows { + if rows[i].RunID != "" { + runIDs = append(runIDs, rows[i].RunID) + } + } + + if len(runIDs) == 0 { + if jsonOut { + data.Workspace = host + return renderEnvelope(ctx, data) + } + cmdio.LogString(ctx, "No active runs found.") + return nil + } + + if !yes { + displayCancelPreview(ctx, rows, host) + confirmed, err := cmdio.AskYesOrNo(ctx, fmt.Sprintf("\nCancel %d run(s) in %s?", len(runIDs), host)) + if err != nil { + return err + } + if !confirmed { + cmdio.LogString(ctx, "Cancellation aborted.") + return root.ErrAlreadyPrinted + } + } + } + + for _, rid := range runIDs { + err := cancelRun(ctx, w, rid) + if err != nil { + data.Failed = append(data.Failed, cancelFailure{RunID: rid, Error: err.Error()}) + if !jsonOut { + if runNotFound(err) { + cmdio.LogString(ctx, fmt.Sprintf("Run %s not found. Please check the run ID and ensure you're using a Job Run ID.", rid)) + } else { + cmdio.LogString(ctx, fmt.Sprintf("Failed to cancel run %s: %s", rid, err)) + } + } + continue + } + data.Cancelled = append(data.Cancelled, rid) + if !jsonOut { + cmdio.LogString(ctx, "Successfully requested cancellation for run "+rid) + } + } + + if jsonOut { + if err := renderEnvelope(ctx, data); err != nil { + return err + } + // Print the envelope, but still exit non-zero on any failure. + if len(data.Failed) > 0 { + return root.ErrAlreadyPrinted + } + return nil + } + + if len(data.Failed) > 0 { + cmdio.LogString(ctx, fmt.Sprintf("%d run(s) failed to cancel.", len(data.Failed))) + return root.ErrAlreadyPrinted + } + if all || len(data.Cancelled) > 1 { + cmdio.LogString(ctx, fmt.Sprintf("Successfully requested cancellation for %d run(s).", len(data.Cancelled))) + } + return nil + } + return cmd } + +// runNotFound reports whether err means the run does not exist. The cancel +// endpoint returns 400 INVALID_PARAMETER_VALUE ("Run <id> does not exist") for +// an unknown run, and the SDK only remaps that to ErrResourceDoesNotExist for +// the runs/get path, not cancel — so we also detect the raw code here. +func runNotFound(err error) bool { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return true + } + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + return apiErr.StatusCode == http.StatusBadRequest && apiErr.ErrorCode == "INVALID_PARAMETER_VALUE" + } + return false +} + +// cancelRun requests cancellation of a single job run. The cancel is async, so +// the returned waiter is ignored. +func cancelRun(ctx context.Context, w *databricks.WorkspaceClient, rid string) error { + runID, err := strconv.ParseInt(rid, 10, 64) + if err != nil || runID <= 0 { + return fmt.Errorf("invalid run ID %q: must be a positive integer", rid) + } + _, err = w.Jobs.CancelRun(ctx, jobs.CancelRun{RunId: runID}) + return err +} + +// displayCancelPreview shows the runs that `cancel --all` is about to terminate. +func displayCancelPreview(ctx context.Context, rows []listRow, host string) { + var sb strings.Builder + fmt.Fprintf(&sb, "\nWorkspace: %s\n", host) + fmt.Fprintf(&sb, "Found %d active run(s) to cancel:\n\n", len(rows)) + + tw := tabwriter.NewWriter(&sb, 0, 0, 2, ' ', 0) + fmt.Fprintln(tw, "Run ID\tExperiment\tStarted") + for i := range rows { + experiment := orNA(rows[i].Experiment) + started := na + if rows[i].StartedAt != nil { + started = *rows[i].StartedAt + } + fmt.Fprintf(tw, "%s\t%s\t%s\n", rows[i].RunID, experiment, started) + } + tw.Flush() + + cmdio.LogString(ctx, strings.TrimRight(sb.String(), "\n")) +} diff --git a/experimental/air/cmd/cancel_test.go b/experimental/air/cmd/cancel_test.go new file mode 100644 index 00000000000..c676567003c --- /dev/null +++ b/experimental/air/cmd/cancel_test.go @@ -0,0 +1,289 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "testing/iotest" + + "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/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +// runCancelAll runs `cancel --all` against w with the given output mode and +// stdin, capturing output into buf. +func runCancelAll(t *testing.T, w *databricks.WorkspaceClient, out flags.Output, in io.Reader, buf *bytes.Buffer) error { + t.Helper() + cmd := withOutput(newCancelCommand(), out) + require.NoError(t, cmd.Flags().Set("all", "true")) + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, in, buf, buf, "", "")) + cmd.SetContext(cmdctx.SetWorkspaceClient(ctx, w)) + return cmd.RunE(cmd, nil) +} + +// cancelEnvelope decodes the air JSON envelope with the cancel payload. +type cancelEnvelope struct { + V int `json:"v"` + Data cancelData `json:"data"` +} + +// runCancel runs the cancel command against w with the given output mode and +// stdin, capturing stdout/stderr into buf. +func runCancel(t *testing.T, w *databricks.WorkspaceClient, out flags.Output, in string, buf *bytes.Buffer, args ...string) (*cobra.Command, error) { + t.Helper() + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, strings.NewReader(in), buf, buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, w) + cmd := withOutput(newCancelCommand(), out) + cmd.SetContext(ctx) + return cmd, cmd.RunE(cmd, args) +} + +func TestCancelArgs(t *testing.T) { + tests := []struct { + name string + all bool + args []string + wantErr string + }{ + {name: "one id", args: []string{"123"}}, + {name: "many ids", args: []string{"123", "456"}}, + {name: "all", all: true}, + {name: "no input", wantErr: "provide at least one JOB_RUN_ID, or use --all"}, + {name: "ids with all", all: true, args: []string{"123"}, wantErr: "cannot combine JOB_RUN_ID arguments with --all"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + cmd := newCancelCommand() + if tc.all { + require.NoError(t, cmd.Flags().Set("all", "true")) + } + err := cmd.Args(cmd, tc.args) + if tc.wantErr == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErr) + }) + } +} + +func TestCancelRunInvalidID(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + for _, id := range []string{"abc", "0", "-1"} { + err := cancelRun(t.Context(), m.WorkspaceClient, id) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid run ID") + } +} + +func TestCancelByIDSuccess(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 456}).Return(nil, nil) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "123", "456") + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, "Successfully requested cancellation for run 123") + assert.Contains(t, out, "Successfully requested cancellation for run 456") + // More than one run cancelled prints the count summary. + assert.Contains(t, out, "Successfully requested cancellation for 2 run(s).") +} + +func TestCancelByIDNotFound(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apierr.ErrResourceDoesNotExist) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "5") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + out := buf.String() + assert.Contains(t, out, "Run 5 not found") + assert.Contains(t, out, "1 run(s) failed to cancel.") +} + +func TestCancelByIDNotFoundInvalidParam(t *testing.T) { + // The cancel endpoint reports an unknown run as 400 INVALID_PARAMETER_VALUE, + // which the SDK does not remap to ErrResourceDoesNotExist for this path. + m := mocks.NewMockWorkspaceClient(t) + apiErr := &apierr.APIError{StatusCode: http.StatusBadRequest, ErrorCode: "INVALID_PARAMETER_VALUE", Message: "Run 5 does not exist."} + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apiErr) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "5") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Run 5 not found") +} + +func TestCancelPartialFailureJSON(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apierr.ErrResourceDoesNotExist) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputJSON, "", &buf, "123", "5") + // The envelope is printed, but a failure still exits non-zero. + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, []string{"123"}, got.Data.Cancelled) + require.Len(t, got.Data.Failed, 1) + assert.Equal(t, "5", got.Data.Failed[0].RunID) + assert.False(t, got.Data.All) +} + +func TestCancelByIDSuccessJSON(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputJSON, "", &buf, "123") + require.NoError(t, err) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, []string{"123"}, got.Data.Cancelled) + assert.Empty(t, got.Data.Failed) +} + +func TestCancelByIDGenericFailure(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 7}).Return(nil, errors.New("boom")) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "7") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Failed to cancel run 7: boom") +} + +func TestCancelAllNoActiveRuns(t *testing.T) { + w := newTestWorkspaceClient(t, runsServer(t, runsListBody(t, "")).URL) + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputText, nil, &buf)) + assert.Contains(t, buf.String(), "No active runs found.") +} + +func TestCancelAllNoActiveRunsJSON(t *testing.T) { + srv := runsServer(t, runsListBody(t, "")) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputJSON, nil, &buf)) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Empty(t, got.Data.Cancelled) + assert.True(t, got.Data.All) + assert.Equal(t, srv.URL, got.Data.Workspace) +} + +func TestCancelAllConfirmYes(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airJobRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + airJobRun(222, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-b"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputText, strings.NewReader("y\n"), &buf)) + out := buf.String() + assert.Contains(t, out, "active run(s) to cancel") + assert.Contains(t, out, "Successfully requested cancellation for run 111") + assert.Contains(t, out, "Successfully requested cancellation for run 222") +} + +func TestCancelAllAbort(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airJobRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, strings.NewReader("n\n"), &buf) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Cancellation aborted.") +} + +func TestCancelAllConfirmReadError(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airJobRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, iotest.ErrReader(errors.New("read failed")), &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "read failed") +} + +func TestCancelAllMeError(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, iam.MeRequest{}).Return(nil, errors.New("nope")) + + var buf bytes.Buffer + err := runCancelAll(t, m.WorkspaceClient, flags.OutputText, nil, &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to resolve current user") +} + +func TestCancelAllListError(t *testing.T) { + // Me succeeds (default empty user), but listing active runs fails. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jobsRunsListPath { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code":"INTERNAL","message":"boom"}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, nil, &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to list active runs") +} + +func TestDisplayCancelPreview(t *testing.T) { + var buf bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, nil, &buf, &buf, "", "")) + + started := "2026-06-05 17:32 UTC" + rows := []listRow{ + {RunID: "111", Experiment: "exp-a", StartedAt: &started}, + {RunID: "222"}, // no experiment or start time -> N/A + } + displayCancelPreview(ctx, rows, "https://my-workspace.cloud.databricks.test") + + out := buf.String() + assert.Contains(t, out, "Workspace: https://my-workspace.cloud.databricks.test") + assert.Contains(t, out, "Found 2 active run(s) to cancel:") + assert.Contains(t, out, "Run ID") + assert.Contains(t, out, "111") + assert.Contains(t, out, "exp-a") + assert.Contains(t, out, "222") + assert.Contains(t, out, na) +} diff --git a/experimental/air/cmd/stubs_test.go b/experimental/air/cmd/stubs_test.go index 4607d7d9eac..e28d7f66730 100644 --- a/experimental/air/cmd/stubs_test.go +++ b/experimental/air/cmd/stubs_test.go @@ -14,7 +14,6 @@ import ( func TestStubCommandsReturnNotImplemented(t *testing.T) { stubs := map[string]*cobra.Command{ "logs": newLogsCommand(), - "cancel": newCancelCommand(), "register-image": newRegisterImageCommand(), } From cb73518660c86351181c53a1ee0146ff46569f55 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Mon, 6 Jul 2026 19:36:09 -0700 Subject: [PATCH 12/18] AIR CLI Integration: fast `air list` via active run index (#5814) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Changes 1. Default to active-only; add --all-status (replaces --active). Plain air list previously scanned every run of every state through Jobs runs/list. It now lists only active runs by default; `--all-status` opts into all states. 2. AiTrainingService index fast path for --all-status scoped to yourself. Instead of scanning the Jobs firehose, it fetches cheap (job_run_id, submit_time) pairs from GET /api/2.0/ai-training/workflows, orders by submit time, keeps the newest --limit, and surfaces only those via concurrent Jobs runs/get. If the index is unavailable, it silently falls back to the Jobs scan so the command never hard-fails. --all-users and other-user filters always use the scan (the index is per-user only). 3. Terminal runs are immutable, so once hydrated, their row is cached; repeat `--all-status` calls skip runs/get + get-output + MLflow for those ids. The runFetcher now wraps a listStrategy (jobsScanStrategy | indexStrategy) behind the same next(want)/exhausted contract, so the interactive table, JSON, and one-shot output paths are unchanged. This PR also fixes a pre-existing recvcheck lint failure and a latent stale-loading guard bug in list_tui.go (fetch helpers converted to value receivers). The index path over-fetches (skips the newest-N truncation) when a --filter on task fields is active, so a filtered-out run can't shrink the result below --limit. ## Why air list in the Go CLI was noticeably slower than the Python AIR CLI — in both plain and --limit modes. The Python CLI's speed comes from three architectural choices, not a faster scan; this PR ports all three to reach parity. ## Testing - Unit tests: index ordering/limit, 403/404-drop vs 500-propagate, parseSubmitTimeMs, cache-hit-skips-network, gate routing, silent fallback, filter over-fetch. - Acceptance: --all-status end-to-end (index → runs/get → get-output) renders fully populated columns. - gofmt, lint-q (0 issues), full air acceptance suite green. --- acceptance/experimental/air/help/output.txt | 6 +- acceptance/experimental/air/list/output.txt | 24 +++ acceptance/experimental/air/list/script | 6 + acceptance/experimental/air/list/test.toml | 34 ++++ experimental/air/cmd/aitraining.go | 109 ++++++++++ experimental/air/cmd/aitraining_test.go | 92 +++++++++ experimental/air/cmd/joblist.go | 80 +++++++- experimental/air/cmd/list.go | 183 +++++++++++------ experimental/air/cmd/list_cache.go | 79 ++++++++ experimental/air/cmd/list_cache_test.go | 56 ++++++ experimental/air/cmd/list_filter.go | 8 + experimental/air/cmd/list_index.go | 141 +++++++++++++ experimental/air/cmd/list_index_test.go | 212 ++++++++++++++++++++ experimental/air/cmd/list_test.go | 10 +- 14 files changed, 973 insertions(+), 67 deletions(-) create mode 100644 experimental/air/cmd/aitraining.go create mode 100644 experimental/air/cmd/aitraining_test.go create mode 100644 experimental/air/cmd/list_cache.go create mode 100644 experimental/air/cmd/list_cache_test.go create mode 100644 experimental/air/cmd/list_index.go create mode 100644 experimental/air/cmd/list_index_test.go diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index 8eac074581d..ee89e778d6b 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -12,7 +12,7 @@ Usage: Available Commands: cancel Cancel one or more runs get Show status, configuration, and timing details for a specific run - list List your recent runs (active and completed) for the current profile + list List your active runs for the current profile (use --all-status for finished runs) logs Stream or fetch logs for a run register-image Mirror a Docker image into the workspace registry run Submit a training workload from a YAML config @@ -30,13 +30,13 @@ Use "databricks experimental air [command] --help" for more information about a === list help >>> [CLI] experimental air list --help -List your recent runs (active and completed) for the current profile +List your active runs for the current profile (use --all-status for finished runs) Usage: databricks experimental air list [flags] Flags: - --active Show only active runs + --all-status Show runs in all states (default: active only) --all-users Show runs from all users --filter stringArray Filter runs, e.g. experiment=foo* (repeatable) -h, --help help for list diff --git a/acceptance/experimental/air/list/output.txt b/acceptance/experimental/air/list/output.txt index 3f84d281ad4..e27cfb0bf39 100644 --- a/acceptance/experimental/air/list/output.txt +++ b/acceptance/experimental/air/list/output.txt @@ -22,3 +22,27 @@ ] } } + +=== list --all-status (text, via AiTrainingService index) +>>> [CLI] experimental air list --all-status + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 + +=== list --all-status (json) +>>> [CLI] experimental air list --all-status -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "runs": [ + { + "run_id": "[NUMID]", + "run_name": "qwen-train", + "user": "[USERNAME]", + "status": "SUCCESS", + "started_at": "[TIMESTAMP]", + "is_sweep": false + } + ] + } +} diff --git a/acceptance/experimental/air/list/script b/acceptance/experimental/air/list/script index 14702b283e7..df547794c6f 100644 --- a/acceptance/experimental/air/list/script +++ b/acceptance/experimental/air/list/script @@ -3,3 +3,9 @@ trace $CLI experimental air list title "list (json)" trace $CLI experimental air list -o json + +title "list --all-status (text, via AiTrainingService index)" +trace $CLI experimental air list --all-status + +title "list --all-status (json)" +trace $CLI experimental air list --all-status -o json diff --git a/acceptance/experimental/air/list/test.toml b/acceptance/experimental/air/list/test.toml index 23dd0f6ba5b..82f1f829d85 100644 --- a/acceptance/experimental/air/list/test.toml +++ b/acceptance/experimental/air/list/test.toml @@ -2,6 +2,10 @@ [EnvMatrix] DATABRICKS_BUNDLE_ENGINE = [] +# Disable the on-disk run cache so --all-status output is deterministic across runs. +[Env] +DATABRICKS_CACHE_ENABLED = "false" + # The SDK occasionally probes host reachability with a HEAD request; stub it so # the test is deterministic. [[Server]] @@ -51,3 +55,33 @@ Pattern = "GET /api/2.2/jobs/runs/get-output" Response.Body = ''' {"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}} ''' + +# `air list --all-status` scoped to the current user is served by the +# AiTrainingService index: it returns cheap (job_run_id, submit_time) pairs, which +# the CLI orders and then hydrates into full runs via runs/get. +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows" +Response.Body = ''' +{"training_workflows": [{"job_run_id": "334747067049496", "submit_time": "2024-06-05T17:32:39Z"}]} +''' + +# runs/get hydrates one index id into the same shape as a runs/list element. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 334747067049496, + "run_name": "qwen-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717608759000, + "end_time": 1717608771000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{ + "run_id": 334747067049497, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/qwen-train", + "deployments": [{"compute": {"accelerator_type": "GPU_8xH100", "accelerator_count": 8}}] + } + }] +} +''' diff --git a/experimental/air/cmd/aitraining.go b/experimental/air/cmd/aitraining.go new file mode 100644 index 00000000000..478f3fbb50a --- /dev/null +++ b/experimental/air/cmd/aitraining.go @@ -0,0 +1,109 @@ +package aircmd + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "time" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" +) + +// aiTrainingWorkflowsPath is the AiTrainingService index of the caller's own AIR +// runs. It returns cheap (job_run_id, submit_time) pairs, letting `air list` +// order and page without scanning the Jobs runs/list firehose. +const aiTrainingWorkflowsPath = "/api/2.0/ai-training/workflows" + +// workflowRef is one run from the index: its Jobs run id and submission time. +type workflowRef struct { + jobRunID int64 + submitTimeMs int64 +} + +type aiTrainingWorkflow struct { + // job_run_id is a Jobs run id; tolerate it arriving as a JSON number or string. + JobRunID json.Number `json:"job_run_id"` + // submit_time is a proto Timestamp, serialized over HTTP as either an RFC3339 + // string or a {seconds, nanos} object. + SubmitTime json.RawMessage `json:"submit_time"` +} + +type aiTrainingWorkflowsResponse struct { + TrainingWorkflows []aiTrainingWorkflow `json:"training_workflows"` + NextPageToken string `json:"next_page_token"` +} + +// listAiTrainingWorkflows pages the index and returns every workflow ref the +// caller owns. Pagination stops at the end or when a page token repeats, which +// guards against a stuck or cycling cursor without an arbitrary page cap. +func listAiTrainingWorkflows(ctx context.Context, w *databricks.WorkspaceClient, activeOnly bool) ([]workflowRef, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var refs []workflowRef + seenTokens := map[string]bool{} + // The index can return the same job_run_id on more than one page; dedupe so + // the newest-`limit` truncation counts unique runs, not repeats. + seenIDs := map[int64]bool{} + var pageToken string + for { + query := map[string]any{} + if activeOnly { + query["active_only"] = true + } + if pageToken != "" { + query["page_token"] = pageToken + } + + var resp aiTrainingWorkflowsResponse + err = apiClient.Do(ctx, http.MethodGet, aiTrainingWorkflowsPath, nil, nil, query, &resp) + if err != nil { + return nil, fmt.Errorf("failed to list training workflows: %w", err) + } + + for _, wf := range resp.TrainingWorkflows { + id, err := wf.JobRunID.Int64() + if err != nil || id == 0 || seenIDs[id] { + continue + } + seenIDs[id] = true + refs = append(refs, workflowRef{jobRunID: id, submitTimeMs: parseSubmitTimeMs(wf.SubmitTime)}) + } + + if resp.NextPageToken == "" || seenTokens[resp.NextPageToken] { + break + } + seenTokens[resp.NextPageToken] = true + pageToken = resp.NextPageToken + } + return refs, nil +} + +// parseSubmitTimeMs converts a proto Timestamp (RFC3339 string or {seconds, nanos} +// object) to epoch milliseconds, or 0 when absent or unparseable (so it sorts last). +func parseSubmitTimeMs(raw json.RawMessage) int64 { + if len(raw) == 0 { + return 0 + } + + var s string + if json.Unmarshal(raw, &s) == nil { + if t, err := time.Parse(time.RFC3339, s); err == nil { + return t.UnixMilli() + } + return 0 + } + + var obj struct { + Seconds int64 `json:"seconds"` + Nanos int64 `json:"nanos"` + } + if json.Unmarshal(raw, &obj) == nil { + return obj.Seconds*1000 + obj.Nanos/1_000_000 + } + return 0 +} diff --git a/experimental/air/cmd/aitraining_test.go b/experimental/air/cmd/aitraining_test.go new file mode 100644 index 00000000000..4c44266d8cb --- /dev/null +++ b/experimental/air/cmd/aitraining_test.go @@ -0,0 +1,92 @@ +package aircmd + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseSubmitTimeMs(t *testing.T) { + cases := []struct { + name string + raw string + want int64 + }{ + {"rfc3339", `"2023-11-14T22:13:20Z"`, 1700000000000}, + {"rfc3339 offset", `"2023-11-14T22:13:20+00:00"`, 1700000000000}, + {"seconds and nanos", `{"seconds": 1700000000, "nanos": 500000000}`, 1700000000500}, + {"seconds only", `{"seconds": 1700000000}`, 1700000000000}, + {"empty", ``, 0}, + {"garbage string", `"not-a-time"`, 0}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, parseSubmitTimeMs(json.RawMessage(tc.raw))) + }) + } +} + +// indexServer serves paginated AiTrainingService responses, one body per call, +// tracking whether the index was hit. +func indexServer(t *testing.T, hit *bool, bodies ...string) *httptest.Server { + t.Helper() + call := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == aiTrainingWorkflowsPath { + *hit = true + body := bodies[min(call, len(bodies)-1)] + call++ + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListAiTrainingWorkflowsPaginates(t *testing.T) { + page1 := `{"training_workflows":[{"job_run_id":"1","submit_time":"2023-11-14T22:13:20Z"}],"next_page_token":"tok"}` + page2 := `{"training_workflows":[{"job_run_id":2,"submit_time":{"seconds":1700000100}}]}` + var hit bool + srv := indexServer(t, &hit, page1, page2) + + refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) + require.NoError(t, err) + require.Len(t, refs, 2) + assert.Equal(t, int64(1), refs[0].jobRunID) + assert.Equal(t, int64(1700000000000), refs[0].submitTimeMs) + assert.Equal(t, int64(2), refs[1].jobRunID) +} + +func TestListAiTrainingWorkflowsStopsOnRepeatedToken(t *testing.T) { + // A cursor that always returns the same token must not loop forever. The + // repeated id is also deduped, so only one ref survives. + page := `{"training_workflows":[{"job_run_id":1,"submit_time":"2023-11-14T22:13:20Z"}],"next_page_token":"tok"}` + var hit bool + srv := indexServer(t, &hit, page) + + refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) + require.NoError(t, err) + require.Len(t, refs, 1) + assert.Equal(t, int64(1), refs[0].jobRunID) +} + +func TestListAiTrainingWorkflowsDedupesIDs(t *testing.T) { + // The same job_run_id on multiple pages must be counted once, so the + // newest-limit truncation doesn't silently return fewer unique runs. + page1 := `{"training_workflows":[{"job_run_id":1,"submit_time":"2023-11-14T22:13:20Z"},{"job_run_id":2,"submit_time":"2023-11-14T22:13:21Z"}],"next_page_token":"tok"}` + page2 := `{"training_workflows":[{"job_run_id":2,"submit_time":"2023-11-14T22:13:21Z"},{"job_run_id":3,"submit_time":"2023-11-14T22:13:22Z"}]}` + var hit bool + srv := indexServer(t, &hit, page1, page2) + + refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) + require.NoError(t, err) + require.Len(t, refs, 3) + got := []int64{refs[0].jobRunID, refs[1].jobRunID, refs[2].jobRunID} + assert.ElementsMatch(t, []int64{1, 2, 3}, got) +} diff --git a/experimental/air/cmd/joblist.go b/experimental/air/cmd/joblist.go index 8a50921264e..c84f87c9df0 100644 --- a/experimental/air/cmd/joblist.go +++ b/experimental/air/cmd/joblist.go @@ -2,17 +2,27 @@ package aircmd import ( "context" + "errors" "fmt" "net/http" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" + "golang.org/x/sync/errgroup" ) -// jobsRunsListPath is the Jobs runs/list endpoint. We call it directly (rather -// than via the typed SDK) because the SDK's RunTask omits ai_runtime_task, the -// task type the AI runtime now submits. -const jobsRunsListPath = "/api/2.2/jobs/runs/list" +// jobsRunsListPath and jobsRunsGetPath are the Jobs endpoints we call directly +// (rather than via the typed SDK) because the SDK's RunTask omits +// ai_runtime_task, the task type the AI runtime now submits. +const ( + jobsRunsListPath = "/api/2.2/jobs/runs/list" + jobsRunsGetPath = "/api/2.2/jobs/runs/get" +) + +// hydrateConcurrency bounds the parallel runs/get calls when hydrating a batch +// of run ids from the AiTrainingService index. +const hydrateConcurrency = 16 type jobsRunsListResponse struct { Runs []jobRun `json:"runs"` @@ -155,6 +165,16 @@ func jobTiming(r *jobRun) (startMillis, endMillis int64) { return startMillis, endMillis } +// isTerminal reports whether a run has finished and its details are immutable, +// so its row is safe to cache. +func isTerminal(r *jobRun) bool { + switch r.State.LifeCycleState { + case "TERMINATED", "INTERNAL_ERROR", "SKIPPED": + return true + } + return false +} + // fetchJobRunsPage fetches one page of Jobs runs/list. query carries the request // params (and page_token across calls). func fetchJobRunsPage(ctx context.Context, w *databricks.WorkspaceClient, query map[string]any) (*jobsRunsListResponse, error) { @@ -170,3 +190,55 @@ func fetchJobRunsPage(ctx context.Context, w *databricks.WorkspaceClient, query } return &resp, nil } + +// fetchJobRun fetches a single run via runs/get. The response mirrors one +// runs/list element, so it deserializes into jobRun directly. +func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (*jobRun, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var run jobRun + query := map[string]any{"run_id": runID, "expand_tasks": true} + err = apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &run) + if err != nil { + return nil, err + } + return &run, nil +} + +// hydrateJobRuns fetches the given run ids concurrently via runs/get, preserving +// input order. runs/get enforces per-run view ACLs, so an id the caller can't +// view (403) or that has been purged (404) is dropped; any other error is +// systemic and fails the whole batch. +func hydrateJobRuns(ctx context.Context, w *databricks.WorkspaceClient, ids []int64) ([]*jobRun, error) { + runs := make([]*jobRun, len(ids)) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(hydrateConcurrency) + for i, id := range ids { + g.Go(func() error { + run, err := fetchJobRun(gctx, w, id) + if err != nil { + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok && + (apiErr.StatusCode == http.StatusForbidden || apiErr.StatusCode == http.StatusNotFound) { + return nil // not viewable or purged: drop this id + } + return fmt.Errorf("failed to get run %d: %w", id, err) + } + runs[i] = run + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, err + } + + hydrated := make([]*jobRun, 0, len(runs)) + for _, run := range runs { + if run != nil { + hydrated = append(hydrated, run) + } + } + return hydrated, nil +} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 8c716f5556a..d5e5e48bcc5 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -60,29 +60,32 @@ type listedRun struct { // listQuery holds the resolved inputs to a runFetcher. type listQuery struct { activeOnly bool + allUsers bool userFilter string + currentUser string filters listFilters fetchMLflow bool + limit int } func newListCommand() *cobra.Command { var ( - limit int - active bool - allUsers bool - filters []string + limit int + allStatus bool + allUsers bool + filters []string ) cmd := &cobra.Command{ Use: "list", Args: root.NoArgs, - Short: "List your recent runs (active and completed) for the current profile", + Short: "List your active runs for the current profile (use --all-status for finished runs)", } cmd.PreRunE = root.MustWorkspaceClient cmd.Flags().IntVar(&limit, "limit", 20, "Maximum number of runs to show") - cmd.Flags().BoolVar(&active, "active", false, "Show only active runs") + cmd.Flags().BoolVar(&allStatus, "all-status", false, "Show runs in all states (default: active only)") cmd.Flags().BoolVar(&allUsers, "all-users", false, "Show runs from all users") cmd.Flags().StringArrayVar(&filters, "filter", nil, "Filter runs, e.g. experiment=foo* (repeatable)") @@ -103,19 +106,24 @@ func newListCommand() *cobra.Command { // unless --all-users is set. runs/list has no creator param, so the // creator is matched while scanning. userFilter := f.User + var currentUser string if userFilter == "" && !allUsers { me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) if err != nil { return fmt.Errorf("failed to resolve current user: %w", err) } - userFilter = me.UserName + currentUser = me.UserName + userFilter = currentUser } fetcher := newRunFetcher(ctx, w, listQuery{ - activeOnly: active, + activeOnly: !allStatus, + allUsers: allUsers, userFilter: userFilter, + currentUser: currentUser, filters: f, fetchMLflow: root.OutputType(cmd) == flags.OutputText, + limit: limit, }) // JSON prints the newest `limit` runs once. Text renders the table: @@ -135,23 +143,96 @@ func newListCommand() *cobra.Command { return cmd } -// runFetcher pages Jobs runs/list on demand, yielding AIR runs that match the -// user and filters. It buffers a page's leftover runs so successive next() calls -// resume where the last stopped — driving both one-shot output and lazy paging. +// listStrategy is a source of matching runs, pulled in batches. Two implement it: +// jobsScanStrategy pages runs/list; indexStrategy hydrates the AiTrainingService +// index. The fetcher wraps whichever is chosen. +type listStrategy interface { + // next returns up to want more matching runs (already row-built + task id). + next(want int) ([]listedRun, error) + // done reports whether the source has no more runs to yield. + done() bool + // truncated reports whether a safety cap stopped the scan short of the end. + truncated() bool +} + +// runFetcher yields matching rows in batches, driving both one-shot output and +// the interactive table's lazy paging. It wraps a listStrategy and adds the +// shared tail: MLflow enrichment (text only) and row projection. type runFetcher struct { ctx context.Context w *databricks.WorkspaceClient - query map[string]any - userFilter string - filters listFilters fetchMLflow bool + strategy listStrategy - pending []jobRun // runs from the last page not yet inspected - scanned int exhausted bool } func newRunFetcher(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *runFetcher { + return &runFetcher{ + ctx: ctx, + w: w, + fetchMLflow: q.fetchMLflow, + strategy: newListStrategy(ctx, w, q), + } +} + +// newListStrategy picks the fetch source. The AiTrainingService index serves only +// the caller's own runs, so it's used for an all-status self-scoped list; if the +// index load fails (e.g. endpoint unavailable in this workspace), we fall back to +// the Jobs scan so the command still returns. Everything else — the default +// active list, --all-users, and --all-status for another user — uses the scan. +func newListStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) listStrategy { + useIndex := !q.activeOnly && !q.allUsers && (q.userFilter == "" || q.userFilter == q.currentUser) + if !useIndex { + return newJobsScanStrategy(ctx, w, q) + } + idx := newIndexStrategy(ctx, w, q, q.limit) + if err := idx.load(); err != nil { + log.Debugf(ctx, "air list: AiTrainingService index unavailable, falling back to Jobs scan: %v", err) + return newJobsScanStrategy(ctx, w, q) + } + return idx +} + +// next pulls the next batch from the strategy, enriches it with MLflow links for +// text output, and projects it to rows. It sets exhausted once the strategy is +// drained so the interactive table knows to stop paging. +func (f *runFetcher) next(want int) ([]listRow, error) { + entries, err := f.strategy.next(want) + if err != nil { + return nil, err + } + f.exhausted = f.strategy.done() + + // MLflow links appear only in the text table, so the per-run get-output + // lookups are skipped for JSON output (which omits the column anyway). + if f.fetchMLflow { + setMLflowLinks(f.ctx, f.w, entries) + } + + rows := make([]listRow, len(entries)) + for i, e := range entries { + rows[i] = e.row + } + return rows, nil +} + +// jobsScanStrategy pages Jobs runs/list, keeping the AIR runs that match the user +// and filters. It buffers a page's leftover runs so successive next() calls +// resume where the last stopped. +type jobsScanStrategy struct { + ctx context.Context + w *databricks.WorkspaceClient + query map[string]any + userFilter string + filters listFilters + + pending []jobRun // runs from the last page not yet inspected + scanned int + drained bool +} + +func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *jobsScanStrategy { query := map[string]any{ "run_type": "SUBMIT_RUN", "expand_tasks": true, @@ -160,84 +241,74 @@ func newRunFetcher(ctx context.Context, w *databricks.WorkspaceClient, q listQue if q.activeOnly { query["active_only"] = true } - return &runFetcher{ - ctx: ctx, - w: w, - query: query, - userFilter: q.userFilter, - filters: q.filters, - fetchMLflow: q.fetchMLflow, + return &jobsScanStrategy{ + ctx: ctx, + w: w, + query: query, + userFilter: q.userFilter, + filters: q.filters, } } -// next returns up to want more matching rows, paging runs/list (and buffering the -// leftover runs of a page) until it has enough, the server has no more pages, or -// it has scanned maxListScan runs. MLflow links are filled in for text output. -func (f *runFetcher) next(want int) ([]listRow, error) { +func (s *jobsScanStrategy) next(want int) ([]listedRun, error) { var entries []listedRun for len(entries) < want { - if len(f.pending) == 0 { - if f.exhausted || f.scanned >= maxListScan { + if len(s.pending) == 0 { + if s.drained || s.scanned >= maxListScan { break } - if err := f.fetchPage(); err != nil { + if err := s.fetchPage(); err != nil { return nil, err } continue } - run := &f.pending[0] - f.pending = f.pending[1:] - f.scanned++ + run := &s.pending[0] + s.pending = s.pending[1:] + s.scanned++ if !isAirRun(run) { continue } - if f.userFilter != "" && run.CreatorUserName != f.userFilter { + if s.userFilter != "" && run.CreatorUserName != s.userFilter { continue } - if !f.filters.matches(run) { + if !s.filters.matches(run) { continue } entries = append(entries, listedRun{row: buildListRow(run), taskRunID: taskRunID(run)}) } - if f.scanned >= maxListScan { - f.exhausted = true - } + return entries, nil +} - // MLflow links appear only in the text table, so the per-run get-output - // lookups are skipped for JSON output (which omits the column anyway). - if f.fetchMLflow { - setMLflowLinks(f.ctx, f.w, entries) - } +func (s *jobsScanStrategy) done() bool { + return (s.drained && len(s.pending) == 0) || s.scanned >= maxListScan +} - rows := make([]listRow, len(entries)) - for i, e := range entries { - rows[i] = e.row - } - return rows, nil +func (s *jobsScanStrategy) truncated() bool { + return s.scanned >= maxListScan } // fetchPage loads the next runs/list page into the pending buffer, marking the -// fetcher exhausted once the server reports no further pages. -func (f *runFetcher) fetchPage() error { - resp, err := fetchJobRunsPage(f.ctx, f.w, f.query) +// strategy drained once the server reports no further pages. +func (s *jobsScanStrategy) fetchPage() error { + resp, err := fetchJobRunsPage(s.ctx, s.w, s.query) if err != nil { return err } - f.pending = resp.Runs + s.pending = resp.Runs if resp.NextPageToken == "" { - f.exhausted = true + s.drained = true } else { - f.query["page_token"] = resp.NextPageToken + s.query["page_token"] = resp.NextPageToken } return nil } -// warnIfTruncated logs when the scan hit maxListScan, so one-shot output signals +// warnIfTruncated logs when a scan hit its safety cap, so one-shot output signals // its results may be incomplete. func warnIfTruncated(ctx context.Context, f *runFetcher) { - if f.scanned >= maxListScan { + if f.strategy.truncated() { log.Warnf(ctx, "air list: stopped after scanning %d runs; results may be incomplete", maxListScan) } } diff --git a/experimental/air/cmd/list_cache.go b/experimental/air/cmd/list_cache.go new file mode 100644 index 00000000000..fb4442a1127 --- /dev/null +++ b/experimental/air/cmd/list_cache.go @@ -0,0 +1,79 @@ +package aircmd + +import ( + "context" + "time" + + "github.com/databricks/cli/libs/cache" +) + +// The AiTrainingService index path caches hydrated terminal runs on disk: +// terminal runs are immutable, so once we've paid for runs/get + get-output + +// MLflow we persist the finished row and skip those round-trips next time. The +// TTL matches AICM's ~60-day retention, after which the run drops out of the +// index anyway. +const ( + listCacheComponent = "air-list-runs" + listCacheTTL = 60 * 24 * time.Hour +) + +// listCacheKey fingerprints a cached run. Host isolates workspaces (a Jobs run +// id is unique only within one), matching how libs/cache namespaces entries. +type listCacheKey struct { + Host string `json:"host"` + RunID int64 `json:"run_id"` +} + +// cachedRun is the persisted value: every listRow field (including the +// table-only columns, which listRow tags json:"-" and so wouldn't survive a +// direct marshal) plus the submit time. +type cachedRun struct { + RunID string `json:"run_id"` + RunName string `json:"run_name"` + User string `json:"user"` + Status string `json:"status"` + StartedAt *string `json:"started_at"` + IsSweep bool `json:"is_sweep"` + Experiment string `json:"experiment"` + Duration string `json:"duration"` + MLflowURL string `json:"mlflow_url"` + Accelerators string `json:"accelerators"` + SubmitTimeMs int64 `json:"submit_time_ms"` +} + +func (c cachedRun) toRow() listRow { + return listRow{ + RunID: c.RunID, RunName: c.RunName, User: c.User, Status: c.Status, + StartedAt: c.StartedAt, IsSweep: c.IsSweep, Experiment: c.Experiment, + Duration: c.Duration, MLflowURL: c.MLflowURL, Accelerators: c.Accelerators, + } +} + +func cachedRunFromRow(r listRow, submitTimeMs int64) cachedRun { + return cachedRun{ + RunID: r.RunID, RunName: r.RunName, User: r.User, Status: r.Status, + StartedAt: r.StartedAt, IsSweep: r.IsSweep, Experiment: r.Experiment, + Duration: r.Duration, MLflowURL: r.MLflowURL, Accelerators: r.Accelerators, + SubmitTimeMs: submitTimeMs, + } +} + +// newListCache builds the cache for the index path. It fails open, so a nil +// return (or any cache error) just means every run is hydrated from the API. +func newListCache(ctx context.Context) *cache.Cache { + return cache.NewCache(ctx, listCacheComponent, listCacheTTL, nil) +} + +// cachedRow returns the cached row for a run, or (zero, false) on miss. +func cachedRow(ctx context.Context, c *cache.Cache, host string, runID int64) (listRow, bool) { + entry, ok := cache.Get[cachedRun](ctx, c, listCacheKey{Host: host, RunID: runID}) + if !ok { + return listRow{}, false + } + return entry.toRow(), true +} + +// putRow caches a terminal run's finished row under its submit time. +func putRow(ctx context.Context, c *cache.Cache, host string, runID, submitTimeMs int64, row listRow) { + cache.Put(ctx, c, listCacheKey{Host: host, RunID: runID}, cachedRunFromRow(row, submitTimeMs)) +} diff --git a/experimental/air/cmd/list_cache_test.go b/experimental/air/cmd/list_cache_test.go new file mode 100644 index 00000000000..3ac695e8d0d --- /dev/null +++ b/experimental/air/cmd/list_cache_test.go @@ -0,0 +1,56 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestListCacheRoundTrip(t *testing.T) { + t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) + ctx := t.Context() + c := newListCache(ctx) + + _, ok := cachedRow(ctx, c, "https://host.test", 42) + require.False(t, ok, "miss before write") + + row := listRow{RunID: "42", Experiment: "exp", Status: "SUCCESS"} + putRow(ctx, c, "https://host.test", 42, 1700000000000, row) + + got, ok := cachedRow(ctx, c, "https://host.test", 42) + require.True(t, ok, "hit after write") + assert.Equal(t, row, got) + + // Different host is a different key. + _, ok = cachedRow(ctx, c, "https://other.test", 42) + assert.False(t, ok) +} + +func TestIndexStrategyServesCachedRowWithoutFetch(t *testing.T) { + t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) + + refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} + srv, hits := indexAndGetServer(t, refs, map[int64]jobRun{7: indexRun(7, 1000_000)}, nil, nil) + host := srv.URL + + // Pre-seed the cache for run 7 so hydration should skip runs/get entirely. + ctx := t.Context() + putRow(ctx, newListCache(ctx), host, 7, 1000_000, listRow{RunID: "7", Status: "SUCCESS"}) + + f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 1) + assert.Equal(t, "7", rows[0].RunID) + assert.Equal(t, 0, hits.get, "cached run must not hit runs/get") +} + +func TestIsTerminal(t *testing.T) { + assert.True(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "TERMINATED"}})) + assert.True(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "INTERNAL_ERROR"}})) + assert.False(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "RUNNING"}})) + assert.False(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "PENDING"}})) +} diff --git a/experimental/air/cmd/list_filter.go b/experimental/air/cmd/list_filter.go index 6ac47c74b0e..272141e80d1 100644 --- a/experimental/air/cmd/list_filter.go +++ b/experimental/air/cmd/list_filter.go @@ -10,6 +10,14 @@ import ( // supportedFilterKeys are the keys accepted by `air list --filter KEY=VALUE`. var supportedFilterKeys = []string{"accelerator_type", "experiment", "num_accelerators", "user"} +// hasTaskFilter reports whether any filter is applied to a run's task fields +// (experiment or accelerators), i.e. matched after a run is fetched rather than +// while scanning. The index path uses this to skip its newest-N truncation, so a +// dropped match doesn't shrink the result below --limit. +func (f listFilters) hasTaskFilter() bool { + return f.Experiment != "" || f.AcceleratorType != "" || f.NumAccelerators != nil +} + // listFilters holds the parsed `--filter` values for `air list`. type listFilters struct { // User is an exact creator-email match diff --git a/experimental/air/cmd/list_index.go b/experimental/air/cmd/list_index.go new file mode 100644 index 00000000000..097859ba402 --- /dev/null +++ b/experimental/air/cmd/list_index.go @@ -0,0 +1,141 @@ +package aircmd + +import ( + "cmp" + "context" + "slices" + + "github.com/databricks/cli/libs/cache" + "github.com/databricks/databricks-sdk-go" +) + +// indexStrategy serves the caller's own runs from the AiTrainingService index: +// it fetches every run id up front (cheap id+timestamp pairs), orders them +// newest-first, keeps the newest `limit`, then hydrates them into full rows in +// want-sized batches via Jobs runs/get. Terminal rows are cached so repeat calls +// skip the network. Unlike the Jobs scan it can't lazy-page (it must sort the +// whole id set first), but it still yields in batches so the table paints early. +type indexStrategy struct { + ctx context.Context + w *databricks.WorkspaceClient + activeOnly bool + filters listFilters + limit int + cache *cache.Cache + + ids []int64 // newest-first run ids to hydrate, resolved on first next() + pos int + loaded bool +} + +func newIndexStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery, limit int) *indexStrategy { + return &indexStrategy{ + ctx: ctx, + w: w, + activeOnly: q.activeOnly, + filters: q.filters, + limit: limit, + cache: newListCache(ctx), + } +} + +// load fetches and orders the index once. It returns an error only when the +// index endpoint itself fails, letting the caller fall back to the Jobs scan. +func (s *indexStrategy) load() error { + refs, err := listAiTrainingWorkflows(s.ctx, s.w, s.activeOnly) + if err != nil { + return err + } + slices.SortFunc(refs, func(a, b workflowRef) int { return cmp.Compare(b.submitTimeMs, a.submitTimeMs) }) + // Keep only the newest `limit` ids so hydration is bounded — but skip that when + // a task filter is active, since it drops matches post-hydration and we'd + // otherwise return fewer than `limit`. The caller stops pulling at `limit`. + if s.limit > 0 && len(refs) > s.limit && !s.filters.hasTaskFilter() { + refs = refs[:s.limit] + } + s.ids = make([]int64, len(refs)) + for i, r := range refs { + s.ids[i] = r.jobRunID + } + s.loaded = true + return nil +} + +func (s *indexStrategy) next(want int) ([]listedRun, error) { + if !s.loaded { + if err := s.load(); err != nil { + return nil, err + } + } + + var entries []listedRun + for len(entries) < want && s.pos < len(s.ids) { + end := min(s.pos+want-len(entries), len(s.ids)) + batch := s.ids[s.pos:end] + s.pos = end + + rows, err := s.hydrate(batch) + if err != nil { + return nil, err + } + entries = append(entries, rows...) + } + return entries, nil +} + +func (s *indexStrategy) done() bool { + return s.loaded && s.pos >= len(s.ids) +} + +// truncated is always false: the index path is bounded by limit, not a scan cap. +func (s *indexStrategy) truncated() bool { return false } + +// hydrate turns a batch of run ids into rows, serving cached terminal rows +// without a network call and fetching the rest via runs/get. Freshly hydrated +// terminal runs are cached. Results keep the input (newest-first) order, then +// the batch is re-sorted by start time since concurrent hydration reorders it. +func (s *indexStrategy) hydrate(ids []int64) ([]listedRun, error) { + host := s.w.Config.Host + + rows := make([]listedRun, 0, len(ids)) + var toFetch []int64 + for _, id := range ids { + if row, ok := cachedRow(s.ctx, s.cache, host, id); ok { + rows = append(rows, listedRun{row: row, taskRunID: id}) + continue + } + toFetch = append(toFetch, id) + } + + runs, err := hydrateJobRuns(s.ctx, s.w, toFetch) + if err != nil { + return nil, err + } + for _, run := range runs { + if !s.filters.matches(run) { + continue + } + row := buildListRow(run) + rows = append(rows, listedRun{row: row, taskRunID: taskRunID(run)}) + if isTerminal(run) { + start, _ := jobTiming(run) + putRow(s.ctx, s.cache, host, run.RunID, start, row) + } + } + + // Concurrent hydration reorders runs, so re-sort the batch newest-first. The + // ISO start timestamp sorts lexicographically; a missing time ("") sorts last. + slices.SortStableFunc(rows, func(a, b listedRun) int { + return cmp.Compare(rowStartKey(b.row), rowStartKey(a.row)) + }) + return rows, nil +} + +// rowStartKey returns a row's ISO start timestamp for ordering, or "" when the +// run hasn't started (which sorts last under descending comparison). +func rowStartKey(r listRow) string { + if r.StartedAt == nil { + return "" + } + return *r.StartedAt +} diff --git a/experimental/air/cmd/list_index_test.go b/experimental/air/cmd/list_index_test.go new file mode 100644 index 00000000000..ec8d082e3a0 --- /dev/null +++ b/experimental/air/cmd/list_index_test.go @@ -0,0 +1,212 @@ +package aircmd + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// indexRun is a terminal AIR run with a start time, for index-path hydration. +func indexRun(id, startMillis int64) jobRun { + r := airJobRun(id, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/exp") + r.State = jobState{LifeCycleState: "TERMINATED", ResultState: "SUCCESS"} + r.Tasks[0].StartTime = startMillis + r.Tasks[0].EndTime = startMillis + 1000 + return r +} + +// indexAndGetServer serves the AiTrainingService index (a single page of the +// given refs) and runs/get for each id, recording hit counts per endpoint. A +// runID in forbidden returns 403; in missing returns 404. +type indexHits struct{ index, get int } + +func indexAndGetServer(t *testing.T, refs []workflowRef, runs map[int64]jobRun, forbidden, missing map[int64]bool) (*httptest.Server, *indexHits) { + t.Helper() + hits := &indexHits{} + wfs := make([]map[string]any, len(refs)) + for i, r := range refs { + wfs[i] = map[string]any{"job_run_id": strconv.FormatInt(r.jobRunID, 10), "submit_time": map[string]any{"seconds": r.submitTimeMs / 1000}} + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case aiTrainingWorkflowsPath: + hits.index++ + _ = json.NewEncoder(w).Encode(map[string]any{"training_workflows": wfs}) + case jobsRunsGetPath: + hits.get++ + id, _ := strconv.ParseInt(r.URL.Query().Get("run_id"), 10, 64) + if forbidden[id] { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"message":"forbidden"}`)) + return + } + if missing[id] { + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"message":"not found"}`)) + return + } + run := runs[id] + _ = json.NewEncoder(w).Encode(run) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv, hits +} + +func TestIndexStrategyOrdersAndLimits(t *testing.T) { + // Three runs, out of submit-time order; newest two should win, newest-first. + refs := []workflowRef{ + {jobRunID: 1, submitTimeMs: 1000_000}, + {jobRunID: 2, submitTimeMs: 3000_000}, + {jobRunID: 3, submitTimeMs: 2000_000}, + } + runs := map[int64]jobRun{ + 1: indexRun(1, 1000_000), + 2: indexRun(2, 3000_000), + 3: indexRun(3, 2000_000), + } + srv, _ := indexAndGetServer(t, refs, runs, nil, nil) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 2, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 2) + assert.Equal(t, "2", rows[0].RunID) // submit 3000 + assert.Equal(t, "3", rows[1].RunID) // submit 2000 + assert.True(t, f.exhausted) +} + +func TestIndexStrategyOverFetchesWithTaskFilter(t *testing.T) { + // With a task filter and limit 1, the newest run doesn't match; the strategy + // must keep hydrating past `limit` to find the match rather than truncating. + refs := []workflowRef{ + {jobRunID: 1, submitTimeMs: 3000_000}, + {jobRunID: 2, submitTimeMs: 2000_000}, + } + run1 := indexRun(1, 3000_000) + run1.Tasks[0].AiRuntimeTask.Experiment = "/Users/me@example.com/llama" + run2 := indexRun(2, 2000_000) + run2.Tasks[0].AiRuntimeTask.Experiment = "/Users/me@example.com/qwen" + srv, _ := indexAndGetServer(t, refs, map[int64]jobRun{1: run1, 2: run2}, nil, nil) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 1, + filters: listFilters{Experiment: "qwen"}, + }) + rows, err := f.next(1) + require.NoError(t, err) + require.Len(t, rows, 1) + assert.Equal(t, "2", rows[0].RunID) // found despite being the older, second id +} + +func TestIndexStrategyDropsForbiddenAndMissing(t *testing.T) { + refs := []workflowRef{ + {jobRunID: 1, submitTimeMs: 3000_000}, + {jobRunID: 2, submitTimeMs: 2000_000}, + {jobRunID: 3, submitTimeMs: 1000_000}, + } + runs := map[int64]jobRun{1: indexRun(1, 3000_000), 3: indexRun(3, 1000_000)} + srv, _ := indexAndGetServer(t, refs, runs, map[int64]bool{2: true}, nil) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 2) // run 2 (403) dropped + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "3", rows[1].RunID) +} + +func TestIndexStrategyPropagatesServerError(t *testing.T) { + refs := []workflowRef{{jobRunID: 1, submitTimeMs: 1000_000}} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case aiTrainingWorkflowsPath: + wfs := []map[string]any{{"job_run_id": "1", "submit_time": map[string]any{"seconds": int64(1000)}}} + _ = json.NewEncoder(w).Encode(map[string]any{"training_workflows": wfs}) + case jobsRunsGetPath: + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"message":"boom"}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + _ = refs + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + }) + _, err := f.next(10) + require.Error(t, err) // 500 is systemic, not an ACL drop +} + +func TestNewListStrategyGate(t *testing.T) { + // --all-users and other-user filters must NOT touch the index endpoint. + cases := []struct { + name string + q listQuery + wantIndex bool + }{ + {"active default → scan", listQuery{activeOnly: true, userFilter: "me@example.com", currentUser: "me@example.com"}, false}, + {"all-status self → index", listQuery{userFilter: "me@example.com", currentUser: "me@example.com", limit: 5}, true}, + {"all-status all-users → scan", listQuery{allUsers: true}, false}, + {"all-status other user → scan", listQuery{userFilter: "other@example.com", currentUser: "me@example.com"}, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var indexHit bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == aiTrainingWorkflowsPath { + indexHit = true + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + newListStrategy(t.Context(), newTestWorkspaceClient(t, srv.URL), tc.q) + assert.Equal(t, tc.wantIndex, indexHit) + }) + } +} + +func TestNewListStrategyFallsBackWhenIndexFails(t *testing.T) { + // Index 500 must silently fall back to the Jobs scan, not fail the command. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == aiTrainingWorkflowsPath { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"message":"boom"}`)) + return + } + if r.URL.Path == jobsRunsListPath { + _, _ = fmt.Fprint(w, runsListBody(t, "", airJobRun(1, "me@example.com", "GPU_1xH100", 1, "exp"))) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 1) // served by the Jobs scan fallback +} diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go index 11185528a2c..8d186eb048d 100644 --- a/experimental/air/cmd/list_test.go +++ b/experimental/air/cmd/list_test.go @@ -67,6 +67,7 @@ func TestListAirRunsFiltersUserAndType(t *testing.T) { srv := runsServer(t, runsListBody(t, "", runs...)) rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + activeOnly: true, userFilter: "me@example.com", }).next(10) require.NoError(t, err) @@ -83,7 +84,8 @@ func TestListAirRunsExperimentFilter(t *testing.T) { srv := runsServer(t, runsListBody(t, "", runs...)) rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - filters: listFilters{Experiment: "qwen*"}, + activeOnly: true, + filters: listFilters{Experiment: "qwen*"}, }).next(10) require.NoError(t, err) require.Len(t, rows, 1) @@ -98,7 +100,7 @@ func TestListAirRunsLimitTruncates(t *testing.T) { } srv := runsServer(t, runsListBody(t, "", runs...)) - rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{}).next(2) + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{activeOnly: true}).next(2) require.NoError(t, err) require.Len(t, rows, 2) assert.Equal(t, "1", rows[0].RunID) @@ -110,7 +112,7 @@ func TestListAirRunsPaginates(t *testing.T) { page2 := runsListBody(t, "", airJobRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) srv := runsServer(t, page1, page2) - rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{}).next(10) + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{activeOnly: true}).next(10) require.NoError(t, err) require.Len(t, rows, 2) assert.Equal(t, "1", rows[0].RunID) @@ -127,7 +129,7 @@ func TestRunFetcherResumesAcrossCalls(t *testing.T) { airJobRun(3, "me@example.com", "GPU_1xH100", 1, "exp-c"), } srv := runsServer(t, runsListBody(t, "", runs...)) - f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{}) + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{activeOnly: true}) first, err := f.next(2) require.NoError(t, err) From b3b7ec8a1f7e5af521ff5c8b6ff2b4c9b1a3630d Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Tue, 7 Jul 2026 21:47:52 -0700 Subject: [PATCH 13/18] AIR CLI: drop max_retries from the ai_runtime_task payload (#5813) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Changes Stop sending `max_retries` / `retry_on_timeout` on the `ai_runtime_task` submit payload. Both fields are removed from the `submitTask` struct and from `buildSubmitPayload`. The `max_retries` YAML config field and its validation stay in the schema; it is simply no longer put on the wire for this submission path. ## Why On the `ai_runtime_task` path, execution retries are driven by the **AI Runtime service (AICM)**, not by the Jobs task `max_retries` field — so setting it on the task had no effect on how many times a failing workload actually retried. Observed directly: a run submitted with `max_retries: 0` still made 4 attempts (3 retries). `air get` on that run confirmed the field round-tripped correctly (`Max Retries 0`) while the workload retried anyway. The Python CLI's native `ai_runtime_task` branch omits `max_retries` for the same reason; this matches that behavior. Leaving the field on the payload was misleading: a user setting `max_retries: 0` to disable retries would still see retries. ## Tests - `TestBuildSubmitPayload` asserts the marshaled task contains neither `max_retries` nor `retry_on_timeout`. - Removed `TestBuildSubmitPayload_NoRetries` (it asserted the now-removed field was sent). - `go test ./experimental/air/...` and `./task lint-q` pass. --- experimental/air/cmd/runsubmit.go | 12 ++++++++++-- experimental/air/cmd/runsubmit_test.go | 21 ++++++++++++++++++--- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index b32ebd622c6..49f51a8aa66 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -66,6 +66,13 @@ type jobEnvironment struct { } // submitTask is the single task air submits: a native ai_runtime_task. +// +// max_retries is always sent (including 0) so the user's YAML value is honored: +// setting it to 0 explicitly disables retries rather than falling back to the +// server default. retry_on_timeout is sent only when retries are allowed, and is +// omitempty so the wire form matches the Python CLI (which never emits a bare +// "false"). Jobs performs the retries — each attempt is a fresh AI Runtime +// workload. type submitTask struct { TaskKey string `json:"task_key"` RunIf string `json:"run_if"` @@ -125,8 +132,9 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string) jobsSubmitR EnvironmentKey: aiRuntimeEnvironmentKey, MaxRetries: cfg.maxRetries(), } - // max_retries 0 (no retries) is sent explicitly; retry_on_timeout only - // applies when retries are allowed. + // retry_on_timeout only makes sense when retries are allowed; otherwise omit + // it (matches Python's native path, which sets retry_on_timeout only under + // the same > 0 gate). st.RetryOnTimeout = st.MaxRetries > 0 return jobsSubmitRun{ diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index bfd92dcd58e..d57b14f80a6 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -61,22 +61,37 @@ func TestBuildSubmitPayload(t *testing.T) { assert.Equal(t, aiRuntimeCompute{AcceleratorType: "GPU_8xH100", AcceleratorCount: 16}, at.Deployments[0].Compute) } -func TestBuildSubmitPayload_NoRetries(t *testing.T) { +func TestBuildSubmitPayloadDefaultRetries(t *testing.T) { + // max_retries unset defaults to 3 (matching the Python native path), so both + // retry fields are sent. cfg := &runConfig{ ExperimentName: "exp", Command: new("x"), Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, - MaxRetries: new(0), } + task := buildSubmitPayload(cfg, "/d/command.sh", "4").Tasks[0] + assert.Equal(t, defaultMaxRetries, task.MaxRetries) + assert.True(t, task.RetryOnTimeout) +} +func TestBuildSubmitPayloadNoRetries(t *testing.T) { + // max_retries: 0 must be sent explicitly so Jobs honors "no retries" instead + // of applying the server default. retry_on_timeout is omitted when retries + // aren't allowed. + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, + MaxRetries: new(0), + } task := buildSubmitPayload(cfg, "/d/command.sh", "4").Tasks[0] assert.Equal(t, 0, task.MaxRetries) assert.False(t, task.RetryOnTimeout) - // max_retries: 0 must be sent, not omitted, so the server honors "no retries". b, err := json.Marshal(task) require.NoError(t, err) assert.Contains(t, string(b), `"max_retries":0`) + assert.NotContains(t, string(b), "retry_on_timeout") } func TestSubmitToken(t *testing.T) { From a3492b801e82ccb65f129e05228dc880215707cd Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Tue, 7 Jul 2026 21:57:16 -0700 Subject: [PATCH 14/18] AIR CLI Integration: authenticate `air get` up front and fail fast (#5729) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Changes Authenticate before any run status or config is fetched or printed: - **PreRunE** maps a `MustWorkspaceClient` failure to an actionable auth error: - no default profile set and `--profile`/`-p` not passed (`config.ErrCannotConfigureDefault`) → *"no default profile is set: pass --profile (-p) or configure a default profile in your .databrickscfg"* - otherwise → *"authentication was not successful: <cause>"* - **RunE** calls `CurrentUser.Me` before fetching/rendering anything, so a credential that resolves locally but is rejected by the workspace also fails fast with the same clear message — and **no partial status/config is shown**. Both errors are permanent (not retryable) and, in `-o json` mode, render as the standard error envelope (`code: UNAUTHENTICATED`). ## Why `air get JOB_RUN_ID` validated authentication only lazily. `MustWorkspaceClient` (PreRunE) calls `Config.Authenticate`, which merely *attaches* credentials (for a PAT it does no server-side check), so an invalid credential or missing profile surfaced as a confusing, generic failure partway through — after a run's config had already started rendering — instead of a clear, up-front error. ## Testing - `TestGetRunAuthFailed` — a rejected `CurrentUser.Me` short-circuits before `GetRun` (no run fetched, nothing rendered). - `TestAuthError` — verifies the no-profile vs generic-auth message mapping (and that the cause is preserved). - Existing not-found tests updated to stub the up-front `Me` success. - `go build ./...`, the air unit tests, the air acceptance suite, and `./task lint-q` (0 issues) all pass. Based on `air-cli` (post-#5685, so the command is `air get JOB_RUN_ID`). This pull request and its description were written by Isaac. --- .../air/get-ai-runtime/out.test.toml | 3 + .../air/get-ai-runtime/output.txt | 52 +++++++ .../experimental/air/get-ai-runtime/script | 9 ++ .../experimental/air/get-ai-runtime/test.toml | 55 +++++++ .../air/get-ai-runtime/training_config.yaml | 8 + experimental/air/cmd/get.go | 68 ++++++++- experimental/air/cmd/get_test.go | 142 +++++++++++++++++- experimental/air/cmd/joblist.go | 45 +++++- experimental/air/cmd/joblist_test.go | 50 ++++++ experimental/air/cmd/render.go | 20 ++- experimental/air/cmd/render_test.go | 27 ++++ 11 files changed, 462 insertions(+), 17 deletions(-) create mode 100644 acceptance/experimental/air/get-ai-runtime/out.test.toml create mode 100644 acceptance/experimental/air/get-ai-runtime/output.txt create mode 100644 acceptance/experimental/air/get-ai-runtime/script create mode 100644 acceptance/experimental/air/get-ai-runtime/test.toml create mode 100644 acceptance/experimental/air/get-ai-runtime/training_config.yaml create mode 100644 experimental/air/cmd/joblist_test.go diff --git a/acceptance/experimental/air/get-ai-runtime/out.test.toml b/acceptance/experimental/air/get-ai-runtime/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/get-ai-runtime/output.txt b/acceptance/experimental/air/get-ai-runtime/output.txt new file mode 100644 index 00000000000..c47790eaa31 --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/output.txt @@ -0,0 +1,52 @@ + +=== get (text) +>>> [CLI] experimental air get 123 + +╭─ Configuration ────────────────────────────────────────────────╮ +│ │ +│ experiment_name: my-exp │ +│ compute: │ +│ accelerator_type: a10 │ +│ num_accelerators: 1 │ +│ command: |- │ +│ for i in $(seq 1 10); do │ +│ echo "step $i" │ +│ done │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +╭─ Metadata ─────────────────────────────────────────────────────╮ +│ │ +│ Run ID 123 │ +│ Status ● SUCCESS │ +│ Submitted 2023-11-14 22:13 UTC │ +│ Retries 0 │ +│ Max Retries 3 │ +│ Duration 12s │ +│ Experiment my-exp │ +│ MLflow Run my-run │ +│ User user@example.com │ +│ Accelerators 1x A10 │ +│ Environment N/A │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +Run URL: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] +MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 + +=== get (json) +>>> [CLI] experimental air get 123 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "run_id": "123", + "status": "SUCCESS", + "started_at": "[TIMESTAMP]", + "duration_seconds": 12, + "attempt_number": 0, + "experiment_name": "my-exp", + "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[NUMID]", + "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0" + } +} diff --git a/acceptance/experimental/air/get-ai-runtime/script b/acceptance/experimental/air/get-ai-runtime/script new file mode 100644 index 00000000000..616e0e3087d --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/script @@ -0,0 +1,9 @@ +# Seed the run's training_config.yaml next to command.sh so `air get` can +# download and render it in the Configuration box. +$CLI workspace import "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc/training_config.yaml" --file training_config.yaml --format AUTO &> LOG.import + +title "get (text)" +trace $CLI experimental air get 123 + +title "get (json)" +trace $CLI experimental air get 123 -o json diff --git a/acceptance/experimental/air/get-ai-runtime/test.toml b/acceptance/experimental/air/get-ai-runtime/test.toml new file mode 100644 index 00000000000..c344dbb7af6 --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/test.toml @@ -0,0 +1,55 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# The typed SDK GetRun response: an ai_runtime_task run has no gen_ai_compute_task, +# so the task comes back empty (the SDK has no field for ai_runtime_task). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "run_page_url": "https://my-workspace.cloud.databricks.test/jobs/runs/123", + "creator_user_name": "user@example.com", + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [ + { + "task_key": "train", + "run_id": 456, + "attempt_number": 0, + "max_retries": 3, + "ai_runtime_task": { + "experiment": "my-exp", + "deployments": [ + { + "command_path": "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc/command.sh", + "compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1} + } + ] + } + } + ] +} +''' + +# MLflow identifiers for the deep-link (runs/get-output is not modeled by the typed SDK). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = ''' +{"gen_ai_compute_output": {"run_info": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}} +''' + +# The MLflow Run cell shows the run's name, fetched from the MLflow REST API. +[[Server]] +Pattern = "GET /api/2.0/mlflow/runs/get" +Response.Body = ''' +{"run": {"info": {"run_name": "my-run"}}} +''' diff --git a/acceptance/experimental/air/get-ai-runtime/training_config.yaml b/acceptance/experimental/air/get-ai-runtime/training_config.yaml new file mode 100644 index 00000000000..5f6060ecbbc --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/training_config.yaml @@ -0,0 +1,8 @@ +experiment_name: my-exp +compute: + accelerator_type: a10 + num_accelerators: 1 +command: |- + for i in $(seq 1 10); do + echo "step $i" + done diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 0da1f082e12..a6b47c977f7 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -1,14 +1,18 @@ package aircmd import ( + "context" "errors" "fmt" + "path" "strconv" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/spf13/cobra" ) @@ -38,6 +42,8 @@ type getData struct { AcceleratorsDisplay string `json:"-"` EnvironmentDisplay string `json:"-"` MaxRetriesDisplay string `json:"-"` + // TrainingConfigPath is the run's config file, downloaded for the config box. + TrainingConfigPath string `json:"-"` // Sweep replaces the single-run view for foreach runs. Sweep *sweepInfo `json:"-"` } @@ -63,6 +69,26 @@ Sweep Tasks: {{- end}} ` +// errNoProfile is the actionable message shown when no credentials are +// configured: no default profile, no --profile (-p), and no auth environment. +var errNoProfile = errors.New("no default profile is set: pass --profile (-p) or configure a default profile in your .databrickscfg") + +// authError classifies a workspace-client or Me() probe failure. Only genuinely +// auth-shaped errors surface as UNAUTHENTICATED/PERMANENT: missing profile, +// SDK auth wrappers, or an API 401/403. Anything else (network blip, 429, 5xx) +// is transient and reported as INTERNAL_ERROR/TRANSIENT so the caller can retry. +func authError(ctx context.Context, cmd *cobra.Command, err error) error { + if errors.Is(err, config.ErrCannotConfigureDefault) { + return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false, errNoProfile) + } + if errors.Is(err, apierr.ErrUnauthenticated) || errors.Is(err, apierr.ErrPermissionDenied) { + return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false, + fmt.Errorf("authentication was not successful: %w", err)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to verify authentication: %w", err)) +} + // newGetCommand returns the `air get JOB_RUN_ID` command, which shows status, // configuration, and timing details for a specific run. func newGetCommand() *cobra.Command { @@ -75,14 +101,16 @@ func newGetCommand() *cobra.Command { }, } - // Match Python: a client/auth failure is a JSON error envelope in -o json mode, - // not a bare error. ErrAlreadyPrinted passes through (it was handled upstream). + // Resolve and authenticate the workspace client up front so an auth failure + // fails fast here, before any run status or config is fetched or printed. + // ErrAlreadyPrinted passes through (it was handled upstream); other failures + // become an actionable auth error (JSON envelope in -o json mode). cmd.PreRunE = func(cmd *cobra.Command, args []string) error { err := root.MustWorkspaceClient(cmd, args) if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { return err } - return renderError(cmd.Context(), cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + return authError(cmd.Context(), cmd, err) } cmd.RunE = func(cmd *cobra.Command, args []string) error { @@ -95,7 +123,19 @@ func newGetCommand() *cobra.Command { fmt.Errorf("invalid JOB_RUN_ID %q: must be a positive integer", args[0])) } - run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + // Validate authentication against the workspace before fetching or + // rendering anything. MustWorkspaceClient's Config.Authenticate only + // attaches credentials (e.g. it does not check a PAT server-side), so + // without this a bad credential would surface as a confusing failure + // mid-render instead of a clear "not authenticated" error here. + if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil { + return authError(ctx, cmd, err) + } + + // Fetch the run once, in both the typed and raw shapes: the typed jobs.Run + // drives the display path, and the raw jobRun preserves the ai_runtime_task + // the typed model drops (used below without a second roundtrip). + run, rawRun, err := fetchRun(ctx, w, runID) if err != nil { // The backend returns this when the run ID is unknown to the user. if errors.Is(err, apierr.ErrResourceDoesNotExist) { @@ -121,6 +161,10 @@ func newGetCommand() *cobra.Command { } if task := findForEachTask(run); task != nil { data.Sweep = buildSweepInfo(ctx, w, task) + } else if genAIComputeTask(run) == nil { + // The typed SDK drops ai_runtime_task, so read it from the raw run we + // already fetched above. + enrichFromRawRun(rawRun, &data) } if root.OutputType(cmd) != flags.OutputText { @@ -142,6 +186,22 @@ func newGetCommand() *cobra.Command { return cmd } +// enrichFromRawRun fills the config path, experiment, and accelerators from the +// raw run (the ai_runtime_task the typed model drops). Best-effort: empty fields +// leave the existing "N/A" fallbacks in place. +func enrichFromRawRun(raw *jobRun, data *getData) { + if cmdPath := raw.commandPath(); cmdPath != "" { + data.TrainingConfigPath = path.Join(path.Dir(cmdPath), trainingConfigName) + } + if exp := jobExperiment(raw); exp != "" { + data.ExperimentName = &exp + data.ExperimentDisplay = exp + } + if a := acceleratorLabel(jobCompute(raw)); a != "" { + data.AcceleratorsDisplay = a + } +} + // buildGetData extracts the fields we display from a run. The text-view cells // are pre-rendered here with their "N/A" fallbacks; the styled renderer adds the // hyperlinks and colors once the dashboard and MLflow identifiers are known. diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go index 8ff864e387c..88a7bdfc4bd 100644 --- a/experimental/air/cmd/get_test.go +++ b/experimental/air/cmd/get_test.go @@ -3,6 +3,9 @@ package aircmd import ( "bytes" "encoding/json" + "errors" + "net/http" + "net/http/httptest" "testing" "text/template" @@ -11,6 +14,7 @@ import ( "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/experimental/mocks" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/stretchr/testify/assert" @@ -53,25 +57,92 @@ func TestGetRunInvalidID(t *testing.T) { assert.Contains(t, err.Error(), "invalid JOB_RUN_ID") } +// notFoundGetServer serves the auth probe plus a runs/get that reports the run +// as missing (400 INVALID_PARAMETER_VALUE, which the SDK maps to +// ErrResourceDoesNotExist for this path). +func notFoundGetServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jobsRunsGetPath { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error_code":"INVALID_PARAMETER_VALUE","message":"Run 5 does not exist."}`)) + return + } + // Me() probe and any other config discovery. + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + })) + t.Cleanup(srv.Close) + return srv +} + func TestGetRunNotFound(t *testing.T) { + srv := notFoundGetServer(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(newGetCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"5"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "run 5 not found") +} + +func TestGetRunAuthFailed(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) - m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return( - nil, apierr.ErrResourceDoesNotExist) + // A genuine auth failure (permission denied) is validated before the run is + // fetched, so GetRun is never reached and nothing is rendered. + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, apierr.ErrPermissionDenied) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) cmd := withOutput(newGetCommand(), flags.OutputText) cmd.SetContext(ctx) err := cmd.RunE(cmd, []string{"5"}) require.Error(t, err) - assert.Contains(t, err.Error(), "run 5 not found") + assert.Contains(t, err.Error(), "authentication was not successful") +} + +func TestGetRunAuthTransient(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + // A transient failure at the auth probe must not be misreported as an auth + // error; it surfaces as a retryable internal error instead. + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, errors.New("connection reset")) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newGetCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"5"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to verify authentication") + assert.NotContains(t, err.Error(), "authentication was not successful") +} + +func TestAuthError(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + cmd := withOutput(newGetCommand(), flags.OutputText) + + // No configurable credentials maps to the missing-profile hint. + noProfile := authError(ctx, cmd, config.ErrCannotConfigureDefault) + require.Error(t, noProfile) + assert.Contains(t, noProfile.Error(), "no default profile is set") + + // A 401 / 403 (via the SDK sentinels) is a real auth failure. + unauth := authError(ctx, cmd, apierr.ErrUnauthenticated) + require.Error(t, unauth) + assert.Contains(t, unauth.Error(), "authentication was not successful") + + denied := authError(ctx, cmd, apierr.ErrPermissionDenied) + require.Error(t, denied) + assert.Contains(t, denied.Error(), "authentication was not successful") + + transient := authError(ctx, cmd, errors.New("connection reset")) + require.Error(t, transient) + assert.Contains(t, transient.Error(), "failed to verify authentication") + assert.Contains(t, transient.Error(), "connection reset") } func TestGetRunNotFoundJSON(t *testing.T) { var buf bytes.Buffer - m := mocks.NewMockWorkspaceClient(t) - m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 5}).Return( - nil, apierr.ErrResourceDoesNotExist) - ctx := cmdctx.SetWorkspaceClient(t.Context(), m.WorkspaceClient) + srv := notFoundGetServer(t) + ctx := cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL)) ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) cmd := withOutput(newGetCommand(), flags.OutputJSON) cmd.SetContext(ctx) @@ -149,6 +220,63 @@ func TestBuildGetData(t *testing.T) { assert.Equal(t, int64(12), *d.DurationSeconds) } +func TestEnrichFromRawRun(t *testing.T) { + rawRun := func(t *testing.T, body string) *jobRun { + t.Helper() + var r jobRun + require.NoError(t, json.Unmarshal([]byte(body), &r)) + return &r + } + + t.Run("fills config path, experiment, and accelerators", func(t *testing.T) { + raw := rawRun(t, `{"run_id":5,"tasks":[{"ai_runtime_task":{ + "experiment":"/Users/me@example.com/my-exp", + "deployments":[{"command_path":"/Workspace/run/command.sh","compute":{"accelerator_type":"GPU_1xA10","accelerator_count":1}}] + }}]}`) + data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na} + enrichFromRawRun(raw, data) + assert.Equal(t, "/Workspace/run/training_config.yaml", data.TrainingConfigPath) + assert.Equal(t, "my-exp", data.ExperimentDisplay) + require.NotNil(t, data.ExperimentName) + assert.Equal(t, "my-exp", *data.ExperimentName) + assert.Equal(t, "1x A10", data.AcceleratorsDisplay) + }) + + t.Run("leaves fallbacks when the run has no ai_runtime_task", func(t *testing.T) { + raw := rawRun(t, `{"run_id":5,"tasks":[{}]}`) + data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na} + enrichFromRawRun(raw, data) + assert.Empty(t, data.TrainingConfigPath) + assert.Equal(t, na, data.ExperimentDisplay) + assert.Equal(t, na, data.AcceleratorsDisplay) + }) +} + +func TestFetchRun(t *testing.T) { + t.Run("parses the run into both the typed and raw shapes from one call", func(t *testing.T) { + body := `{"run_id":5,"creator_user_name":"me@example.com","tasks":[{"ai_runtime_task":{ + "experiment":"/Users/me@example.com/my-exp", + "deployments":[{"command_path":"/Workspace/run/command.sh","compute":{"accelerator_type":"GPU_1xA10","accelerator_count":1}}] + }}]}` + srv := runGetServer(t, body) + typed, raw, err := fetchRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5) + require.NoError(t, err) + // Typed jobs.Run carries the common fields; raw jobRun keeps ai_runtime_task. + assert.Equal(t, int64(5), typed.RunId) + assert.Equal(t, "me@example.com", typed.CreatorUserName) + assert.Equal(t, "/Workspace/run/command.sh", raw.commandPath()) + }) + + t.Run("propagates a fetch failure", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + _, _, err := fetchRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5) + require.Error(t, err) + }) +} + func TestBuildGetDataEmpty(t *testing.T) { // A run with no tasks, creator, or timing renders every text cell as "N/A". d := buildGetData(&jobs.Run{RunId: 7}) diff --git a/experimental/air/cmd/joblist.go b/experimental/air/cmd/joblist.go index c84f87c9df0..c6f94e1a69f 100644 --- a/experimental/air/cmd/joblist.go +++ b/experimental/air/cmd/joblist.go @@ -2,6 +2,7 @@ package aircmd import ( "context" + "encoding/json" "errors" "fmt" "net/http" @@ -9,6 +10,7 @@ import ( "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" "golang.org/x/sync/errgroup" ) @@ -66,6 +68,8 @@ type jobAiRuntimeTask struct { type aiRuntimeDeploy struct { Compute airCompute `json:"compute"` + // CommandPath is command.sh; training_config.yaml sits beside it. + CommandPath string `json:"command_path"` } type airCompute struct { @@ -192,7 +196,8 @@ func fetchJobRunsPage(ctx context.Context, w *databricks.WorkspaceClient, query } // fetchJobRun fetches a single run via runs/get. The response mirrors one -// runs/list element, so it deserializes into jobRun directly. +// runs/list element, so it deserializes into jobRun directly; expand_tasks +// pulls the ai_runtime_task the typed SDK omits. func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (*jobRun, error) { apiClient, err := client.New(w.Config) if err != nil { @@ -203,11 +208,38 @@ func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64 query := map[string]any{"run_id": runID, "expand_tasks": true} err = apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &run) if err != nil { - return nil, err + return nil, fmt.Errorf("failed to get run %d: %w", runID, err) } return &run, nil } +// fetchRun fetches a run once and returns it in both shapes: the typed SDK +// jobs.Run the display path consumes, and the raw jobRun that preserves the +// ai_runtime_task the typed model drops. The single runs/get body is unmarshalled +// into both, avoiding a second identical roundtrip. +func fetchRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (*jobs.Run, *jobRun, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, nil, fmt.Errorf("failed to create API client: %w", err) + } + + var body json.RawMessage + query := map[string]any{"run_id": runID, "expand_tasks": true} + if err := apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &body); err != nil { + return nil, nil, err + } + + var typed jobs.Run + if err := json.Unmarshal(body, &typed); err != nil { + return nil, nil, fmt.Errorf("failed to parse run %d: %w", runID, err) + } + var raw jobRun + if err := json.Unmarshal(body, &raw); err != nil { + return nil, nil, fmt.Errorf("failed to parse run %d: %w", runID, err) + } + return &typed, &raw, nil +} + // hydrateJobRuns fetches the given run ids concurrently via runs/get, preserving // input order. runs/get enforces per-run view ACLs, so an id the caller can't // view (403) or that has been purged (404) is dropped; any other error is @@ -242,3 +274,12 @@ func hydrateJobRuns(ctx context.Context, w *databricks.WorkspaceClient, ids []in } return hydrated, nil } + +// commandPath returns the run's command.sh path, or "" when it has no deployment. +func (r *jobRun) commandPath() string { + ai, _ := r.airTask() + if ai == nil || len(ai.Deployments) == 0 { + return "" + } + return ai.Deployments[0].CommandPath +} diff --git a/experimental/air/cmd/joblist_test.go b/experimental/air/cmd/joblist_test.go new file mode 100644 index 00000000000..d8f03fc60ba --- /dev/null +++ b/experimental/air/cmd/joblist_test.go @@ -0,0 +1,50 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// runGetServer serves one runs/get body and a stub for the SDK config probe. +func runGetServer(t *testing.T, body string) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jobsRunsGetPath { + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestFetchJobRunParsesAiRuntimeTask(t *testing.T) { + body := `{ + "run_id": 5, + "tasks": [{ + "ai_runtime_task": { + "experiment": "my-exp", + "deployments": [{ + "command_path": "/Workspace/Users/me/.air/cli_launch/my-exp/my-exp_abc/command.sh", + "compute": {"accelerator_count": 1, "accelerator_type": "GPU_1xA10"} + }] + } + }] + }` + srv := runGetServer(t, body) + + run, err := fetchJobRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5) + require.NoError(t, err) + assert.Equal(t, "/Workspace/Users/me/.air/cli_launch/my-exp/my-exp_abc/command.sh", run.commandPath()) +} + +func TestCommandPathEmpty(t *testing.T) { + // No ai_runtime_task deployment: no command path. + assert.Empty(t, (&jobRun{Tasks: []jobTask{{}}}).commandPath()) + assert.Empty(t, (&jobRun{}).commandPath()) +} diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go index ff7c0298671..3d1fe9637c5 100644 --- a/experimental/air/cmd/render.go +++ b/experimental/air/cmd/render.go @@ -115,10 +115,8 @@ func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceCl } var sections []string - if task := genAIComputeTask(run); task != nil { - if body := colorizeConfig(p, configYAML(ctx, w, task)); body != "" { - sections = append(sections, renderBox(p, configBoxTitle, body)) - } + if body := colorizeConfig(p, resolveConfigYAML(ctx, w, run, data)); body != "" { + sections = append(sections, renderBox(p, configBoxTitle, body)) } sections = append(sections, renderBox(p, metadataBoxTitle, renderFields(p, colorOn, view))) @@ -148,6 +146,20 @@ func genAIComputeTask(run *jobs.Run) *jobs.GenAiComputeTask { return run.Tasks[0].GenAiComputeTask } +// resolveConfigYAML returns the config box body: from the downloaded config file +// when we have its path, else from the legacy task. +func resolveConfigYAML(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run, data *getData) string { + if data.TrainingConfigPath != "" { + if raw := downloadConfig(ctx, w, data.TrainingConfigPath); len(raw) > 0 { + return strings.TrimRight(reformatYAMLForDisplay(raw), "\n") + } + } + if task := genAIComputeTask(run); task != nil { + return configYAML(ctx, w, task) + } + return "" +} + // configYAML returns the run's resolved training config as YAML for the box. The // full config (including the command/script) lives in the run's parameters, not // the structured task fields, so we prefer the inline parameters, then the diff --git a/experimental/air/cmd/render_test.go b/experimental/air/cmd/render_test.go index 71e2732fe0f..e950bb22331 100644 --- a/experimental/air/cmd/render_test.go +++ b/experimental/air/cmd/render_test.go @@ -55,6 +55,33 @@ func TestConfigYAML(t *testing.T) { }) } +func TestResolveConfigYAML(t *testing.T) { + ctx := t.Context() + + t.Run("downloads the training config path for new runs", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockWorkspaceAPI().EXPECT(). + Download(mock.Anything, "/Workspace/run/training_config.yaml"). + Return(io.NopCloser(strings.NewReader("experiment_name: from-air\n")), nil) + run := &jobs.Run{Tasks: []jobs.RunTask{{}}} + data := &getData{TrainingConfigPath: "/Workspace/run/training_config.yaml"} + assert.Equal(t, "experiment_name: from-air", resolveConfigYAML(ctx, m.WorkspaceClient, run, data)) + }) + + t.Run("falls back to the legacy task when no path is set", func(t *testing.T) { + run := &jobs.Run{Tasks: []jobs.RunTask{{GenAiComputeTask: &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/exp", + }}}} + got := resolveConfigYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, run, &getData{}) + assert.Contains(t, got, "experiment_name: exp") + }) + + t.Run("empty when neither source is present", func(t *testing.T) { + run := &jobs.Run{Tasks: []jobs.RunTask{{}}} + assert.Empty(t, resolveConfigYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, run, &getData{})) + }) +} + func TestSynthConfigYAML(t *testing.T) { task := &jobs.GenAiComputeTask{ MlflowExperimentName: "/Users/me@example.com/stream-latency-test", From cff51063788ec57affd62d54acb8c6eb4890c5c6 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Thu, 9 Jul 2026 17:59:39 +0000 Subject: [PATCH 15/18] AIR CLI: mkdirs parent before workspace import in get-ai-runtime test The acceptance testserver's /workspace/import handler now returns 404 when the parent folder does not exist (matching the real API, which does not mkdir -p). The get-ai-runtime script imported training_config.yaml to a deep path without creating the parent, so `air get` exited 1 once merged with main. Create the containing folder first, matching the convention in other tests. Also drop the stale "Milestone 0" scaffolding comment on air's New(). Co-authored-by: Isaac --- acceptance/experimental/air/get-ai-runtime/script | 4 +++- experimental/air/cmd/air.go | 3 --- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/acceptance/experimental/air/get-ai-runtime/script b/acceptance/experimental/air/get-ai-runtime/script index 616e0e3087d..3f41b089cdd 100644 --- a/acceptance/experimental/air/get-ai-runtime/script +++ b/acceptance/experimental/air/get-ai-runtime/script @@ -1,5 +1,7 @@ # Seed the run's training_config.yaml next to command.sh so `air get` can -# download and render it in the Configuration box. +# download and render it in the Configuration box. The import API does not +# create missing parents, so mkdirs the containing folder first. +$CLI workspace mkdirs "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc" &> LOG.mkdirs $CLI workspace import "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc/training_config.yaml" --file training_config.yaml --format AUTO &> LOG.import title "get (text)" diff --git a/experimental/air/cmd/air.go b/experimental/air/cmd/air.go index 81ffb2dd346..fbf40a34b52 100644 --- a/experimental/air/cmd/air.go +++ b/experimental/air/cmd/air.go @@ -7,9 +7,6 @@ import ( ) // New returns the root command for the experimental AI runtime CLI. -// -// Milestone 0: scaffolds the command group with every subcommand registered as a -// stub (not yet implemented), pending the port from the Python `air` CLI. func New() *cobra.Command { cmd := &cobra.Command{ Use: "air", From 131a0f772edbfd311988e829bb5b7c7360f30fc7 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Thu, 9 Jul 2026 20:36:07 +0000 Subject: [PATCH 16/18] AIR CLI: filter cached rows in air list; fix runs/submit test handler shadowing The list cache key is {host, run id} and excludes the filter, so a row cached by one query would bypass a different query's task filter (e.g. an unfiltered list caches a run, then --filter experiment=foo returns it anyway). Persist the raw filter inputs alongside the row and run every cache hit through a shared matchesFields comparator, so cached rows are filtered without re-fetching and the live and cache paths can't drift. Also fix TestSubmitWorkload: main's testserver now registers a default jobs/runs/submit handler, and the router is first-wins, so the test's custom handler must be registered before AddDefaultHandlers to claim the route. Co-authored-by: Isaac --- experimental/air/cmd/list_cache.go | 44 +++++++++++++------------ experimental/air/cmd/list_cache_test.go | 12 ++++--- experimental/air/cmd/list_filter.go | 40 ++++++++++++++++------ experimental/air/cmd/list_index.go | 13 +++++--- experimental/air/cmd/runsubmit_test.go | 4 ++- 5 files changed, 72 insertions(+), 41 deletions(-) diff --git a/experimental/air/cmd/list_cache.go b/experimental/air/cmd/list_cache.go index fb4442a1127..6f0b95a97d3 100644 --- a/experimental/air/cmd/list_cache.go +++ b/experimental/air/cmd/list_cache.go @@ -26,19 +26,20 @@ type listCacheKey struct { // cachedRun is the persisted value: every listRow field (including the // table-only columns, which listRow tags json:"-" and so wouldn't survive a -// direct marshal) plus the submit time. +// direct marshal), the filter inputs, and the submit time. type cachedRun struct { - RunID string `json:"run_id"` - RunName string `json:"run_name"` - User string `json:"user"` - Status string `json:"status"` - StartedAt *string `json:"started_at"` - IsSweep bool `json:"is_sweep"` - Experiment string `json:"experiment"` - Duration string `json:"duration"` - MLflowURL string `json:"mlflow_url"` - Accelerators string `json:"accelerators"` - SubmitTimeMs int64 `json:"submit_time_ms"` + RunID string `json:"run_id"` + RunName string `json:"run_name"` + User string `json:"user"` + Status string `json:"status"` + StartedAt *string `json:"started_at"` + IsSweep bool `json:"is_sweep"` + Experiment string `json:"experiment"` + Duration string `json:"duration"` + MLflowURL string `json:"mlflow_url"` + Accelerators string `json:"accelerators"` + Fields filterFields `json:"filter_fields"` + SubmitTimeMs int64 `json:"submit_time_ms"` } func (c cachedRun) toRow() listRow { @@ -49,12 +50,12 @@ func (c cachedRun) toRow() listRow { } } -func cachedRunFromRow(r listRow, submitTimeMs int64) cachedRun { +func cachedRunFromRow(r listRow, fields filterFields, submitTimeMs int64) cachedRun { return cachedRun{ RunID: r.RunID, RunName: r.RunName, User: r.User, Status: r.Status, StartedAt: r.StartedAt, IsSweep: r.IsSweep, Experiment: r.Experiment, Duration: r.Duration, MLflowURL: r.MLflowURL, Accelerators: r.Accelerators, - SubmitTimeMs: submitTimeMs, + Fields: fields, SubmitTimeMs: submitTimeMs, } } @@ -64,16 +65,17 @@ func newListCache(ctx context.Context) *cache.Cache { return cache.NewCache(ctx, listCacheComponent, listCacheTTL, nil) } -// cachedRow returns the cached row for a run, or (zero, false) on miss. -func cachedRow(ctx context.Context, c *cache.Cache, host string, runID int64) (listRow, bool) { +// cachedRow returns the cached row and its filter fields for a run, or +// (zero, zero, false) on miss. +func cachedRow(ctx context.Context, c *cache.Cache, host string, runID int64) (listRow, filterFields, bool) { entry, ok := cache.Get[cachedRun](ctx, c, listCacheKey{Host: host, RunID: runID}) if !ok { - return listRow{}, false + return listRow{}, filterFields{}, false } - return entry.toRow(), true + return entry.toRow(), entry.Fields, true } -// putRow caches a terminal run's finished row under its submit time. -func putRow(ctx context.Context, c *cache.Cache, host string, runID, submitTimeMs int64, row listRow) { - cache.Put(ctx, c, listCacheKey{Host: host, RunID: runID}, cachedRunFromRow(row, submitTimeMs)) +// putRow caches a terminal run's finished row and filter fields under its submit time. +func putRow(ctx context.Context, c *cache.Cache, host string, runID, submitTimeMs int64, row listRow, fields filterFields) { + cache.Put(ctx, c, listCacheKey{Host: host, RunID: runID}, cachedRunFromRow(row, fields, submitTimeMs)) } diff --git a/experimental/air/cmd/list_cache_test.go b/experimental/air/cmd/list_cache_test.go index 3ac695e8d0d..f6b50d4cd09 100644 --- a/experimental/air/cmd/list_cache_test.go +++ b/experimental/air/cmd/list_cache_test.go @@ -12,18 +12,20 @@ func TestListCacheRoundTrip(t *testing.T) { ctx := t.Context() c := newListCache(ctx) - _, ok := cachedRow(ctx, c, "https://host.test", 42) + _, _, ok := cachedRow(ctx, c, "https://host.test", 42) require.False(t, ok, "miss before write") row := listRow{RunID: "42", Experiment: "exp", Status: "SUCCESS"} - putRow(ctx, c, "https://host.test", 42, 1700000000000, row) + fields := filterFields{Experiment: "exp", GPUType: "GPU_1xA10", GPUCount: 1} + putRow(ctx, c, "https://host.test", 42, 1700000000000, row, fields) - got, ok := cachedRow(ctx, c, "https://host.test", 42) + got, gotFields, ok := cachedRow(ctx, c, "https://host.test", 42) require.True(t, ok, "hit after write") assert.Equal(t, row, got) + assert.Equal(t, fields, gotFields) // Different host is a different key. - _, ok = cachedRow(ctx, c, "https://other.test", 42) + _, _, ok = cachedRow(ctx, c, "https://other.test", 42) assert.False(t, ok) } @@ -36,7 +38,7 @@ func TestIndexStrategyServesCachedRowWithoutFetch(t *testing.T) { // Pre-seed the cache for run 7 so hydration should skip runs/get entirely. ctx := t.Context() - putRow(ctx, newListCache(ctx), host, 7, 1000_000, listRow{RunID: "7", Status: "SUCCESS"}) + putRow(ctx, newListCache(ctx), host, 7, 1000_000, listRow{RunID: "7", Status: "SUCCESS"}, filterFields{}) f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, diff --git a/experimental/air/cmd/list_filter.go b/experimental/air/cmd/list_filter.go index 272141e80d1..b1cedcb0b99 100644 --- a/experimental/air/cmd/list_filter.go +++ b/experimental/air/cmd/list_filter.go @@ -58,29 +58,49 @@ func parseListFilters(raw []string) (listFilters, error) { return f, nil } +// filterFields are the task-filter inputs, cached alongside the row so a cache +// hit can be re-filtered without re-fetching the run. +type filterFields struct { + Experiment string `json:"experiment"` + GPUType string `json:"gpu_type"` + GPUCount int `json:"gpu_count"` +} + +func filterFieldsFromRun(run *jobRun) filterFields { + gpuType, count := jobCompute(run) + return filterFields{ + Experiment: jobExperiment(run), + GPUType: gpuType, + GPUCount: count, + } +} + // matches reports whether a run satisfies the experiment, accelerator-type and // accelerator-count filters. The user filter is applied separately while // scanning, since it maps onto the run's creator rather than its task. func (f listFilters) matches(run *jobRun) bool { + return f.matchesFields(filterFieldsFromRun(run)) +} + +// matchesFields is the shared comparator for live runs and cached rows, so the +// two paths can't drift. +func (f listFilters) matchesFields(fields filterFields) bool { if f.Experiment != "" { - matched, err := path.Match(strings.ToLower(f.Experiment), strings.ToLower(jobExperiment(run))) + matched, err := path.Match(strings.ToLower(f.Experiment), strings.ToLower(fields.Experiment)) if err != nil || !matched { return false } } - if f.AcceleratorType != "" || f.NumAccelerators != nil { - gpuType, count := jobCompute(run) - if f.AcceleratorType != "" { - display := strings.ToLower(gpuDisplayName(gpuType)) - if !strings.Contains(display, strings.ToLower(f.AcceleratorType)) { - return false - } - } - if f.NumAccelerators != nil && count != *f.NumAccelerators { + if f.AcceleratorType != "" { + display := strings.ToLower(gpuDisplayName(fields.GPUType)) + if !strings.Contains(display, strings.ToLower(f.AcceleratorType)) { return false } } + if f.NumAccelerators != nil && fields.GPUCount != *f.NumAccelerators { + return false + } return true } diff --git a/experimental/air/cmd/list_index.go b/experimental/air/cmd/list_index.go index 097859ba402..d6c6cf6958f 100644 --- a/experimental/air/cmd/list_index.go +++ b/experimental/air/cmd/list_index.go @@ -99,9 +99,13 @@ func (s *indexStrategy) hydrate(ids []int64) ([]listedRun, error) { rows := make([]listedRun, 0, len(ids)) var toFetch []int64 + // The cache key excludes the filter, so a cached row is still run through the + // active filter; a non-matching hit is dropped, not re-fetched. for _, id := range ids { - if row, ok := cachedRow(s.ctx, s.cache, host, id); ok { - rows = append(rows, listedRun{row: row, taskRunID: id}) + if row, fields, ok := cachedRow(s.ctx, s.cache, host, id); ok { + if s.filters.matchesFields(fields) { + rows = append(rows, listedRun{row: row, taskRunID: id}) + } continue } toFetch = append(toFetch, id) @@ -112,14 +116,15 @@ func (s *indexStrategy) hydrate(ids []int64) ([]listedRun, error) { return nil, err } for _, run := range runs { - if !s.filters.matches(run) { + fields := filterFieldsFromRun(run) + if !s.filters.matchesFields(fields) { continue } row := buildListRow(run) rows = append(rows, listedRun{row: row, taskRunID: taskRunID(run)}) if isTerminal(run) { start, _ := jobTiming(run) - putRow(s.ctx, s.cache, host, run.RunID, start, row) + putRow(s.ctx, s.cache, host, run.RunID, start, row, fields) } } diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index d57b14f80a6..374b584bc63 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -137,13 +137,15 @@ func TestJobsSubmitClient(t *testing.T) { func TestSubmitWorkload(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) - testserver.AddDefaultHandlers(server) + // Register before AddDefaultHandlers: the router is first-wins, so this must claim the route ahead of the default handler. var got jobsSubmitRun server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { require.NoError(t, json.Unmarshal(req.Body, &got)) return submitRunResponse{RunID: 777} }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) require.NoError(t, err) From 78f916d35a19183377d60ba8636a55064e94808a Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Thu, 9 Jul 2026 22:02:34 +0000 Subject: [PATCH 17/18] AIR CLI: test that air list filters cached rows Add coverage for the cache-filter fix: a cached row that doesn't match the active task filter is dropped (not returned) and not re-fetched, and a matching cached row is still served from cache. The first test fails against the prior unfiltered cache-hit path, guarding the regression. Co-authored-by: Isaac --- experimental/air/cmd/list_cache_test.go | 46 +++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/experimental/air/cmd/list_cache_test.go b/experimental/air/cmd/list_cache_test.go index f6b50d4cd09..f03dd7bc29f 100644 --- a/experimental/air/cmd/list_cache_test.go +++ b/experimental/air/cmd/list_cache_test.go @@ -50,6 +50,52 @@ func TestIndexStrategyServesCachedRowWithoutFetch(t *testing.T) { assert.Equal(t, 0, hits.get, "cached run must not hit runs/get") } +func TestIndexStrategyFiltersCachedRow(t *testing.T) { + t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) + + refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} + srv, hits := indexAndGetServer(t, refs, map[int64]jobRun{7: indexRun(7, 1000_000)}, nil, nil) + host := srv.URL + + // Cache hit under experiment "bar" must be filtered out by an experiment=foo query. + ctx := t.Context() + putRow(ctx, newListCache(ctx), host, 7, 1000_000, + listRow{RunID: "7", Status: "SUCCESS", Experiment: "bar"}, + filterFields{Experiment: "bar"}) + + f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + filters: listFilters{Experiment: "foo"}, + }) + rows, err := f.next(10) + require.NoError(t, err) + assert.Empty(t, rows, "cached row not matching the filter must be dropped") + assert.Equal(t, 0, hits.get, "non-matching cached row must not hit runs/get") +} + +func TestIndexStrategyServesMatchingCachedRow(t *testing.T) { + t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) + + refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} + srv, hits := indexAndGetServer(t, refs, map[int64]jobRun{7: indexRun(7, 1000_000)}, nil, nil) + host := srv.URL + + ctx := t.Context() + putRow(ctx, newListCache(ctx), host, 7, 1000_000, + listRow{RunID: "7", Status: "SUCCESS", Experiment: "foo"}, + filterFields{Experiment: "foo"}) + + f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + filters: listFilters{Experiment: "foo"}, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 1) + assert.Equal(t, "7", rows[0].RunID) + assert.Equal(t, 0, hits.get, "matching cached row must not hit runs/get") +} + func TestIsTerminal(t *testing.T) { assert.True(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "TERMINATED"}})) assert.True(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "INTERNAL_ERROR"}})) From eaf24a85acee2d5188df603dafe577e5f08b9588 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db <riddhi.bhagwat@databricks.com> Date: Thu, 9 Jul 2026 22:35:09 +0000 Subject: [PATCH 18/18] AIR CLI: disable MSYS path conversion in get-ai-runtime test On Windows, Git Bash rewrites the leading-/ workspace paths passed to `workspace mkdirs` and `workspace import` into C:/... paths, so mkdirs creates a mangled directory and the import 404s on the missing parent. Set MSYS_NO_PATHCONV=1 for this test, matching the export-dir precedent. Co-authored-by: Isaac --- acceptance/experimental/air/get-ai-runtime/test.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/acceptance/experimental/air/get-ai-runtime/test.toml b/acceptance/experimental/air/get-ai-runtime/test.toml index c344dbb7af6..ad0c25bb7f2 100644 --- a/acceptance/experimental/air/get-ai-runtime/test.toml +++ b/acceptance/experimental/air/get-ai-runtime/test.toml @@ -2,6 +2,11 @@ [EnvMatrix] DATABRICKS_BUNDLE_ENGINE = [] +# On Windows, Git Bash rewrites the leading-/ workspace paths passed to +# `workspace mkdirs`/`import` into C:/... paths; disable that conversion. +[Env] +MSYS_NO_PATHCONV = "1" + # The SDK occasionally probes host reachability with a HEAD request; stub it so # the test is deterministic. [[Server]]