[air] Add GPU_8xB300 support to DABs

.nextchanges/bundles/b300.md

@@ -0,0 +1 @@
+* Add `GPU_8xB300` support to AI Runtime tasks. ([#6651](https://github.com/databricks/cli/pull/6651))

acceptance/bundle/validate/enum/databricks.yml

@@ -57,3 +57,19 @@ resources:
               ebs_volume_type: "INVALID_VOLUME_TYPE"
             node_type_id: "i3.xlarge"
             num_workers: 1
+
+    b300_valid:
+      tasks:
+        - task_key: "b300"
+          environment_key: "default"
+          ai_runtime_task:
+            experiment: "b300"
+            deployments:
+              - command_path: "/Workspace/train.sh"
+                compute:
+                  accelerator_type: "GPU_8xB300"
+                  accelerator_count: 8
+      environments:
+        - environment_key: "default"
+          spec:
+            environment_version: "6"

bundle/internal/schema/annotations.yml

@@ -1149,6 +1149,17 @@ resources:
                 The email of an active workspace user. Non-admin users can only set this field to their own email.
         "tasks":
           "$fields":
+            "ai_runtime_task":
+              "$fields":
+                "deployments":
+                  "$fields":
+                    "compute":
+                      "$fields":
+                        "accelerator_type":
+                          "$type":
+                            "enum":
+                              - |-
+                                GPU_8xB300
             "alert_task":
               "$fields":
                 "subscribers":

bundle/internal/validation/enum.go

@@ -17,6 +17,7 @@ import (
 	"github.com/databricks/cli/libs/structs/structpath"
 	"github.com/databricks/cli/libs/structs/structtag"
 	"github.com/databricks/cli/libs/structs/structwalk"
+	"github.com/databricks/databricks-sdk-go/service/jobs"
 )
 
 type EnumPatternInfo struct {
@@ -29,6 +30,12 @@ type EnumPatternInfo struct {
 	Values string
 }
 
+// additionalEnumValues contains values supported by the CLI before they are
+// published in the corresponding SDK enum.
+var additionalEnumValues = map[reflect.Type][]string{
+	reflect.TypeFor[jobs.ComputeSpecAcceleratorType](): {"GPU_8xB300"},
+}
+
 // isEnumType checks if a type is an enum (string type with a Values() method)
 func isEnumType(typ reflect.Type) bool {
 	// Must be a string type
@@ -105,6 +112,11 @@ func getEnumValues(typ reflect.Type) ([]string, error) {
 		value := valuesSlice.Index(i)
 		enumStrings = append(enumStrings, value.String())
 	}
+	for _, value := range additionalEnumValues[typ] {
+		if !slices.Contains(enumStrings, value) {
+			enumStrings = append(enumStrings, value)
+		}
+	}
 
 	return enumStrings, nil
 }

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.
 	//

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 TestConvertToDabsB300(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