@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 100 lines 2.84 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import MagicMock, patch
5
6 import pytest
7
8 from vllm.v1.executor.ray_utils import get_bundles_sorted_by_node
9
10 NODE_A = "node_a"
11 NODE_B = "node_b"
12 NODE_C = "node_c"
13
14 IP_A = "10.0.0.1"
15 IP_B = "10.0.0.2"
16 IP_C = "10.0.0.3"
17
18 NODE_ID_TO_IP = {NODE_A: IP_A, NODE_B: IP_B, NODE_C: IP_C}
19
20 MOCK_RAY_NODES = [
21 {"NodeID": NODE_A, "NodeManagerAddress": IP_A, "Alive": True},
22 {"NodeID": NODE_B, "NodeManagerAddress": IP_B, "Alive": True},
23 {"NodeID": NODE_C, "NodeManagerAddress": IP_C, "Alive": True},
24 ]
25
26
27 @pytest.mark.parametrize(
28 "bundles_to_node_id,bundle_specs,expected",
29 [
30 pytest.param(
31 {0: NODE_C, 1: NODE_A, 2: NODE_B, 3: NODE_C, 4: NODE_A, 5: NODE_B},
32 [{"GPU": 1}] * 6,
33 [
34 (1, NODE_A, IP_A),
35 (4, NODE_A, IP_A),
36 (2, NODE_B, IP_B),
37 (5, NODE_B, IP_B),
38 (0, NODE_C, IP_C),
39 (3, NODE_C, IP_C),
40 ],
41 ),
42 pytest.param(
43 {0: NODE_B, 1: NODE_B, 2: NODE_A, 3: NODE_A},
44 [{"GPU": 1}] * 4,
45 [
46 (2, NODE_A, IP_A),
47 (3, NODE_A, IP_A),
48 (0, NODE_B, IP_B),
49 (1, NODE_B, IP_B),
50 ],
51 ),
52 pytest.param(
53 {0: NODE_C, 1: NODE_B, 2: NODE_C, 3: NODE_B},
54 [{"GPU": 1}] * 4,
55 [
56 (1, NODE_B, IP_B),
57 (3, NODE_B, IP_B),
58 (0, NODE_C, IP_C),
59 (2, NODE_C, IP_C),
60 ],
61 ),
62 pytest.param(
63 {0: NODE_A, 1: NODE_A, 2: NODE_A},
64 [{"GPU": 1}] * 3,
65 [(0, NODE_A, IP_A), (1, NODE_A, IP_A), (2, NODE_A, IP_A)],
66 ),
67 pytest.param(
68 {},
69 [],
70 [],
71 ),
72 pytest.param(
73 {0: NODE_A, 1: NODE_B, 2: NODE_A},
74 [{"CPU": 1}, {"GPU": 1}, {"GPU": 1}],
75 [(2, NODE_A, IP_A), (1, NODE_B, IP_B)],
76 ),
77 ],
78 )
79 def test_get_bundles_sorted_by_node(bundles_to_node_id, bundle_specs, expected):
80 mock_pg = MagicMock()
81 mock_pg.bundle_specs = bundle_specs
82
83 mock_ctx = MagicMock()
84 mock_ctx.get_node_id.return_value = NODE_A
85
86 with (
87 patch(
88 "vllm.v1.executor.ray_utils.placement_group_table",
89 return_value={"bundles_to_node_id": bundles_to_node_id},
90 ),
91 patch("vllm.v1.executor.ray_utils.ray") as mock_ray,
92 patch("vllm.v1.executor.ray_utils.current_platform") as mock_platform,
93 ):
94 mock_ray.get_runtime_context.return_value = mock_ctx
95 mock_ray.nodes.return_value = MOCK_RAY_NODES
96 mock_platform.ray_device_key = "GPU"
97
98 result = get_bundles_sorted_by_node(mock_pg)
99
100 assert result == expected