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 4, 00:00 UTC
[Frontend][Bugfix] Abort ASR engine requests on cancellation (#41266)
Signed-off-by: abdulrahman-cohere <abdulrahman.abdulrazzag@cohere.com> Signed-off-by: <> Co-authored-by: Cursor Agent <cursor-agent@cursor.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
abdulrahman-cohere committed
May 10, 2026 at 02:51 UTC
a5d0a5afba521e875454a9cd920734bef73d404f
2 files changed
+284
-54
tests/entrypoints/openai/speech_to_text/test_speech_to_text_cancellation.py
new
+191
@@ -0,0 +1,191 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import asyncio
5
+from types import SimpleNamespace
6
+from unittest.mock import AsyncMock, Mock
7
+
8
+import pytest
9
+
10
+from vllm.entrypoints.openai.speech_to_text.protocol import TranscriptionResponse
11
+from vllm.entrypoints.openai.speech_to_text.speech_to_text import OpenAISpeechToText
12
+
13
+
14
+async def _never_finishes():
15
+ await asyncio.Event().wait()
16
+ yield
17
+
18
+
19
+async def _records_start_then_never_finishes(started_request_ids, request_id):
20
+ started_request_ids.append(request_id)
21
+ await asyncio.Event().wait()
22
+ yield
23
+
24
+
25
+@pytest.mark.asyncio
26
+@pytest.mark.parametrize(
27
+ ("engine_inputs", "expected_request_ids"),
28
+ [
29
+ ([{"prompt": "chunk"}], ["transcribe-outer-request"]),
30
+ (
31
+ [{"prompt": "chunk-0"}, {"prompt": "chunk-1"}],
32
+ ["transcribe-outer-request-0", "transcribe-outer-request-1"],
33
+ ),
34
+ ],
35
+)
36
+async def test_non_streaming_cancel_aborts_engine_requests(
37
+ engine_inputs, expected_request_ids
38
+):
39
+ engine_client = SimpleNamespace(
40
+ errored=False,
41
+ generate=Mock(side_effect=lambda *_args, **_kwargs: _never_finishes()),
42
+ abort=AsyncMock(),
43
+ is_tracing_enabled=AsyncMock(return_value=False),
44
+ )
45
+
46
+ server = OpenAISpeechToText.__new__(OpenAISpeechToText)
47
+ server.engine_client = engine_client
48
+ server.task_type = "transcribe"
49
+ server.models = SimpleNamespace(model_name=lambda: "audio")
50
+ server.model_config = SimpleNamespace(max_model_len=1024)
51
+ server.model_cls = SimpleNamespace(no_space_languages=set())
52
+ server.default_sampling_params = {}
53
+ server.asr_config = SimpleNamespace(max_audio_clip_s=30)
54
+ server._check_model = AsyncMock(return_value=None)
55
+ server._maybe_get_adapters = Mock(return_value=None)
56
+ server._preprocess_speech_to_text = AsyncMock(return_value=(engine_inputs, 40.0))
57
+ server._log_inputs = Mock()
58
+
59
+ request = SimpleNamespace(
60
+ model="audio",
61
+ response_format="json",
62
+ stream=False,
63
+ use_beam_search=False,
64
+ max_completion_tokens=None,
65
+ language="en",
66
+ prompt="",
67
+ to_sampling_params=Mock(return_value=object()),
68
+ )
69
+ raw_request = SimpleNamespace(
70
+ headers={"X-Request-Id": "outer-request"},
71
+ state=SimpleNamespace(),
72
+ )
73
+
74
+ task = asyncio.create_task(
75
+ server._create_speech_to_text(
76
+ audio_data=b"audio",
77
+ request=request,
78
+ raw_request=raw_request,
79
+ response_class=TranscriptionResponse,
80
+ stream_generator_method=Mock(),
81
+ )
82
+ )
83
+ await asyncio.sleep(0)
84
+
85
+ task.cancel()
86
+ with pytest.raises(asyncio.CancelledError):
87
+ await task
88
+
89
+ generated_request_ids = [
90
+ call.args[2] for call in engine_client.generate.call_args_list
91
+ ]
92
+ assert generated_request_ids == expected_request_ids
93
+ engine_client.abort.assert_awaited_once_with(expected_request_ids)
94
+
95
+
96
+@pytest.mark.asyncio
97
+async def test_non_streaming_cancel_advances_all_chunk_generators():
98
+ started_request_ids: list[str] = []
99
+ engine_client = SimpleNamespace(
100
+ errored=False,
101
+ generate=Mock(
102
+ side_effect=lambda *_args, **_kwargs: (
103
+ _records_start_then_never_finishes(started_request_ids, _args[2])
104
+ )
105
+ ),
106
+ abort=AsyncMock(),
107
+ is_tracing_enabled=AsyncMock(return_value=False),
108
+ )
109
+
110
+ engine_inputs = [
111
+ {"prompt": "chunk-0"},
112
+ {"prompt": "chunk-1"},
113
+ {"prompt": "chunk-2"},
114
+ ]
115
+ server = OpenAISpeechToText.__new__(OpenAISpeechToText)
116
+ server.engine_client = engine_client
117
+ server.task_type = "transcribe"
118
+ server.models = SimpleNamespace(model_name=lambda: "audio")
119
+ server.model_config = SimpleNamespace(max_model_len=1024)
120
+ server.model_cls = SimpleNamespace(no_space_languages=set())
121
+ server.default_sampling_params = {}
122
+ server.asr_config = SimpleNamespace(max_audio_clip_s=30)
123
+ server._check_model = AsyncMock(return_value=None)
124
+ server._maybe_get_adapters = Mock(return_value=None)
125
+ server._preprocess_speech_to_text = AsyncMock(return_value=(engine_inputs, 90.0))
126
+ server._log_inputs = Mock()
127
+
128
+ request = SimpleNamespace(
129
+ model="audio",
130
+ response_format="json",
131
+ stream=False,
132
+ use_beam_search=False,
133
+ max_completion_tokens=None,
134
+ language="en",
135
+ prompt="",
136
+ to_sampling_params=Mock(return_value=object()),
137
+ )
138
+ raw_request = SimpleNamespace(
139
+ headers={"X-Request-Id": "outer-request"},
140
+ state=SimpleNamespace(),
141
+ )
142
+
143
+ task = asyncio.create_task(
144
+ server._create_speech_to_text(
145
+ audio_data=b"audio",
146
+ request=request,
147
+ raw_request=raw_request,
148
+ response_class=TranscriptionResponse,
149
+ stream_generator_method=Mock(),
150
+ )
151
+ )
152
+ await asyncio.sleep(0.01)
153
+
154
+ expected_request_ids = [
155
+ "transcribe-outer-request-0",
156
+ "transcribe-outer-request-1",
157
+ "transcribe-outer-request-2",
158
+ ]
159
+ assert set(started_request_ids) == set(expected_request_ids)
160
+
161
+ task.cancel()
162
+ with pytest.raises(asyncio.CancelledError):
163
+ await task
164
+
165
+
166
+@pytest.mark.asyncio
167
+async def test_language_detection_cancel_aborts_engine_request():
168
+ engine_client = SimpleNamespace(
169
+ generate=Mock(return_value=_never_finishes()),
170
+ abort=AsyncMock(),
171
+ )
172
+
173
+ server = OpenAISpeechToText.__new__(OpenAISpeechToText)
174
+ server.engine_client = engine_client
175
+ server.asr_config = SimpleNamespace()
176
+ server.tokenizer = Mock()
177
+ server.model_cls = SimpleNamespace(
178
+ get_language_detection_prompt=Mock(return_value={"prompt": "detect"}),
179
+ get_language_token_ids=Mock(return_value=[1]),
180
+ parse_language_detection_output=Mock(),
181
+ )
182
+
183
+ request_id = "transcribe-outer-request-lang_detect"
184
+ task = asyncio.create_task(server._detect_language(Mock(), request_id))
185
+ await asyncio.sleep(0)
186
+
187
+ task.cancel()
188
+ with pytest.raises(asyncio.CancelledError):
189
+ await task
190
+
191
+ engine_client.abort.assert_awaited_once_with(request_id)
vllm/entrypoints/openai/speech_to_text/speech_to_text.py
+93
-54
@@ -49,6 +49,7 @@ from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt
49
from vllm.renderers.inputs.preprocess import parse_enc_dec_prompt, parse_model_prompt
50
from vllm.sampling_params import BeamSearchParams, SamplingParams
51
from vllm.tokenizers import get_tokenizer
52
+from vllm.utils.async_utils import merge_async_iterators
53
54
SpeechToTextResponse: TypeAlias = TranscriptionResponse | TranslationResponse
55
SpeechToTextResponseVerbose: TypeAlias = (
@@ -163,10 +164,17 @@ class OpenAISpeechToText(OpenAIServing):
164
request_id,
165
)
166
166
- final_output: RequestOutput
167
- async for final_output in result_generator:
168
- if final_output.finished:
169
- break
167
+ try:
168
+ final_output: RequestOutput
169
+ async for final_output in result_generator:
170
+ if final_output.finished:
171
+ break
172
+ except asyncio.CancelledError:
173
+ await asyncio.gather(
174
+ self.engine_client.abort(request_id),
175
+ return_exceptions=True,
176
+ )
177
+ raise
178
179
token_ids = list(final_output.outputs[0].token_ids)
180
lang = self.model_cls.parse_language_detection_output(
@@ -458,41 +466,55 @@ class OpenAISpeechToText(OpenAIServing):
466
if request.response_format == "verbose_json":
467
sampling_params.logprobs = 1
468
469
+ engine_request_ids = [
470
+ request_id if len(engine_inputs) == 1 else f"{request_id}-{idx}"
471
+ for idx in range(len(engine_inputs))
472
+ ]
473
list_result_generator = []
462
- for i, engine_input in enumerate(engine_inputs):
463
- request_id_item = f"{request_id}_{i}"
464
-
465
- self._log_inputs(
466
- request_id_item,
467
- engine_input,
468
- params=sampling_params,
469
- lora_request=lora_request,
470
- )
471
-
472
- trace_headers = (
473
- None
474
- if raw_request is None
475
- else await self._get_trace_headers(raw_request.headers)
476
- )
477
-
478
- if isinstance(sampling_params, BeamSearchParams):
479
- generator = self.beam_search(
480
- prompt=engine_input,
474
+ try:
475
+ for request_id_item, engine_input in zip(engine_request_ids, engine_inputs):
476
+ self._log_inputs(
477
+ request_id_item,
478
+ engine_input,
479
params=sampling_params,
482
- request_id=request_id_item,
480
lora_request=lora_request,
484
- trace_headers=trace_headers,
481
)
486
- else:
487
- generator = self.engine_client.generate(
488
- engine_input,
489
- sampling_params,
490
- request_id_item,
491
- lora_request=lora_request,
492
- trace_headers=trace_headers,
482
+
483
+ trace_headers = (
484
+ None
485
+ if raw_request is None
486
+ else await self._get_trace_headers(raw_request.headers)
487
)
488
495
- list_result_generator.append(generator)
489
+ if isinstance(sampling_params, BeamSearchParams):
490
+ generator = self.beam_search(
491
+ prompt=engine_input,
492
+ params=sampling_params,
493
+ request_id=request_id_item,
494
+ lora_request=lora_request,
495
+ trace_headers=trace_headers,
496
+ )
497
+ else:
498
+ generator = self.engine_client.generate(
499
+ engine_input,
500
+ sampling_params,
501
+ request_id_item,
502
+ lora_request=lora_request,
503
+ trace_headers=trace_headers,
504
+ )
505
+
506
+ list_result_generator.append(generator)
507
+ except asyncio.CancelledError:
508
+ logger.info(
509
+ "Request %s cancelled; aborting %d transcription engine request(s).",
510
+ request_id,
511
+ len(engine_request_ids),
512
+ )
513
+ await asyncio.gather(
514
+ self.engine_client.abort(engine_request_ids),
515
+ return_exceptions=True,
516
+ )
517
+ raise
518
519
separator = asr_inter_chunk_separator(
520
request.language, self.model_cls.no_space_languages
@@ -508,10 +530,12 @@ class OpenAISpeechToText(OpenAIServing):
530
separator,
531
)
532
# Non-streaming response.
511
- total_segments = []
512
- text_parts = []
533
try:
534
assert list_result_generator is not None
535
+ chunk_segment_parts: list[list[SpeechToTextSegment]] = [
536
+ [] for _ in list_result_generator
537
+ ]
538
+ chunk_text_parts: list[list[str]] = [[] for _ in list_result_generator]
539
segments_types: dict[str, type[SpeechToTextSegment]] = {
540
"transcribe": TranscriptionSegment,
541
"translate": TranslationSegment,
@@ -522,28 +546,34 @@ class OpenAISpeechToText(OpenAIServing):
546
assert len(list_result_generator) == 1, (
547
"`max_audio_clip_s` is set to None, audio cannot be chunked"
548
)
525
- for idx, result_generator in enumerate(list_result_generator):
549
+ result_generator = merge_async_iterators(*list_result_generator)
550
+ async for idx, op in result_generator:
551
start_time = (
552
float(idx * chunk_size_in_s) if chunk_size_in_s is not None else 0.0
553
)
529
- async for op in result_generator:
530
- if request.response_format == "verbose_json":
531
- assert op.outputs[0].logprobs
532
- segments: list[SpeechToTextSegment] = (
533
- self._get_verbose_segments(
534
- tokens=tuple(op.outputs[0].token_ids),
535
- segment_class=segment_class,
536
- request=request,
537
- start_time=start_time,
538
- log_probs=op.outputs[0].logprobs,
539
- )
540
- )
554
+ if request.response_format == "verbose_json":
555
+ assert op.outputs[0].logprobs
556
+ segments: list[SpeechToTextSegment] = self._get_verbose_segments(
557
+ tokens=tuple(op.outputs[0].token_ids),
558
+ segment_class=segment_class,
559
+ request=request,
560
+ start_time=start_time,
561
+ log_probs=op.outputs[0].logprobs,
562
+ )
563
542
- total_segments.extend(segments)
543
- text_parts.extend([seg.text for seg in segments])
544
- else:
545
- raw_text = op.outputs[0].text
546
- text_parts.append(self.model_cls.post_process_output(raw_text))
564
+ chunk_segment_parts[idx].extend(segments)
565
+ chunk_text_parts[idx].extend([seg.text for seg in segments])
566
+ else:
567
+ raw_text = op.outputs[0].text
568
+ chunk_text_parts[idx].append(
569
+ self.model_cls.post_process_output(raw_text)
570
+ )
571
+ total_segments = [
572
+ segment
573
+ for segment_parts in chunk_segment_parts
574
+ for segment in segment_parts
575
+ ]
576
+ text_parts = [text for text_part in chunk_text_parts for text in text_part]
577
text = separator.join(text_parts)
578
if self.task_type == "transcribe":
579
final_response: ResponseType
@@ -583,7 +613,16 @@ class OpenAISpeechToText(OpenAIServing):
613
)
614
return final_response
615
except asyncio.CancelledError:
586
- return self.create_error_response("Client disconnected")
616
+ logger.info(
617
+ "Request %s cancelled; aborting %d transcription engine request(s).",
618
+ request_id,
619
+ len(engine_request_ids),
620
+ )
621
+ await asyncio.gather(
622
+ self.engine_client.abort(engine_request_ids),
623
+ return_exceptions=True,
624
+ )
625
+ raise
626
627
async def _speech_to_text_stream_generator(
628
self,