@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 206 lines 6.28 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Orchestration-level integration tests for RayExecutorV2."""
4
5 import gc
6 import os
7 import pathlib
8
9 import pytest
10 import ray
11
12 pytestmark = pytest.mark.usefixtures("enable_ray_v2_backend")
13
14 MODEL = "facebook/opt-125m"
15
16
17 def _get_env_var(worker, name):
18 return os.environ.get(name)
19
20
21 def _ray_init():
22 """Start Ray with the project root on workers' PYTHONPATH.
23
24 Without this, workers cannot unpickle actor classes defined in the
25 ``tests`` package, causing FunctionActorManager to fall back to
26 TemporaryActor which drops async method signatures."""
27 project_root = str(pathlib.Path(__file__).resolve().parents[2])
28 ray.init(
29 ignore_reinit_error=True,
30 runtime_env={"env_vars": {"PYTHONPATH": project_root}},
31 )
32
33
34 @pytest.fixture
35 def ray_init():
36 _ray_init()
37
38
39 class _AsyncLLMActor:
40 def start(self, pg, bundle_indices=None, ray_runtime_env=None):
41 os.environ["VLLM_USE_RAY_V2_EXECUTOR_BACKEND"] = "1"
42 # Needed so collective_rpc can pickle _get_env_var over the
43 # AsyncLLM -> EngineCore ZMQ boundary.
44 os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
45 if bundle_indices is not None:
46 os.environ["VLLM_RAY_BUNDLE_INDICES"] = bundle_indices
47 else:
48 os.environ.pop("VLLM_RAY_BUNDLE_INDICES", None)
49
50 from vllm.engine.arg_utils import AsyncEngineArgs
51 from vllm.v1.engine.async_llm import AsyncLLM
52 from vllm.v1.executor.abstract import Executor
53
54 engine_args = AsyncEngineArgs(
55 model=MODEL,
56 tensor_parallel_size=2,
57 distributed_executor_backend="ray",
58 enforce_eager=True,
59 max_model_len=256,
60 gpu_memory_utilization=0.8,
61 )
62 vllm_config = engine_args.create_engine_config()
63 vllm_config.parallel_config.placement_group = pg
64 if ray_runtime_env is not None:
65 vllm_config.parallel_config.ray_runtime_env = ray_runtime_env
66
67 executor_class = Executor.get_class(vllm_config)
68 self.engine = AsyncLLM(
69 vllm_config=vllm_config,
70 executor_class=executor_class,
71 log_stats=False,
72 log_requests=False,
73 )
74
75 async def generate(self, prompt):
76 from vllm.sampling_params import SamplingParams
77
78 params = SamplingParams(max_tokens=16)
79 result = None
80 async for output in self.engine.generate(
81 prompt, params, request_id="test_request_id"
82 ):
83 result = output
84 assert result is not None
85 return result.outputs[0].text
86
87 async def generate_and_get_worker_envs(self, prompt, env_names):
88 from vllm.sampling_params import SamplingParams
89
90 params = SamplingParams(max_tokens=16)
91 result = None
92 async for output in self.engine.generate(
93 prompt, params, request_id="test_request_id"
94 ):
95 result = output
96 assert result is not None
97 text = result.outputs[0].text
98
99 env_results = {}
100 for name in env_names:
101 vals = await self.engine.collective_rpc(
102 _get_env_var, timeout=10, args=(name,)
103 )
104 env_results[name] = vals
105 return text, env_results
106
107 def shutdown(self):
108 if engine := getattr(self, "engine", None):
109 engine.shutdown()
110 del self.engine
111 gc.collect()
112
113
114 AsyncLLMActor = ray.remote(num_cpus=0, max_concurrency=1)(_AsyncLLMActor)
115
116
117 def test_multi_replicas(ray_init):
118 pg1 = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
119 pg2 = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
120 ray.get([pg1.ready(), pg2.ready()])
121
122 actor1 = AsyncLLMActor.remote()
123 actor2 = AsyncLLMActor.remote()
124
125 ray.get(actor1.start.remote(pg1))
126 ray.get(actor2.start.remote(pg2))
127
128 out1, out2 = ray.get(
129 [
130 actor1.generate.remote("Hello world"),
131 actor2.generate.remote("Hello world"),
132 ]
133 )
134 assert len(out1) > 0
135 assert len(out2) > 0
136
137
138 def test_multi_replicas_with_bundle_indices(ray_init):
139 pg = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 4, strategy="PACK")
140 ray.get(pg.ready())
141
142 actor1 = AsyncLLMActor.remote()
143 actor2 = AsyncLLMActor.remote()
144
145 ray.get(actor1.start.remote(pg, bundle_indices="2,1"))
146 ray.get(actor2.start.remote(pg, bundle_indices="0,3"))
147
148 out1, out2 = ray.get(
149 [
150 actor1.generate.remote("Hello world"),
151 actor2.generate.remote("Hello world"),
152 ]
153 )
154 assert len(out1) > 0
155 assert len(out2) > 0
156
157
158 def test_env_var_and_runtime_env_propagation():
159 """Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env
160 propagate to RayWorkerProc actors.
161 """
162 sentinel_vars = {
163 "NCCL_DEBUG": "INFO",
164 "HF_TOKEN": "test_sentinel_token",
165 }
166 for k, v in sentinel_vars.items():
167 os.environ[k] = v
168
169 try:
170 # Called directly (not via the ray_init fixture) because sentinel
171 # env vars must be in os.environ before ray.init() so that Ray
172 # worker processes inherit them.
173 _ray_init()
174
175 pg = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
176 ray.get(pg.ready())
177
178 # Include the project root so that RayWorkerProc actors can
179 # unpickle _get_env_var.
180 project_root = str(pathlib.Path(__file__).resolve().parents[2])
181 ray_runtime_env = {
182 "env_vars": {
183 "RAY_RUNTIME_ENV_TEST": "ray_runtime_env",
184 "PYTHONPATH": project_root,
185 },
186 }
187
188 actor = AsyncLLMActor.remote()
189 ray.get(actor.start.remote(pg, ray_runtime_env=ray_runtime_env))
190
191 all_env_names = list(sentinel_vars) + ["RAY_RUNTIME_ENV_TEST"]
192 text, env_results = ray.get(
193 actor.generate_and_get_worker_envs.remote("Hello world", all_env_names)
194 )
195 assert len(text) > 0
196
197 for name, expected in sentinel_vars.items():
198 for val in env_results[name]:
199 assert val == expected
200
201 for val in env_results["RAY_RUNTIME_ENV_TEST"]:
202 assert val == "ray_runtime_env"
203
204 finally:
205 for k in sentinel_vars:
206 os.environ.pop(k, None)