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
[Misc] Support passing multiple request ids at once to `AsyncLLM.abort()` (#22944)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Aug 15, 2025 at 17:00 UTC
ad0297d1139f55cbd602652afd54276a8ae217ce
6 files changed
+105
-14
tests/v1/engine/test_async_llm.py
+76
-1
@@ -212,6 +212,79 @@ async def test_abort(
212
assert not engine.output_processor.has_unfinished_requests()
213
214
215
+@pytest.mark.parametrize(
216
+ "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
217
+@pytest.mark.asyncio
218
+async def test_multi_abort(
219
+ monkeypatch: pytest.MonkeyPatch,
220
+ output_kind: RequestOutputKind,
221
+):
222
+
223
+ with monkeypatch.context() as m, ExitStack() as after:
224
+ m.setenv("VLLM_USE_V1", "1")
225
+
226
+ with set_default_torch_num_threads(1):
227
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
228
+ after.callback(engine.shutdown)
229
+
230
+ NUM_REQUESTS = 50
231
+ NUM_EXPECTED_TOKENS = 100
232
+ NUM_EXPECTED_TOKENS_LONG = 50000
233
+ REQUEST_IDS_TO_ABORT = [5, 10, 15, 20, 25]
234
+ PARALLEL_SAMPLE_REQ_IDS = [5, 15, 30, 35]
235
+
236
+ request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
237
+
238
+ # Create concurrent requests.
239
+ tasks: list[asyncio.Task] = []
240
+ for idx, request_id in enumerate(request_ids):
241
+ max_tokens = (NUM_EXPECTED_TOKENS_LONG if
242
+ (idx
243
+ in REQUEST_IDS_TO_ABORT) else NUM_EXPECTED_TOKENS)
244
+ n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1
245
+ tasks.append(
246
+ asyncio.create_task(
247
+ generate(engine, request_id, TEXT_PROMPT, output_kind,
248
+ max_tokens, n)))
249
+
250
+ # Let requests start
251
+ await asyncio.sleep(0.5)
252
+
253
+ # Use multi-abort to abort multiple requests at once
254
+ abort_request_ids = [request_ids[i] for i in REQUEST_IDS_TO_ABORT]
255
+ await engine.abort(abort_request_ids)
256
+
257
+ # Wait for all tasks to complete
258
+ results = await asyncio.gather(*tasks, return_exceptions=True)
259
+
260
+ # Verify results
261
+ for idx, result in enumerate(results):
262
+ if idx in REQUEST_IDS_TO_ABORT:
263
+ # Aborted requests should return partial results
264
+ assert isinstance(
265
+ result, tuple
266
+ ), f"Request {idx} should have completed with partial results"
267
+ num_generated_tokens, request_id = result
268
+ # Should have generated some tokens before abort
269
+ assert num_generated_tokens > 0, (
270
+ f"Aborted request "
271
+ f"{request_id} should have generated some tokens")
272
+ else:
273
+ # Non-aborted requests should complete normally
274
+ assert isinstance(
275
+ result,
276
+ tuple), f"Request {idx} should have completed successfully"
277
+ num_generated_tokens, request_id = result
278
+ n = 3 if idx in PARALLEL_SAMPLE_REQ_IDS else 1
279
+ expected_tokens = NUM_EXPECTED_TOKENS * n
280
+ assert num_generated_tokens == expected_tokens, (
281
+ f"{request_id} generated {num_generated_tokens} but "
282
+ f"expected {expected_tokens}")
283
+
284
+ # Make sure all aborted requests were cleaned up
285
+ assert not engine.output_processor.has_unfinished_requests()
286
+
287
+
288
@pytest.mark.parametrize("n", [1, 3])
289
@pytest.mark.parametrize(
290
"engine_args,prompt",
@@ -460,7 +533,9 @@ async def test_abort_final_output(
533
token_count = sum(
534
len(output.outputs[0].token_ids) for output in outputs)
535
assert token_count > 0
463
- assert len(final_output.outputs[0].token_ids) == 0
536
+ # This would ordinarily be 0, but could end up > 0 if the
537
+ # final abort is coalesced with another chunk in the output queue.
538
+ assert len(final_output.outputs[0].token_ids) >= 0
539
else:
540
# For FINAL_ONLY, we should only get the final output
541
assert len(outputs) == 0
vllm/engine/async_llm_engine.py
+4
-1
@@ -998,7 +998,7 @@ class AsyncLLMEngine(EngineClient):
998
await self.abort(request_id)
999
raise
1000
1001
- async def abort(self, request_id: str) -> None:
1001
+ async def abort(self, request_id: Union[str, Iterable[str]]) -> None:
1002
"""Abort a request.
1003
1004
Abort a submitted request. If the request is finished or not found,
@@ -1007,6 +1007,9 @@ class AsyncLLMEngine(EngineClient):
1007
Args:
1008
request_id: The unique id of the request.
1009
"""
1010
+ if not isinstance(request_id, str):
1011
+ raise RuntimeError("Only single-request abort supported in"
1012
+ " deprecated V0")
1013
if not self.is_running:
1014
raise AsyncEngineDeadError(
1015
"Background loop is not running. If it was running, "
vllm/engine/multiprocessing/client.py
+7
-3
@@ -5,8 +5,8 @@ import asyncio
5
import copy
6
import pickle
7
from contextlib import contextmanager, suppress
8
-from typing import (Any, AsyncGenerator, Dict, Iterator, List, Mapping,
9
- Optional, Union, cast)
8
+from typing import (Any, AsyncGenerator, Dict, Iterable, Iterator, List,
9
+ Mapping, Optional, Union, cast)
10
11
import cloudpickle
12
import psutil
@@ -404,9 +404,13 @@ class MQLLMEngineClient(EngineClient):
404
error_message="Unable to start RPC Server",
405
socket=socket)
406
407
- async def abort(self, request_id: str):
407
+ async def abort(self, request_id: Union[str, Iterable[str]]):
408
"""Send an ABORT_REQUEST signal to the RPC Server"""
409
410
+ if not isinstance(request_id, str):
411
+ raise RuntimeError("Only single-request abort supported in"
412
+ " deprecated V0")
413
+
414
with suppress(MQClientClosedError):
415
await self._send_one_way_rpc_request(
416
request=RPCAbortRequest(request_id), socket=self.input_socket)
vllm/engine/protocol.py
+4
-3
@@ -3,7 +3,7 @@
3
4
import asyncio
5
from abc import ABC, abstractmethod
6
-from typing import AsyncGenerator, Mapping, Optional
6
+from typing import AsyncGenerator, Iterable, Mapping, Optional, Union
7
8
from vllm.beam_search import BeamSearchSequence, create_sort_beams_key_function
9
from vllm.config import DecodingConfig, ModelConfig, VllmConfig
@@ -229,11 +229,12 @@ class EngineClient(ABC):
229
...
230
231
@abstractmethod
232
- async def abort(self, request_id: str) -> None:
232
+ async def abort(self, request_id: Union[str, Iterable[str]]) -> None:
233
"""Abort a request.
234
235
Args:
236
- request_id: The unique id of the request.
236
+ request_id: The unique id of the request,
237
+ or an iterable of such ids.
238
"""
239
...
240
vllm/utils/__init__.py
+5
@@ -1315,6 +1315,11 @@ def common_broadcastable_dtype(dtypes: Collection[torch.dtype]):
1315
)
1316
1317
1318
+def as_list(maybe_list: Iterable[T]) -> list[T]:
1319
+ """Convert iterable to list, unless it's already a list."""
1320
+ return maybe_list if isinstance(maybe_list, list) else list(maybe_list)
1321
+
1322
+
1323
# `collections` helpers
1324
def is_list_of(
1325
value: object,
vllm/v1/engine/async_llm.py
+9
-6
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import asyncio
4
import time
5
-from collections.abc import AsyncGenerator, Mapping
5
+from collections.abc import AsyncGenerator, Iterable, Mapping
6
from copy import copy
7
from typing import Any, Optional, Union
8
@@ -27,7 +27,8 @@ from vllm.transformers_utils.config import (
27
from vllm.transformers_utils.tokenizer import AnyTokenizer
28
from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
29
from vllm.usage.usage_lib import UsageContext
30
-from vllm.utils import Device, cancel_task_threadsafe, cdiv, deprecate_kwargs
30
+from vllm.utils import (Device, as_list, cancel_task_threadsafe, cdiv,
31
+ deprecate_kwargs)
32
from vllm.v1.engine import EngineCoreRequest
33
from vllm.v1.engine.core_client import EngineCoreClient
34
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
@@ -431,14 +432,16 @@ class AsyncLLM(EngineClient):
432
433
self.output_handler = asyncio.create_task(output_handler())
434
434
- async def abort(self, request_id: str) -> None:
435
+ async def abort(self, request_id: Union[str, Iterable[str]]) -> None:
436
"""Abort RequestId in OutputProcessor and EngineCore."""
437
437
- request_ids = self.output_processor.abort_requests((request_id, ))
438
- await self.engine_core.abort_requests_async(request_ids)
438
+ request_ids = (request_id, ) if isinstance(
439
+ request_id, str) else as_list(request_id)
440
+ all_request_ids = self.output_processor.abort_requests(request_ids)
441
+ await self.engine_core.abort_requests_async(all_request_ids)
442
443
if self.log_requests:
441
- logger.info("Aborted request %s.", request_id)
444
+ logger.info("Aborted request(s) %s.", ",".join(request_ids))
445
446
async def encode(
447
self,