From ff3869d7a0c2aa2db2650a5c6d68f1def4b1bcd8 Mon Sep 17 00:00:00 2001 From: Lin Yuan Date: Wed, 28 Oct 2026 03:54:06 +0000 Subject: [PATCH 1/4] air: support RTX6000 accelerators --- .nextchanges/air/rtx6000.md | 1 + acceptance/air/run-submit/output.txt | 49 ++++++++++++++++++++++++++ acceptance/air/run-submit/rtx6000.yaml | 9 +++++ acceptance/air/run-submit/script | 6 ++++ cmd/air/compute.go | 15 ++++---- cmd/air/compute_test.go | 7 +++- cmd/air/format.go | 13 +++---- cmd/air/format_test.go | 2 ++ cmd/air/logdownload_test.go | 1 + cmd/air/runconfig_test.go | 3 +- cmd/air/runsubmit_test.go | 19 ++++++++++ 11 files changed, 110 insertions(+), 15 deletions(-) create mode 100644 .nextchanges/air/rtx6000.md create mode 100644 acceptance/air/run-submit/rtx6000.yaml diff --git a/.nextchanges/air/rtx6000.md b/.nextchanges/air/rtx6000.md new file mode 100644 index 00000000000..1bc2b190c9e --- /dev/null +++ b/.nextchanges/air/rtx6000.md @@ -0,0 +1 @@ +* Add support for the GPU_1xRTX6000 accelerator in AI Runtime commands. diff --git a/acceptance/air/run-submit/output.txt b/acceptance/air/run-submit/output.txt index 516c4f122d7..83bb8c0dbee 100644 --- a/acceptance/air/run-submit/output.txt +++ b/acceptance/air/run-submit/output.txt @@ -153,3 +153,52 @@ Stream logs after submission using: ] } } + +=== submit an RTX6000 workload with AI Runtime 6 +>>> [CLI] air run -f rtx6000.yaml +Submitting experiment: rtx6000-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 air logs 555 + +=== the RTX6000 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": "rtx6000-submit-smoke", + "tasks": [ + { + "ai_runtime_task": { + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/rtx6000-submit-smoke/rtx6000-submit-smoke_[RUN_ID]/command.sh", + "compute": { + "accelerator_count": 1, + "accelerator_type": "GPU_1xRTX6000" + } + } + ], + "experiment": "rtx6000-submit-smoke" + }, + "environment_key": "default", + "max_retries": 0, + "run_if": "ALL_SUCCESS", + "task_key": "rtx6000-submit-smoke" + } + ] + } +} diff --git a/acceptance/air/run-submit/rtx6000.yaml b/acceptance/air/run-submit/rtx6000.yaml new file mode 100644 index 00000000000..2b0caa622a9 --- /dev/null +++ b/acceptance/air/run-submit/rtx6000.yaml @@ -0,0 +1,9 @@ +experiment_name: rtx6000-submit-smoke +command: nvidia-smi +compute: + accelerator_type: GPU_1xRTX6000 + num_accelerators: 1 +environment: + version: 6 + dependencies: [] +max_retries: 0 diff --git a/acceptance/air/run-submit/script b/acceptance/air/run-submit/script index 50627f98b21..da5178173af 100644 --- a/acceptance/air/run-submit/script +++ b/acceptance/air/run-submit/script @@ -24,4 +24,10 @@ trace $CLI 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 +title "submit an RTX6000 workload with AI Runtime 6" +trace $CLI air run -f rtx6000.yaml + +title "the RTX6000 accelerator and runtime reach the submit payload" +trace print_requests.py //api/2.2/jobs/runs/submit + rm -fr .git diff --git a/cmd/air/compute.go b/cmd/air/compute.go index 1977eb24441..20a3895c09b 100644 --- a/cmd/air/compute.go +++ b/cmd/air/compute.go @@ -11,15 +11,16 @@ import ( type gpuType string const ( - gpuType1xA10 gpuType = "GPU_1xA10" - gpuType8xH100 gpuType = "GPU_8xH100" - gpuType1xH100 gpuType = "GPU_1xH100" - gpuType8xB300 gpuType = "GPU_8xB300" + gpuType1xA10 gpuType = "GPU_1xA10" + gpuType8xH100 gpuType = "GPU_8xH100" + gpuType1xH100 gpuType = "GPU_1xH100" + gpuType8xB300 gpuType = "GPU_8xB300" + gpuType1xRTX6000 gpuType = "GPU_1xRTX6000" ) // gpuTypes lists every accelerator type understood by this CLI. // Workspace availability is enforced by the server. -var gpuTypes = []gpuType{gpuType1xA10, gpuType1xH100, gpuType8xH100, gpuType8xB300} +var gpuTypes = []gpuType{gpuType1xA10, gpuType1xH100, gpuType8xH100, gpuType8xB300, gpuType1xRTX6000} func validGPUTypesHint() string { names := make([]string, len(gpuTypes)) @@ -33,7 +34,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, gpuType8xB300: + case gpuType1xA10, gpuType8xH100, gpuType1xH100, gpuType8xB300, gpuType1xRTX6000: return gpuType(value), nil } return "", fmt.Errorf("invalid GPU type %q: %s", value, validGPUTypesHint()) @@ -78,7 +79,7 @@ func parsePriorityClass(value string) (priorityClass, error) { // round multiple of this since accelerators are allocated in whole nodes. func gpusPerNode(g gpuType) (int, error) { switch g { - case gpuType1xA10, gpuType1xH100: + case gpuType1xA10, gpuType1xH100, gpuType1xRTX6000: return 1, nil case gpuType8xH100, gpuType8xB300: return 8, nil diff --git a/cmd/air/compute_test.go b/cmd/air/compute_test.go index 5a099c5adf5..a8d472c3f70 100644 --- a/cmd/air/compute_test.go +++ b/cmd/air/compute_test.go @@ -17,6 +17,7 @@ func TestParseGPUType(t *testing.T) { {"GPU_8xH100", gpuType8xH100}, {"GPU_1xH100", gpuType1xH100}, {"GPU_8xB300", gpuType8xB300}, + {"GPU_1xRTX6000", gpuType1xRTX6000}, } for _, tt := range tests { t.Run(tt.in, func(t *testing.T) { @@ -30,7 +31,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_8xb300", "GPU_2xH100", "h100_80gb", "a10", "b200", ""} { + for _, in := range []string{"gpu_1xa10", "GPU_1XA10", "gpu_8xb300", "gpu_1xrtx6000", "GPU_1XRTX6000", "RTX6000", "GPU_8xRTX6000", "GPU_2xH100", "h100_80gb", "a10", "b200", ""} { t.Run(in, func(t *testing.T) { _, err := parseGPUType(in) require.Error(t, err) @@ -48,6 +49,7 @@ func TestGPUsPerNode(t *testing.T) { {gpuType1xH100, 1}, {gpuType8xH100, 8}, {gpuType8xB300, 8}, + {gpuType1xRTX6000, 1}, } for _, tt := range tests { t.Run(string(tt.in), func(t *testing.T) { @@ -101,6 +103,9 @@ 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"}, + {"RTX6000 single node", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xRTX6000"}, ""}, + {"RTX6000 multiple nodes", computeConfig{NumAccelerators: 3, AcceleratorType: "GPU_1xRTX6000"}, ""}, + {"RTX6000 zero", computeConfig{NumAccelerators: 0, AcceleratorType: "GPU_1xRTX6000"}, "must be positive"}, {"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"}, ""}, diff --git a/cmd/air/format.go b/cmd/air/format.go index 1271f89b4ae..89e2420fd3b 100644 --- a/cmd/air/format.go +++ b/cmd/air/format.go @@ -77,12 +77,13 @@ func forceLiteralBlockStrings(node *yaml.Node) { // gpuDisplayNames maps the GPU identifiers returned by the backend to the short // names we show to users. Unknown identifiers are shown unchanged. var gpuDisplayNames = map[string]string{ - "h100_80gb": "H100", - "a10": "A10", - "GPU_1xA10": "A10", - "GPU_8xH100": "H100", - "GPU_1xH100": "H100", - "GPU_8xB300": "B300", + "h100_80gb": "H100", + "a10": "A10", + "GPU_1xA10": "A10", + "GPU_8xH100": "H100", + "GPU_1xH100": "H100", + "GPU_8xB300": "B300", + "GPU_1xRTX6000": "RTX6000", } // isRunning reports whether a run is currently executing (lifecycle RUNNING), as diff --git a/cmd/air/format_test.go b/cmd/air/format_test.go index 842806e6113..7b4bbc80154 100644 --- a/cmd/air/format_test.go +++ b/cmd/air/format_test.go @@ -70,6 +70,7 @@ func TestGpuDisplayName(t *testing.T) { assert.Equal(t, "H100", gpuDisplayName("GPU_8xH100")) assert.Equal(t, "H100", gpuDisplayName("GPU_1xH100")) assert.Equal(t, "B300", gpuDisplayName("GPU_8xB300")) + assert.Equal(t, "RTX6000", gpuDisplayName("GPU_1xRTX6000")) // Unknown identifiers pass through unchanged. assert.Equal(t, "b200", gpuDisplayName("b200")) assert.Empty(t, gpuDisplayName("")) @@ -228,6 +229,7 @@ func TestAcceleratorLabel(t *testing.T) { 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)) + assert.Equal(t, "1x RTX6000", acceleratorLabel("GPU_1xRTX6000", 1)) // The RPC may report a count without a recognized type. assert.Equal(t, "8x", acceleratorLabel("", 8)) } diff --git a/cmd/air/logdownload_test.go b/cmd/air/logdownload_test.go index 57b6982bddf..9a065cbf0b0 100644 --- a/cmd/air/logdownload_test.go +++ b/cmd/air/logdownload_test.go @@ -43,6 +43,7 @@ func TestResolveNodeCount(t *testing.T) { }{ {"GPU_1xA10", 2, 2}, {"GPU_1xH100", 4, 4}, + {"GPU_1xRTX6000", 3, 3}, {"GPU_8xH100", 16, 2}, {"GPU_8xB300", 16, 2}, } diff --git a/cmd/air/runconfig_test.go b/cmd/air/runconfig_test.go index 622c8408b51..6115715bf0b 100644 --- a/cmd/air/runconfig_test.go +++ b/cmd/air/runconfig_test.go @@ -808,7 +808,8 @@ func TestRunCommandHelpComputeOptions(t *testing.T) { builtInOptions := ` GPU_1xA10 - 1 accelerator per node GPU_1xH100 - 1 accelerator per node GPU_8xH100 - 8 accelerators per node - GPU_8xB300 - 8 accelerators per node` + GPU_8xB300 - 8 accelerators per node + GPU_1xRTX6000 - 1 accelerator per node` tests := []struct { name string statusCode int diff --git a/cmd/air/runsubmit_test.go b/cmd/air/runsubmit_test.go index 5c5a2e7f23a..bbd6718c264 100644 --- a/cmd/air/runsubmit_test.go +++ b/cmd/air/runsubmit_test.go @@ -113,6 +113,25 @@ func TestBuildSubmitPayloadB300(t *testing.T) { assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorType("GPU_8xB300"), AcceleratorCount: 8}, p.Tasks[0].AiRuntimeTask.Deployments[0].Compute) } +func TestBuildSubmitPayloadRTX6000(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, ` +experiment_name: rtx6000 +command: nvidia-smi +compute: + accelerator_type: GPU_1xRTX6000 + num_accelerators: 1 +`)) + require.NoError(t, err) + + p := buildSubmitPayload(cfg, "/d/command.sh", "6", "", snapshotResult{}, nil) + require.Len(t, p.Tasks, 1) + require.NotNil(t, p.Tasks[0].AiRuntimeTask) + require.Len(t, p.Tasks[0].AiRuntimeTask.Deployments, 1) + computeJSON, err := json.Marshal(p.Tasks[0].AiRuntimeTask.Deployments[0].Compute) + require.NoError(t, err) + assert.JSONEq(t, `{"accelerator_type":"GPU_1xRTX6000","accelerator_count":1}`, string(computeJSON)) +} + func TestSubmitRun(t *testing.T) { for _, tc := range []struct { name string From 5b6fe4271d94114b07812a031a7416bda8be2f24 Mon Sep 17 00:00:00 2001 From: Lin Yuan Date: Wed, 28 Oct 2026 06:10:32 +0000 Subject: [PATCH 2/4] air: link RTX6000 release note to pull request --- .nextchanges/air/rtx6000.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.nextchanges/air/rtx6000.md b/.nextchanges/air/rtx6000.md index 1bc2b190c9e..6f218cfbd6d 100644 --- a/.nextchanges/air/rtx6000.md +++ b/.nextchanges/air/rtx6000.md @@ -1 +1 @@ -* Add support for the GPU_1xRTX6000 accelerator in AI Runtime commands. +* Add support for the GPU_1xRTX6000 accelerator in AI Runtime commands. ([#7214](https://github.com/databricks/cli/pull/7214)) From 58bfd2928d245c1253cec150e351cc78668c5afb Mon Sep 17 00:00:00 2001 From: Lin Yuan Date: Wed, 28 Oct 2026 06:21:18 +0000 Subject: [PATCH 3/4] air: require workspace confirmation for RTX6000 help --- acceptance/air/config-help/output.txt | 1 + acceptance/air/config-help/test.toml | 3 ++- cmd/air/runconfig.go | 4 ++++ cmd/air/runconfig_test.go | 8 ++++++-- 4 files changed, 13 insertions(+), 3 deletions(-) diff --git a/acceptance/air/config-help/output.txt b/acceptance/air/config-help/output.txt index be048a1e392..16fedb9ffa9 100644 --- a/acceptance/air/config-help/output.txt +++ b/acceptance/air/config-help/output.txt @@ -77,6 +77,7 @@ config.compute Supported accelerator types: GPU_1xH100 - 1 accelerator per node [Private Preview] GPU_8xH100 - 8 accelerators per node + GPU_1xRTX6000 - 1 accelerator per node Fields: num_accelerators Total number of GPUs to allocate. diff --git a/acceptance/air/config-help/test.toml b/acceptance/air/config-help/test.toml index 965694857c9..79cb5772cdf 100644 --- a/acceptance/air/config-help/test.toml +++ b/acceptance/air/config-help/test.toml @@ -10,7 +10,8 @@ Response.Body = ''' { "compute_options": [ {"hardware_accelerator": "GPU_1xH100", "per_node_accelerator_count": 1, "launch_stage": "PRIVATE_PREVIEW"}, - {"hardware_accelerator": "GPU_8xH100", "per_node_accelerator_count": 8, "launch_stage": "GA"} + {"hardware_accelerator": "GPU_8xH100", "per_node_accelerator_count": 8, "launch_stage": "GA"}, + {"hardware_accelerator": "GPU_1xRTX6000", "per_node_accelerator_count": 1, "launch_stage": "GA"} ] } ''' diff --git a/cmd/air/runconfig.go b/cmd/air/runconfig.go index d7e4dc0f100..31c701cdc12 100644 --- a/cmd/air/runconfig.go +++ b/cmd/air/runconfig.go @@ -902,6 +902,10 @@ func renderConfigField(w io.Writer, f configField, computeOptions []computeOptio func fallbackComputeOptions() []computeOption { options := make([]computeOption, 0, len(gpuTypes)) for _, acceleratorType := range gpuTypes { + // RTX6000 availability must be confirmed by the workspace API during rollout. + if acceleratorType == gpuType1xRTX6000 { + continue + } perNode, err := gpusPerNode(acceleratorType) if err != nil { continue diff --git a/cmd/air/runconfig_test.go b/cmd/air/runconfig_test.go index 6115715bf0b..e3c2f34520a 100644 --- a/cmd/air/runconfig_test.go +++ b/cmd/air/runconfig_test.go @@ -808,8 +808,7 @@ func TestRunCommandHelpComputeOptions(t *testing.T) { builtInOptions := ` GPU_1xA10 - 1 accelerator per node GPU_1xH100 - 1 accelerator per node GPU_8xH100 - 8 accelerators per node - GPU_8xB300 - 8 accelerators per node - GPU_1xRTX6000 - 1 accelerator per node` + GPU_8xB300 - 8 accelerators per node` tests := []struct { name string statusCode int @@ -822,6 +821,11 @@ func TestRunCommandHelpComputeOptions(t *testing.T) { response: `{"compute_options":[{"hardware_accelerator":"GPU_8xH100","per_node_accelerator_count":4,"launch_stage":"PUBLIC_PREVIEW"},{"hardware_accelerator":"GPU_FUTURE","per_node_accelerator_count":16},{"hardware_accelerator":"GPU_1xA10"}]}`, wantOptions: " GPU_8xH100 - 4 accelerators per node [Public Preview]\n GPU_FUTURE - 16 accelerators per node\n GPU_1xA10 - 1 accelerator per node", }, + { + name: "workspace confirms RTX6000 availability", + response: `{"compute_options":[{"hardware_accelerator":"GPU_1xRTX6000","per_node_accelerator_count":1,"launch_stage":"GA"}]}`, + wantOptions: " GPU_1xRTX6000 - 1 accelerator per node", + }, { name: "empty response", response: `{}`, From 509dc4464aacaac31f3444b237a3f46c3f9b48d6 Mon Sep 17 00:00:00 2001 From: Lin Yuan Date: Wed, 28 Oct 2026 08:48:09 -0700 Subject: [PATCH 4/4] Update .nextchanges/air/rtx6000.md Co-authored-by: caroline-db --- .nextchanges/air/rtx6000.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.nextchanges/air/rtx6000.md b/.nextchanges/air/rtx6000.md index 6f218cfbd6d..6b6732bb293 100644 --- a/.nextchanges/air/rtx6000.md +++ b/.nextchanges/air/rtx6000.md @@ -1 +1 @@ -* Add support for the GPU_1xRTX6000 accelerator in AI Runtime commands. ([#7214](https://github.com/databricks/cli/pull/7214)) +* Add workload support for the GPU_1xRTX6000 accelerator type in enabled workspaces. ([#7214](https://github.com/databricks/cli/pull/7214))