@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 176 lines 5.93 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """V2 ModelRunner + pipeline parallel + data parallel integration tests.
4
5 Covers the interaction between the V2 model runner's PP sampled-token
6 broadcast and the DP per-step all-reduce across a few concurrency
7 regimes. Requires 4 GPUs (DP=2, PP=2, TP=1) on CUDA.
8 """
9
10 import asyncio
11 import contextlib
12 import os
13 from contextlib import ExitStack
14
15 import pytest
16
17 from vllm import SamplingParams
18 from vllm.engine.arg_utils import AsyncEngineArgs
19 from vllm.platforms import current_platform
20 from vllm.sampling_params import RequestOutputKind
21 from vllm.v1.engine.async_llm import AsyncLLM
22
23 PP_DP_MODEL = "ibm-research/PowerMoE-3b" # smallest cached MoE that supports PP
24 PROMPT = "This is a test of data parallel and pipeline parallel together"
25
26
27 def _gpu_skip_reason() -> str | None:
28 if not current_platform.is_cuda():
29 return "requires CUDA"
30 n = current_platform.device_count()
31 if n < 4:
32 return f"requires 4 GPUs, got {n}"
33 return None
34
35
36 _GPU_SKIP = _gpu_skip_reason()
37
38 pytestmark = [
39 pytest.mark.skipif(
40 os.environ.get("VLLM_USE_V2_MODEL_RUNNER", "1") == "0",
41 reason="V2 model runner disabled via VLLM_USE_V2_MODEL_RUNNER=0",
42 ),
43 pytest.mark.skipif(_GPU_SKIP is not None, reason=_GPU_SKIP or ""),
44 ]
45
46
47 def _engine_args(async_scheduling: bool) -> AsyncEngineArgs:
48 return AsyncEngineArgs(
49 model=PP_DP_MODEL,
50 pipeline_parallel_size=2,
51 data_parallel_size=2,
52 data_parallel_backend="mp",
53 tensor_parallel_size=1,
54 max_model_len=4096,
55 max_num_batched_tokens=2048,
56 max_num_seqs=256,
57 async_scheduling=async_scheduling,
58 enable_prefix_caching=False,
59 enforce_eager=False,
60 enable_expert_parallel=False,
61 )
62
63
64 async def _generate(engine: AsyncLLM, prompt: str, max_tokens: int) -> int:
65 """Run one streaming completion and return the number of tokens it yielded."""
66 sampling_params = SamplingParams(
67 max_tokens=max_tokens,
68 ignore_eos=True,
69 output_kind=RequestOutputKind.DELTA,
70 temperature=0.0,
71 )
72 request_id = f"req-{id(prompt):x}-{max_tokens}"
73 total = 0
74 async for out in engine.generate(
75 request_id=request_id, prompt=prompt, sampling_params=sampling_params
76 ):
77 total += len(out.outputs[0].token_ids)
78 return total
79
80
81 @pytest.mark.asyncio
82 @pytest.mark.parametrize("async_scheduling", [True, False])
83 async def test_pp_dp_v2_low_concurrency(async_scheduling: bool):
84 """A single in-flight request at a time, repeated, to exercise the
85 PP slot ring under empty batches between decodes."""
86 with ExitStack() as after:
87 engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling))
88 after.callback(engine.shutdown)
89
90 for _ in range(4):
91 n = await _generate(engine, PROMPT, max_tokens=16)
92 assert n == 16
93
94
95 @pytest.mark.asyncio
96 @pytest.mark.parametrize("async_scheduling", [True, False])
97 async def test_pp_dp_v2_mid_concurrency(async_scheduling: bool):
98 """64 concurrent requests, staggered, to exercise the steady-state
99 DP all-reduce + PP slot-ring path."""
100 with ExitStack() as after:
101 engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling))
102 after.callback(engine.shutdown)
103
104 async def _one(i: int) -> int:
105 await asyncio.sleep(0.01 * i) # stagger so DP load-balances
106 return await _generate(engine, f"{PROMPT} {i}", max_tokens=64)
107
108 results = await asyncio.gather(*[_one(i) for i in range(64)])
109 assert all(n == 64 for n in results), results
110
111
112 @pytest.mark.asyncio
113 async def test_pp_dp_v2_abort_mid_decode():
114 """Cancel half the in-flight requests mid-stream and confirm the
115 engine survives the abort storm."""
116 with ExitStack() as after:
117 engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True))
118 after.callback(engine.shutdown)
119
120 async def _maybe_cancel(i: int):
121 sampling_params = SamplingParams(
122 max_tokens=64,
123 ignore_eos=True,
124 output_kind=RequestOutputKind.DELTA,
125 temperature=0.0,
126 )
127 request_id = f"abort-req-{i}"
128 count = 0
129 cancel_at = 4 if i % 2 == 0 else 64
130 async for out in engine.generate(
131 request_id=request_id,
132 prompt=f"{PROMPT} {i}",
133 sampling_params=sampling_params,
134 ):
135 count += len(out.outputs[0].token_ids)
136 if count >= cancel_at:
137 break
138 return count, i
139
140 results = await asyncio.gather(*[_maybe_cancel(i) for i in range(32)])
141 for count, i in results:
142 if i % 2 == 0:
143 assert count >= 4
144 else:
145 assert count == 64
146
147 # Engine must still serve after the abort storm.
148 final = await _generate(engine, "post-abort warmup", max_tokens=8)
149 assert final == 8
150
151
152 @pytest.mark.asyncio
153 async def test_pp_dp_v2_pause_resume():
154 """Pause an engine with a request in flight, then resume and confirm
155 new requests still work."""
156 with ExitStack() as after:
157 engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True))
158 after.callback(engine.shutdown)
159
160 # Start a long-running generation, let some decoding happen, then
161 # pause (abort mode) and confirm the in-flight task terminates.
162 inflight = asyncio.create_task(_generate(engine, PROMPT, max_tokens=128))
163 await asyncio.sleep(0.5)
164
165 assert not await engine.is_paused()
166 await engine.pause_generation(mode="abort")
167 assert await engine.is_paused()
168
169 with contextlib.suppress(Exception):
170 await inflight
171
172 await engine.resume_generation()
173 assert not await engine.is_paused()
174
175 n = await _generate(engine, PROMPT, max_tokens=8)
176 assert n == 8