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 1, 20:00 UTC
[Core][Engine] allow DP ray placement groups to be set on specific nodes (#44669)
Signed-off-by: walterbm <walter.beller.morales@gmail.com>
Walter Beller-Morales committed
Jun 5, 2026 at 16:07 UTC
c73b0d0db97b90e32b7c8dc2a92756eaf0ec237c
3 files changed
+173
-11
tests/v1/engine/test_dp_placement_node_allowlist.py
new
+119
@@ -0,0 +1,119 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for VLLM_RAY_DP_PLACEMENT_NODE_IPS."""
4
+
5
+from types import SimpleNamespace
6
+from unittest.mock import patch
7
+
8
+import pytest
9
+
10
+import vllm.v1.engine.utils as utils
11
+from vllm.v1.engine.utils import CoreEngineActorManager
12
+
13
+
14
+def _vllm_config(
15
+ *, dp_size, dp_local, master_ip, world_size=1, all2all_backend="naive"
16
+):
17
+ parallel = SimpleNamespace(
18
+ data_parallel_master_ip=master_ip,
19
+ data_parallel_size=dp_size,
20
+ data_parallel_size_local=dp_local,
21
+ world_size=world_size,
22
+ all2all_backend=all2all_backend,
23
+ )
24
+ return SimpleNamespace(parallel_config=parallel)
25
+
26
+
27
+def _resources(node_gpus: dict[str, int]):
28
+ # node_gpus: {ip: gpu_count}; plus a CPU-only head node.
29
+ res = {
30
+ f"id-{ip}": {"GPU": float(g), f"node:{ip}": 1.0} for ip, g in node_gpus.items()
31
+ }
32
+ res["id-head"] = {
33
+ "CPU": 8.0,
34
+ "node:__internal_head__": 1.0,
35
+ "node:10.9.9.9": 1.0,
36
+ }
37
+ return res
38
+
39
+
40
+def _run(cfg, resources):
41
+ created = []
42
+
43
+ def fake_pg(name, strategy, bundles):
44
+ created.append({"name": name, "strategy": strategy, "bundles": bundles})
45
+ return object()
46
+
47
+ with (
48
+ patch(
49
+ "ray._private.state.available_resources_per_node",
50
+ return_value=resources,
51
+ ),
52
+ patch.object(utils, "current_platform", SimpleNamespace(ray_device_key="GPU")),
53
+ patch("ray.util.placement_group", side_effect=fake_pg),
54
+ ):
55
+ pgs, local_ranks = CoreEngineActorManager.create_dp_placement_groups(cfg)
56
+ return pgs, local_ranks, created
57
+
58
+
59
+def _pinned_ips(created):
60
+ return {
61
+ key.split(":", 1)[1]
62
+ for pg in created
63
+ for bundle in pg["bundles"]
64
+ for key in bundle
65
+ if key.startswith("node:")
66
+ }
67
+
68
+
69
+def test_allowlist_confines_dp_to_listed_nodes(monkeypatch):
70
+ monkeypatch.setenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", "10.0.0.1,10.0.0.3")
71
+ resources = _resources({"10.0.0.1": 8, "10.0.0.2": 8, "10.0.0.3": 8, "10.0.0.4": 8})
72
+ cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.1")
73
+
74
+ pgs, _, created = _run(cfg, resources)
75
+
76
+ assert len(pgs) == 16 # 8 on .1 (master) + 8 on .3
77
+ assert _pinned_ips(created) <= {"10.0.0.1", "10.0.0.3"}
78
+
79
+
80
+def test_empty_allowlist_is_noop(monkeypatch):
81
+ monkeypatch.delenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", raising=False)
82
+ resources = _resources({"10.0.0.1": 8, "10.0.0.2": 8})
83
+ cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.1")
84
+
85
+ pgs, _, created = _run(cfg, resources)
86
+
87
+ assert len(pgs) == 16
88
+ assert _pinned_ips(created) == {"10.0.0.1", "10.0.0.2"} # all nodes used
89
+
90
+
91
+def test_master_auto_added_with_warning(monkeypatch):
92
+ # Allowlist omits the master; vLLM must still keep it and warn.
93
+ monkeypatch.setenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", "10.0.0.3")
94
+ resources = _resources({"10.0.0.1": 8, "10.0.0.3": 8})
95
+ cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.1")
96
+ _, _, created = _run(cfg, resources)
97
+
98
+ assert _pinned_ips(created) == {"10.0.0.1", "10.0.0.3"}
99
+
100
+
101
+def test_allowlist_isolates_two_engines(monkeypatch):
102
+ # Engine B is confined to .2/.4, so it can never touch engine A's master .1.
103
+ monkeypatch.setenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", "10.0.0.2,10.0.0.4")
104
+ resources = _resources({"10.0.0.1": 8, "10.0.0.2": 8, "10.0.0.3": 8, "10.0.0.4": 8})
105
+ cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.2")
106
+
107
+ _, _, created = _run(cfg, resources)
108
+
109
+ assert _pinned_ips(created) <= {"10.0.0.2", "10.0.0.4"}
110
+
111
+
112
+def test_allowlist_too_small_raises(monkeypatch):
113
+ # Master alone can't hold all ranks and no other node is allowed.
114
+ monkeypatch.setenv("VLLM_RAY_DP_PLACEMENT_NODE_IPS", "10.0.0.1")
115
+ resources = _resources({"10.0.0.1": 8, "10.0.0.2": 8})
116
+ cfg = _vllm_config(dp_size=16, dp_local=8, master_ip="10.0.0.1")
117
+
118
+ with pytest.raises(ValueError): # not enough placement groups created
119
+ _run(cfg, resources)
vllm/envs.py
+8
@@ -157,6 +157,7 @@ if TYPE_CHECKING:
157
VLLM_DP_MASTER_PORT: int = 0
158
VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
159
VLLM_RAY_DP_PACK_STRATEGY: Literal["strict", "fill", "span"] = "strict"
160
+ VLLM_RAY_DP_PLACEMENT_NODE_IPS: str = ""
161
VLLM_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY: str = ""
162
VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: str = ""
163
VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
@@ -1324,6 +1325,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
1325
"VLLM_RAY_DP_PACK_STRATEGY": lambda: os.getenv(
1326
"VLLM_RAY_DP_PACK_STRATEGY", "strict"
1327
),
1328
+ # Optional comma-separated list of node IPs that Ray data-parallel
1329
+ # placement groups may use. When set, create_dp_placement_groups only
1330
+ # considers these nodes (the DP master node is always included).
1331
+ # This environment variable is ignored if data-parallel-backend is not Ray.
1332
+ "VLLM_RAY_DP_PLACEMENT_NODE_IPS": lambda: os.getenv(
1333
+ "VLLM_RAY_DP_PLACEMENT_NODE_IPS", ""
1334
+ ),
1335
# Comma-separated *additional* prefixes of env vars to copy from the
1336
# driver to Ray workers. These are merged with the built-in defaults
1337
# defined in ``vllm.ray.ray_env`` (VLLM_, etc.). Example: "MYLIB_,OTHER_"
vllm/v1/engine/utils.py
+46
-11
@@ -101,6 +101,23 @@ def _get_bundle_node_ip(bundle: dict[str, float]) -> str:
101
raise ValueError(f"Missing node affinity in placement bundle: {bundle}")
102
103
104
+def _node_ip_from_resources(node_resources: dict) -> str | None:
105
+ """Return the node IP encoded in a Ray per-node resource dict, or None.
106
+
107
+ Ray advertises each node's IP as a ``node:<ip>`` resource key. The head node
108
+ also carries ``node:__internal_head__``, and placement groups add
109
+ ``..._group_...`` keys; both are ignored.
110
+ """
111
+ for key in node_resources:
112
+ if (
113
+ key.startswith("node:")
114
+ and key != "node:__internal_head__"
115
+ and "_group_" not in key
116
+ ):
117
+ return key.split(":", 1)[1]
118
+ return None
119
+
120
+
121
class CoreEngineProcManager:
122
"""
123
Utility class to handle creation, readiness, and shutdown
@@ -506,6 +523,32 @@ class CoreEngineActorManager:
523
assert dp_master_ip_key in nodes[0], (
524
f"The DP master node (ip: {dp_master_ip}) is missing or dead"
525
)
526
+
527
+ # optionally restrict DP placement to a caller-provided node set.
528
+ requested_node_ips = {
529
+ ip.strip()
530
+ for ip in envs.VLLM_RAY_DP_PLACEMENT_NODE_IPS.split(",")
531
+ if ip.strip()
532
+ }
533
+ if requested_node_ips:
534
+ allowed_node_ips = set(requested_node_ips)
535
+ # The master node must host the local ranks, so it has to be allowed.
536
+ if dp_master_ip not in allowed_node_ips:
537
+ allowed_node_ips.add(dp_master_ip)
538
+ filtered_nodes = [
539
+ node_resources
540
+ for node_resources in nodes
541
+ if _node_ip_from_resources(node_resources) in allowed_node_ips
542
+ ]
543
+ logger.info(
544
+ "VLLM_RAY_DP_PLACEMENT_NODE_IPS set; restricting DP placement "
545
+ "from %d to %d node(s): %s",
546
+ len(nodes),
547
+ len(filtered_nodes),
548
+ sorted(allowed_node_ips),
549
+ )
550
+ nodes = filtered_nodes
551
+
552
device_str = current_platform.ray_device_key
553
n_node_devices: list[int] = [
554
int(node_resources[device_str])
@@ -572,18 +615,10 @@ class CoreEngineActorManager:
615
# for "span" pack strategy
616
collected_bundles = []
617
for node_resources in nodes:
575
- node_ip_keys = [
576
- key
577
- for key in node_resources
578
- if key != "node:__internal_head__"
579
- and key.startswith("node:")
580
- and "_group_" not in key
581
- ]
582
- assert len(node_ip_keys) == 1, (
583
- f"Zero or multiple node IP keys found in node resources: {node_ip_keys}"
618
+ node_ip = _node_ip_from_resources(node_resources)
619
+ assert node_ip is not None, (
620
+ f"No node IP key found in node resources: {node_resources}"
621
)
585
- node_ip_key = node_ip_keys[0]
586
- node_ip = node_ip_key.split(":")[1]
622
623
n_device_on_node = int(node_resources.get(device_str, 0))
624
if pack_strategy == "span" and n_device_on_node != 0: