[air] Add GPU_8xB300 support to the experimental AIR CLI

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: []

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

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)

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"]

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

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"
+      }
+    ]
+  }
+}

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

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

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"},

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) {

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

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))
 }

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))

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)