Turn provisioned-capacities -> pools

acceptance/experimental/air/config-help/output.txt

@@ -63,10 +63,10 @@ config.compute
   Which accelerators to run on and how many.
 
   Fields:
-    num_accelerators         Total number of GPUs to allocate.
-    accelerator_type         Which accelerator to run on, e.g. GPU_1xA10.
-    provisioned_capacity_id  Pre-provisioned AIR capacity reservation id.
-    priority_class           Scheduling priority within the reservation: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest).
+    num_accelerators  Total number of GPUs to allocate.
+    accelerator_type  Which accelerator to run on, e.g. GPU_1xA10.
+    pool_id           GPU pool to run on, by id.
+    priority_class    Scheduling priority within the pool: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest).
 
 Use "-h config.compute.<field>" for details on a field.
 
@@ -78,16 +78,16 @@ config.mlflow_artifact_location
   Type:     string
   Required: no
 
->>> [CLI] experimental air run -h config.compute.provisioned_capacity_id
-config.compute.provisioned_capacity_id
-  Pre-provisioned AIR capacity reservation id. Must be 1-255 characters. Contact your Databricks account team to provision capacity.
+>>> [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'.
 
   Type:     string
   Required: no
 
 >>> [CLI] experimental air run -h config.compute.priority_class
 config.compute.priority_class
-  Scheduling priority within the reservation: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest). Requires provisioned_capacity_id.
+  Scheduling priority within the pool: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest). Requires pool_id.
 
   Type:     string
   Required: no
@@ -143,7 +143,7 @@ config.compute.num_accelerators
 >>> [CLI] experimental air run -h config.compute.acclerator_type
 Error: unknown config field "config.compute.acclerator_type"; did you mean "accelerator_type"?
 
-fields under "config.compute" are: accelerator_type, num_accelerators, priority_class, provisioned_capacity_id
+fields under "config.compute" are: accelerator_type, num_accelerators, pool_id, priority_class
 
 === free-form map keys are not schema fields
 >>> [CLI] experimental air run -h config.parameters.learning_rate

acceptance/experimental/air/config-help/script

@@ -14,7 +14,7 @@ trace $CLI experimental air run -h config.compute
 
 title "new submission fields are documented"
 trace $CLI experimental air run -h config.mlflow_artifact_location
-trace $CLI experimental air run -h config.compute.provisioned_capacity_id
+trace $CLI experimental air run -h config.compute.pool_id
 trace $CLI experimental air run -h config.compute.priority_class
 trace $CLI experimental air run -h config.environment.dependencies
 

acceptance/experimental/air/run-submit-deps/output.txt

@@ -12,7 +12,7 @@ Stream logs after submission using:
 === only config + command are uploaded; no requirements.yaml
 >>> print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep
 {"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/command.sh", "q": {"overwrite": "true"}, "raw_body": "python train.py"}
-{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/training_config.yaml", "q": {"overwrite": "true"}, "raw_body": "experiment_name: deps-smoke\ncommand: python train.py\ncompute:\n    accelerator_type: GPU_1xH100\n    num_accelerators: 1\n    provisioned_capacity_id: capacity-123\nenvironment:\n    version: 5\n    dependencies:\n        - numpy\n        - torch==2.3.0\nmlflow_artifact_location: /Volumes/main/default/air-artifacts\n"}
+{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/training_config.yaml", "q": {"overwrite": "true"}, "raw_body": "experiment_name: deps-smoke\ncommand: python train.py\ncompute:\n    accelerator_type: GPU_1xH100\n    num_accelerators: 1\n    pool_id: capacity-123\nenvironment:\n    version: 5\n    dependencies:\n        - numpy\n        - torch==2.3.0\nmlflow_artifact_location: /Volumes/main/default/air-artifacts\n"}
 
 === inline deps, artifact location, and capacity id reach the submit payload
 >>> print_requests.py //api/2.2/jobs/runs/submit

acceptance/experimental/air/run-submit-deps/run.yaml

@@ -3,7 +3,7 @@ command: python train.py
 compute:
   accelerator_type: GPU_1xH100
   num_accelerators: 1
-  provisioned_capacity_id: capacity-123
+  pool_id: capacity-123
 environment:
   version: 5
   dependencies:

experimental/air/cmd/compute.go

@@ -89,10 +89,10 @@ func gpusPerNode(g gpuType) (int, error) {
 // computeConfig is the `compute` block of the run YAML: which accelerators to
 // use and how many.
 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."`
-	ProvisionedCapacityID *string `yaml:"provisioned_capacity_id" help:"Pre-provisioned AIR capacity reservation id. Must be 1-255 characters. Contact your Databricks account team to provision capacity."`
-	PriorityClass         *string `yaml:"priority_class" help:"Scheduling priority within the reservation: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest). Requires provisioned_capacity_id."`
+	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'."`
+	PriorityClass   *string `yaml:"priority_class" help:"Scheduling priority within the pool: BEST_EFFORT (lowest, preemptable), NORMAL, or CRITICAL (highest). Requires pool_id."`
 }
 
 // validate checks the compute block against the backend's constraints.
@@ -114,26 +114,26 @@ func (c *computeConfig) validate() error {
 		return fmt.Errorf("compute.num_accelerators for %s must be a multiple of %d, got %d", c.AcceleratorType, perNode, c.NumAccelerators)
 	}
 
-	if c.ProvisionedCapacityID != nil {
-		v := strings.TrimSpace(*c.ProvisionedCapacityID)
+	if c.PoolID != nil {
+		v := strings.TrimSpace(*c.PoolID)
 		if v == "" {
-			return errors.New("compute.provisioned_capacity_id cannot be empty")
+			return errors.New("compute.pool_id cannot be empty")
 		}
 		if len(v) > 255 {
-			return fmt.Errorf("compute.provisioned_capacity_id must be 255 characters or less, got %d", len(v))
+			return fmt.Errorf("compute.pool_id must be 255 characters or less, got %d", len(v))
 		}
-		*c.ProvisionedCapacityID = v
+		*c.PoolID = v
 	}
 
 	if c.PriorityClass != nil {
 		p, err := parsePriorityClass(*c.PriorityClass)
 		if err != nil {
 			return fmt.Errorf("compute.priority_class: %w", err)
 		}
-		// A priority class only ranks pending work within a reservation, so it
-		// requires one.
-		if c.ProvisionedCapacityID == nil {
-			return errors.New("compute.priority_class requires compute.provisioned_capacity_id — priority applies only to a pre-provisioned capacity reservation")
+		// A priority class only ranks pending work within a pool, so it requires
+		// one.
+		if c.PoolID == nil {
+			return errors.New("compute.priority_class requires compute.pool_id — priority applies only within a GPU pool")
 		}
 		*c.PriorityClass = string(p)
 	}

experimental/air/cmd/compute_test.go

@@ -78,7 +78,7 @@ func TestParsePriorityClass(t *testing.T) {
 func TestComputeConfigValidateNormalizesPriorityClass(t *testing.T) {
 	// A lower-case value is upper-cased in place so the submitted payload carries
 	// the enum's contract value.
-	cfg := computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new("cap"), PriorityClass: new("critical")}
+	cfg := computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PoolID: new("cap"), PriorityClass: new("critical")}
 	require.NoError(t, cfg.validate())
 	assert.Equal(t, "CRITICAL", *cfg.PriorityClass)
 }
@@ -92,16 +92,16 @@ func TestComputeConfigValidate(t *testing.T) {
 		{"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"}, ""},
-		{"capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(" capacity ")}, ""},
-		{"empty capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(" ")}, "cannot be empty"},
-		{"long capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(strings.Repeat("a", 256))}, "255 characters or less"},
+		{"pool id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PoolID: new(" capacity ")}, ""},
+		{"empty pool id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PoolID: new(" ")}, "cannot be empty"},
+		{"long pool id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PoolID: new(strings.Repeat("a", 256))}, "255 characters or less"},
 		{"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"},
-		{"priority class", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new("cap"), PriorityClass: new("critical")}, ""},
-		{"priority class requires reservation", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PriorityClass: new("NORMAL")}, "requires compute.provisioned_capacity_id"},
-		{"invalid priority class", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new("cap"), PriorityClass: new("urgent")}, "invalid priority_class"},
+		{"priority class", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PoolID: new("cap"), PriorityClass: new("critical")}, ""},
+		{"priority class requires reservation", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PriorityClass: new("NORMAL")}, "requires compute.pool_id"},
+		{"invalid priority class", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", PoolID: new("cap"), PriorityClass: new("urgent")}, "invalid priority_class"},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

experimental/air/cmd/runsubmit.go

@@ -119,12 +119,12 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri
 	}
 }
 
-func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.SubmitRun, provisionedCapacityID, priorityClass, unityCatalogImagePath string) (int64, error) {
+func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.SubmitRun, poolID, priorityClass, unityCatalogImagePath string) (int64, error) {
 	// None of these fields are modeled by the SDK's AiRuntimeTask, so a run that
 	// sets any of them has to go through the raw /api/2.2 body. priority_class only
-	// ever appears alongside a reservation (validation enforces it), but route on
+	// ever appears alongside a pool (validation enforces it), but route on
 	// all of them so none can be silently dropped.
-	if provisionedCapacityID == "" && priorityClass == "" && unityCatalogImagePath == "" {
+	if poolID == "" && priorityClass == "" && unityCatalogImagePath == "" {
 		wait, err := w.Jobs.Submit(ctx, payload)
 		if err != nil {
 			return 0, err
@@ -142,7 +142,7 @@ func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.
 	if err := decoder.Decode(&body); err != nil {
 		return 0, fmt.Errorf("failed to decode AIR submit payload: %w", err)
 	}
-	if err := injectReservationFields(body, provisionedCapacityID, priorityClass); err != nil {
+	if err := injectPoolFields(body, poolID, priorityClass); err != nil {
 		return 0, err
 	}
 	if unityCatalogImagePath != "" {
@@ -165,19 +165,19 @@ func submitRun(ctx context.Context, w *databricks.WorkspaceClient, payload jobs.
 	return response.RunId, nil
 }
 
-// injectReservationFields sets the reservation-only fields the SDK does not
-// model onto the decoded submit body: priority_class rides directly on the
-// ai_runtime_task, while provisioned_capacity_id rides on the deployment's
-// compute spec. Each is set only when non-empty.
-func injectReservationFields(body map[string]any, provisionedCapacityID, priorityClass string) error {
+// injectPoolFields sets the pool-only fields the SDK does not model onto the
+// decoded submit body: priority_class rides directly on the ai_runtime_task,
+// while provisioned_capacity_id (the wire name for the pool) rides on the
+// deployment's compute spec. Each is set only when non-empty.
+func injectPoolFields(body map[string]any, poolID, priorityClass string) error {
 	aiRuntimeTask, err := aiRuntimeTaskFromSubmitBody(body)
 	if err != nil {
 		return err
 	}
 	if priorityClass != "" {
 		aiRuntimeTask["priority_class"] = priorityClass
 	}
-	if provisionedCapacityID != "" {
+	if poolID != "" {
 		deployments, ok := aiRuntimeTask["deployments"].([]any)
 		if !ok || len(deployments) != 1 {
 			return errors.New("AIR submit payload must contain exactly one deployment")
@@ -190,7 +190,7 @@ func injectReservationFields(body map[string]any, provisionedCapacityID, priorit
 		if !ok {
 			return errors.New("AIR submit payload is missing deployment compute")
 		}
-		computeSpec["provisioned_capacity_id"] = provisionedCapacityID
+		computeSpec["provisioned_capacity_id"] = poolID
 	}
 	return nil
 }
@@ -364,16 +364,18 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 	payload := buildSubmitPayload(cfg, commandPath, dlRuntimeImage(ctx, runtimeVersion), usagePolicyID, snap, deps)
 	payload.IdempotencyToken = token
 
-	provisionedCapacityID := ""
-	if cfg.Compute.ProvisionedCapacityID != nil {
-		provisionedCapacityID = *cfg.Compute.ProvisionedCapacityID
+	// The pool id is sent on the wire as provisioned_capacity_id (the backend's
+	// name for a GPU pool); only the user-facing YAML field is pool_id.
+	poolID := ""
+	if cfg.Compute.PoolID != nil {
+		poolID = *cfg.Compute.PoolID
 	}
 	priorityClass := ""
 	if cfg.Compute.PriorityClass != nil {
 		priorityClass = *cfg.Compute.PriorityClass
 	}
 	// Submit returns as soon as the run is created; we don't wait for it to finish.
-	runID, err := submitRun(ctx, w, payload, provisionedCapacityID, priorityClass, cfg.unityCatalogImagePath())
+	runID, err := submitRun(ctx, w, payload, poolID, priorityClass, cfg.unityCatalogImagePath())
 	if err != nil {
 		return 0, "", err
 	}

experimental/air/cmd/runsubmit_test.go

@@ -85,7 +85,7 @@ func TestBuildSubmitPayload(t *testing.T) {
 	assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu8xH100, AcceleratorCount: 16}, at.Deployments[0].Compute)
 }
 
-func TestSubmitRunInjectsProvisionedCapacityID(t *testing.T) {
+func TestSubmitRunInjectsPoolID(t *testing.T) {
 	server := testserver.New(t)
 	t.Cleanup(server.Close)
 	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {

experimental/air/cmd/validateconfig.go

@@ -66,7 +66,9 @@ func validateConfigRequest(cfg *runConfig, commandPath string) map[string]any {
 	if cfg.Compute != nil {
 		compute["accelerator_type"] = cfg.Compute.AcceleratorType
 		compute["accelerator_count"] = cfg.Compute.NumAccelerators
-		putOpt(compute, "provisioned_capacity_id", cfg.Compute.ProvisionedCapacityID)
+		// Wire field stays provisioned_capacity_id (the backend name); the YAML
+		// field is pool_id.
+		putOpt(compute, "provisioned_capacity_id", cfg.Compute.PoolID)
 	}
 	task := map[string]any{
 		"experiment":  cfg.ExperimentName,

experimental/air/cmd/validateconfig_test.go

@@ -115,7 +115,7 @@ func TestValidateConfigRequestCarriesPriorityClass(t *testing.T) {
 	srv := validateServer(t, http.StatusOK, `{}`, &gotReq)
 
 	cfg := baseRunConfig()
-	cfg.Compute.ProvisionedCapacityID = new("cap-8xh100-res")
+	cfg.Compute.PoolID = new("cap-8xh100-res")
 	cfg.Compute.PriorityClass = new("CRITICAL")
 	err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), cfg, "/Workspace/Users/me/cmd.sh")
 	require.NoError(t, err)