[air] Change pool commands to be resource-first

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

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,21 +72,38 @@ 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
   -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 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

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
 

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,15 +27,15 @@ 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
 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]",

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

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())

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)
 	}

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."`
 }
 

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
 }
 

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{}))

experimental/air/cmd/list.go

@@ -175,8 +175,6 @@ Filters:
 		return renderListText(cmd, fetcher, limit)
 	}
 
-	cmd.AddCommand(newListPoolsCommand())
-
 	return cmd
 }
 

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)

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})