[CNXT-2446][AIR] support RTX6000 accelerators

.nextchanges/air/rtx6000.md

@@ -0,0 +1 @@
+* Add workload support for the GPU_1xRTX6000 accelerator type in enabled workspaces. ([#7214](https://github.com/databricks/cli/pull/7214))

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.

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

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

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

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

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

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

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

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

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

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

cmd/air/runconfig_test.go

@@ -821,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:    `{}`,

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