Skip to content

[air] Add air list/get provisioned_capacity commands - #6595

Open
vinchenzo-db wants to merge 1 commit into
mainfrom
air-provisioned-capacity
Open

[air] Add air list/get provisioned_capacity commands#6595
vinchenzo-db wants to merge 1 commit into
mainfrom
air-provisioned-capacity

Conversation

@vinchenzo-db

Copy link
Copy Markdown
Contributor

Adds two read commands for pre-provisioned AI Runtime capacity reservations (the same reservation a run targets via compute.provisioned_capacity_id):

  • air list provisioned_capacity — the workspace's reservations (id, accelerator type, reserved accelerators).
  • air get provisioned_capacity <id> — one reservation plus its accelerator usage summary (used / idle).

Both call AiWorkflowService's purpose-built public read API (GET /api/2.0/ai-training/provisioned-capacities[/{id}]) via a raw client.Do, matching the sibling air list workflows call. The platform tracks these internally as "guaranteed capacity"; the public surface (and this CLI) calls them "provisioned capacity". Usage is populated only by Get, so the list view omits it. Output uses the same JSON envelope + text pattern as the other air commands, and the SAFE-gate the server returns before rollout is surfaced as a clear FEATURE_DISABLED error rather than a retryable one.

Tests: unit tests for the projections, resource-name handling, int64-as-string decoding, pagination, and the JSON/text/not-found/feature-disabled paths; an acceptance test drives the real binary against a mocked endpoint (text + JSON goldens). Verified live against e2-dogfood: the CLI reaches the deployed handler and classifies its responses correctly (the API is SAFE-gated off there, so the successful list/get shape is covered by the goldens).

Changes

Why

Tests

Adds two read commands for pre-provisioned AI Runtime capacity reservations
(the same reservation a run targets via compute.provisioned_capacity_id):

- `air list provisioned_capacity` — the workspace's reservations
  (id, accelerator type, reserved accelerators).
- `air get provisioned_capacity <id>` — one reservation plus its accelerator
  usage summary (used / idle).

Both call AiWorkflowService's purpose-built public read API
(GET /api/2.0/ai-training/provisioned-capacities[/{id}]) via a raw client.Do,
matching the sibling `air list` workflows call. The platform tracks these
internally as "guaranteed capacity"; the public surface (and this CLI) calls
them "provisioned capacity". Usage is populated only by Get, so the list view
omits it. Output uses the same JSON envelope + text pattern as the other air
commands, and the SAFE-gate the server returns before rollout is surfaced as a
clear FEATURE_DISABLED error rather than a retryable one.

Tests: unit tests for the projections, resource-name handling, int64-as-string
decoding, pagination, and the JSON/text/not-found/feature-disabled paths; an
acceptance test drives the real binary against a mocked endpoint (text + JSON
goldens). Verified live against e2-dogfood: the CLI reaches the deployed
handler and classifies its responses correctly (the API is SAFE-gated off there,
so the successful list/get shape is covered by the goldens).

Co-authored-by: Isaac <no-reply@databricks.com>
@github-actions

github-actions Bot commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

Approval status: pending

/acceptance/experimental/air/ - needs approval

5 files changed
Suggested: @ben-hansen-db
Also eligible: @apeforest, @bfontain, @panchalhp-db, @maggiewang-db, @pardis-beikzadeh-db, @caroline-db

/experimental/air/ - needs approval

5 files changed
Suggested: @ben-hansen-db
Also eligible: @apeforest, @bfontain, @panchalhp-db, @maggiewang-db, @pardis-beikzadeh-db, @caroline-db

Any maintainer (@andrewnester, @anton-107, @denik, @pietern, @shreyas-goenka, @simonfaltum, @renaudhartert-db, @janniklasrose, @lennartkats-db, @rugpanov, @rclarey) can approve all areas.
See OWNERS for ownership rules.

@eng-dev-ecosystem-bot

Copy link
Copy Markdown
Collaborator

Integration test report

Commit: 07094fc

Run: 34389279991

Env 💚​RECOVERED ✅​pass 🙈​skip Time
💚​ aws linux 1 275 17 7:20
💚​ aws windows 1 277 15 7:18
💚​ azure linux 1 274 17 10:59
💚​ azure windows 1 276 15 7:33
💚​ gcp linux 1 275 17 10:34
💚​ gcp windows 1 277 15 6:54
Test Name aws linux aws windows azure linux azure windows gcp linux gcp windows
💚​ TestAccept 💚​R 💚​R 💚​R 💚​R 💚​R 💚​R
Top 13 slowest tests (at least 2 minutes):
duration env testname
4:38 azure windows TestAccept
4:12 aws windows TestAccept
3:44 gcp windows TestAccept
2:53 gcp linux TestFilerWorkspaceFilesExtensionsStat
2:33 azure linux TestImportDir
2:07 aws linux TestAccept
2:07 aws windows TestFilerWorkspaceFilesExtensionsRead
2:05 azure linux TestAccept
2:05 gcp linux TestAccept
2:05 gcp windows TestImportDirDoesNotOverwrite
2:05 azure linux TestExportDir
2:00 azure linux TestFilerWorkspaceFilesExtensionsStat
2:00 aws windows TestFilerRecursiveDelete/workspace_files

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants