From c927fb35c52d353580b94a032261b98a07b6d02f Mon Sep 17 00:00:00 2001 From: Caroline Chen <324939130+caroline-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 19:32:31 +0000 Subject: [PATCH] air: make pool commands resource-first --- .../experimental/air/config-help/output.txt | 2 +- acceptance/experimental/air/help/output.txt | 32 +++++++++++++------ acceptance/experimental/air/help/script | 3 ++ acceptance/experimental/air/pools/output.txt | 8 ++--- acceptance/experimental/air/pools/script | 8 ++--- experimental/air/cmd/air.go | 1 + experimental/air/cmd/air_test.go | 2 +- experimental/air/cmd/compute.go | 2 +- experimental/air/cmd/get.go | 2 -- experimental/air/cmd/get_test.go | 7 ++-- experimental/air/cmd/list.go | 2 -- experimental/air/cmd/pools.go | 29 ++++++++++++----- experimental/air/cmd/pools_test.go | 17 ++++++---- 13 files changed, 71 insertions(+), 44 deletions(-) diff --git a/acceptance/experimental/air/config-help/output.txt b/acceptance/experimental/air/config-help/output.txt index dc313f32eed..a22e7cfd8de 100644 --- a/acceptance/experimental/air/config-help/output.txt +++ b/acceptance/experimental/air/config-help/output.txt @@ -80,7 +80,7 @@ config.mlflow_artifact_location >>> [CLI] experimental air run -h config.compute.pool_id config.compute.pool_id - GPU pool to run on, by id. A GPU pool is a pre-provisioned accelerator reservation; contact your Databricks account team to provision one. Must be 1-255 characters. List available pools with 'air list pools'. + GPU pool to run on, by id. A GPU pool is a pre-provisioned accelerator reservation; contact your Databricks account team to provision one. Must be 1-255 characters. List available pools with 'air pools list'. Type: string Required: no diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index 5928ff9e6b6..9f6561dc20f 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -14,6 +14,7 @@ Available Commands: get Show status, configuration, and timing details for a specified run list List active runs for the current user logs Stream or fetch logs for a run + pools View provisioned GPU pools run Submit a training workload from a YAML config Flags: @@ -41,7 +42,6 @@ Filters: Usage: databricks experimental air list [flags] - databricks experimental air list [command] Examples: databricks experimental air list --filter accelerator_type=H100 @@ -51,9 +51,6 @@ Examples: --filter num_accelerators=8 \ --filter user=person@example.com -Available Commands: - pools List the GPU pools available to the current workspace - Flags: --all-status Show runs in all states (default: active only) --all-users Show runs from all users @@ -67,8 +64,6 @@ Global Flags: -p, --profile string ~/.databrickscfg profile -t, --target string bundle target to use (if applicable) -Use "databricks experimental air list [command] --help" for more information about a command. - === get help >>> [CLI] experimental air get --help Show status, configuration, and timing details for a specified run. @@ -77,13 +72,30 @@ JOB_RUN_ID is the job run ID returned by "air run". Usage: databricks experimental air get JOB_RUN_ID [flags] - databricks experimental air get [command] + +Flags: + -h, --help help for get + +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) + +=== pools help +>>> [CLI] experimental air pools --help +View provisioned GPU pools + +Usage: + databricks experimental air pools [flags] + databricks experimental air pools [command] Available Commands: - pool Show a GPU pool, including its accelerator usage (POOL_ID may be omitted when the workspace has exactly one pool) + get Show a GPU pool, including its accelerator usage (POOL_ID may be omitted when the workspace has exactly one pool) + list List the GPU pools available to the current workspace Flags: - -h, --help help for get + -h, --help help for pools Global Flags: --debug enable debug logging @@ -91,7 +103,7 @@ Global Flags: -p, --profile string ~/.databrickscfg profile -t, --target string bundle target to use (if applicable) -Use "databricks experimental air get [command] --help" for more information about a command. +Use "databricks experimental air pools [command] --help" for more information about a command. === logs help >>> [CLI] experimental air logs --help diff --git a/acceptance/experimental/air/help/script b/acceptance/experimental/air/help/script index 3c088d278d0..a3706e6cd96 100644 --- a/acceptance/experimental/air/help/script +++ b/acceptance/experimental/air/help/script @@ -10,6 +10,9 @@ trace $CLI experimental air list --help title "get help" trace $CLI experimental air get --help +title "pools help" +trace $CLI experimental air pools --help + title "logs help" trace $CLI experimental air logs --help diff --git a/acceptance/experimental/air/pools/output.txt b/acceptance/experimental/air/pools/output.txt index 4df5949628c..b545ac1355e 100644 --- a/acceptance/experimental/air/pools/output.txt +++ b/acceptance/experimental/air/pools/output.txt @@ -1,12 +1,12 @@ === list pools (text) ->>> [CLI] experimental air list pools +>>> [CLI] experimental air pools list ID ACCELERATOR RESERVED pool-8xh100-alpha GPU_8xH100 64 pool-1xh100-beta GPU_1xH100 8 === list pools (json) ->>> [CLI] experimental air list pools -o json +>>> [CLI] experimental air pools list -o json { "v": 1, "ts": "[TIMESTAMP]", @@ -27,7 +27,7 @@ pool-1xh100-beta GPU_1xH100 8 } === get pool (text) ->>> [CLI] experimental air get pool pool-8xh100-alpha +>>> [CLI] experimental air pools get pool-8xh100-alpha Pool ID: pool-8xh100-alpha Accelerator Type: GPU_8xH100 Reserved Accelerators: 64 @@ -35,7 +35,7 @@ Used Accelerators: 40 Idle Accelerators: 24 === get pool (json) ->>> [CLI] experimental air get pool pool-8xh100-alpha -o json +>>> [CLI] experimental air pools get pool-8xh100-alpha -o json { "v": 1, "ts": "[TIMESTAMP]", diff --git a/acceptance/experimental/air/pools/script b/acceptance/experimental/air/pools/script index f7424db52ae..f0dfb0390dc 100644 --- a/acceptance/experimental/air/pools/script +++ b/acceptance/experimental/air/pools/script @@ -1,11 +1,11 @@ title "list pools (text)" -trace $CLI experimental air list pools +trace $CLI experimental air pools list title "list pools (json)" -trace $CLI experimental air list pools -o json +trace $CLI experimental air pools list -o json title "get pool (text)" -trace $CLI experimental air get pool pool-8xh100-alpha +trace $CLI experimental air pools get pool-8xh100-alpha title "get pool (json)" -trace $CLI experimental air get pool pool-8xh100-alpha -o json +trace $CLI experimental air pools get pool-8xh100-alpha -o json diff --git a/experimental/air/cmd/air.go b/experimental/air/cmd/air.go index 2bafd6150a9..bdfc2cc417f 100644 --- a/experimental/air/cmd/air.go +++ b/experimental/air/cmd/air.go @@ -22,6 +22,7 @@ These commands are experimental and may change in future versions.`, cmd.AddCommand(runCommand) cmd.AddCommand(newGetCommand()) cmd.AddCommand(newListCommand()) + cmd.AddCommand(newPoolsCommand()) cmd.AddCommand(newLogsCommand()) cmd.AddCommand(newCancelCommand()) cmd.AddCommand(newConvertToDabsCommand()) diff --git a/experimental/air/cmd/air_test.go b/experimental/air/cmd/air_test.go index fe602693b7d..754b13d43fc 100644 --- a/experimental/air/cmd/air_test.go +++ b/experimental/air/cmd/air_test.go @@ -16,7 +16,7 @@ func TestNewRegistersAllSubcommands(t *testing.T) { registered[c.Name()] = true } - want := []string{"run", "get", "list", "logs", "cancel", "convert-to-dabs"} + want := []string{"run", "get", "list", "pools", "logs", "cancel", "convert-to-dabs"} for _, name := range want { assert.True(t, registered[name], "subcommand %q is not registered", name) } diff --git a/experimental/air/cmd/compute.go b/experimental/air/cmd/compute.go index 0f302c4aaa2..1977eb24441 100644 --- a/experimental/air/cmd/compute.go +++ b/experimental/air/cmd/compute.go @@ -93,7 +93,7 @@ func gpusPerNode(g gpuType) (int, error) { type computeConfig struct { NumAccelerators int `yaml:"num_accelerators" help:"Total number of GPUs to allocate. Must be a positive multiple of the accelerator type's per-node GPU count. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for supported GPU types."` AcceleratorType string `yaml:"accelerator_type" help:"Which accelerator to run on, e.g. GPU_1xA10. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for the current list of supported GPU types. Matched case-sensitively."` - PoolID *string `yaml:"pool_id" help:"GPU pool to run on, by id. A GPU pool is a pre-provisioned accelerator reservation; contact your Databricks account team to provision one. Must be 1-255 characters. List available pools with 'air list pools'."` + PoolID *string `yaml:"pool_id" help:"GPU pool to run on, by id. A GPU pool is a pre-provisioned accelerator reservation; contact your Databricks account team to provision one. Must be 1-255 characters. List available pools with 'air pools list'."` PriorityClass *string `yaml:"priority_class" help:"Scheduling priority within the pool: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest). Requires pool_id."` } diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 956911b8ae4..9f55f0f60e8 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -208,8 +208,6 @@ JOB_RUN_ID is the job run ID returned by "air run".`, return nil } - cmd.AddCommand(newGetPoolCommand()) - return cmd } diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go index 3c1a97074d4..7df2cf5e613 100644 --- a/experimental/air/cmd/get_test.go +++ b/experimental/air/cmd/get_test.go @@ -35,14 +35,11 @@ func renderGet(t *testing.T, data getData) string { // TestGetCommandShape locks in that `get` takes the run id directly as // `air get JOB_RUN_ID` (there is no `run` subcommand — it was collapsed back -// into `get`). Its only subcommand is the `pool` noun. The -// acceptance test exercises the happy path end to end. +// 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) - subs := cmd.Commands() - require.Len(t, subs, 1) - assert.Equal(t, "pool", subs[0].Name()) + assert.Empty(t, cmd.Commands()) // ExactArgs(1): exactly one run id is required. assert.NoError(t, cmd.Args(cmd, []string{"123"})) assert.Error(t, cmd.Args(cmd, []string{})) diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index d88b421f608..b4b7c13916f 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -175,8 +175,6 @@ Filters: return renderListText(cmd, fetcher, limit) } - cmd.AddCommand(newListPoolsCommand()) - return cmd } diff --git a/experimental/air/cmd/pools.go b/experimental/air/cmd/pools.go index 1af7c60fbb4..c6bf6c7d9cb 100644 --- a/experimental/air/cmd/pools.go +++ b/experimental/air/cmd/pools.go @@ -84,8 +84,8 @@ type listProvisionedCapacitiesResponse struct { NextPageToken string `json:"next_page_token"` } -// poolListData is the `air list pools` payload. Usage counts are intentionally -// absent: the list endpoint does not populate them (they come from `air get`). +// poolListData is the `air pools list` payload. Usage counts are intentionally +// absent: the list endpoint does not populate them (they come from `air pools get`). type poolListData struct { Rows []poolRow `json:"pools"` } @@ -96,7 +96,7 @@ type poolRow struct { ReservedAccelerators int64 `json:"reserved_accelerators"` } -// poolDetailData is the `air get pool` payload. Usage is a pointer because it is +// poolDetailData is the `air pools get` payload. Usage is a pointer because it is // populated only when the pool reports it. type poolDetailData struct { ID string `json:"pool_id"` @@ -181,9 +181,22 @@ func getPool(ctx context.Context, w *databricks.WorkspaceClient, id string) (*pr return &pc, nil } -func newListPoolsCommand() *cobra.Command { +func newPoolsCommand() *cobra.Command { cmd := &cobra.Command{ Use: "pools", + Short: "View provisioned GPU pools", + RunE: root.ReportUnknownSubcommand, + } + + cmd.AddCommand(newGetPoolCommand()) + cmd.AddCommand(newListPoolsCommand()) + + return cmd +} + +func newListPoolsCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "list", Args: root.NoArgs, Short: "List the GPU pools available to the current workspace", } @@ -222,7 +235,7 @@ func newListPoolsCommand() *cobra.Command { func newGetPoolCommand() *cobra.Command { cmd := &cobra.Command{ - Use: "pool [POOL_ID]", + Use: "get [POOL_ID]", Args: root.MaximumNArgs(1), Short: "Show a GPU pool, including its accelerator usage (POOL_ID may be omitted when the workspace has exactly one pool)", } @@ -261,7 +274,7 @@ func newGetPoolCommand() *cobra.Command { // ErrNotFound covers a plain 404 as well as RESOURCE_DOES_NOT_EXIST. if errors.Is(err, apierr.ErrNotFound) { return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, - fmt.Errorf("GPU pool %q not found: check the id with `air list pools`", id)) + fmt.Errorf("GPU pool %q not found: check the id with `air pools list`", id)) } return poolAPIError(ctx, cmd, fmt.Sprintf("get GPU pool %q", id), err) } @@ -277,8 +290,8 @@ func newGetPoolCommand() *cobra.Command { return cmd } -// resolveSolePoolID returns the id of the workspace's only pool, for `air get -// pool` with no argument. It errors (with the ids to choose from) when there is +// resolveSolePoolID returns the id of the workspace's only pool, for `air pools +// get` with no argument. It errors (with the ids to choose from) when there is // not exactly one, so the convenience never silently picks among several. func resolveSolePoolID(ctx context.Context, cmd *cobra.Command, w *databricks.WorkspaceClient) (string, error) { pools, err := listPools(ctx, w) diff --git a/experimental/air/cmd/pools_test.go b/experimental/air/cmd/pools_test.go index 0eba5f8b161..0062c48f61a 100644 --- a/experimental/air/cmd/pools_test.go +++ b/experimental/air/cmd/pools_test.go @@ -20,20 +20,25 @@ import ( const poolsBasePath = "/api/2.0/ai-training/provisioned-capacities" func TestPoolsCommandShape(t *testing.T) { + pools := newPoolsCommand() + assert.Equal(t, "pools", pools.Use) + assert.True(t, hasSubcommand(pools, "get")) + assert.True(t, hasSubcommand(pools, "list")) + list := newListPoolsCommand() - assert.Equal(t, "pools", list.Use) + assert.Equal(t, "list", list.Use) assert.NoError(t, list.Args(list, []string{})) assert.Error(t, list.Args(list, []string{"x"})) get := newGetPoolCommand() - assert.Equal(t, "pool [POOL_ID]", get.Use) + assert.Equal(t, "get [POOL_ID]", get.Use) assert.NoError(t, get.Args(get, []string{"pool-1"})) assert.NoError(t, get.Args(get, []string{})) // id optional: resolved when there's one pool assert.Error(t, get.Args(get, []string{"a", "b"})) - // The subcommands are wired under `air list` and `air get`. - assert.True(t, hasSubcommand(newListCommand(), "pools")) - assert.True(t, hasSubcommand(newGetCommand(), "pool")) + assert.True(t, hasSubcommand(New(), "pools")) + assert.False(t, hasSubcommand(newListCommand(), "pools")) + assert.False(t, hasSubcommand(newGetCommand(), "pool")) } func hasSubcommand(parent *cobra.Command, name string) bool { @@ -232,7 +237,7 @@ func TestGetPoolText(t *testing.T) { } func TestGetPoolOmitIDResolvesSolePool(t *testing.T) { - // With exactly one pool, `air get pool` (no id) resolves it and shows usage. + // With exactly one pool, `air pools get` (no id) resolves it and shows usage. list := `{"provisioned_capacities":[{"name":"provisioned-capacities/pool-only","spec":{"accelerator_type":"GPU_8xH100","accelerator_count":64}}]}` detail := `{"name":"provisioned-capacities/pool-only","spec":{"accelerator_type":"GPU_8xH100","accelerator_count":64},"status":{"usage":{"used_accelerator_count":40,"idle_accelerator_count":24}}}` srv := poolsServer(t, []string{list}, map[string]string{"pool-only": detail})