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 14, 05:30 UTC
[Core] Return final response for aborted requests from `AsyncLLM.generate` (#22283)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Aug 14, 2025 at 14:49 UTC
ebcce2cd36a75effd10556942f0467f5f670a080
2 files changed
+107
-13
tests/v1/engine/test_async_llm.py
+87
@@ -13,6 +13,7 @@ from vllm.assets.image import ImageAsset
13
from vllm.config import VllmConfig
14
from vllm.engine.arg_utils import AsyncEngineArgs
15
from vllm.inputs import PromptType
16
+from vllm.outputs import RequestOutput
17
from vllm.platforms import current_platform
18
from vllm.sampling_params import RequestOutputKind
19
from vllm.utils import set_default_torch_num_threads
@@ -398,3 +399,89 @@ async def test_check_health(monkeypatch: pytest.MonkeyPatch):
399
400
# Test 3: Verify healthy engine still works after mock
401
await engine.check_health()
402
+
403
+
404
+@pytest.mark.parametrize(
405
+ "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
406
+@pytest.mark.asyncio
407
+async def test_abort_final_output(
408
+ monkeypatch: pytest.MonkeyPatch,
409
+ output_kind: RequestOutputKind,
410
+):
411
+ """Test that abort() returns a final output with correct information."""
412
+
413
+ with monkeypatch.context() as m, ExitStack() as after:
414
+ m.setenv("VLLM_USE_V1", "1")
415
+
416
+ with set_default_torch_num_threads(1):
417
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
418
+ after.callback(engine.shutdown)
419
+
420
+ request_id = "test-abort-final-output"
421
+
422
+ # Start a long-running request
423
+ sampling_params = SamplingParams(
424
+ max_tokens=3000, # Long enough to allow abort
425
+ ignore_eos=True,
426
+ output_kind=output_kind,
427
+ temperature=0.5,
428
+ seed=42,
429
+ )
430
+
431
+ outputs: list[RequestOutput] = []
432
+ generated = asyncio.create_task(
433
+ collect_outputs(engine, request_id, TEXT_PROMPT, sampling_params,
434
+ outputs))
435
+
436
+ # Let it generate some tokens
437
+ await asyncio.sleep(0.5)
438
+
439
+ # Abort the request
440
+ await engine.abort(request_id)
441
+
442
+ # Wait for generation to complete and return final output
443
+ final_output = await generated
444
+
445
+ # Verify we got a final output
446
+ assert final_output is not None
447
+ assert final_output.finished
448
+ assert len(final_output.outputs) == 1
449
+
450
+ assert final_output.outputs[0].finish_reason == "abort"
451
+ assert final_output.outputs[0].stop_reason is None
452
+
453
+ # Verify num_cached_tokens is set correctly
454
+ assert hasattr(final_output, 'num_cached_tokens')
455
+ assert final_output.num_cached_tokens >= 0
456
+
457
+ # If we got intermediate outputs, verify they are consistent
458
+ if output_kind == RequestOutputKind.DELTA:
459
+ # For DELTA, sum all intermediate tokens should <= final tokens
460
+ token_count = sum(
461
+ len(output.outputs[0].token_ids) for output in outputs)
462
+ assert token_count > 0
463
+ assert len(final_output.outputs[0].token_ids) == 0
464
+ else:
465
+ # For FINAL_ONLY, we should only get the final output
466
+ assert len(outputs) == 0
467
+ assert len(final_output.outputs[0].token_ids) > 0
468
+
469
+ assert not engine.output_processor.has_unfinished_requests()
470
+
471
+
472
+async def collect_outputs(
473
+ engine: AsyncLLM,
474
+ request_id: str,
475
+ prompt: PromptType,
476
+ sampling_params: SamplingParams,
477
+ outputs_list: list[RequestOutput],
478
+) -> Optional[RequestOutput]:
479
+ """Helper to collect outputs and return the final one."""
480
+ final_output: Optional[RequestOutput] = None
481
+ async for output in engine.generate(request_id=request_id,
482
+ prompt=prompt,
483
+ sampling_params=sampling_params):
484
+ if not output.finished:
485
+ outputs_list.append(output)
486
+ final_output = output
487
+ return final_output
vllm/v1/engine/output_processor.py
+20
-13
@@ -107,6 +107,7 @@ class RequestState:
107
self.max_tokens_param = max_tokens_param
108
self.is_prefilling = True
109
self.queue = queue
110
+ self.num_cached_tokens = 0
111
112
self.stats = RequestStateStats(
113
arrival_time=arrival_time) if log_stats else None
@@ -167,7 +168,6 @@ class RequestState:
168
finish_reason: Optional[FinishReason],
169
stop_reason: Union[int, str, None],
170
kv_transfer_params: Optional[dict[str, Any]] = None,
170
- num_cached_tokens: int = 0,
171
) -> Optional[Union[RequestOutput, PoolingRequestOutput]]:
172
173
finished = finish_reason is not None
@@ -195,7 +195,7 @@ class RequestState:
195
return None
196
197
return self._new_request_output(request_id, outputs, finished,
198
- kv_transfer_params, num_cached_tokens)
198
+ kv_transfer_params)
199
200
def _new_request_output(
201
self,
@@ -203,14 +203,14 @@ class RequestState:
203
outputs: Union[list[CompletionOutput], list[PoolingOutput]],
204
finished: bool,
205
kv_transfer_params: Optional[dict[str, Any]] = None,
206
- num_cached_tokens: int = 0,
206
) -> Union[RequestOutput, PoolingRequestOutput]:
207
209
- if isinstance(outputs[0], PoolingOutput):
208
+ first_output = outputs[0]
209
+ if isinstance(first_output, PoolingOutput):
210
assert len(outputs) == 1
211
return PoolingRequestOutput(
212
request_id=request_id,
213
- outputs=outputs[0],
213
+ outputs=first_output,
214
prompt_token_ids=self.prompt_token_ids,
215
finished=finished,
216
)
@@ -229,7 +229,7 @@ class RequestState:
229
outputs=cast(list[CompletionOutput], outputs),
230
finished=finished,
231
kv_transfer_params=kv_transfer_params,
232
- num_cached_tokens=num_cached_tokens,
232
+ num_cached_tokens=self.num_cached_tokens,
233
)
234
235
def _new_completion_output(
@@ -308,11 +308,18 @@ class OutputProcessor:
308
if req_state is not None:
309
self.lora_states.abort_request(req_state)
310
request_ids_to_abort.append(request_id)
311
- else:
312
- parent = self.parent_requests.pop(request_id, None)
313
- if parent and parent.child_requests:
314
- self.abort_requests(parent.child_requests)
315
- request_ids_to_abort.extend(parent.child_requests)
311
+ # Produce final abort output.
312
+ if req_state.queue is not None and (
313
+ request_output := req_state.make_request_output(
314
+ [], None, FinishReason.ABORT, None, None)):
315
+ req_state.queue.put(request_output)
316
+ elif parent := self.parent_requests.get(request_id):
317
+ # Abort children prior to removing the parent.
318
+ if parent.child_requests:
319
+ child_reqs = list(parent.child_requests)
320
+ child_reqs = self.abort_requests(child_reqs)
321
+ request_ids_to_abort.extend(child_reqs)
322
+ self.parent_requests.pop(request_id, None)
323
return request_ids_to_abort
324
325
def add_request(
@@ -390,7 +397,7 @@ class OutputProcessor:
397
finish_reason = engine_core_output.finish_reason
398
stop_reason = engine_core_output.stop_reason
399
kv_transfer_params = engine_core_output.kv_transfer_params
393
- num_cached_tokens = engine_core_output.num_cached_tokens
400
+ req_state.num_cached_tokens = engine_core_output.num_cached_tokens
401
req_state.is_prefilling = False
402
403
if pooling_output is None:
@@ -411,7 +418,7 @@ class OutputProcessor:
418
# 4) Create and handle RequestOutput objects.
419
if request_output := req_state.make_request_output(
420
new_token_ids, pooling_output, finish_reason, stop_reason,
414
- kv_transfer_params, num_cached_tokens):
421
+ kv_transfer_params):
422
if req_state.queue is not None:
423
# AsyncLLM: put into queue for handling by generate().
424
req_state.queue.put(request_output)