diff --git a/acceptance/experimental/air/convert-to-dabs/b300.yaml b/acceptance/experimental/air/convert-to-dabs/b300.yaml new file mode 100644 index 0000000000..7a91f98c75 --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/b300.yaml @@ -0,0 +1,8 @@ +experiment_name: b300-smoke-test +command: nvidia-smi +compute: + accelerator_type: GPU_8xB300 + num_accelerators: 8 +environment: + version: 6 + dependencies: [] diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index 0922646a0d..308b412043 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -101,3 +101,69 @@ Unlike `air run` (which submits an ephemeral run), bundle deploy creates a persistent job that is not garbage-collected. When you are done, remove the job and its uploaded files with: [CLI] bundle destroy + +=== convert a B300 run using the current SDK +>>> [CLI] experimental air convert-to-dabs b300.yaml --output-dir generated-b300 +Wrote a Databricks Asset Bundle to generated-b300: + databricks.yml + generated_artifacts/training_config.yaml + generated_artifacts/command.sh + +To deploy and run this workload as a bundle: + 1. cd generated-b300 + 2. [CLI] bundle validate + 3. [CLI] bundle deploy + 4. [CLI] bundle run b300-smoke-test --no-wait + +bundle deploy uploads the code source and launch scripts automatically. +To see what it deployed and where: [CLI] bundle summary + +Unlike `air run` (which submits an ephemeral run), bundle deploy creates a +persistent job that is not garbage-collected. When you are done, remove the +job and its uploaded files with: + [CLI] bundle destroy + +=== the B300 bundle preserves its accelerator and AI Runtime version +>>> cat generated-b300/databricks.yml +bundle: + name: b300-smoke-test +sync: + paths: + - generated_artifacts +targets: + dev: + mode: development + default: true +resources: + jobs: + b300-smoke-test: + name: b300-smoke-test + tasks: + - task_key: b300-smoke-test + environment_key: default + max_retries: 3 + ai_runtime_task: + experiment: b300-smoke-test + deployments: + - command_path: ./generated_artifacts/command.sh + compute: + accelerator_type: GPU_8xB300 + accelerator_count: 8 + environments: + - environment_key: default + spec: + environment_version: "6" + +=== the current SDK warns about B300 but does not block bundle validation +>>> [CLI] bundle validate +Warning: invalid value "GPU_8xB300" for enum field. Valid values are [GPU_1xA10 GPU_1xH100 GPU_8xH100] + at resources.jobs.b300-smoke-test.tasks[0].ai_runtime_task.deployments[0].compute.accelerator_type + in databricks.yml:23:37 + +Name: b300-smoke-test +Target: dev +Workspace: + User: [USERNAME] + Path: /Workspace/Users/[USERNAME]/.bundle/b300-smoke-test/dev + +Found 1 warning diff --git a/acceptance/experimental/air/convert-to-dabs/script b/acceptance/experimental/air/convert-to-dabs/script index fd39053c92..cd9dff9f3d 100644 --- a/acceptance/experimental/air/convert-to-dabs/script +++ b/acceptance/experimental/air/convert-to-dabs/script @@ -15,3 +15,12 @@ errcode trace $CLI experimental air convert-to-dabs train.yaml title "--force overwrites it" trace $CLI experimental air convert-to-dabs train.yaml --force + +title "convert a B300 run using the current SDK" +trace $CLI experimental air convert-to-dabs b300.yaml --output-dir generated-b300 + +title "the B300 bundle preserves its accelerator and AI Runtime version" +trace cat generated-b300/databricks.yml + +title "the current SDK warns about B300 but does not block bundle validation" +(cd generated-b300 && trace $CLI bundle validate) diff --git a/acceptance/experimental/air/convert-to-dabs/test.toml b/acceptance/experimental/air/convert-to-dabs/test.toml index aa531fbd10..956c8efc16 100644 --- a/acceptance/experimental/air/convert-to-dabs/test.toml +++ b/acceptance/experimental/air/convert-to-dabs/test.toml @@ -1,3 +1,3 @@ # convert-to-dabs writes the bundle in place (next to train.yaml + src/). These are # generated artifacts, not committed inputs, so exclude them from the repo-diff check. -Ignore = ["databricks.yml", "generated_artifacts", "generated-docker"] +Ignore = ["databricks.yml", "generated_artifacts", "generated-b300", "generated-docker"] diff --git a/acceptance/experimental/air/run-submit/b300.yaml b/acceptance/experimental/air/run-submit/b300.yaml new file mode 100644 index 0000000000..c69924cbcd --- /dev/null +++ b/acceptance/experimental/air/run-submit/b300.yaml @@ -0,0 +1,9 @@ +experiment_name: b300-submit-smoke +command: nvidia-smi +compute: + accelerator_type: GPU_8xB300 + num_accelerators: 8 +environment: + version: 6 + dependencies: [] +max_retries: 0 diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt index c20c8467f1..52aa612190 100644 --- a/acceptance/experimental/air/run-submit/output.txt +++ b/acceptance/experimental/air/run-submit/output.txt @@ -77,3 +77,52 @@ Stream logs after submission using: ] } } + +=== submit a B300 workload with AI Runtime 6 +>>> [CLI] experimental air run -f b300.yaml +Submitting experiment: b300-submit-smoke +Submitted workload with Job Run ID: 555 +View job run at: [DATABRICKS_URL]/jobs/runs/555 + +Tip: use --watch when submitting a run to stream logs to your terminal. +Stream logs after submission using: + databricks experimental air logs 555 + +=== the B300 accelerator and runtime reach the submit payload +>>> print_requests.py //api/2.2/jobs/runs/submit +{ + "method": "POST", + "path": "/api/2.2/jobs/runs/submit", + "body": { + "environments": [ + { + "environment_key": "default", + "spec": { + "environment_version": "6" + } + } + ], + "idempotency_token": "[UUID]", + "run_name": "b300-submit-smoke", + "tasks": [ + { + "ai_runtime_task": { + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/b300-submit-smoke/b300-submit-smoke_[RUN_ID]/command.sh", + "compute": { + "accelerator_count": 8, + "accelerator_type": "GPU_8xB300" + } + } + ], + "experiment": "b300-submit-smoke" + }, + "environment_key": "default", + "max_retries": 0, + "run_if": "ALL_SUCCESS", + "task_key": "b300-submit-smoke" + } + ] + } +} diff --git a/acceptance/experimental/air/run-submit/script b/acceptance/experimental/air/run-submit/script index 26161d63ad..e2fa42eea6 100644 --- a/acceptance/experimental/air/run-submit/script +++ b/acceptance/experimental/air/run-submit/script @@ -16,4 +16,10 @@ title "the ai_runtime_task and additive job permission grants" # Include MLflow paths so an accidental experiment request appears in the golden output. trace print_requests.py --get //api/2.2/jobs/runs/submit //api/2.2/jobs/runs/get //api/2.0/permissions //api/2.0/mlflow --sort --unique +title "submit a B300 workload with AI Runtime 6" +trace $CLI experimental air run -f b300.yaml + +title "the B300 accelerator and runtime reach the submit payload" +trace print_requests.py //api/2.2/jobs/runs/submit + rm -fr .git diff --git a/experimental/air/cmd/compute.go b/experimental/air/cmd/compute.go index c3357eaf40..0f302c4aaa 100644 --- a/experimental/air/cmd/compute.go +++ b/experimental/air/cmd/compute.go @@ -14,10 +14,12 @@ const ( gpuType1xA10 gpuType = "GPU_1xA10" gpuType8xH100 gpuType = "GPU_8xH100" gpuType1xH100 gpuType = "GPU_1xH100" + gpuType8xB300 gpuType = "GPU_8xB300" ) -// gpuTypes lists every valid type. Used for validation error messages. -var gpuTypes = []gpuType{gpuType1xA10, gpuType1xH100, gpuType8xH100} +// gpuTypes lists every accelerator type understood by this CLI. +// Workspace availability is enforced by the server. +var gpuTypes = []gpuType{gpuType1xA10, gpuType1xH100, gpuType8xH100, gpuType8xB300} func validGPUTypesHint() string { names := make([]string, len(gpuTypes)) @@ -31,7 +33,7 @@ func validGPUTypesHint() string { // exact: the server's lookup is case-sensitive. func parseGPUType(value string) (gpuType, error) { switch gpuType(value) { - case gpuType1xA10, gpuType8xH100, gpuType1xH100: + case gpuType1xA10, gpuType8xH100, gpuType1xH100, gpuType8xB300: return gpuType(value), nil } return "", fmt.Errorf("invalid GPU type %q: %s", value, validGPUTypesHint()) @@ -78,7 +80,7 @@ func gpusPerNode(g gpuType) (int, error) { switch g { case gpuType1xA10, gpuType1xH100: return 1, nil - case gpuType8xH100: + case gpuType8xH100, gpuType8xB300: return 8, nil } // Unreachable: callers resolve g through parseGPUType first, which rejects diff --git a/experimental/air/cmd/compute_test.go b/experimental/air/cmd/compute_test.go index 467e2acfa9..5a099c5adf 100644 --- a/experimental/air/cmd/compute_test.go +++ b/experimental/air/cmd/compute_test.go @@ -16,6 +16,7 @@ func TestParseGPUType(t *testing.T) { {"GPU_1xA10", gpuType1xA10}, {"GPU_8xH100", gpuType8xH100}, {"GPU_1xH100", gpuType1xH100}, + {"GPU_8xB300", gpuType8xB300}, } for _, tt := range tests { t.Run(tt.in, func(t *testing.T) { @@ -29,7 +30,7 @@ func TestParseGPUType(t *testing.T) { func TestParseGPUTypeInvalid(t *testing.T) { // Wrong casing is rejected rather than fixed up; legacy types (h100_80gb, a10) // can no longer be submitted; unknown types are rejected. - for _, in := range []string{"gpu_1xa10", "GPU_1XA10", "GPU_2xH100", "h100_80gb", "a10", "b200", ""} { + for _, in := range []string{"gpu_1xa10", "GPU_1XA10", "gpu_8xb300", "GPU_2xH100", "h100_80gb", "a10", "b200", ""} { t.Run(in, func(t *testing.T) { _, err := parseGPUType(in) require.Error(t, err) @@ -46,6 +47,7 @@ func TestGPUsPerNode(t *testing.T) { {gpuType1xA10, 1}, {gpuType1xH100, 1}, {gpuType8xH100, 8}, + {gpuType8xB300, 8}, } for _, tt := range tests { t.Run(string(tt.in), func(t *testing.T) { @@ -99,6 +101,14 @@ func TestComputeConfigValidate(t *testing.T) { {"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"}, + {"B300 single node", computeConfig{NumAccelerators: 8, AcceleratorType: "GPU_8xB300"}, ""}, + {"B300 two nodes", computeConfig{NumAccelerators: 16, AcceleratorType: "GPU_8xB300"}, ""}, + {"B300 three nodes", computeConfig{NumAccelerators: 24, AcceleratorType: "GPU_8xB300"}, ""}, + {"B300 count one", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_8xB300"}, "multiple of 8"}, + {"B300 count four", computeConfig{NumAccelerators: 4, AcceleratorType: "GPU_8xB300"}, "multiple of 8"}, + {"B300 count twelve", computeConfig{NumAccelerators: 12, AcceleratorType: "GPU_8xB300"}, "multiple of 8"}, + {"B300 zero", computeConfig{NumAccelerators: 0, AcceleratorType: "GPU_8xB300"}, "must be positive"}, + {"B300 negative", computeConfig{NumAccelerators: -8, AcceleratorType: "GPU_8xB300"}, "must be positive"}, {"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"}, diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index eadffe1c49..fe96574237 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -244,6 +244,29 @@ func TestConvertToDabsRuntimeVersionEnvOverride(t *testing.T) { assert.Equal(t, "7", get(t, root, env+".spec.environment_version").MustString()) } +func TestConvertToDabsSupportsB300(t *testing.T) { + path := writeConfigFile(t, "run.yaml", ` +experiment_name: b300 +command: python train.py +compute: + accelerator_type: GPU_8xB300 + num_accelerators: 16 +environment: + version: 6 + dependencies: [] +`) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) + require.NoError(t, err) + + deployment := "resources.jobs.b300.tasks[0].ai_runtime_task.deployments[0]" + assert.Equal(t, "GPU_8xB300", get(t, root, deployment+".compute.accelerator_type").MustString()) + assert.Equal(t, int64(16), get(t, root, deployment+".compute.accelerator_count").MustInt()) + assert.Equal(t, "6", get(t, root, "resources.jobs.b300.environments[0].spec.environment_version").MustString()) +} + // remote_volume can't be honored by a converted bundle (bundle deploy owns the // artifact upload location), so it is rejected rather than silently ignored. func TestConvertToDabsRejectsRemoteVolume(t *testing.T) { diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index e496a1e89c..1271f89b4a 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -82,6 +82,7 @@ var gpuDisplayNames = map[string]string{ "GPU_1xA10": "A10", "GPU_8xH100": "H100", "GPU_1xH100": "H100", + "GPU_8xB300": "B300", } // isRunning reports whether a run is currently executing (lifecycle RUNNING), as diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go index 4cc6dbf909..842806e611 100644 --- a/experimental/air/cmd/format_test.go +++ b/experimental/air/cmd/format_test.go @@ -69,6 +69,7 @@ func TestGpuDisplayName(t *testing.T) { assert.Equal(t, "A10", gpuDisplayName("a10")) assert.Equal(t, "H100", gpuDisplayName("GPU_8xH100")) assert.Equal(t, "H100", gpuDisplayName("GPU_1xH100")) + assert.Equal(t, "B300", gpuDisplayName("GPU_8xB300")) // Unknown identifiers pass through unchanged. assert.Equal(t, "b200", gpuDisplayName("b200")) assert.Empty(t, gpuDisplayName("")) @@ -226,6 +227,7 @@ func TestAcceleratorLabel(t *testing.T) { assert.Empty(t, acceleratorLabel("GPU_8xH100", 0)) assert.Equal(t, "8x H100", acceleratorLabel("GPU_8xH100", 8)) assert.Equal(t, "1x A10", acceleratorLabel("GPU_1xA10", 1)) + assert.Equal(t, "16x B300", acceleratorLabel("GPU_8xB300", 16)) // The RPC may report a count without a recognized type. assert.Equal(t, "8x", acceleratorLabel("", 8)) } diff --git a/experimental/air/cmd/logdownload_test.go b/experimental/air/cmd/logdownload_test.go index b7717980f4..57b6982bdd 100644 --- a/experimental/air/cmd/logdownload_test.go +++ b/experimental/air/cmd/logdownload_test.go @@ -44,6 +44,7 @@ func TestResolveNodeCount(t *testing.T) { {"GPU_1xA10", 2, 2}, {"GPU_1xH100", 4, 4}, {"GPU_8xH100", 16, 2}, + {"GPU_8xB300", 16, 2}, } for _, tt := range tests { n, err := resolveNodeCount(airRunWithCompute(tt.accelType, tt.count)) diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index ba10c1b61d..7a85190743 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -88,6 +88,24 @@ func TestBuildSubmitPayload(t *testing.T) { assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu8xH100, AcceleratorCount: 16}, at.Deployments[0].Compute) } +func TestBuildSubmitPayloadB300(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "b300", + Command: new("nvidia-smi"), + Compute: &computeConfig{AcceleratorType: "GPU_8xB300", NumAccelerators: 8}, + } + + p := buildSubmitPayload(cfg, "/d/command.sh", "databricks_ai_v6", "", snapshotResult{}, nil) + require.Len(t, p.Environments, 1) + require.NotNil(t, p.Environments[0].Spec) + assert.Equal(t, "workspace-base-environments/databricks_ai_v6", p.Environments[0].Spec.BaseEnvironment) + assert.Empty(t, p.Environments[0].Spec.EnvironmentVersion) + require.Len(t, p.Tasks, 1) + require.NotNil(t, p.Tasks[0].AiRuntimeTask) + require.Len(t, p.Tasks[0].AiRuntimeTask.Deployments, 1) + assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorType("GPU_8xB300"), AcceleratorCount: 8}, p.Tasks[0].AiRuntimeTask.Deployments[0].Compute) +} + func TestSubmitRunInjectsPoolID(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close)