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 4, 17:30 UTC
[DP][Ray] Pin DP control bundle to same node as first GPU bundle (#39167)
Signed-off-by: Shahar Mor <smor@nvidia.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
shaharmor98 committed
Apr 23, 2026 at 20:21 UTC
e9ba519f450fd0c3eea5cda44e73eec3ad34f654
1 file changed
+27
-2
vllm/v1/engine/utils.py
+27
-2
@@ -80,6 +80,21 @@ class EngineHandshakeMetadata:
80
parallel_config: dict[str, int | str | list[int]]
81
82
83
+def _make_control_bundle(node_ip: str) -> dict[str, float]:
84
+ # The engine actor is scheduled on the final CPU-only bundle. Keep that
85
+ # bundle colocated with the group's first GPU bundle so the actor does not
86
+ # float to an unrelated node and reorder worker ranks away from the
87
+ # advertised DP bootstrap host.
88
+ return {"CPU": 1.0, "node:" + node_ip: 0.001}
89
+
90
+
91
+def _get_bundle_node_ip(bundle: dict[str, float]) -> str:
92
+ for key in bundle:
93
+ if key.startswith("node:"):
94
+ return key.split(":", 1)[1]
95
+ raise ValueError(f"Missing node affinity in placement bundle: {bundle}")
96
+
97
+
98
class CoreEngineProcManager:
99
"""
100
Utility class to handle creation, readiness, and shutdown
@@ -597,10 +612,20 @@ class CoreEngineActorManager:
612
if len(collected_bundles) < world_size:
613
continue
614
600
- bundles = collected_bundles + [{"CPU": 1.0}]
615
+ control_node_ip = _get_bundle_node_ip(collected_bundles[0])
616
+ bundles = collected_bundles + [
617
+ _make_control_bundle(control_node_ip)
618
+ ]
619
collected_bundles = []
620
else:
603
- bundles = device_bundle * world_size + [{"CPU": 1.0}]
621
+ # STRICT_PACK already keeps every bundle in the placement
622
+ # group on one node, so the explicit node affinity on the
623
+ # control bundle is redundant for correctness here. Keep it
624
+ # anyway for consistency with the span path and to preserve
625
+ # intent if this scheduling strategy changes later.
626
+ bundles = device_bundle * world_size + [
627
+ _make_control_bundle(node_ip)
628
+ ]
629
630
pg = ray.util.placement_group(
631
name=f"dp_rank_{len(placement_groups)}",