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 15, 07:15 UTC
[Bugfix][Core] Update cancellation logic in `generate()` to handle Generator exits (#19225)
Co-authored-by: Adolfo Victoria <adovi@meta.com>
Adolfo Victoria committed
Jun 6, 2025 at 13:17 UTC
ca27f0f9c1452a0e73126be2b1666c3067cf6290
2 files changed
+129
-45
tests/v1/engine/test_async_llm.py
+126
-43
@@ -22,9 +22,11 @@ if not current_platform.is_cuda():
22
pytest.skip(reason="V1 currently only supported on CUDA.",
23
allow_module_level=True)
24
25
-TEXT_ENGINE_ARGS = AsyncEngineArgs(model="meta-llama/Llama-3.2-1B-Instruct",
26
- enforce_eager=True,
27
- disable_log_requests=True)
25
+TEXT_ENGINE_ARGS = AsyncEngineArgs(
26
+ model="meta-llama/Llama-3.2-1B-Instruct",
27
+ enforce_eager=True,
28
+ disable_log_requests=True,
29
+)
30
31
VISION_ENGINE_ARGS = AsyncEngineArgs(model="Qwen/Qwen2-VL-2B-Instruct",
32
enforce_eager=True,
@@ -41,28 +43,33 @@ VISION_PROMPT = {
43
"prompt": VISION_PROMPT_TEMPLATE,
44
"multi_modal_data": {
45
"image": ImageAsset("stop_sign").pil_image
44
- }
46
+ },
47
}
48
49
48
-async def generate(engine: AsyncLLM,
49
- request_id: str,
50
- prompt: PromptType,
51
- output_kind: RequestOutputKind,
52
- max_tokens: int,
53
- n: int = 1,
54
- prompt_logprobs: Optional[int] = None) -> tuple[int, str]:
50
+async def generate(
51
+ engine: AsyncLLM,
52
+ request_id: str,
53
+ prompt: PromptType,
54
+ output_kind: RequestOutputKind,
55
+ max_tokens: int,
56
+ n: int = 1,
57
+ prompt_logprobs: Optional[int] = None,
58
+ cancel_after: Optional[int] = None,
59
+) -> tuple[int, str]:
60
# Ensure generate doesn't complete too fast for cancellation test.
61
await asyncio.sleep(0.2)
62
63
count = 0
59
- sampling_params = SamplingParams(max_tokens=max_tokens,
60
- ignore_eos=True,
61
- output_kind=output_kind,
62
- temperature=0.5,
63
- seed=33,
64
- n=n,
65
- prompt_logprobs=prompt_logprobs)
64
+ sampling_params = SamplingParams(
65
+ max_tokens=max_tokens,
66
+ ignore_eos=True,
67
+ output_kind=output_kind,
68
+ temperature=0.5,
69
+ seed=33,
70
+ n=n,
71
+ prompt_logprobs=prompt_logprobs,
72
+ )
73
async for out in engine.generate(request_id=request_id,
74
prompt=prompt,
75
sampling_params=sampling_params):
@@ -73,20 +80,27 @@ async def generate(engine: AsyncLLM,
80
else:
81
count = num_tokens
82
76
- await asyncio.sleep(0.)
83
+ if cancel_after is not None and count >= cancel_after:
84
+ return count, request_id
85
+
86
+ await asyncio.sleep(0.0)
87
88
return count, request_id
89
90
91
@pytest.mark.parametrize(
92
"output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
83
-@pytest.mark.parametrize("engine_args,prompt",
84
- [(TEXT_ENGINE_ARGS, TEXT_PROMPT),
85
- (VISION_ENGINE_ARGS, VISION_PROMPT)])
93
+@pytest.mark.parametrize(
94
+ "engine_args,prompt",
95
+ [(TEXT_ENGINE_ARGS, TEXT_PROMPT), (VISION_ENGINE_ARGS, VISION_PROMPT)],
96
+)
97
@pytest.mark.asyncio
87
-async def test_load(monkeypatch: pytest.MonkeyPatch,
88
- output_kind: RequestOutputKind,
89
- engine_args: AsyncEngineArgs, prompt: PromptType):
98
+async def test_load(
99
+ monkeypatch: pytest.MonkeyPatch,
100
+ output_kind: RequestOutputKind,
101
+ engine_args: AsyncEngineArgs,
102
+ prompt: PromptType,
103
+):
104
# TODO(rickyx): Remove monkeypatch once we have a better way to test V1
105
# so that in the future when we switch, we don't have to change all the
106
# tests.
@@ -125,13 +139,17 @@ async def test_load(monkeypatch: pytest.MonkeyPatch,
139
140
@pytest.mark.parametrize(
141
"output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
128
-@pytest.mark.parametrize("engine_args,prompt",
129
- [(TEXT_ENGINE_ARGS, TEXT_PROMPT),
130
- (VISION_ENGINE_ARGS, VISION_PROMPT)])
142
+@pytest.mark.parametrize(
143
+ "engine_args,prompt",
144
+ [(TEXT_ENGINE_ARGS, TEXT_PROMPT), (VISION_ENGINE_ARGS, VISION_PROMPT)],
145
+)
146
@pytest.mark.asyncio
132
-async def test_abort(monkeypatch: pytest.MonkeyPatch,
133
- output_kind: RequestOutputKind,
134
- engine_args: AsyncEngineArgs, prompt: PromptType):
147
+async def test_abort(
148
+ monkeypatch: pytest.MonkeyPatch,
149
+ output_kind: RequestOutputKind,
150
+ engine_args: AsyncEngineArgs,
151
+ prompt: PromptType,
152
+):
153
154
with monkeypatch.context() as m, ExitStack() as after:
155
m.setenv("VLLM_USE_V1", "1")
@@ -150,8 +168,9 @@ async def test_abort(monkeypatch: pytest.MonkeyPatch,
168
# Create concurrent requests.
169
tasks: list[asyncio.Task] = []
170
for idx, request_id in enumerate(request_ids):
153
- max_tokens = NUM_EXPECTED_TOKENS_LONG if (
154
- idx in REQUEST_IDS_TO_ABORT) else NUM_EXPECTED_TOKENS
171
+ max_tokens = (NUM_EXPECTED_TOKENS_LONG if
172
+ (idx
173
+ in REQUEST_IDS_TO_ABORT) else NUM_EXPECTED_TOKENS)
174
n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1
175
tasks.append(
176
asyncio.create_task(
@@ -192,12 +211,17 @@ async def test_abort(monkeypatch: pytest.MonkeyPatch,
211
212
213
@pytest.mark.parametrize("n", [1, 3])
195
-@pytest.mark.parametrize("engine_args,prompt",
196
- [(TEXT_ENGINE_ARGS, TEXT_PROMPT),
197
- (VISION_ENGINE_ARGS, VISION_PROMPT)])
214
+@pytest.mark.parametrize(
215
+ "engine_args,prompt",
216
+ [(TEXT_ENGINE_ARGS, TEXT_PROMPT), (VISION_ENGINE_ARGS, VISION_PROMPT)],
217
+)
218
@pytest.mark.asyncio
199
-async def test_finished_flag(monkeypatch: pytest.MonkeyPatch, n: int,
200
- engine_args: AsyncEngineArgs, prompt: PromptType):
219
+async def test_finished_flag(
220
+ monkeypatch: pytest.MonkeyPatch,
221
+ n: int,
222
+ engine_args: AsyncEngineArgs,
223
+ prompt: PromptType,
224
+):
225
226
with monkeypatch.context() as m, ExitStack() as after:
227
m.setenv("VLLM_USE_V1", "1")
@@ -205,11 +229,13 @@ async def test_finished_flag(monkeypatch: pytest.MonkeyPatch, n: int,
229
engine = AsyncLLM.from_engine_args(engine_args)
230
after.callback(engine.shutdown)
231
208
- sampling_params = SamplingParams(max_tokens=100,
209
- output_kind=RequestOutputKind.DELTA,
210
- temperature=1.0,
211
- seed=33,
212
- n=n)
232
+ sampling_params = SamplingParams(
233
+ max_tokens=100,
234
+ output_kind=RequestOutputKind.DELTA,
235
+ temperature=1.0,
236
+ seed=33,
237
+ n=n,
238
+ )
239
outputs = [
240
out
241
async for out in engine.generate(request_id="request-33",
@@ -222,6 +248,63 @@ async def test_finished_flag(monkeypatch: pytest.MonkeyPatch, n: int,
248
assert outputs[-1].finished
249
250
251
+@pytest.mark.parametrize(
252
+ "engine_args,prompt",
253
+ [(TEXT_ENGINE_ARGS, TEXT_PROMPT), (VISION_ENGINE_ARGS, VISION_PROMPT)],
254
+)
255
+@pytest.mark.asyncio
256
+async def test_mid_stream_cancellation(monkeypatch: pytest.MonkeyPatch,
257
+ engine_args: AsyncEngineArgs,
258
+ prompt: PromptType):
259
+ """Test that requests can be cancelled mid-stream."""
260
+ with monkeypatch.context() as m, ExitStack() as after:
261
+ m.setenv("VLLM_USE_V1", "1")
262
+
263
+ engine = AsyncLLM.from_engine_args(engine_args)
264
+ after.callback(engine.shutdown)
265
+
266
+ NUM_REQUESTS = 100
267
+ NUM_TOKENS = 1000
268
+ NUM_EXPECTED_TOKENS = 20
269
+
270
+ request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
271
+
272
+ # Create concurrent requests that will be cancelled mid-stream
273
+ tasks = []
274
+ for request_id in request_ids:
275
+ tasks.append(
276
+ asyncio.create_task(
277
+ generate(
278
+ engine,
279
+ request_id,
280
+ prompt,
281
+ RequestOutputKind.DELTA,
282
+ NUM_TOKENS,
283
+ cancel_after=NUM_EXPECTED_TOKENS,
284
+ )))
285
+
286
+ # Wait for all tasks to complete
287
+ results = await asyncio.gather(*tasks)
288
+
289
+ # Verify all tasks were cancelled at the expected point
290
+ for num_generated_tokens, request_id in results:
291
+ assert num_generated_tokens == NUM_EXPECTED_TOKENS, (
292
+ f"{request_id} generated {num_generated_tokens} tokens but "
293
+ f"expected to cancel after {NUM_EXPECTED_TOKENS}")
294
+
295
+ # Make sure no requests are left hanging
296
+ assert not engine.output_processor.has_unfinished_requests()
297
+
298
+ # Confirm we can reuse the request id after the cancellations.
299
+ request_id = request_ids[0]
300
+ task = asyncio.create_task(
301
+ generate(engine, request_id, prompt, RequestOutputKind.DELTA,
302
+ NUM_EXPECTED_TOKENS))
303
+ num_generated_tokens, request_id = await task
304
+ assert num_generated_tokens == NUM_EXPECTED_TOKENS
305
+ assert not engine.output_processor.has_unfinished_requests()
306
+
307
+
308
class MockLoggingStatLogger(LoggingStatLogger):
309
310
def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
vllm/v1/engine/async_llm.py
+3
-2
@@ -332,8 +332,9 @@ class AsyncLLM(EngineClient):
332
yield out
333
334
# If the request is disconnected by the client, generate()
335
- # is cancelled. So, we abort the request if we end up here.
336
- except asyncio.CancelledError:
335
+ # is cancelled or the generator is garbage collected. So,
336
+ # we abort the request if we end up here.
337
+ except (asyncio.CancelledError, GeneratorExit):
338
await self.abort(request_id)
339
if self.log_requests:
340
logger.info("Request %s aborted.", request_id)