From e5f1ef84c9225a20688a1151bc186a4c8ac83d2e Mon Sep 17 00:00:00 2001 From: Caroline Chen <324939130+caroline-db@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:01:58 +0000 Subject: [PATCH 1/3] Add B300 support to experimental AIR CLI and DABs --- .../experimental/air/run-submit/b300.yaml | 9 ++++ .../experimental/air/run-submit/output.txt | 49 +++++++++++++++++++ acceptance/experimental/air/run-submit/script | 6 +++ experimental/air/cmd/compute.go | 7 +-- experimental/air/cmd/compute_test.go | 12 ++++- experimental/air/cmd/convert_to_dabs.go | 2 +- experimental/air/cmd/convert_to_dabs_test.go | 23 +++++++++ experimental/air/cmd/format.go | 1 + experimental/air/cmd/format_test.go | 2 + experimental/air/cmd/logdownload_test.go | 1 + experimental/air/cmd/runsubmit_test.go | 17 +++++++ 11 files changed, 124 insertions(+), 5 deletions(-) create mode 100644 acceptance/experimental/air/run-submit/b300.yaml diff --git a/acceptance/experimental/air/run-submit/b300.yaml b/acceptance/experimental/air/run-submit/b300.yaml new file mode 100644 index 00000000000..c69924cbcde --- /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 c20c8467f1e..52aa612190f 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 26161d63ad4..e2fa42eea67 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 c3357eaf401..36854932e2a 100644 --- a/experimental/air/cmd/compute.go +++ b/experimental/air/cmd/compute.go @@ -14,10 +14,11 @@ 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} +var gpuTypes = []gpuType{gpuType1xA10, gpuType1xH100, gpuType8xH100, gpuType8xB300} func validGPUTypesHint() string { names := make([]string, len(gpuTypes)) @@ -31,7 +32,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 +79,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 467e2acfa9a..5a099c5adf5 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.go b/experimental/air/cmd/convert_to_dabs.go index ad6ab819ca1..951c11b4879 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -82,7 +82,7 @@ does not contact the workspace.`, ctx := cmd.Context() yamlPath := args[0] - cfg, err := loadRunConfig(yamlPath) + cfg, err := loadRunConfigWithOverrides(ctx, yamlPath, nil) if err != nil { return err } diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index eadffe1c491..c8f35e7640c 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 TestConvertToDabsB300(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 := loadRunConfigWithOverrides(t.Context(), path, nil) + 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 e496a1e89cc..1271f89b4ae 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 4cc6dbf909f..842806e6113 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 b7717980f42..57b6982bddf 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 ba10c1b61dd..9f95b8da7f7 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -88,6 +88,23 @@ 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", "6", "", snapshotResult{}, nil) + require.Len(t, p.Environments, 1) + require.NotNil(t, p.Environments[0].Spec) + assert.Equal(t, "6", 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) From 96a12dcbd2f5b830766b8fab11b7ed01646f6d24 Mon Sep 17 00:00:00 2001 From: Caroline Chen <324939130+caroline-db@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:32:16 +0000 Subject: [PATCH 2/3] [air] Split DAB B300 support into follow-up --- experimental/air/cmd/convert_to_dabs.go | 6 +++++- experimental/air/cmd/convert_to_dabs_test.go | 13 ++++--------- experimental/air/cmd/runsubmit_test.go | 5 +++-- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 951c11b4879..a8ef5e830f2 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -82,7 +82,7 @@ does not contact the workspace.`, ctx := cmd.Context() yamlPath := args[0] - cfg, err := loadRunConfigWithOverrides(ctx, yamlPath, nil) + cfg, err := loadRunConfig(yamlPath) if err != nil { return err } @@ -115,6 +115,10 @@ does not contact the workspace.`, // env/secret/param sidecars) to write under generated_artifacts/. It does not touch the // 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") + } + // idempotency_token is intentionally not mapped: it dedups a single runs/submit // call, which has no analogue for a persistent, repeatedly-runnable bundle job. // diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index c8f35e7640c..bd18e438a5e 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -244,7 +244,7 @@ func TestConvertToDabsRuntimeVersionEnvOverride(t *testing.T) { assert.Equal(t, "7", get(t, root, env+".spec.environment_version").MustString()) } -func TestConvertToDabsB300(t *testing.T) { +func TestConvertToDabsRejectsB300(t *testing.T) { path := writeConfigFile(t, "run.yaml", ` experiment_name: b300 command: python train.py @@ -255,16 +255,11 @@ environment: version: 6 dependencies: [] `) - loaded, err := loadRunConfigWithOverrides(t.Context(), path, nil) - require.NoError(t, err) - - root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) + loaded, err := loadRunConfig(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()) + _, _, err = convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) + require.EqualError(t, err, "GPU_8xB300 is not yet supported by convert-to-dabs; use air run until DAB support is available") } // remote_volume can't be honored by a converted bundle (bundle deploy owns the diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 9f95b8da7f7..7a851907439 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -95,10 +95,11 @@ func TestBuildSubmitPayloadB300(t *testing.T) { Compute: &computeConfig{AcceleratorType: "GPU_8xB300", NumAccelerators: 8}, } - p := buildSubmitPayload(cfg, "/d/command.sh", "6", "", snapshotResult{}, nil) + 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, "6", p.Environments[0].Spec.EnvironmentVersion) + 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) From 95e3389dc6358febfab0f15436e72c09a3ec589b Mon Sep 17 00:00:00 2001 From: Caroline Chen <324939130+caroline-db@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:50:08 +0000 Subject: [PATCH 3/3] [air] Allow B300 conversion to DABs --- .../air/convert-to-dabs/b300.yaml | 8 +++ .../air/convert-to-dabs/output.txt | 66 +++++++++++++++++++ .../experimental/air/convert-to-dabs/script | 9 +++ .../air/convert-to-dabs/test.toml | 2 +- experimental/air/cmd/compute.go | 3 +- experimental/air/cmd/convert_to_dabs.go | 4 -- experimental/air/cmd/convert_to_dabs_test.go | 11 +++- 7 files changed, 94 insertions(+), 9 deletions(-) create mode 100644 acceptance/experimental/air/convert-to-dabs/b300.yaml 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 00000000000..7a91f98c757 --- /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 0922646a0d0..308b4120436 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 fd39053c927..cd9dff9f3d8 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 aa531fbd106..956c8efc16e 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/experimental/air/cmd/compute.go b/experimental/air/cmd/compute.go index 36854932e2a..0f302c4aaa2 100644 --- a/experimental/air/cmd/compute.go +++ b/experimental/air/cmd/compute.go @@ -17,7 +17,8 @@ const ( gpuType8xB300 gpuType = "GPU_8xB300" ) -// gpuTypes lists every valid type. Used for validation error messages. +// 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 { diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index a8ef5e830f2..ad6ab819ca1 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -115,10 +115,6 @@ does not contact the workspace.`, // env/secret/param sidecars) to write under generated_artifacts/. It does not touch the // 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") - } - // idempotency_token is intentionally not mapped: it dedups a single runs/submit // call, which has no analogue for a persistent, repeatedly-runnable bundle job. // diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index bd18e438a5e..fe96574237a 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -244,7 +244,7 @@ func TestConvertToDabsRuntimeVersionEnvOverride(t *testing.T) { assert.Equal(t, "7", get(t, root, env+".spec.environment_version").MustString()) } -func TestConvertToDabsRejectsB300(t *testing.T) { +func TestConvertToDabsSupportsB300(t *testing.T) { path := writeConfigFile(t, "run.yaml", ` experiment_name: b300 command: python train.py @@ -258,8 +258,13 @@ environment: loaded, err := loadRunConfig(path) require.NoError(t, err) - _, _, err = convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) - require.EqualError(t, err, "GPU_8xB300 is not yet supported by convert-to-dabs; use air run until DAB support is available") + 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