From 48dc9f56368f2876bbfed5f1cfc23a540ae95937 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Tue, 9 Jun 2026 22:09:39 +0000 Subject: [PATCH 1/2] tools: make list_embeds.py compatible with Python 3.6 subprocess's text= keyword was added in 3.7; universal_newlines is the exact equivalent and also works on older interpreters. Unblocks running the acceptance suite on a Python 3.6 host. Co-authored-by: Isaac --- tools/list_embeds.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/list_embeds.py b/tools/list_embeds.py index ba534aa47e5..aa672be3afb 100755 --- a/tools/list_embeds.py +++ b/tools/list_embeds.py @@ -20,7 +20,7 @@ def main(): # files are silently missed until the file is staged or committed. out = subprocess.check_output( ["git", "grep", "--no-color", "-E", "^" + EMBED, "--", "*.go"], - text=True, + universal_newlines=True, ) paths = set() for line in out.splitlines(): From 0474207e14a8348d4cc8c92f8acb065f63fe1ce1 Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Thu, 11 Jun 2026 18:52:04 +0000 Subject: [PATCH 2/2] cmd/experimental/air: scaffold AI runtime CLI command package Add the experimental `air` command group as the Go port surface for the Python `air` CLI. Every subcommand (run, status, list, logs, cancel, register-image) is registered as a stub that returns a not-implemented error; the real implementations land in later milestones. Includes unit tests for the command-tree wiring and the not-implemented stubs, plus an acceptance test exercising the stubs end-to-end. Co-authored-by: Isaac --- .../air/unimplemented/out.test.toml | 3 ++ .../experimental/air/unimplemented/output.txt | 36 +++++++++++++++++++ .../experimental/air/unimplemented/script | 19 ++++++++++ .../experimental/air/unimplemented/test.toml | 6 ++++ cmd/experimental/air/air.go | 36 +++++++++++++++++++ cmd/experimental/air/air_test.go | 22 ++++++++++++ cmd/experimental/air/cancel.go | 27 ++++++++++++++ cmd/experimental/air/list.go | 31 ++++++++++++++++ cmd/experimental/air/logs.go | 34 ++++++++++++++++++ cmd/experimental/air/register_image.go | 33 +++++++++++++++++ cmd/experimental/air/run.go | 36 +++++++++++++++++++ cmd/experimental/air/status.go | 19 ++++++++++ cmd/experimental/air/stubs_test.go | 30 ++++++++++++++++ cmd/experimental/experimental.go | 2 ++ 14 files changed, 334 insertions(+) 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 cmd/experimental/air/air.go create mode 100644 cmd/experimental/air/air_test.go create mode 100644 cmd/experimental/air/cancel.go create mode 100644 cmd/experimental/air/list.go create mode 100644 cmd/experimental/air/logs.go create mode 100644 cmd/experimental/air/register_image.go create mode 100644 cmd/experimental/air/run.go create mode 100644 cmd/experimental/air/status.go create mode 100644 cmd/experimental/air/stubs_test.go 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..3dc88de3b77 --- /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 + +=== status +>>> [CLI] experimental air status 123 +Error: `air status` 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..83397b4b741 --- /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 "status" +errcode trace $CLI experimental air status 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..832b783573b --- /dev/null +++ b/acceptance/experimental/air/unimplemented/test.toml @@ -0,0 +1,6 @@ +Local = true +Cloud = false + +# Stubs fail locally before any API call, so no server stubs needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] diff --git a/cmd/experimental/air/air.go b/cmd/experimental/air/air.go new file mode 100644 index 00000000000..00283e95ecd --- /dev/null +++ b/cmd/experimental/air/air.go @@ -0,0 +1,36 @@ +package air + +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(newStatusCommand()) + 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/cmd/experimental/air/air_test.go b/cmd/experimental/air/air_test.go new file mode 100644 index 00000000000..7f0ecbb127a --- /dev/null +++ b/cmd/experimental/air/air_test.go @@ -0,0 +1,22 @@ +package air + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestNewRegistersAllSubcommands asserts the `ai` 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", "status", "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/cmd/experimental/air/cancel.go b/cmd/experimental/air/cancel.go new file mode 100644 index 00000000000..43e8a752c7d --- /dev/null +++ b/cmd/experimental/air/cancel.go @@ -0,0 +1,27 @@ +package air + +import ( + "github.com/spf13/cobra" +) + +func newCancelCommand() *cobra.Command { + var ( + all bool + yes bool + ) + + cmd := &cobra.Command{ + Use: "cancel [RUN_ID...]", + Args: cobra.ArbitraryArgs, + 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") + + return cmd +} diff --git a/cmd/experimental/air/list.go b/cmd/experimental/air/list.go new file mode 100644 index 00000000000..1310abbf5af --- /dev/null +++ b/cmd/experimental/air/list.go @@ -0,0 +1,31 @@ +package air + +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/cmd/experimental/air/logs.go b/cmd/experimental/air/logs.go new file mode 100644 index 00000000000..710b8dd515e --- /dev/null +++ b/cmd/experimental/air/logs.go @@ -0,0 +1,34 @@ +package air + +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 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 (default: 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") + + return cmd +} diff --git a/cmd/experimental/air/register_image.go b/cmd/experimental/air/register_image.go new file mode 100644 index 00000000000..36257e747e3 --- /dev/null +++ b/cmd/experimental/air/register_image.go @@ -0,0 +1,33 @@ +package air + +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().BoolVar(&interactiveAuth, "interactive-authenticate", 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/cmd/experimental/air/run.go b/cmd/experimental/air/run.go new file mode 100644 index 00000000000..b3e56dbca82 --- /dev/null +++ b/cmd/experimental/air/run.go @@ -0,0 +1,36 @@ +package air + +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/cmd/experimental/air/status.go b/cmd/experimental/air/status.go new file mode 100644 index 00000000000..25767e13103 --- /dev/null +++ b/cmd/experimental/air/status.go @@ -0,0 +1,19 @@ +package air + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newStatusCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "status RUN_ID", + Args: root.ExactArgs(1), + Short: "Show status and configuration for a run", + RunE: func(cmd *cobra.Command, args []string) error { + return notImplemented("status") + }, + } + + return cmd +} diff --git a/cmd/experimental/air/stubs_test.go b/cmd/experimental/air/stubs_test.go new file mode 100644 index 00000000000..1eec39fc3a2 --- /dev/null +++ b/cmd/experimental/air/stubs_test.go @@ -0,0 +1,30 @@ +package air + +import ( + "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(), + "status": newStatusCommand(), + "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.ErrorContains(t, err, "not implemented") + }) + } +} diff --git a/cmd/experimental/experimental.go b/cmd/experimental/experimental.go index 8d9827c5c94..bc39c9ab946 100644 --- a/cmd/experimental/experimental.go +++ b/cmd/experimental/experimental.go @@ -1,6 +1,7 @@ package experimental import ( + "github.com/databricks/cli/cmd/experimental/air" 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(air.New()) cmd.AddCommand(aitoolscmd.NewAitoolsCmd()) cmd.AddCommand(geniecmd.NewGenieCmd()) cmd.AddCommand(postgrescmd.New())