From ed3e3f82b5ab545f1ae27494b96668b804f2466e Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Tue, 15 Sep 2026 22:27:00 +0000 Subject: [PATCH 1/2] [air] Rename run YAML field compute.provisioned_capacity_id to pool_id Product renamed the concept to "GPU pool", so the run YAML field that selects a reservation becomes compute.pool_id (hard rename, no back-compat alias). priority_class now reads "requires pool_id". The value is still sent on the wire as provisioned_capacity_id (the backend's name for a pool); only the user-facing YAML field changed. Updated help text, validation messages, tests, and the config-help / run-submit-deps acceptance goldens. Co-authored-by: Isaac --- .../experimental/air/config-help/output.txt | 18 ++++++------- .../experimental/air/config-help/script | 2 +- .../air/run-submit-deps/output.txt | 2 +- .../experimental/air/run-submit-deps/run.yaml | 2 +- experimental/air/cmd/compute.go | 26 +++++++++---------- experimental/air/cmd/compute_test.go | 14 +++++----- experimental/air/cmd/runsubmit.go | 6 +++-- experimental/air/cmd/validateconfig.go | 4 ++- experimental/air/cmd/validateconfig_test.go | 2 +- 9 files changed, 40 insertions(+), 36 deletions(-) diff --git a/acceptance/experimental/air/config-help/output.txt b/acceptance/experimental/air/config-help/output.txt index a17388562ad..c747b4ea6d5 100644 --- a/acceptance/experimental/air/config-help/output.txt +++ b/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." 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 diff --git a/acceptance/experimental/air/config-help/script b/acceptance/experimental/air/config-help/script index da83063b64f..8a9f2b302b3 100644 --- a/acceptance/experimental/air/config-help/script +++ b/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 diff --git a/acceptance/experimental/air/run-submit-deps/output.txt b/acceptance/experimental/air/run-submit-deps/output.txt index 0445557d826..618148f4a6f 100644 --- a/acceptance/experimental/air/run-submit-deps/output.txt +++ b/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 diff --git a/acceptance/experimental/air/run-submit-deps/run.yaml b/acceptance/experimental/air/run-submit-deps/run.yaml index e9f6b58bff1..deb17acf080 100644 --- a/acceptance/experimental/air/run-submit-deps/run.yaml +++ b/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: diff --git a/experimental/air/cmd/compute.go b/experimental/air/cmd/compute.go index 01b307b22bc..c3357eaf401 100644 --- a/experimental/air/cmd/compute.go +++ b/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,15 +114,15 @@ 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 { @@ -130,10 +130,10 @@ func (c *computeConfig) validate() error { 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) } diff --git a/experimental/air/cmd/compute_test.go b/experimental/air/cmd/compute_test.go index c402b1e595b..467e2acfa9a 100644 --- a/experimental/air/cmd/compute_test.go +++ b/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) { diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index c580fec8c47..4cb98bf1a47 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -364,9 +364,11 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run payload := buildSubmitPayload(cfg, commandPath, dlRuntimeImage(ctx, runtimeVersion), usagePolicyID, snap, deps) payload.IdempotencyToken = token + // 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. provisionedCapacityID := "" - if cfg.Compute.ProvisionedCapacityID != nil { - provisionedCapacityID = *cfg.Compute.ProvisionedCapacityID + if cfg.Compute.PoolID != nil { + provisionedCapacityID = *cfg.Compute.PoolID } priorityClass := "" if cfg.Compute.PriorityClass != nil { diff --git a/experimental/air/cmd/validateconfig.go b/experimental/air/cmd/validateconfig.go index 8646311c84b..5004a6dbc70 100644 --- a/experimental/air/cmd/validateconfig.go +++ b/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, diff --git a/experimental/air/cmd/validateconfig_test.go b/experimental/air/cmd/validateconfig_test.go index 2c91dfb4a7e..8278074650f 100644 --- a/experimental/air/cmd/validateconfig_test.go +++ b/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) From cce565572a310357462dfe92a643790898ee13b9 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 16 Sep 2026 17:52:58 +0000 Subject: [PATCH 2/2] [air] Rename internal provisioned-capacity identifiers to pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Carry the user-facing pool rename through to the Go internals of the submit path: provisionedCapacityID (variable and params) becomes poolID, and injectReservationFields becomes injectPoolFields. The on-wire key stays provisioned_capacity_id — that is the backend's contract, so only the Go identifiers and the concept wording in comments change, not any string sent to or parsed from the server. Co-authored-by: Isaac --- experimental/air/cmd/runsubmit.go | 28 +++++++++++++------------- experimental/air/cmd/runsubmit_test.go | 2 +- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 4cb98bf1a47..9be8da1d527 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/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,11 +165,11 @@ 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 @@ -177,7 +177,7 @@ func injectReservationFields(body map[string]any, provisionedCapacityID, priorit 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 } @@ -366,16 +366,16 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run // 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. - provisionedCapacityID := "" + poolID := "" if cfg.Compute.PoolID != nil { - provisionedCapacityID = *cfg.Compute.PoolID + 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 } diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 73c9f72ae59..c4be517d82a 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/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 {