[air] Add GPU_8xB300 support to the experimental AIR CLI - #6649
Conversation
Integration test reportCommit: 95e3389
Top 3 slowest tests (at least 2 minutes):
|
689bdff to
c51a90b
Compare
ben-hansen-db
left a comment
ben-hansen-db
left a comment
There was a problem hiding this comment.
lgtm, note Pardis' work on using our internal list api, we should do that here as well but this is fine for now
maggiewang-db
left a comment
maggiewang-db
left a comment
There was a problem hiding this comment.
Code looks good. Should we have some backend validation to reject client image version < 6 when specified B300?
| // code_source; the emitted `tgz` artifact packages it at deploy. | ||
| func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir string) (map[string]dyn.Value, []uploadItem, error) { | ||
| if gpuType(cfg.Compute.AcceleratorType) == gpuType8xB300 { | ||
| return nil, nil, errors.New("GPU_8xB300 is not yet supported by convert-to-dabs; use air run until DAB support is available") |
There was a problem hiding this comment.
is it hard to support that? why is it blocked?
note that I made a b300 dab successfully today
There was a problem hiding this comment.
go sdk does not yet have the b300 accelerator type enum and this will print a warning so I just added this guard.
but you're right, I confirmed it will still work and run successfully, just with a printed warning. removed this guard, and the warning will be removed when the go sdk is updated
| gpuType1xA10 gpuType = "GPU_1xA10" | ||
| gpuType8xH100 gpuType = "GPU_8xH100" | ||
| gpuType1xH100 gpuType = "GPU_1xH100" | ||
| gpuType8xB300 gpuType = "GPU_8xB300" |
There was a problem hiding this comment.
rather than adding b300 directly like this, we should switch the impl to match python side what I am doing here: https://github.com/databricks-eng/universe/pull/2598845/changes
reason: B300 OD is a gated pupr feature. Backend knows which customers have access and which don't. we should call the backend for a given profile and let the backend tell us the dynamically generated list of gpu types that are available for a given workspace for a given user
There was a problem hiding this comment.
adding list compute options in a separate PR as a follow up and will request your review. this will require setting up an ai-training proxy to the aicm api btw.
8b450bc to
cc0186a
Compare
Backend already has validation to reject client image < 6 for B300s. I want --dry-run to also be able to handle this validation but that will be a separate decision / PR. |
cc0186a to
4b9cdc7
Compare
4b9cdc7 to
95e3389
Compare
5 participants
Changes
GPU_8xB300support to the experimental AIR CLI with eight GPUs per node and positive accelerator counts that are multiples of eight.databricks_ai_v6environment routing.convert-to-dabsuntil DAB support lands in a separate PR.Why
Enable direct B300 workloads without coupling the experimental AIR change to Bundle-owned schema and validation changes. Jobs remains authoritative for supported environment versions, including the B300 minimum, for now. A follow-up may incorporate that validation into
--dry-runonce its final behavior is defined.Tests
NVIDIA B300 SXM6 AC.This PR was written with Codex.