diff --git a/.nextchanges/bundles/b300.md b/.nextchanges/bundles/b300.md new file mode 100644 index 00000000000..292a4aeb60f --- /dev/null +++ b/.nextchanges/bundles/b300.md @@ -0,0 +1 @@ +* Add `GPU_8xB300` support to AI Runtime tasks. ([#6651](https://github.com/databricks/cli/pull/6651)) diff --git a/acceptance/bundle/validate/enum/databricks.yml b/acceptance/bundle/validate/enum/databricks.yml index 8bff766951c..0c0a2e598a6 100644 --- a/acceptance/bundle/validate/enum/databricks.yml +++ b/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" diff --git a/bundle/internal/schema/annotations.yml b/bundle/internal/schema/annotations.yml index b35ce0e4ae0..a54da9008cd 100644 --- a/bundle/internal/schema/annotations.yml +++ b/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": diff --git a/bundle/internal/validation/enum.go b/bundle/internal/validation/enum.go index df58193c898..202c3ee566c 100644 --- a/bundle/internal/validation/enum.go +++ b/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 } diff --git a/bundle/internal/validation/generated/enum_fields.go b/bundle/internal/validation/generated/enum_fields.go index 2ac14ca4ebe..5dd1dd9b557 100644 --- a/bundle/internal/validation/generated/enum_fields.go +++ b/bundle/internal/validation/generated/enum_fields.go @@ -99,11 +99,11 @@ var EnumFields = map[string][]string{ "resources.jobs.*.permissions[*].level": {"CAN_MANAGE", "CAN_MANAGE_RUN", "CAN_VIEW", "IS_OWNER"}, "resources.jobs.*.schedule.pause_status": {"PAUSED", "UNPAUSED"}, "resources.jobs.*.schedule.sql_condition.trigger_mode": {"QUERY_RETURNS_ROWS", "RESULT_VALUE_CHANGES"}, - "resources.jobs.*.tasks[*].ai_runtime_task.deployments[*].compute.accelerator_type": {"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"}, + "resources.jobs.*.tasks[*].ai_runtime_task.deployments[*].compute.accelerator_type": {"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"}, "resources.jobs.*.tasks[*].compute.hardware_accelerator": {"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"}, "resources.jobs.*.tasks[*].condition_task.op": {"EQUAL_TO", "GREATER_THAN", "GREATER_THAN_OR_EQUAL", "LESS_THAN", "LESS_THAN_OR_EQUAL", "NOT_EQUAL"}, "resources.jobs.*.tasks[*].dbt_task.source": {"GIT", "WORKSPACE"}, - "resources.jobs.*.tasks[*].for_each_task.task.ai_runtime_task.deployments[*].compute.accelerator_type": {"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"}, + "resources.jobs.*.tasks[*].for_each_task.task.ai_runtime_task.deployments[*].compute.accelerator_type": {"GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"}, "resources.jobs.*.tasks[*].for_each_task.task.compute.hardware_accelerator": {"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"}, "resources.jobs.*.tasks[*].for_each_task.task.condition_task.op": {"EQUAL_TO", "GREATER_THAN", "GREATER_THAN_OR_EQUAL", "LESS_THAN", "LESS_THAN_OR_EQUAL", "NOT_EQUAL"}, "resources.jobs.*.tasks[*].for_each_task.task.dbt_task.source": {"GIT", "WORKSPACE"}, diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index cb9f9d337f6..69166a4eee5 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -8508,12 +8508,14 @@ "enum": [ "GPU_1xA10", "GPU_1xH100", - "GPU_8xH100" + "GPU_8xH100", + "GPU_8xB300" ], "enumDescriptions": [ "[Public Preview]", "[Public Preview]", - "[Public Preview]" + "[Public Preview]", + "" ] }, { diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 8c2a9b6ecc6..06d2e0311ac 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 f557065c087..f6d229fccb1 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 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 diff --git a/python/databricks/bundles/jobs/_models/compute_spec_accelerator_type.py b/python/databricks/bundles/jobs/_models/compute_spec_accelerator_type.py index c3428f40d43..218cce4766f 100644 --- a/python/databricks/bundles/jobs/_models/compute_spec_accelerator_type.py +++ b/python/databricks/bundles/jobs/_models/compute_spec_accelerator_type.py @@ -14,8 +14,10 @@ class ComputeSpecAcceleratorType(Enum): GPU_1X_A10 = "GPU_1xA10" GPU_1X_H100 = "GPU_1xH100" GPU_8X_H100 = "GPU_8xH100" + GPU_8X_B300 = "GPU_8xB300" ComputeSpecAcceleratorTypeParam = ( - Literal["GPU_1xA10", "GPU_1xH100", "GPU_8xH100"] | ComputeSpecAcceleratorType + Literal["GPU_1xA10", "GPU_1xH100", "GPU_8xH100", "GPU_8xB300"] + | ComputeSpecAcceleratorType )