@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 119 lines 3.96 KB
Raw
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)