A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[DP] Create placement groups by ray_device_key (#25026)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Xinyu Chen committed
Sep 17, 2025 at 16:57 UTC
bb58dc8c20315038ea5e14007de7269dfaec1ce4
1 file changed
+11
-8
vllm/v1/engine/utils.py
+11
-8
@@ -334,20 +334,22 @@ class CoreEngineActorManager:
334
"No nodes with resources found in Ray cluster.")
335
assert dp_master_ip_key in nodes[0], (
336
"The DP master node (ip: %s) is missing or dead", dp_master_ip)
337
+ device_str = current_platform.ray_device_key
338
for node_resources in nodes:
338
- if "GPU" not in node_resources:
339
+ if device_str not in node_resources:
340
continue
341
# For now, each DP rank can only be assigned to one node
342
# TODO(rui): support allocating a single DP rank
343
# to multiple nodes
343
- available_engine_count = int(node_resources["GPU"]) // world_size
344
+ available_engine_count = int(
345
+ node_resources[device_str]) // world_size
346
if dp_master_ip_key in node_resources:
347
assert available_engine_count >= local_engine_count, (
348
"Not enough resources to allocate DP ranks "
349
f"on DP master node {dp_master_ip}")
350
for i in range(local_engine_count):
351
bundles = [{
350
- "GPU": 1.0,
352
+ device_str: 1.0,
353
"node:" + dp_master_ip: 0.001
354
}] * world_size + [{
355
"CPU": 1.0
@@ -363,7 +365,7 @@ class CoreEngineActorManager:
365
for i in range(available_engine_count):
366
if len(placement_groups) == num_pg_to_create:
367
break
366
- bundles = [{"GPU": 1.0}] * world_size + [{"CPU": 1.0}]
368
+ bundles = [{device_str: 1.0}] * world_size + [{"CPU": 1.0}]
369
pg = ray.util.placement_group(
370
name=f"dp_rank_{len(placement_groups)}",
371
strategy="STRICT_PACK",
@@ -415,17 +417,18 @@ class CoreEngineActorManager:
417
local_dp_ranks = []
418
num_pg_created = 0
419
420
+ device_str = current_platform.ray_device_key
421
for node in nodes:
422
if num_pg_created >= num_pg_to_create:
423
break
424
425
node_ip = node.node_ip
426
node_id = node.node_id
424
- available_gpus = int(available_resources[node_id]["GPU"])
427
+ available_gpus = int(available_resources[node_id][device_str])
428
429
# Get total GPUs on this node from the node's resources
430
# Ray stores node resources with node ID as key
428
- total_gpus = int(total_resources[node_id]["GPU"])
431
+ total_gpus = int(total_resources[node_id][device_str])
432
433
# Calculate used GPUs and used engines on this node
434
used_gpus = max(0, total_gpus - available_gpus)
@@ -444,13 +447,13 @@ class CoreEngineActorManager:
447
# Create bundles with node constraint for master node
448
if node_ip == dp_master_ip:
449
bundles = [{
447
- "GPU": 1.0,
450
+ device_str: 1.0,
451
"node:" + dp_master_ip: 0.001
452
}] * world_size + [{
453
"CPU": 1.0
454
}]
455
else:
453
- bundles = [{"GPU": 1.0}] * world_size + [{"CPU": 1.0}]
456
+ bundles = [{device_str: 1.0}] * world_size + [{"CPU": 1.0}]
457
458
pg = ray.util.placement_group(
459
name=f"dp_rank_{rank}",