Skip to content

Commit f9d4c56

Browse files
kupratyu-splunkkupratyu-splunk
authored
Merge pull request #142 from splunk/fix/k0s-installer-node-mutations
feat(AIP-4445): Add RTX pro 6000 support for ai tier
2 parents 7ac4050 + 3eb4134 commit f9d4c56

8 files changed

Lines changed: 602 additions & 48 deletions

File tree

‎config/configs/applications.yaml‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,17 @@ applications:
5555
max_ongoing_requests: 6
5656
ray_actor_options:
5757
num_gpus: 2
58+
RTX_PRO_6000_BLACKWELL:
59+
autoscaling_config:
60+
max_replicas: {{.Replicas.Gemma431bIt}}
61+
min_replicas: {{.Replicas.Gemma431bIt}}
62+
target_ongoing_requests: 4
63+
max_ongoing_requests: 10
64+
ray_actor_options:
65+
num_gpus: 1
66+
resources:
67+
"gpu_count:2": 0.001
68+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
5869
options:
5970
autoscaling_config:
6071
max_replicas: {{.Replicas.Gemma431bIt}}
@@ -78,6 +89,14 @@ applications:
7889
max_num_batched_tokens: 4096
7990
max_num_seqs: 2
8091
tensor_parallel_size: 2
92+
RTX_PRO_6000_BLACKWELL:
93+
engine_args:
94+
dtype: bfloat16
95+
gpu_memory_utilization: 0.9
96+
max_model_len: 120000
97+
max_num_batched_tokens: 4096
98+
max_num_seqs: 1
99+
tensor_parallel_size: 1
81100
model_config:
82101
openai_serving_config:
83102
chat:
@@ -162,6 +181,17 @@ applications:
162181
L40S:
163182
ray_actor_options:
164183
num_gpus: 1
184+
RTX_PRO_6000_BLACKWELL:
185+
autoscaling_config:
186+
max_replicas: {{.Replicas.GptOss20b}}
187+
min_replicas: {{.Replicas.GptOss20b}}
188+
target_ongoing_requests: 4
189+
max_ongoing_requests: 8
190+
ray_actor_options:
191+
num_gpus: 0.5
192+
resources:
193+
"gpu_count:2": 0.001
194+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
165195
options:
166196
autoscaling_config:
167197
max_replicas: {{.Replicas.GptOss20b}}
@@ -178,6 +208,10 @@ applications:
178208
engine_args:
179209
gpu_memory_utilization: 0.95
180210
tensor_parallel_size: 1
211+
RTX_PRO_6000_BLACKWELL:
212+
engine_args:
213+
gpu_memory_utilization: 0.5
214+
tensor_parallel_size: 1
181215
model_config:
182216
openai_serving_config:
183217
chat:
@@ -251,6 +285,12 @@ applications:
251285
L40S:
252286
ray_actor_options:
253287
num_gpus: 0.075
288+
RTX_PRO_6000_BLACKWELL:
289+
ray_actor_options:
290+
num_gpus: 0.031
291+
resources:
292+
"gpu_count:2": 0.001
293+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
254294
options:
255295
autoscaling_config:
256296
max_replicas: {{.Replicas.UaeLarge}}
@@ -266,6 +306,9 @@ applications:
266306
L40S:
267307
engine_args:
268308
gpu_memory_utilization: 0.075
309+
RTX_PRO_6000_BLACKWELL:
310+
engine_args:
311+
gpu_memory_utilization: 0.031
269312
model_config:
270313
engine_args:
271314
gpu_memory_utilization: 0.15
@@ -317,6 +360,12 @@ applications:
317360
H100:
318361
ray_actor_options:
319362
num_gpus: 0.005
363+
RTX_PRO_6000_BLACKWELL:
364+
ray_actor_options:
365+
num_gpus: 0.004
366+
resources:
367+
"gpu_count:2": 0.001
368+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
320369
options:
321370
autoscaling_config:
322371
max_replicas: {{.Replicas.AllMinilmL6V2}}
@@ -329,6 +378,9 @@ applications:
329378
H100:
330379
engine_args:
331380
gpu_memory_utilization: 0.005
381+
RTX_PRO_6000_BLACKWELL:
382+
engine_args:
383+
gpu_memory_utilization: 0.004
332384
model_config:
333385
engine_args:
334386
gpu_memory_utilization: 0.01
@@ -380,6 +432,12 @@ applications:
380432
H100:
381433
ray_actor_options:
382434
num_gpus: 0.005
435+
RTX_PRO_6000_BLACKWELL:
436+
ray_actor_options:
437+
num_gpus: 0.004
438+
resources:
439+
"gpu_count:2": 0.001
440+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
383441
options:
384442
autoscaling_config:
385443
max_replicas: {{.Replicas.BiEncoder}}
@@ -392,6 +450,9 @@ applications:
392450
H100:
393451
engine_args:
394452
gpu_memory_utilization: 0.005
453+
RTX_PRO_6000_BLACKWELL:
454+
engine_args:
455+
gpu_memory_utilization: 0.004
395456
model_config:
396457
engine_args:
397458
gpu_memory_utilization: 0.01
@@ -447,6 +508,12 @@ applications:
447508
L40S:
448509
ray_actor_options:
449510
num_gpus: 0.1
511+
RTX_PRO_6000_BLACKWELL:
512+
ray_actor_options:
513+
num_gpus: 0.05
514+
resources:
515+
"gpu_count:2": 0.001
516+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
450517
options:
451518
autoscaling_config:
452519
max_replicas: {{.Replicas.MbartTranslator}}
@@ -500,6 +567,12 @@ applications:
500567
L40S:
501568
ray_actor_options:
502569
num_gpus: 0.05
570+
RTX_PRO_6000_BLACKWELL:
571+
ray_actor_options:
572+
num_gpus: 0.021
573+
resources:
574+
"gpu_count:2": 0.001
575+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
503576
options:
504577
autoscaling_config:
505578
max_replicas: {{.Replicas.XlmRobertaLanguageClassifier}}
@@ -515,6 +588,9 @@ applications:
515588
L40S:
516589
engine_args:
517590
gpu_memory_utilization: 0.05
591+
RTX_PRO_6000_BLACKWELL:
592+
engine_args:
593+
gpu_memory_utilization: 0.021
518594
model_config:
519595
engine_args:
520596
gpu_memory_utilization: 0.1
@@ -605,6 +681,12 @@ applications:
605681
H100:
606682
ray_actor_options:
607683
num_gpus: 0.005
684+
RTX_PRO_6000_BLACKWELL:
685+
ray_actor_options:
686+
num_gpus: 0.004
687+
resources:
688+
"gpu_count:2": 0.001
689+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
608690
options:
609691
autoscaling_config:
610692
max_replicas: {{.Replicas.CrossEncoder}}
@@ -617,6 +699,9 @@ applications:
617699
H100:
618700
engine_args:
619701
gpu_memory_utilization: 0.005
702+
RTX_PRO_6000_BLACKWELL:
703+
engine_args:
704+
gpu_memory_utilization: 0.004
620705
model_config:
621706
engine_args:
622707
gpu_memory_utilization: 0.01
@@ -672,6 +757,12 @@ applications:
672757
L40S:
673758
ray_actor_options:
674759
num_gpus: 0.05
760+
RTX_PRO_6000_BLACKWELL:
761+
ray_actor_options:
762+
num_gpus: 0.021
763+
resources:
764+
"gpu_count:2": 0.001
765+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
675766
options:
676767
autoscaling_config:
677768
max_replicas: {{.Replicas.E5LanguageClassifier}}
@@ -687,6 +778,9 @@ applications:
687778
L40S:
688779
engine_args:
689780
gpu_memory_utilization: 0.05
781+
RTX_PRO_6000_BLACKWELL:
782+
engine_args:
783+
gpu_memory_utilization: 0.021
690784
model_config:
691785
engine_args:
692786
gpu_memory_utilization: 0.1
@@ -741,6 +835,12 @@ applications:
741835
L40S:
742836
ray_actor_options:
743837
num_gpus: 0.025
838+
RTX_PRO_6000_BLACKWELL:
839+
ray_actor_options:
840+
num_gpus: 0.013
841+
resources:
842+
"gpu_count:2": 0.001
843+
"accelerator_type:RTX_PRO_6000_BLACKWELL": 0.001
744844
options:
745845
autoscaling_config:
746846
max_replicas: {{.Replicas.PromptInjectionCrossEncoder}}

‎config/configs/features/saia.yaml‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,8 @@ instanceScale:
2323
h100-1-gpu: 2
2424
H100_NVL:
2525
h100-nvl-0-gpu: 1
26-
h100-nvl-1-gpu: 2
26+
h100-nvl-1-gpu: 2
27+
RTX_PRO_6000_BLACKWELL:
28+
rtx-pro-6000-blackwell-0-gpu: 1
29+
rtx-pro-6000-blackwell-1-gpu: 0
30+
rtx-pro-6000-blackwell-2-gpu: 1

‎config/configs/instance.yaml‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,4 +85,37 @@ H100_NVL:
8585
cpu: "16"
8686
memory: "48Gi"
8787
ephemeral-storage: "100Gi"
88-
nvidia.com/gpu: "1"
88+
nvidia.com/gpu: "1"
89+
RTX_PRO_6000_BLACKWELL:
90+
- tier: rtx-pro-6000-blackwell-0-gpu
91+
gpusPerPod: 0
92+
env:
93+
NVIDIA_VISIBLE_DEVICES: void
94+
resources:
95+
limits:
96+
cpu: "16"
97+
memory: "24Gi"
98+
ephemeral-storage: "50Gi"
99+
nvidia.com/gpu: "0"
100+
requests:
101+
cpu: "4"
102+
- tier: rtx-pro-6000-blackwell-1-gpu
103+
gpusPerPod: 1
104+
resources:
105+
requests:
106+
cpu: "4"
107+
limits:
108+
cpu: "16"
109+
memory: "48Gi"
110+
ephemeral-storage: "150Gi"
111+
nvidia.com/gpu: "1"
112+
- tier: rtx-pro-6000-blackwell-2-gpu
113+
gpusPerPod: 2
114+
resources:
115+
requests:
116+
cpu: "4"
117+
limits:
118+
cpu: "8"
119+
memory: "96Gi"
120+
ephemeral-storage: "350Gi"
121+
nvidia.com/gpu: "2"

‎pkg/ai/raybuilder/configmap_apps_test.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,60 @@ func Test_ApplicationsYAML_IsWellFormed(t *testing.T) {
151151
}
152152
}
153153

154+
// Test_ApplicationsYAML_BiEncoderRTXMemoryMatchesRayAllocation prevents Ray's
155+
// fractional GPU reservation from drifting below the model server's actual
156+
// memory allocation for RTX Pro 6000 Blackwell.
157+
func Test_ApplicationsYAML_BiEncoderRTXMemoryMatchesRayAllocation(t *testing.T) {
158+
type rayActorOptions struct {
159+
NumGPUs float64 `yaml:"num_gpus"`
160+
}
161+
type gpuTypeOptions struct {
162+
RayActorOptions rayActorOptions `yaml:"ray_actor_options"`
163+
}
164+
type deploymentConfig struct {
165+
GPUTypeOptionsOverride map[string]gpuTypeOptions `yaml:"gpu_type_options_override"`
166+
}
167+
type engineArgs struct {
168+
GPUMemoryUtilization float64 `yaml:"gpu_memory_utilization"`
169+
}
170+
type modelConfigOverride struct {
171+
EngineArgs engineArgs `yaml:"engine_args"`
172+
}
173+
type app struct {
174+
Name string `yaml:"name"`
175+
Args struct {
176+
DeploymentConfigs map[string]deploymentConfig `yaml:"deployment_configs"`
177+
ModelDefinition struct {
178+
GPUTypeModelConfigOverride map[string]modelConfigOverride `yaml:"gpu_type_model_config_override"`
179+
} `yaml:"model_definition"`
180+
} `yaml:"args"`
181+
}
182+
var doc struct {
183+
Applications []app `yaml:"applications"`
184+
}
185+
require.NoError(t, yaml.Unmarshal([]byte(maskGoTemplates(readApplicationsYAMLFromRepo(t))), &doc))
186+
187+
const accelerator = "RTX_PRO_6000_BLACKWELL"
188+
for _, application := range doc.Applications {
189+
if application.Name != "BiEncoder" {
190+
continue
191+
}
192+
deployment, ok := application.Args.DeploymentConfigs["EmbeddingModelDeployment"]
193+
require.True(t, ok, "BiEncoder is missing EmbeddingModelDeployment")
194+
rayOptions, ok := deployment.GPUTypeOptionsOverride[accelerator]
195+
require.True(t, ok, "BiEncoder is missing the %s Ray allocation", accelerator)
196+
modelOverride, ok := application.Args.ModelDefinition.GPUTypeModelConfigOverride[accelerator]
197+
require.True(t, ok, "BiEncoder is missing the %s model memory override", accelerator)
198+
199+
require.InDelta(t, 0.004, rayOptions.RayActorOptions.NumGPUs, 0.000001)
200+
require.InDelta(t, rayOptions.RayActorOptions.NumGPUs,
201+
modelOverride.EngineArgs.GPUMemoryUtilization, 0.000001,
202+
"BiEncoder's RTX model memory utilization must match its Ray GPU reservation")
203+
return
204+
}
205+
t.Fatal("BiEncoder application not found")
206+
}
207+
154208
func keys(m map[string]any) []string {
155209
out := make([]string, 0, len(m))
156210
for k := range m {

‎tools/artifacts_download_upload_scripts/download_from_huggingface.sh‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@ Usage: $(basename "$0") [--accelerator <type>] [--skip-if-staged] [--help]
1616
1717
Options:
1818
-a, --accelerator <type> GPU accelerator type for which to download models.
19-
Supported: l40s (default), h100
19+
Supported: l40s (default), h100, rtx_pro_6000_blackwell
20+
(rtx_pro_6000_blackwell uses the same quantized
21+
artifacts as h100 — i.e. the w4a16 Gemma variant)
2022
--skip-if-staged Check the configured object store first; skip
2123
downloading (and uploading) any artifact that is
2224
already fully staged there.
@@ -66,21 +68,23 @@ if [[ -z "$ACCEL_FLAG" && -z "$ACCELERATOR" ]]; then
6668
echo "Select GPU type:"
6769
echo " 1) l40s"
6870
echo " 2) h100"
69-
read -rp "Enter 1 or 2: " GPU_CHOICE
71+
echo " 3) rtx_pro_6000_blackwell"
72+
read -rp "Enter 1, 2 or 3: " GPU_CHOICE
7073
case "$GPU_CHOICE" in
7174
1) ACCEL_FLAG="l40s" ;;
7275
2) ACCEL_FLAG="h100" ;;
73-
*) echo "Error: invalid choice '${GPU_CHOICE}'. Please enter 1 or 2." >&2; exit 1 ;;
76+
3) ACCEL_FLAG="rtx_pro_6000_blackwell" ;;
77+
*) echo "Error: invalid choice '${GPU_CHOICE}'. Please enter 1, 2 or 3." >&2; exit 1 ;;
7478
esac
7579
fi
7680

7781
ACCEL="$(printf '%s' "${ACCEL_FLAG:-${ACCELERATOR}}" | tr '[:upper:]' '[:lower:]')"
7882

7983
case "${ACCEL}" in
8084
l40s|"") CONFIG_FILE="./model_artifacts_configs.yaml" ;;
81-
h100) CONFIG_FILE="./model_artifacts_configs_h100.yaml" ;;
85+
h100|rtx_pro_6000_blackwell) CONFIG_FILE="./model_artifacts_configs_h100.yaml" ;;
8286
*)
83-
echo "Error: unsupported accelerator '${ACCEL}'. Supported values: l40s, h100" >&2
87+
echo "Error: unsupported accelerator '${ACCEL}'. Supported values: l40s, h100, rtx_pro_6000_blackwell" >&2
8488
exit 1
8589
;;
8690
esac

0 commit comments

Comments
 (0)