diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index eaa45e3320..c07f747a27 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -42,6 +42,7 @@ from vllm.v1.engine.output_processor import ( OutputProcessor, RequestOutputCollector, RequestState, + StreamingInputOutputCollector, ) from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.metrics.loggers import PrometheusStatLogger @@ -1642,6 +1643,53 @@ async def test_cumulative_output_collector_n(): assert third[0].text == "c" +def _make_boundary_output( + token_id: int, finish_reason: str | None, finished: bool = False +) -> RequestOutput: + return RequestOutput( + request_id="my-request-id", + prompt=None, + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="a", + token_ids=[token_id], + cumulative_logprob=None, + logprobs=None, + finish_reason=finish_reason, + ) + ], + finished=finished, + ) + + +@pytest.mark.asyncio +async def test_streaming_input_collector_preserves_chunk_boundary(): + """A pending chunk-boundary marker (finish_reason on a non-finished + output) must not be erased by merging the next chunk's output into it.""" + collector = StreamingInputOutputCollector( + RequestOutputKind.DELTA, request_id="my-request-id-int" + ) + + collector.put(_make_boundary_output(0, "stop")) + collector.put(_make_boundary_output(1, None)) + collector.put(_make_boundary_output(2, None)) + + boundary = await collector.get() + assert boundary.outputs[0].finish_reason == "stop" + assert boundary.outputs[0].token_ids == [0] + + following = await collector.get() + assert following.outputs[0].finish_reason is None + assert following.outputs[0].token_ids == [1, 2] + + assert collector.output is None + assert not collector._backlog + assert not collector.ready.is_set() + + def _make_streaming_input_request( request_id: str, prompt_token_ids: list[int] ) -> EngineCoreRequest: diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index f282c1860b..4ef942d0b0 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -53,7 +53,11 @@ from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.core_client import EngineCoreClient from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm.v1.engine.input_processor import InputProcessor -from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector +from vllm.v1.engine.output_processor import ( + OutputProcessor, + RequestOutputCollector, + StreamingInputOutputCollector, +) from vllm.v1.engine.parallel_sampling import ParentRequest from vllm.v1.executor import Executor from vllm.v1.fault_tolerance.utils import FaultToleranceRequest, FaultToleranceResult @@ -620,7 +624,9 @@ class AsyncLLM(EngineClient): self.input_processor.assign_request_id(final_req) internal_req_id = final_req.request_id - queue = RequestOutputCollector(sampling_params.output_kind, internal_req_id) + queue = StreamingInputOutputCollector( + sampling_params.output_kind, internal_req_id + ) async def handle_inputs(): cancelled = False diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index c5f25e4276..6008dd0d12 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -113,6 +113,53 @@ class RequestOutputCollector: self._input_stream_task = None +def _is_chunk_boundary( + output: RequestOutput | PoolingRequestOutput | Exception | None, +) -> bool: + """Whether the output marks the end of a streaming-input chunk: a + finish_reason on a request that has not finished.""" + return ( + isinstance(output, RequestOutput) + and not output.finished + and any(completion.finished() for completion in output.outputs) + ) + + +class StreamingInputOutputCollector(RequestOutputCollector): + """Collector for streaming-input requests. Doesn't merge across input + chunk boundaries (a finish_reason on a non-finished output).""" + + def __init__(self, output_kind: RequestOutputKind, request_id: str): + super().__init__(output_kind, request_id) + self._backlog: deque[RequestOutput] = deque() + + def put(self, output: RequestOutput | PoolingRequestOutput | Exception) -> None: + if isinstance(output, Exception): + self._backlog.clear() + elif self._backlog or _is_chunk_boundary(self.output): + assert isinstance(output, RequestOutput) + if self._backlog and not _is_chunk_boundary(self._backlog[-1]): + self._backlog[-1].add(output, aggregate=self.aggregate) + else: + self._backlog.append(output) + return + super().put(output) + + async def get(self) -> RequestOutput | PoolingRequestOutput: + output = await super().get() + if self._backlog: + self.output = self._backlog.popleft() + self.ready.set() + return output + + def get_nowait(self) -> RequestOutput | PoolingRequestOutput | None: + output = super().get_nowait() + if self._backlog: + self.output = self._backlog.popleft() + self.ready.set() + return output + + @dataclass class OutputProcessorOutput: request_outputs: list[RequestOutput | PoolingRequestOutput]