From 9bdc2c1383ac59ec609c873a9f45674c060d83e3 Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Sun, 9 Aug 2026 08:27:22 +0000 Subject: [PATCH 1/7] feat(chat): support n>1 via server-side fan-out Fans a single n>1 request into N independent engine.generate() calls with distinct random_seeds, collating into N choices. Works for both pytorch and turbomind (engine-agnostic handler-layer approach). n==1 keeps the original single-generator fast path. Co-Authored-By: Claude --- .../serve/openai/chat_completions/serving.py | 457 +++++++++++++++++- .../openai/chat_completions/validation.py | 13 + .../serve/openai/chat_completions/conftest.py | 142 ++++++ .../chat_completions/test_n_completions.py | 241 +++++++++ 4 files changed, 833 insertions(+), 20 deletions(-) create mode 100644 tests/test_lmdeploy/serve/openai/chat_completions/conftest.py create mode 100644 tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py diff --git a/lmdeploy/serve/openai/chat_completions/serving.py b/lmdeploy/serve/openai/chat_completions/serving.py index b8190fc9c1..97a4bea89c 100644 --- a/lmdeploy/serve/openai/chat_completions/serving.py +++ b/lmdeploy/serve/openai/chat_completions/serving.py @@ -1,10 +1,13 @@ # Copyright (c) OpenMMLab. All rights reserved. from __future__ import annotations +import asyncio import json import time -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import aclosing +from copy import deepcopy +from dataclasses import dataclass, field from http import HTTPStatus import shortuuid @@ -42,6 +45,273 @@ logger = get_logger('lmdeploy') +@dataclass +class _FanoutResult: + """Per-choice aggregated output of one fan-out generator. + + The handler post-processes these (parser, logprobs, tool calls) to build + the final ``ChatCompletionResponseChoice`` list. Carrying the raw collected + state keeps ``_fanout_generate_collect`` a pure, engine-agnostic helper + that can be unit-tested with fake generators. + """ + index: int + final_res: object + text: str + token_ids: list = field(default_factory=list) + logprobs: list = field(default_factory=list) + cache_block_ids: list = field(default_factory=list) + remote_token_ids: list = field(default_factory=list) + + +class _ClientDisconnected(Exception): + """Raised inside fan-out consumption when the client disconnects.""" + + +async def _fanout_generate_collect( + generators: list[tuple[int, AsyncGenerator]], + prompt_tokens: int | None = None, + *, + disconnect_check: Callable[[], Awaitable[bool]] | None = None, +) -> tuple[list[_FanoutResult], dict[str, int]]: + """Consume N independent engine generators concurrently and aggregate usage. + + Each generator is treated as a black box yielding ``GenOut``-like objects; + the fan-out is therefore engine-agnostic and works for both pytorch and + turbomind. If any generator raises, the whole request fails (OpenAI-style: + a single n>1 request is all-or-nothing). ``completion_tokens`` is the sum + across choices; ``prompt_tokens`` is counted once (taken from + ``prompt_tokens`` if provided, else from the first choice's + ``input_token_len`` since all choices share the same prompt). + + Args: + generators: list of ``(index, async_generator)`` pairs. + prompt_tokens: explicit prompt token count; if ``None`` it is derived + from the first result's ``input_token_len``. + disconnect_check: optional async callback returning ``True`` when the + client has disconnected; the generator is then closed and + ``_ClientDisconnected`` is raised. + + Returns: + ``(results, usage)`` where ``results`` is a list of ``_FanoutResult`` + ordered by index, and ``usage`` is a dict with ``prompt_tokens``, + ``completion_tokens`` and ``cached_tokens``. + """ + if not generators: + return [], { + 'prompt_tokens': prompt_tokens or 0, + 'completion_tokens': 0, + 'cached_tokens': 0, + } + + async def _consume(index: int, gen: AsyncGenerator) -> _FanoutResult: + final_res = None + text = '' + token_ids: list = [] + logprobs: list = [] + cache_block_ids: list = [] + remote_token_ids: list = [] + async for res in gen: + if disconnect_check is not None and await disconnect_check(): + await gen.aclose() + raise _ClientDisconnected( + f'client disconnected during fan-out choice {index}') + final_res = res + text += res.response + if res.token_ids: + token_ids.extend(res.token_ids) + if res.logprobs: + logprobs.extend(res.logprobs) + cache_block_ids.append(res.cache_block_ids) + remote_token_ids.append(res.token_ids) + if final_res is None: + raise RuntimeError( + f'fan-out choice {index} produced no output') + return _FanoutResult( + index=index, + final_res=final_res, + text=text, + token_ids=token_ids, + logprobs=logprobs, + cache_block_ids=cache_block_ids, + remote_token_ids=remote_token_ids, + ) + + results = await asyncio.gather( + *[_consume(idx, gen) for idx, gen in generators]) + results.sort(key=lambda r: r.index) + total_completion = sum(r.final_res.generate_token_len for r in results) + total_cached = sum(getattr(r.final_res, 'cached_tokens', 0) for r in results) + resolved_prompt = (prompt_tokens if prompt_tokens is not None + else results[0].final_res.input_token_len) + usage = { + 'prompt_tokens': resolved_prompt, + 'completion_tokens': total_completion, + 'cached_tokens': total_cached, + } + return results, usage + + +async def _fanout_generate_stream( + generators: list[tuple[int, AsyncGenerator]], + parsers: list, + request: ChatCompletionRequest, + request_id: str, + created_time: int, + model_name: str, + tokenizer: str, + include_usage: bool, +) -> AsyncGenerator[str, None]: + """Interleave N fan-out generators into a single SSE stream. + + Each choice is processed with its own stateful ``response_parser`` (parsers + hold incremental tag-buffering state). Deltas are emitted as they arrive + from any generator, each tagged with its choice ``index``. After all + choices finish, a final aggregated usage chunk is emitted (when + ``include_usage``), followed by ``[DONE]``. Errors propagate to the whole + request. + """ + n = len(generators) + queue: asyncio.Queue = asyncio.Queue() + _DONE = 'done' + _DELTA = 'delta' + _ERROR = 'error' + + async def consume(index: int, gen: AsyncGenerator, parser) -> None: + streaming_tools = False + final_usage: UsageInfo | None = None + try: + async for res in gen: + logprobs = None + output_token_logprobs = None + if request.logprobs and res.logprobs: + logprobs = _create_chat_completion_logprobs( + tokenizer, res.token_ids, res.logprobs) + if request.return_logprob: + output_token_logprobs = _create_output_token_logprobs( + res.token_ids, res.logprobs) + if res.finish_reason and include_usage: + final_usage = UsageInfo.build( + prompt_tokens=res.input_token_len, + completion_tokens=res.generate_token_len, + cached_tokens=res.cached_tokens, + ) + delta_token_ids = (res.token_ids + if res.token_ids is not None else []) + stream_deltas = parser.stream_chunk(res.response, + delta_token_ids) + if not stream_deltas: + if res.finish_reason is None and not delta_token_ids: + continue + stream_deltas = [(DeltaMessage(role='assistant', + content=''), False)] + should_validate_complete = ( + res.finish_reason in ('stop', 'length') and + (request.return_token_ids + or request.return_routed_experts)) + if should_validate_complete and not parser.validate_complete(): + res.finish_reason = 'parse_error' + for delta_index, (delta_message, + tool_emitted) in enumerate(stream_deltas): + if tool_emitted: + streaming_tools = True + is_last_delta = delta_index == len(stream_deltas) - 1 + finish_reason = res.finish_reason if is_last_delta else None + chunk_logprobs = logprobs if is_last_delta else None + chunk_output_token_logprobs = (output_token_logprobs + if is_last_delta else None) + if (request.tool_choice != 'none' + and parser.tool_parser is not None): + if finish_reason == 'stop' and streaming_tools is True: + finish_reason = 'tool_calls' + routed_experts = (res.routed_experts + if finish_reason is not None else None) + stream_output_ids = delta_token_ids if ( + request.return_token_ids + and is_last_delta) else None + choice_data = ChatCompletionResponseStreamChoice( + index=index, + delta=delta_message, + finish_reason=finish_reason, + logprobs=chunk_logprobs, + output_token_logprobs=chunk_output_token_logprobs, + output_ids=stream_output_ids, + routed_experts=routed_experts, + ) + choice_data = maybe_filter_parallel_tool_calls( + choice_data, request) + response = ChatCompletionStreamResponse( + id=request_id, + created=created_time, + model=model_name, + choices=[choice_data], + usage=None, + ) + response_dict = response.model_dump(mode='json', + exclude_none=True) + if include_usage: + response_dict['usage'] = None + if res.cache_block_ids is not None and is_last_delta: + response_dict['cache_block_ids'] = res.cache_block_ids + response_dict['remote_token_ids'] = res.token_ids + await queue.put((_DELTA, response_dict)) + await queue.put((_DONE, index, final_usage)) + except Exception as e: # noqa: BLE001 + await queue.put((_ERROR, e)) + raise + + tasks = [ + asyncio.create_task(consume(idx, gen, parsers[idx])) + for idx, gen in generators + ] + pending_usages: dict[int, UsageInfo] = {} + done_count = 0 + try: + while done_count < n: + item = await queue.get() + kind = item[0] + if kind == _DELTA: + yield f'data: {json.dumps(item[1])}\n\n' + elif kind == _DONE: + done_count += 1 + _, idx, final_usage = item + if final_usage is not None: + pending_usages[idx] = final_usage + elif kind == _ERROR: + raise item[1] + if include_usage and pending_usages: + prompt_tokens = next(iter( + pending_usages.values())).prompt_tokens + completion_tokens = sum( + u.completion_tokens for u in pending_usages.values()) + cached_tokens = sum( + (u.prompt_tokens_details.cached_tokens + if u.prompt_tokens_details is not None else 0) + for u in pending_usages.values()) + agg_usage = UsageInfo.build( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cached_tokens=cached_tokens, + ) + usage_resp = ChatCompletionStreamResponse( + id=request_id, + created=created_time, + model=model_name, + choices=[], + usage=agg_usage, + ) + yield f'data: {usage_resp.model_dump_json(exclude_none=True)}\n\n' + yield 'data: [DONE]\n\n' + finally: + for t in tasks: + if not t.done(): + t.cancel() + for t in tasks: + try: + await t + except (asyncio.CancelledError, Exception): + pass + + def register(router: APIRouter, server_context) -> None: @router.post('/v1/chat/completions', @@ -231,28 +501,175 @@ async def chat_completions_v1(request: ChatCompletionRequest, '`enable_thinking` in `chat_template_kwargs` will override the value in request.' ) - session = server_context.create_session(request.session_id) - try: - preprocessed = await server_context.async_engine.preprocess( - request.messages, - session, - gen_config=gen_config, - tools=request.tools, - reasoning_effort=request.reasoning_effort, - do_preprocess=do_preprocess, - adapter_name=adapter_name, - chat_template_kwargs=chat_template_kwargs or None, - input_ids=resolved_input_ids, - media_io_kwargs=request.media_io_kwargs, - mm_processor_kwargs=request.mm_processor_kwargs) - except RequestError as error: - return create_request_error_response(error) - result_generator = server_context.async_engine.generate( - preprocessed, - stream_response=True) # always use stream to enable batching include_usage = bool(request.stream_options and request.stream_options.include_usage) + # ------------------------------------------------------------------ + # n > 1: server-side fan-out. + # A single n>1 request becomes N independent engine.generate() calls + # with distinct random seeds, collated into N choices. This is + # engine-agnostic handler-layer logic, so both pytorch and turbomind + # are covered without per-engine changes: each inner generate() still + # runs with n=1 (the engine-level n>1 fallback warning in + # async_engine.py is intentionally kept). n==1 keeps the original + # single-generator fast path below (no overhead). + # ------------------------------------------------------------------ + if request.n and request.n > 1: + n_choices = request.n + # The single session created earlier is unused on the fan-out + # path (each choice gets its own session below); release it. + server_context.session_manager.remove(session) + fanout_sessions = [ + server_context.create_session(request.session_id) + for _ in range(n_choices) + ] + fanout_parsers = [parser_cls(request) for _ in range(n_choices)] + fanout_generators: list[tuple[int, AsyncGenerator]] = [] + for i in range(n_choices): + sub_gen_config = deepcopy(gen_config) + # Per-choice seed: derive seed+i when request.seed is set, else + # leave None so the engine randomizes each choice independently + # (handled in AsyncEngine._determine_gen_config). + sub_gen_config.random_seed = ((request.seed + i) + if request.seed is not None + else None) + gen = server_context.async_engine.generate( + request.messages, + fanout_sessions[i], + gen_config=sub_gen_config, + tools=request.tools, + reasoning_effort=request.reasoning_effort, + stream_response=True, # always stream to enable batching + do_preprocess=do_preprocess, + adapter_name=adapter_name, + chat_template_kwargs=chat_template_kwargs or None, + input_ids=resolved_input_ids, + media_io_kwargs=request.media_io_kwargs, + mm_processor_kwargs=request.mm_processor_kwargs, + ) + fanout_generators.append((i, gen)) + + gen_list = [g for _, g in fanout_generators] + + # Streaming fan-out: interleave deltas from all N generators. + if request.stream: + stream_gen = _fanout_generate_stream( + fanout_generators, + fanout_parsers, + request, + request_id, + created_time, + model_name, + tokenizer, + include_usage, + ) + stream_generator = with_request_cleanup( + stream_gen, gen_list, fanout_sessions, + server_context.session_manager) + return StreamingResponse(stream_generator, + media_type='text/event-stream') + + # Non-streaming fan-out: consume all N generators concurrently. + async def _fanout_nonstream(): + try: + results, usage_dict = await _fanout_generate_collect( + fanout_generators, + prompt_tokens=None, + disconnect_check=raw_request.is_disconnected, + ) + except _ClientDisconnected: + for s in fanout_sessions: + await s.async_abort() + return create_error_response( + HTTPStatus.BAD_REQUEST, 'Client disconnected') + + choices = [] + for res in results: + sub_parser = fanout_parsers[res.index] + tool_calls = None + reasoning_content = None + try: + raw_text = res.text + text, tool_calls, reasoning_content = \ + sub_parser.parse_complete( + res.text, res.token_ids) + should_validate_complete = ( + res.final_res.finish_reason in ('stop', 'length') + and (request.return_token_ids + or request.return_routed_experts)) + if should_validate_complete and not \ + sub_parser.validate_complete(raw_text): + res.final_res.finish_reason = 'parse_error' + if isinstance(tool_calls, list) and len(tool_calls): + if res.final_res.finish_reason == 'stop': + res.final_res.finish_reason = 'tool_calls' + except Exception as e: # noqa: BLE001 + logger.error( + f'Failed to parse {res.text}. Exception: {e}.') + return create_error_response( + HTTPStatus.BAD_REQUEST, + 'Failed to parse fc related info to json format!') + + message = ChatMessage( + role='assistant', + content=text, + tool_calls=tool_calls, + reasoning_content=reasoning_content, + ) + choice_logprobs = None + if request.logprobs and len(res.logprobs): + choice_logprobs = _create_chat_completion_logprobs( + tokenizer, res.token_ids, res.logprobs) + choice_output_token_logprobs = None + if request.return_logprob and len(res.logprobs): + choice_output_token_logprobs = \ + _create_output_token_logprobs( + res.token_ids, res.logprobs) + choice_data = ChatCompletionResponseChoice( + index=res.index, + message=message, + logprobs=choice_logprobs, + output_token_logprobs=choice_output_token_logprobs, + finish_reason=res.final_res.finish_reason, + output_ids=(res.token_ids + if request.return_token_ids else None), + routed_experts=(res.final_res.routed_experts + if request.return_routed_experts + else None), + ) + choice_data = maybe_filter_parallel_tool_calls( + choice_data, request) + choices.append(choice_data) + + usage = UsageInfo.build( + prompt_tokens=usage_dict['prompt_tokens'], + completion_tokens=usage_dict['completion_tokens'], + cached_tokens=usage_dict['cached_tokens'], + ) + return ChatCompletionResponse( + id=request_id, + created=created_time, + model=model_name, + choices=choices, + usage=usage, + ).model_dump() + + return await _fanout_nonstream() + + result_generator = server_context.async_engine.generate( + request.messages, + session, + gen_config=gen_config, + tools=request.tools, + reasoning_effort=request.reasoning_effort, + stream_response=True, # always use stream to enable batching + do_preprocess=do_preprocess, + adapter_name=adapter_name, + chat_template_kwargs=chat_template_kwargs or None, + input_ids=resolved_input_ids, + media_io_kwargs=request.media_io_kwargs, + mm_processor_kwargs=request.mm_processor_kwargs) + def create_stream_response_json( index: int, delta_message: DeltaMessage, diff --git a/lmdeploy/serve/openai/chat_completions/validation.py b/lmdeploy/serve/openai/chat_completions/validation.py index 13d696b7c9..9c6d9418ff 100644 --- a/lmdeploy/serve/openai/chat_completions/validation.py +++ b/lmdeploy/serve/openai/chat_completions/validation.py @@ -4,6 +4,10 @@ from lmdeploy.serve.openai.protocol import ChatCompletionRequest +# Upper bound for `n` (number of choices). Each choice is a separate +# engine.generate() call on the fan-out path, so cap to protect resources. +_MAX_FANOUT_N = 128 + def check_request(request: ChatCompletionRequest, server_context) -> str: engine_config = server_context.engine_config @@ -37,12 +41,21 @@ def check_request(request: ChatCompletionRequest, server_context) -> str: # check sampling settings if request.n <= 0: return f'The n {request.n!r} must be a positive int.' + # n > 1 is implemented as server-side fan-out (N independent engine + # generate() calls). Cap it to prevent unbounded resource use. + if request.n > _MAX_FANOUT_N: + return (f'The n {request.n!r} exceeds the maximum supported ' + f'choices ({_MAX_FANOUT_N}).') if request.top_p is not None and not (0 < request.top_p <= 1): return f'The top_p {request.top_p!r} must be in (0, 1].' if request.top_k is not None and request.top_k < 0: return f'The top_k {request.top_k!r} cannot be a negative integer.' if request.temperature is not None and not (0 <= request.temperature <= 2): return f'The temperature {request.temperature!r} must be in [0, 2]' + # seed validation: per-choice seeds are derived as `seed + i` for n > 1, + # so a negative seed could collide with engine internals; reject it. + if request.seed is not None and request.seed < 0: + return f'The seed {request.seed!r} must be a non-negative int.' # Validate input_ids and image_data constraints. # messages has higher priority. input_ids and image_data are only used when diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py b/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py new file mode 100644 index 0000000000..edc6fa6b22 --- /dev/null +++ b/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py @@ -0,0 +1,142 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Shared fakes for ``/v1/chat/completions`` handler tests.""" +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from lmdeploy.serve.openai.endpoints.chat_completions import register +from lmdeploy.serve.openai.protocol import DeltaMessage + + +class FakeTokenizer: + model = SimpleNamespace(model='fake-tokenizer') + + +class FakeAsyncEngine: + """Engine fake whose ``generate`` returns distinct outputs per call. + + Each call yields a ``GenOut``-like stream whose text encodes the call + index, so fan-out tests can assert the N choices are distinct. + """ + + model_name = 'fake-model' + backend_config = SimpleNamespace(adapters=[], logprobs_mode=None) + + def __init__(self): + self.session_mgr = FakeSessionManager() + self.tokenizer = SimpleNamespace(model=FakeTokenizer()) + self.call_count = 0 + self.gen_configs = [] + + def generate(self, prompt, session, **kwargs): + self.call_count += 1 + self.gen_configs.append(kwargs.get('gen_config')) + call_index = self.call_count + + async def _generator(): + yield SimpleNamespace( + response=f'choice-{call_index}', + token_ids=[call_index], + input_token_len=4, + generate_token_len=call_index, + finish_reason='stop', + logprobs=None, + cached_tokens=0, + routed_experts=None, + cache_block_ids=None, + ) + + return _generator() + + +class PassthroughResponseParser: + """Stateful passthrough parser mirroring the real ResponseParser API.""" + + tool_parser_cls = None + + def __init__(self, request): + self.request = request + self.tool_parser = None + self._chunks = [] + + def stream_chunk(self, delta_text, delta_token_ids, **kwargs): + if not delta_text: + return [] + return [(DeltaMessage(content=delta_text), False)] + + def parse_complete(self, text, token_ids=None, **kwargs): + return text, None, None + + def validate_complete(self, raw_text=None): + return True + + +class FakeSessionManager: + + def __init__(self): + self.removed = [] + self._ids = set() + + def has(self, session_id): + return session_id in self._ids + + def remove(self, session): + self.removed.append(session) + + +class FakeSession: + + def __init__(self, session_id): + self.session_id = session_id + self.epoch = 0 + self.aborted = False + + async def async_abort(self): + self.aborted = True + + +class FakeServerContext: + response_parser_cls = PassthroughResponseParser + + def __init__(self): + self.async_engine = FakeAsyncEngine() + self.default_gen_config = {} + + @property + def engine_config(self): + return self.async_engine.backend_config + + @property + def session_manager(self): + return self.async_engine.session_mgr + + def create_session(self, session_id): + return FakeSession(session_id) + + +class FakeRawRequest: + + def __init__(self, payload=None): + self._payload = payload or {} + + async def json(self): + return self._payload + + async def is_disconnected(self): + return False + + +@pytest.fixture +def chat_endpoint(): + context = FakeServerContext() + from fastapi import APIRouter + r = APIRouter() + register(r, context) + return r.routes[0].endpoint, context + + +@pytest.fixture +def fake_raw_request(): + return FakeRawRequest() diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py new file mode 100644 index 0000000000..7338d206e4 --- /dev/null +++ b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py @@ -0,0 +1,241 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Unit tests for ``n > 1`` server-side fan-out in the chat completions handler. + +The fan-out is engine-agnostic handler-layer logic: a single ``n > 1`` request +becomes N independent ``engine.generate()`` calls with distinct random seeds, +collated into N choices. These tests cover the pure aggregation helper +``_fanout_generate_collect`` using fake async generators that mimic the engine's +``GenOut`` yields. Both pytorch and turbomind engines are covered because the +fan-out lives entirely in the handler and treats the engine as a black box. + +Note: the repo uses ``asyncio.run`` (no ``pytest-asyncio`` dependency), so async +test bodies are driven through ``asyncio.run``. +""" +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from lmdeploy.serve.openai.endpoints.chat_completions.serving import ( + _fanout_generate_collect, +) + + +def _fake_gen(outputs): + """Build an async generator yielding ``GenOut``-like objects.""" + + async def _gen(): + for o in outputs: + yield o + + return _gen() + + +def _genout(text, completion_tokens, *, prompt_tokens=5, finish_reason='stop'): + return SimpleNamespace( + response=text, + input_token_len=prompt_tokens, + generate_token_len=completion_tokens, + finish_reason=finish_reason, + token_ids=[], + logprobs=None, + cached_tokens=0, + routed_experts=None, + cache_block_ids=None, + ) + + +def test_fanout_assigns_distinct_indices_and_aggregates_completion_tokens(): + """Three generators producing 2/3/1 completion tokens are collated into + three choices with distinct indices; prompt_tokens counted once, + completion_tokens summed.""" + gens = [ + (0, _fake_gen([_genout('a', 2)])), + (1, _fake_gen([_genout('b', 3)])), + (2, _fake_gen([_genout('c', 1)])), + ] + choices, usage = asyncio.run(_fanout_generate_collect(gens, prompt_tokens=5)) + assert len(choices) == 3 + assert {c.index for c in choices} == {0, 1, 2} + assert usage['prompt_tokens'] == 5 # counted once + assert usage['completion_tokens'] == 6 # 2 + 3 + 1 + + +def test_fanout_prompt_tokens_from_generator_overrides_when_unspecified(): + """When prompt_tokens is passed explicitly it is used as the single + prompt-token count (never summed across choices).""" + gens = [(0, _fake_gen([_genout('a', 2, prompt_tokens=99)]))] + choices, usage = asyncio.run(_fanout_generate_collect(gens, prompt_tokens=7)) + assert usage['prompt_tokens'] == 7 + assert usage['completion_tokens'] == 2 + + +def test_fanout_propagates_error_to_whole_request(): + """If any inner generator raises, the whole fan-out request fails.""" + + async def _boom(): + raise RuntimeError('choice 1 failed') + yield # noqa: unreachable, makes it an async generator + + with pytest.raises(RuntimeError, match='choice 1 failed'): + asyncio.run(_fanout_generate_collect([(0, _boom())], prompt_tokens=1)) + + +# --------------------------------------------------------------------------- +# Handler-level integration: exercise the n>1 branch end-to-end with a fake +# engine. Validates wiring (N sessions, N parsers, distinct seeds, aggregated +# usage, N choices) for both streaming and non-streaming. +# --------------------------------------------------------------------------- + +from lmdeploy.serve.openai.endpoints.chat_completions.protocol import ( # noqa: E402 + ChatCompletionRequest, +) + + +def _sse_payloads(text): + import json + payloads = [] + for line in text.splitlines(): + if line.startswith('data: '): + data = line.removeprefix('data: ') + if data == '[DONE]': + continue + payloads.append(json.loads(data)) + return payloads + + +def test_handler_n3_nonstream_returns_three_choices_with_aggregated_usage( + chat_endpoint, fake_raw_request): + """n=3 non-streaming: 3 distinct choices, prompt counted once, + completion_tokens summed; engine called 3 times with distinct seeds.""" + endpoint, context = chat_endpoint + request = ChatCompletionRequest(model='fake-model', + messages=[{'role': 'user', + 'content': 'hi'}], + n=3, + seed=42, + stream=False) + response = asyncio.run(endpoint(request, fake_raw_request)) + + assert response['object'] == 'chat.completion' + assert len(response['choices']) == 3 + assert {c['index'] for c in response['choices']} == {0, 1, 2} + # Each choice got distinct text from a distinct generate() call. + assert {c['message']['content'] + for c in response['choices']} == {'choice-1', 'choice-2', + 'choice-3'} + # prompt_tokens counted once (4), completion_tokens = 1 + 2 + 3 = 6. + assert response['usage']['prompt_tokens'] == 4 + assert response['usage']['completion_tokens'] == 6 + # Engine was invoked 3 times with derived seeds 42, 43, 44. + assert context.async_engine.call_count == 3 + seeds = [gc.random_seed for gc in context.async_engine.gen_configs] + assert seeds == [42, 43, 44] + + +def test_handler_n3_stream_interleaves_three_indices_and_aggregates_usage( + chat_endpoint, fake_raw_request): + """n=3 streaming: deltas carry indices 0/1/2, final usage chunk sums + completion tokens across choices.""" + endpoint, context = chat_endpoint + request = ChatCompletionRequest(model='fake-model', + messages=[{'role': 'user', + 'content': 'hi'}], + n=3, + stream=True, + stream_options={'include_usage': True}) + response = asyncio.run(endpoint(request, fake_raw_request)) + + # StreamingResponse.body is an async iterable; collect it. + body_iterator = response.body_iterator + + async def _collect(): + chunks = [] + async for chunk in body_iterator: + chunks.append(chunk.decode() + if isinstance(chunk, bytes) else chunk) + return ''.join(chunks) + + text = asyncio.run(_collect()) + payloads = _sse_payloads(text) + + choice_indices = set() + for p in payloads: + for c in p.get('choices', []): + choice_indices.add(c['index']) + assert choice_indices == {0, 1, 2} + + # The final chunk carries aggregated usage (prompt once, completion sum). + usage_chunks = [p for p in payloads if p.get('usage') is not None] + assert usage_chunks, 'expected a final usage chunk' + final_usage = usage_chunks[-1]['usage'] + assert final_usage['prompt_tokens'] == 4 + assert final_usage['completion_tokens'] == 6 # 1 + 2 + 3 + assert text.rstrip().endswith('data: [DONE]') + + +def test_handler_n1_keeps_single_generator_fast_path(chat_endpoint, + fake_raw_request): + """n=1 (default) must not fan out: exactly one engine.generate() call.""" + endpoint, context = chat_endpoint + request = ChatCompletionRequest(model='fake-model', + messages=[{'role': 'user', + 'content': 'hi'}], + stream=False) + response = asyncio.run(endpoint(request, fake_raw_request)) + assert len(response['choices']) == 1 + assert context.async_engine.call_count == 1 + + +def test_handler_n3_unseeded_leaves_random_seed_none(chat_endpoint, + fake_raw_request): + """When request.seed is unset, each sub gen_config keeps random_seed=None + so the engine randomizes each choice independently.""" + endpoint, context = chat_endpoint + request = ChatCompletionRequest(model='fake-model', + messages=[{'role': 'user', + 'content': 'hi'}], + n=3, + stream=False) + asyncio.run(endpoint(request, fake_raw_request)) + seeds = [gc.random_seed for gc in context.async_engine.gen_configs] + assert seeds == [None, None, None] + + +def test_validation_rejects_oversized_n(): + """Fan-out resource cap: n above _MAX_FANOUT_N is rejected.""" + from lmdeploy.serve.openai.endpoints.chat_completions.validation import \ + _MAX_FANOUT_N, check_request + from types import SimpleNamespace + + request = ChatCompletionRequest(model='fake-model', + messages=[{'role': 'user', + 'content': 'hi'}], + n=_MAX_FANOUT_N + 1) + ctx = SimpleNamespace( + engine_config=SimpleNamespace(logprobs_mode=None, adapters=[]), + session_manager=SimpleNamespace(has=lambda sid: False), + response_parser_cls=None, + ) + msg = check_request(request, ctx) + assert 'exceeds the maximum' in msg + + +def test_validation_rejects_negative_seed(): + from lmdeploy.serve.openai.endpoints.chat_completions.validation import \ + check_request + from types import SimpleNamespace + + request = ChatCompletionRequest(model='fake-model', + messages=[{'role': 'user', + 'content': 'hi'}], + seed=-7) + ctx = SimpleNamespace( + engine_config=SimpleNamespace(logprobs_mode=None, adapters=[]), + session_manager=SimpleNamespace(has=lambda sid: False), + response_parser_cls=None, + ) + msg = check_request(request, ctx) + assert 'non-negative' in msg From ffcadc4fcc75b53607015e339e13a2d28a369a07 Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Sun, 9 Aug 2026 09:04:38 +0000 Subject: [PATCH 2/7] fix(chat): fan-out session leak, explicit session_id crash, sibling cancellation Fix round 1 (code-review findings): - Non-streaming fan-out now wraps _fanout_nonstream in try/finally calling cleanup_result_generators so N fan-out sessions are removed on every exit path (success, parse-error, disconnect, generator-error). Previously leaked N sessions per non-streaming n>1 request. - Fan-out sub-sessions are auto-generated (create_session(None)) instead of reusing request.session_id N times, which collided in SessionManager.map_user_session_id on the 2nd call for explicit session_ids. - _fanout_generate_collect now runs _consume as explicit Tasks and cancels pending siblings on first exception (asyncio.gather does not cancel siblings by default), then awaits cancellations so engine generators close. - _consume wraps the generator in aclosing() for prompt closure on cancel. - Non-streaming fan-out now propagates with_cache cache_block_ids / remote_token_ids response fields (mirrors n==1 path). Tests: added explicit-session-id, sibling-cancellation, multi-chunk stream, and session-cleanup assertions. All 12 n_completions tests pass; 81 serve tests green. Co-Authored-By: Claude --- .../serve/openai/chat_completions/serving.py | 245 +++++++++++------- .../serve/openai/chat_completions/conftest.py | 55 +++- .../chat_completions/test_n_completions.py | 197 +++++++++++++- 3 files changed, 385 insertions(+), 112 deletions(-) diff --git a/lmdeploy/serve/openai/chat_completions/serving.py b/lmdeploy/serve/openai/chat_completions/serving.py index 97a4bea89c..721c61b4f0 100644 --- a/lmdeploy/serve/openai/chat_completions/serving.py +++ b/lmdeploy/serve/openai/chat_completions/serving.py @@ -34,7 +34,7 @@ UsageInfo, ) from lmdeploy.serve.openai.utils import maybe_filter_parallel_tool_calls -from lmdeploy.serve.utils.request_cleanup import with_request_cleanup +from lmdeploy.serve.utils.request_cleanup import cleanup_result_generators, with_request_cleanup from lmdeploy.serve.utils.server_utils import validate_json_request from lmdeploy.utils import get_logger @@ -73,7 +73,8 @@ async def _fanout_generate_collect( *, disconnect_check: Callable[[], Awaitable[bool]] | None = None, ) -> tuple[list[_FanoutResult], dict[str, int]]: - """Consume N independent engine generators concurrently and aggregate usage. + """Consume N independent engine generators concurrently and aggregate + usage. Each generator is treated as a black box yielding ``GenOut``-like objects; the fan-out is therefore engine-agnostic and works for both pytorch and @@ -110,19 +111,19 @@ async def _consume(index: int, gen: AsyncGenerator) -> _FanoutResult: logprobs: list = [] cache_block_ids: list = [] remote_token_ids: list = [] - async for res in gen: - if disconnect_check is not None and await disconnect_check(): - await gen.aclose() - raise _ClientDisconnected( - f'client disconnected during fan-out choice {index}') - final_res = res - text += res.response - if res.token_ids: - token_ids.extend(res.token_ids) - if res.logprobs: - logprobs.extend(res.logprobs) - cache_block_ids.append(res.cache_block_ids) - remote_token_ids.append(res.token_ids) + async with aclosing(gen): + async for res in gen: + if disconnect_check is not None and await disconnect_check(): + raise _ClientDisconnected( + f'client disconnected during fan-out choice {index}') + final_res = res + text += res.response + if res.token_ids: + token_ids.extend(res.token_ids) + if res.logprobs: + logprobs.extend(res.logprobs) + cache_block_ids.append(res.cache_block_ids) + remote_token_ids.append(res.token_ids) if final_res is None: raise RuntimeError( f'fan-out choice {index} produced no output') @@ -136,8 +137,28 @@ async def _consume(index: int, gen: AsyncGenerator) -> _FanoutResult: remote_token_ids=remote_token_ids, ) - results = await asyncio.gather( - *[_consume(idx, gen) for idx, gen in generators]) + # Run all _consume coroutines as explicit Tasks so that on first exception + # we can cancel the still-running siblings and close their engine + # generators. asyncio.gather(return_exceptions=False) does NOT cancel + # sibling coroutines on error — they would keep producing and hold engine + # resources open after the request has already failed. + tasks = [ + asyncio.ensure_future(_consume(idx, gen)) + for idx, gen in generators + ] + try: + results = await asyncio.gather(*tasks) + except BaseException: + for t in tasks: + if not t.done(): + t.cancel() + # Await cancellations to ensure generators are closed before propagating. + for t in tasks: + try: + await t + except (asyncio.CancelledError, Exception): + pass + raise results.sort(key=lambda r: r.index) total_completion = sum(r.final_res.generate_token_len for r in results) total_cached = sum(getattr(r.final_res, 'cached_tokens', 0) for r in results) @@ -519,8 +540,14 @@ async def chat_completions_v1(request: ChatCompletionRequest, # The single session created earlier is unused on the fan-out # path (each choice gets its own session below); release it. server_context.session_manager.remove(session) + # Sub-sessions are internal: always auto-generate distinct ids + # (passing None) rather than reusing request.session_id. Reusing an + # explicit user session_id N times would collide in + # SessionManager.map_user_session_id on the 2nd call. Auto-gen + # avoids both the collision and any user-facing ambiguity about + # which choice owns the user-visible session id. fanout_sessions = [ - server_context.create_session(request.session_id) + server_context.create_session(None) for _ in range(n_choices) ] fanout_parsers = [parser_cls(request) for _ in range(n_choices)] @@ -571,88 +598,116 @@ async def chat_completions_v1(request: ChatCompletionRequest, # Non-streaming fan-out: consume all N generators concurrently. async def _fanout_nonstream(): + # Mirror the n==1 non-streaming path: ensure engine generators + # are closed and fanout sessions removed on EVERY exit path + # (success, parse-error, disconnect, generator-error). Without + # this, every non-streaming n>1 request would leak N sessions. try: - results, usage_dict = await _fanout_generate_collect( - fanout_generators, - prompt_tokens=None, - disconnect_check=raw_request.is_disconnected, - ) - except _ClientDisconnected: - for s in fanout_sessions: - await s.async_abort() - return create_error_response( - HTTPStatus.BAD_REQUEST, 'Client disconnected') - - choices = [] - for res in results: - sub_parser = fanout_parsers[res.index] - tool_calls = None - reasoning_content = None try: - raw_text = res.text - text, tool_calls, reasoning_content = \ - sub_parser.parse_complete( - res.text, res.token_ids) - should_validate_complete = ( - res.final_res.finish_reason in ('stop', 'length') - and (request.return_token_ids - or request.return_routed_experts)) - if should_validate_complete and not \ - sub_parser.validate_complete(raw_text): - res.final_res.finish_reason = 'parse_error' - if isinstance(tool_calls, list) and len(tool_calls): - if res.final_res.finish_reason == 'stop': - res.final_res.finish_reason = 'tool_calls' - except Exception as e: # noqa: BLE001 - logger.error( - f'Failed to parse {res.text}. Exception: {e}.') + results, usage_dict = await _fanout_generate_collect( + fanout_generators, + prompt_tokens=None, + disconnect_check=raw_request.is_disconnected, + ) + except _ClientDisconnected: + for s in fanout_sessions: + await s.async_abort() return create_error_response( - HTTPStatus.BAD_REQUEST, - 'Failed to parse fc related info to json format!') - - message = ChatMessage( - role='assistant', - content=text, - tool_calls=tool_calls, - reasoning_content=reasoning_content, - ) - choice_logprobs = None - if request.logprobs and len(res.logprobs): - choice_logprobs = _create_chat_completion_logprobs( - tokenizer, res.token_ids, res.logprobs) - choice_output_token_logprobs = None - if request.return_logprob and len(res.logprobs): - choice_output_token_logprobs = \ - _create_output_token_logprobs( - res.token_ids, res.logprobs) - choice_data = ChatCompletionResponseChoice( - index=res.index, - message=message, - logprobs=choice_logprobs, - output_token_logprobs=choice_output_token_logprobs, - finish_reason=res.final_res.finish_reason, - output_ids=(res.token_ids - if request.return_token_ids else None), - routed_experts=(res.final_res.routed_experts - if request.return_routed_experts + HTTPStatus.BAD_REQUEST, 'Client disconnected') + + choices = [] + for res in results: + sub_parser = fanout_parsers[res.index] + tool_calls = None + reasoning_content = None + try: + raw_text = res.text + text, tool_calls, reasoning_content = \ + sub_parser.parse_complete( + res.text, res.token_ids) + should_validate_complete = ( + res.final_res.finish_reason in ('stop', + 'length') + and (request.return_token_ids + or request.return_routed_experts)) + if should_validate_complete and not \ + sub_parser.validate_complete(raw_text): + res.final_res.finish_reason = 'parse_error' + if isinstance(tool_calls, list) and len(tool_calls): + if res.final_res.finish_reason == 'stop': + res.final_res.finish_reason = 'tool_calls' + except Exception as e: # noqa: BLE001 + logger.error( + f'Failed to parse {res.text}. ' + f'Exception: {e}.') + return create_error_response( + HTTPStatus.BAD_REQUEST, + 'Failed to parse fc related info to ' + 'json format!') + + message = ChatMessage( + role='assistant', + content=text, + tool_calls=tool_calls, + reasoning_content=reasoning_content, + ) + choice_logprobs = None + if request.logprobs and len(res.logprobs): + choice_logprobs = _create_chat_completion_logprobs( + tokenizer, res.token_ids, res.logprobs) + choice_output_token_logprobs = None + if request.return_logprob and len(res.logprobs): + choice_output_token_logprobs = \ + _create_output_token_logprobs( + res.token_ids, res.logprobs) + choice_data = ChatCompletionResponseChoice( + index=res.index, + message=message, + logprobs=choice_logprobs, + output_token_logprobs=choice_output_token_logprobs, + finish_reason=res.final_res.finish_reason, + output_ids=(res.token_ids + if request.return_token_ids else None), + routed_experts=(res.final_res.routed_experts + if request.return_routed_experts + else None), + ) + choice_data = maybe_filter_parallel_tool_calls( + choice_data, request) + choices.append(choice_data) + + usage = UsageInfo.build( + prompt_tokens=usage_dict['prompt_tokens'], + completion_tokens=usage_dict['completion_tokens'], + cached_tokens=usage_dict['cached_tokens'], ) - choice_data = maybe_filter_parallel_tool_calls( - choice_data, request) - choices.append(choice_data) - - usage = UsageInfo.build( - prompt_tokens=usage_dict['prompt_tokens'], - completion_tokens=usage_dict['completion_tokens'], - cached_tokens=usage_dict['cached_tokens'], - ) - return ChatCompletionResponse( - id=request_id, - created=created_time, - model=model_name, - choices=choices, - usage=usage, - ).model_dump() + response = ChatCompletionResponse( + id=request_id, + created=created_time, + model=model_name, + choices=choices, + usage=usage, + ).model_dump() + + # Disaggregation cache metadata (mirrors n==1 path). For + # fan-out the per-choice block-id lists are flattened: the + # first choice's first block id and the last choice's last + # remote token ids, matching the n==1 single-block shape. + if with_cache and results: + first = results[0] + last = results[-1] + response['cache_block_ids'] = ( + first.cache_block_ids[0] + if first.cache_block_ids else None) + response['remote_token_ids'] = [ + last.remote_token_ids[-1] + ] if last.remote_token_ids else [] + return response + finally: + await cleanup_result_generators( + gen_list, fanout_sessions, + server_context.session_manager) return await _fanout_nonstream() diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py b/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py index edc6fa6b22..521a9d6a48 100644 --- a/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py +++ b/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py @@ -6,7 +6,7 @@ import pytest -from lmdeploy.serve.openai.endpoints.chat_completions import register +from lmdeploy.serve.openai.chat_completions import register from lmdeploy.serve.openai.protocol import DeltaMessage @@ -74,15 +74,51 @@ def validate_complete(self, raw_text=None): class FakeSessionManager: + """Mimics the real SessionManager's id/mapping semantics closely enough + to surface fan-out session bugs: explicit user_session_ids are mapped + one-to-one and a duplicate raises (like map_user_session_id), while + None/-1 auto-generates a fresh internal id.""" def __init__(self): self.removed = [] - self._ids = set() + self.sessions = {} + self.user_session_id_map = {} + self._next_id = 0 + + def map_user_session_id(self, user_session_id): + if user_session_id in self.user_session_id_map: + raise ValueError( + f'User session id {user_session_id} already exists') + session_id = self._next_id + self._next_id += 1 + self.user_session_id_map[user_session_id] = session_id + return session_id + + def get(self, session_id=None, create_if_not_exists=True, **kwargs): + if not create_if_not_exists: + return self.sessions.get(session_id, None) + if session_id is None: + session_id = self._next_id + self._next_id += 1 + if session_id in self.sessions: + return self.sessions[session_id] + session = FakeSession(session_id) + self.sessions[session_id] = session + return session def has(self, session_id): - return session_id in self._ids + return session_id in self.sessions def remove(self, session): + if session is None: + return + session_id = (session if isinstance(session, int) + else session.session_id) + self.sessions.pop(session_id, None) + # also drop any user mapping pointing at this session_id + for uid, sid in list(self.user_session_id_map.items()): + if sid == session_id: + self.user_session_id_map.pop(uid, None) self.removed.append(session) @@ -112,8 +148,17 @@ def engine_config(self): def session_manager(self): return self.async_engine.session_mgr - def create_session(self, session_id): - return FakeSession(session_id) + def create_session(self, user_session_id): + # Mirror ServerContext.create_session: None/-1 auto-generates; an + # explicit id maps one-to-one and collides on a second use. + if user_session_id is None or user_session_id == -1: + session = self.session_manager.get() + else: + session_id = self.session_manager.map_user_session_id( + user_session_id) + session = self.session_manager.get(session_id) + session.epoch = 0 + return session class FakeRawRequest: diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py index 7338d206e4..50b54e9de6 100644 --- a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py +++ b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py @@ -1,5 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. -"""Unit tests for ``n > 1`` server-side fan-out in the chat completions handler. +"""Unit tests for ``n > 1`` server-side fan-out in the chat completions +handler. The fan-out is engine-agnostic handler-layer logic: a single ``n > 1`` request becomes N independent ``engine.generate()`` calls with distinct random seeds, @@ -18,7 +19,7 @@ import pytest -from lmdeploy.serve.openai.endpoints.chat_completions.serving import ( +from lmdeploy.serve.openai.chat_completions.serving import ( _fanout_generate_collect, ) @@ -64,8 +65,8 @@ def test_fanout_assigns_distinct_indices_and_aggregates_completion_tokens(): def test_fanout_prompt_tokens_from_generator_overrides_when_unspecified(): - """When prompt_tokens is passed explicitly it is used as the single - prompt-token count (never summed across choices).""" + """When prompt_tokens is passed explicitly it is used as the single prompt- + token count (never summed across choices).""" gens = [(0, _fake_gen([_genout('a', 2, prompt_tokens=99)]))] choices, usage = asyncio.run(_fanout_generate_collect(gens, prompt_tokens=7)) assert usage['prompt_tokens'] == 7 @@ -89,7 +90,7 @@ async def _boom(): # usage, N choices) for both streaming and non-streaming. # --------------------------------------------------------------------------- -from lmdeploy.serve.openai.endpoints.chat_completions.protocol import ( # noqa: E402 +from lmdeploy.serve.openai.protocol import ( # noqa: E402 ChatCompletionRequest, ) @@ -108,7 +109,7 @@ def _sse_payloads(text): def test_handler_n3_nonstream_returns_three_choices_with_aggregated_usage( chat_endpoint, fake_raw_request): - """n=3 non-streaming: 3 distinct choices, prompt counted once, + """N=3 non-streaming: 3 distinct choices, prompt counted once, completion_tokens summed; engine called 3 times with distinct seeds.""" endpoint, context = chat_endpoint request = ChatCompletionRequest(model='fake-model', @@ -133,11 +134,16 @@ def test_handler_n3_nonstream_returns_three_choices_with_aggregated_usage( assert context.async_engine.call_count == 3 seeds = [gc.random_seed for gc in context.async_engine.gen_configs] assert seeds == [42, 43, 44] + # All N fan-out sessions (plus the single pre-fan-out session) are removed + # after the request — no session leak on the non-streaming path. + assert len(context.session_manager.removed) == 3 + 1 + # And no sessions remain live in the manager. + assert context.session_manager.sessions == {} def test_handler_n3_stream_interleaves_three_indices_and_aggregates_usage( chat_endpoint, fake_raw_request): - """n=3 streaming: deltas carry indices 0/1/2, final usage chunk sums + """N=3 streaming: deltas carry indices 0/1/2, final usage chunk sums completion tokens across choices.""" endpoint, context = chat_endpoint request = ChatCompletionRequest(model='fake-model', @@ -174,11 +180,15 @@ async def _collect(): assert final_usage['prompt_tokens'] == 4 assert final_usage['completion_tokens'] == 6 # 1 + 2 + 3 assert text.rstrip().endswith('data: [DONE]') + # Streaming fan-out must also clean up all N fan-out sessions (plus the + # pre-fan-out single session) once the stream completes. + assert len(context.session_manager.removed) == 3 + 1 + assert context.session_manager.sessions == {} def test_handler_n1_keeps_single_generator_fast_path(chat_endpoint, fake_raw_request): - """n=1 (default) must not fan out: exactly one engine.generate() call.""" + """N=1 (default) must not fan out: exactly one engine.generate() call.""" endpoint, context = chat_endpoint request = ChatCompletionRequest(model='fake-model', messages=[{'role': 'user', @@ -206,10 +216,10 @@ def test_handler_n3_unseeded_leaves_random_seed_none(chat_endpoint, def test_validation_rejects_oversized_n(): """Fan-out resource cap: n above _MAX_FANOUT_N is rejected.""" - from lmdeploy.serve.openai.endpoints.chat_completions.validation import \ - _MAX_FANOUT_N, check_request from types import SimpleNamespace + from lmdeploy.serve.openai.chat_completions.validation import _MAX_FANOUT_N, check_request + request = ChatCompletionRequest(model='fake-model', messages=[{'role': 'user', 'content': 'hi'}], @@ -224,10 +234,10 @@ def test_validation_rejects_oversized_n(): def test_validation_rejects_negative_seed(): - from lmdeploy.serve.openai.endpoints.chat_completions.validation import \ - check_request from types import SimpleNamespace + from lmdeploy.serve.openai.chat_completions.validation import check_request + request = ChatCompletionRequest(model='fake-model', messages=[{'role': 'user', 'content': 'hi'}], @@ -239,3 +249,166 @@ def test_validation_rejects_negative_seed(): ) msg = check_request(request, ctx) assert 'non-negative' in msg + + +# --------------------------------------------------------------------------- +# Fix-round-1 regression tests: session-id collision, session cleanup, sibling +# cancellation, multi-chunk interleaving. +# --------------------------------------------------------------------------- + + +def test_handler_n3_with_explicit_session_id_does_not_crash(chat_endpoint, + fake_raw_request): + """An explicit user session_id + n>1 must not collide in + SessionManager.map_user_session_id. + + Fan-out sub-sessions are auto-generated (None), so the user id is mapped at most once and N distinct internal + sessions are created. Regression for the crash + leaked-session bug. + """ + endpoint, context = chat_endpoint + request = ChatCompletionRequest(model='fake-model', + messages=[{'role': 'user', + 'content': 'hi'}], + n=3, + session_id=777, + stream=False) + response = asyncio.run(endpoint(request, fake_raw_request)) + + assert len(response['choices']) == 3 + assert {c['index'] for c in response['choices']} == {0, 1, 2} + # The user session_id was mapped exactly once (to the pre-fan-out single + # session, which is then removed). + assert 777 not in context.session_manager.user_session_id_map + # N distinct internal fan-out sessions were created and all cleaned up. + assert context.async_engine.call_count == 3 + assert context.session_manager.sessions == {} + + +def test_fanout_cancels_sibling_generators_on_error(): + """When one fan-out generator raises, the still-running siblings are + cancelled and their generators closed BEFORE the error propagates out of + _fanout_generate_collect (not only at event-loop shutdown). + + Regression for the asyncio.gather-doesn't-cancel-siblings bug. + """ + from lmdeploy.serve.openai.chat_completions.serving import _fanout_generate_collect + + sibling_closed_before_error = {'value': False} + + async def _boom(): + raise RuntimeError('choice 0 failed') + yield # noqa: unreachable + + async def _long_running(): + try: + # Pretend to produce forever; should be cancelled before done. + while True: + yield _genout('x', 1) + await asyncio.sleep(0.01) + except (asyncio.CancelledError, GeneratorExit): + sibling_closed_before_error['value'] = True + raise + + async def _run_and_record_order(): + # The sibling must be cancelled BEFORE _fanout_generate_collect raises. + # We record the closure state synchronously in the except block, while + # still inside the event loop (before asyncio.run tears it down). + with pytest.raises(RuntimeError, match='choice 0 failed'): + await _fanout_generate_collect( + [(0, _boom()), (1, _long_running())], prompt_tokens=1) + return sibling_closed_before_error['value'] + + closed_before = asyncio.run(_run_and_record_order()) + assert closed_before, \ + 'sibling generator was not cancelled/closed before the error propagated' + + +def test_handler_n2_stream_interleaves_multi_chunk_per_choice(chat_endpoint, + fake_raw_request): + """Streaming fan-out where each generator yields multiple chunks: deltas + from both choices are interleaved and each choice's index appears with its + full text content across chunks.""" + + class MultiChunkEngine: + model_name = 'fake-model' + backend_config = SimpleNamespace(adapters=[], logprobs_mode=None) + + def __init__(self): + self.session_mgr = None # wired from the existing context below + self.tokenizer = SimpleNamespace( + model=SimpleNamespace(model='fake-tokenizer')) + self.call_count = 0 + self.gen_configs = [] + + def generate(self, prompt, session, **kwargs): + self.call_count += 1 + self.gen_configs.append(kwargs.get('gen_config')) + idx = self.call_count + + async def _gen(): + for piece in (f'{idx}-a', f'{idx}-b', f'{idx}-c'): + yield SimpleNamespace( + response=piece, + token_ids=[len(piece)], + input_token_len=3, + generate_token_len=len(piece), + finish_reason=None, + logprobs=None, + cached_tokens=0, + routed_experts=None, + cache_block_ids=None, + ) + yield SimpleNamespace( + response='', + token_ids=[], + input_token_len=3, + generate_token_len=0, + finish_reason='stop', + logprobs=None, + cached_tokens=0, + routed_experts=None, + cache_block_ids=None, + ) + + return _gen() + + endpoint, context = chat_endpoint + # Swap in a multi-chunk engine while reusing the context's session manager. + original_engine = context.async_engine + multi_engine = MultiChunkEngine() + multi_engine.session_mgr = original_engine.session_mgr + context.async_engine = multi_engine + try: + request = ChatCompletionRequest( + model='fake-model', + messages=[{'role': 'user', 'content': 'hi'}], + n=2, + stream=True, + stream_options={'include_usage': True}) + response = asyncio.run(endpoint(request, fake_raw_request)) + + async def _collect(): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() + if isinstance(chunk, bytes) else chunk) + return ''.join(chunks) + + text = asyncio.run(_collect()) + finally: + context.async_engine = original_engine + + payloads = _sse_payloads(text) + # Both choices appear, and the concatenated content per index reconstructs + # the full multi-chunk text for that choice. + per_index = {} + for p in payloads: + for c in p.get('choices', []): + per_index.setdefault(c['index'], '') + content = c['delta'].get('content') if c.get('delta') else None + if content: + per_index[c['index']] += content + assert set(per_index) == {0, 1} + assert per_index[0] == '1-a1-b1-c' + assert per_index[1] == '2-a2-b2-c' + assert text.rstrip().endswith('data: [DONE]') From 4db830dc0ff6e8d4de3c00edae42b443691c049e Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Tue, 11 Aug 2026 10:09:41 +0000 Subject: [PATCH 3/7] refactor(chat): isolate multi-choice fan-out lifecycle --- lmdeploy/serve/openai/api_client.py | 2 +- .../serve/openai/chat_completions/fanout.py | 308 ++++++++ .../serve/openai/chat_completions/serving.py | 552 ++------------ .../chat_completions/streaming_response.py | 114 +++ .../openai/chat_completions/validation.py | 22 +- lmdeploy/serve/openai/endpoints/common.py | 6 +- lmdeploy/serve/proxy/proxy.py | 7 +- .../serve/openai/chat_completions/conftest.py | 8 +- .../chat_completions/test_n_completions.py | 719 ++++++++++-------- 9 files changed, 899 insertions(+), 839 deletions(-) create mode 100644 lmdeploy/serve/openai/chat_completions/fanout.py create mode 100644 lmdeploy/serve/openai/chat_completions/streaming_response.py diff --git a/lmdeploy/serve/openai/api_client.py b/lmdeploy/serve/openai/api_client.py index 753b7260e9..445705e350 100644 --- a/lmdeploy/serve/openai/api_client.py +++ b/lmdeploy/serve/openai/api_client.py @@ -125,7 +125,7 @@ def chat_completions_v1( probable tokens with probabilities that add up to top_p or higher are kept for generation. n (int): How many chat completion choices to generate for each - input message. Only support one here. + input message. Accepts values from 1 to 128. stream: whether to stream the results or not. Default to false. max_completion_tokens (int | None): output token nums. Default to None. max_tokens (int | None): output token nums. Default to None. diff --git a/lmdeploy/serve/openai/chat_completions/fanout.py b/lmdeploy/serve/openai/chat_completions/fanout.py new file mode 100644 index 0000000000..1caef1c0bc --- /dev/null +++ b/lmdeploy/serve/openai/chat_completions/fanout.py @@ -0,0 +1,308 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Multi-choice collation for the chat completions endpoint.""" +from __future__ import annotations + +import asyncio +import json +import time +from collections.abc import AsyncGenerator, Awaitable, Callable +from copy import deepcopy +from dataclasses import dataclass + +import shortuuid +from fastapi import Request +from fastapi.responses import Response, StreamingResponse + +from lmdeploy.serve.openai.protocol import ChatCompletionRequest, UsageInfo + +from .streaming_response import ManagedStreamingResponse + +ChatEndpoint = Callable[[ChatCompletionRequest, Request], + Awaitable[dict | Response]] + + +@dataclass +class _FanoutResponseError(Exception): + response: Response + + +class _FanoutRequest: + """Give each recursive endpoint call an isolated JSON payload.""" + + def __init__(self, request: Request, payload: dict): + self._request = request + self._payload = payload + + def __getattr__(self, name): + return getattr(self._request, name) + + async def json(self) -> dict: + return deepcopy(self._payload) + + async def is_disconnected(self) -> bool: + return await self._request.is_disconnected() + + +def _choice_request(request: ChatCompletionRequest, + index: int) -> ChatCompletionRequest: + choice_request = request.model_copy(deep=True) + choice_request.n = 1 + choice_request.session_id = -1 + if request.seed is not None: + choice_request.seed = (request.seed + index) % (1 << 64) + return choice_request + + +async def _cancel_tasks(tasks: list[asyncio.Task]) -> list: + for task in tasks: + if not task.done(): + task.cancel() + return await asyncio.gather(*tasks, return_exceptions=True) + + +async def _close_streaming_response(response: StreamingResponse) -> None: + close_response = getattr(response, 'close', None) + if close_response is not None: + await close_response() + return + close_iterator = getattr(response.body_iterator, 'aclose', None) + if close_iterator is not None: + await close_iterator() + + +async def _close_responses(responses) -> None: + await asyncio.gather(*( + _close_streaming_response(response) + for response in responses + if isinstance(response, StreamingResponse) + ), return_exceptions=True) + + +async def _cleanup_invocations(tasks: list[asyncio.Task]) -> None: + results = await _cancel_tasks(tasks) + await _close_responses(results) + + +def _consume_cleanup_result(task: asyncio.Task) -> None: + try: + task.result() + except BaseException: # cleanup is best-effort after caller cancellation + pass + + +async def _shield_cleanup(awaitable, name: str) -> None: + cleanup_task = asyncio.create_task(awaitable, name=name) + try: + await asyncio.shield(cleanup_task) + except (asyncio.CancelledError, GeneratorExit): + cleanup_task.add_done_callback(_consume_cleanup_result) + raise + + +async def _invoke_choices( + endpoint: ChatEndpoint, + request: ChatCompletionRequest, + raw_request: Request, + payload: dict, +) -> list[dict | StreamingResponse] | Response: + + async def invoke(index: int): + choice_request = _choice_request(request, index) + choice_payload = deepcopy(payload) + choice_payload.update(n=1, session_id=-1, seed=choice_request.seed) + response = await endpoint( + choice_request, + _FanoutRequest(raw_request, choice_payload), + ) + if isinstance(response, Response) and not isinstance( + response, StreamingResponse): + raise _FanoutResponseError(response) + return response + + tasks = [asyncio.create_task(invoke(index)) for index in range(request.n)] + try: + return await asyncio.gather(*tasks) + except _FanoutResponseError as error: + await _shield_cleanup( + _cleanup_invocations(tasks), 'fanout_invocation_cleanup') + return error.response + except BaseException: + await _shield_cleanup( + _cleanup_invocations(tasks), 'fanout_invocation_cleanup') + raise + + +def _cached_tokens(usage: dict) -> int: + details = usage.get('prompt_tokens_details') or {} + return details.get('cached_tokens', 0) + + +def _aggregate_usage(usages: list[dict]) -> UsageInfo: + first_usage = usages[0] + return UsageInfo.build( + prompt_tokens=first_usage.get('prompt_tokens', 0), + completion_tokens=sum( + usage.get('completion_tokens') or 0 for usage in usages), + cached_tokens=_cached_tokens(first_usage), + ) + + +def _collate_responses( + responses: list[dict], + request_id: str, + created_time: int, +) -> dict: + response = deepcopy(responses[0]) + response['id'] = request_id + response['created'] = created_time + response['choices'] = [] + usages = [] + for index, choice_response in enumerate(responses): + choices = choice_response.get('choices') or [] + if len(choices) != 1: + raise RuntimeError( + f'Expected one choice from fan-out request, got {len(choices)}' + ) + choice = deepcopy(choices[0]) + choice['index'] = index + response['choices'].append(choice) + usages.append(choice_response.get('usage') or {}) + response['usage'] = _aggregate_usage(usages).model_dump() + return response + + +async def _stream_choice( + index: int, + response: StreamingResponse, + queue: asyncio.Queue, + request_id: str, + created_time: int, + model_name: str, + stopping: asyncio.Event, +) -> None: + buffer = '' + try: + async for chunk in response.body_iterator: + buffer += chunk.decode() if isinstance(chunk, bytes) else chunk + while '\n\n' in buffer: + event, buffer = buffer.split('\n\n', 1) + for line in event.splitlines(): + if not line.startswith('data: '): + continue + data = line.removeprefix('data: ') + if data == '[DONE]': + continue + payload = json.loads(data) + if payload.get( + 'usage' + ) is not None and not payload.get('choices'): + await queue.put(('usage', index, payload['usage'])) + continue + choices = payload.get('choices') or [] + if len(choices) != 1: + raise RuntimeError( + 'Expected one streaming choice from fan-out request, ' + f'got {len(choices)}') + choices[0]['index'] = index + payload.update( + id=request_id, + created=created_time, + model=model_name, + ) + await queue.put(('data', payload)) + await queue.put(('done', index)) + except asyncio.CancelledError: + if not stopping.is_set(): + await queue.put(( + 'error', + RuntimeError(f'Fan-out choice {index} was cancelled.'), + )) + raise + except Exception as error: # noqa: BLE001 + await queue.put(('error', error)) + raise + finally: + await _close_streaming_response(response) + + +async def _collate_streams( + responses: list[StreamingResponse], + request: ChatCompletionRequest, + request_id: str, + created_time: int, +) -> AsyncGenerator[str, None]: + queue: asyncio.Queue = asyncio.Queue(maxsize=max(1, len(responses) * 2)) + stopping = asyncio.Event() + # produce the streaming responses for each fan-out request + tasks = [ + asyncio.create_task( + _stream_choice(index, response, queue, request_id, created_time, + request.model, stopping)) + for index, response in enumerate(responses) + ] + usages: dict[int, dict] = {} + completed = 0 + include_usage = bool(request.stream_options + and request.stream_options.include_usage) + # consume the streaming responses of each fan-out request, and yield to the client + try: + while completed < len(tasks): + item = await queue.get() + if item[0] == 'data': + yield f'data: {json.dumps(item[1])}\n\n' + elif item[0] == 'usage': + # item[1]: index, iterm[2]: usage payload + usages[item[1]] = item[2] + elif item[0] == 'done': + completed += 1 + else: + raise item[1] + if include_usage and len(usages) == len(tasks): + ordered_usages = [usages[index] for index in range(len(tasks))] + usage_response = { + 'id': request_id, + 'object': 'chat.completion.chunk', + 'created': created_time, + 'model': request.model, + 'choices': [], + 'usage': _aggregate_usage(ordered_usages).model_dump(), + } + yield f'data: {json.dumps(usage_response)}\n\n' + yield 'data: [DONE]\n\n' + finally: + stopping.set() + await _shield_cleanup(_cancel_tasks(tasks), + 'fanout_stream_cleanup') + + +async def fanout_chat_completions( + endpoint: ChatEndpoint, + request: ChatCompletionRequest, + raw_request: Request, + payload: dict, +) -> dict | Response: + """Run the established single-choice endpoint once per requested choice.""" + request_id = f'chatcmpl-{shortuuid.random()}' + created_time = int(time.time()) + responses = await _invoke_choices(endpoint, request, raw_request, payload) + if isinstance(responses, Response): + return responses + if request.stream: + if not all( + isinstance(response, StreamingResponse) + for response in responses): + await _close_responses(responses) + raise RuntimeError( + 'Expected streaming responses from fan-out requests') + stream = _collate_streams(responses, request, request_id, created_time) + return ManagedStreamingResponse( + stream, + cleanup_callbacks=[ + lambda response=response: _close_streaming_response(response) + for response in responses + ], + media_type='text/event-stream') + if not all(isinstance(response, dict) for response in responses): + await _close_responses(responses) + raise RuntimeError('Expected JSON objects from fan-out requests') + return _collate_responses(responses, request_id, created_time) diff --git a/lmdeploy/serve/openai/chat_completions/serving.py b/lmdeploy/serve/openai/chat_completions/serving.py index 721c61b4f0..32c4a6b812 100644 --- a/lmdeploy/serve/openai/chat_completions/serving.py +++ b/lmdeploy/serve/openai/chat_completions/serving.py @@ -1,18 +1,14 @@ # Copyright (c) OpenMMLab. All rights reserved. from __future__ import annotations -import asyncio import json import time -from collections.abc import AsyncGenerator, Awaitable, Callable +from collections.abc import AsyncGenerator from contextlib import aclosing -from copy import deepcopy -from dataclasses import dataclass, field from http import HTTPStatus import shortuuid from fastapi import APIRouter, Depends, Request -from fastapi.responses import StreamingResponse from lmdeploy.pytorch.disagg.conn.protocol import MigrationRequest from lmdeploy.serve.core.exceptions import RequestError @@ -34,305 +30,19 @@ UsageInfo, ) from lmdeploy.serve.openai.utils import maybe_filter_parallel_tool_calls -from lmdeploy.serve.utils.request_cleanup import cleanup_result_generators, with_request_cleanup +from lmdeploy.serve.utils.request_cleanup import with_request_cleanup from lmdeploy.serve.utils.server_utils import validate_json_request from lmdeploy.utils import get_logger +from .fanout import fanout_chat_completions from .logits_processors import logit_bias_logits_processor from .logprobs import _create_chat_completion_logprobs, _create_output_token_logprobs +from .streaming_response import ManagedStreamingResponse from .validation import check_request logger = get_logger('lmdeploy') -@dataclass -class _FanoutResult: - """Per-choice aggregated output of one fan-out generator. - - The handler post-processes these (parser, logprobs, tool calls) to build - the final ``ChatCompletionResponseChoice`` list. Carrying the raw collected - state keeps ``_fanout_generate_collect`` a pure, engine-agnostic helper - that can be unit-tested with fake generators. - """ - index: int - final_res: object - text: str - token_ids: list = field(default_factory=list) - logprobs: list = field(default_factory=list) - cache_block_ids: list = field(default_factory=list) - remote_token_ids: list = field(default_factory=list) - - -class _ClientDisconnected(Exception): - """Raised inside fan-out consumption when the client disconnects.""" - - -async def _fanout_generate_collect( - generators: list[tuple[int, AsyncGenerator]], - prompt_tokens: int | None = None, - *, - disconnect_check: Callable[[], Awaitable[bool]] | None = None, -) -> tuple[list[_FanoutResult], dict[str, int]]: - """Consume N independent engine generators concurrently and aggregate - usage. - - Each generator is treated as a black box yielding ``GenOut``-like objects; - the fan-out is therefore engine-agnostic and works for both pytorch and - turbomind. If any generator raises, the whole request fails (OpenAI-style: - a single n>1 request is all-or-nothing). ``completion_tokens`` is the sum - across choices; ``prompt_tokens`` is counted once (taken from - ``prompt_tokens`` if provided, else from the first choice's - ``input_token_len`` since all choices share the same prompt). - - Args: - generators: list of ``(index, async_generator)`` pairs. - prompt_tokens: explicit prompt token count; if ``None`` it is derived - from the first result's ``input_token_len``. - disconnect_check: optional async callback returning ``True`` when the - client has disconnected; the generator is then closed and - ``_ClientDisconnected`` is raised. - - Returns: - ``(results, usage)`` where ``results`` is a list of ``_FanoutResult`` - ordered by index, and ``usage`` is a dict with ``prompt_tokens``, - ``completion_tokens`` and ``cached_tokens``. - """ - if not generators: - return [], { - 'prompt_tokens': prompt_tokens or 0, - 'completion_tokens': 0, - 'cached_tokens': 0, - } - - async def _consume(index: int, gen: AsyncGenerator) -> _FanoutResult: - final_res = None - text = '' - token_ids: list = [] - logprobs: list = [] - cache_block_ids: list = [] - remote_token_ids: list = [] - async with aclosing(gen): - async for res in gen: - if disconnect_check is not None and await disconnect_check(): - raise _ClientDisconnected( - f'client disconnected during fan-out choice {index}') - final_res = res - text += res.response - if res.token_ids: - token_ids.extend(res.token_ids) - if res.logprobs: - logprobs.extend(res.logprobs) - cache_block_ids.append(res.cache_block_ids) - remote_token_ids.append(res.token_ids) - if final_res is None: - raise RuntimeError( - f'fan-out choice {index} produced no output') - return _FanoutResult( - index=index, - final_res=final_res, - text=text, - token_ids=token_ids, - logprobs=logprobs, - cache_block_ids=cache_block_ids, - remote_token_ids=remote_token_ids, - ) - - # Run all _consume coroutines as explicit Tasks so that on first exception - # we can cancel the still-running siblings and close their engine - # generators. asyncio.gather(return_exceptions=False) does NOT cancel - # sibling coroutines on error — they would keep producing and hold engine - # resources open after the request has already failed. - tasks = [ - asyncio.ensure_future(_consume(idx, gen)) - for idx, gen in generators - ] - try: - results = await asyncio.gather(*tasks) - except BaseException: - for t in tasks: - if not t.done(): - t.cancel() - # Await cancellations to ensure generators are closed before propagating. - for t in tasks: - try: - await t - except (asyncio.CancelledError, Exception): - pass - raise - results.sort(key=lambda r: r.index) - total_completion = sum(r.final_res.generate_token_len for r in results) - total_cached = sum(getattr(r.final_res, 'cached_tokens', 0) for r in results) - resolved_prompt = (prompt_tokens if prompt_tokens is not None - else results[0].final_res.input_token_len) - usage = { - 'prompt_tokens': resolved_prompt, - 'completion_tokens': total_completion, - 'cached_tokens': total_cached, - } - return results, usage - - -async def _fanout_generate_stream( - generators: list[tuple[int, AsyncGenerator]], - parsers: list, - request: ChatCompletionRequest, - request_id: str, - created_time: int, - model_name: str, - tokenizer: str, - include_usage: bool, -) -> AsyncGenerator[str, None]: - """Interleave N fan-out generators into a single SSE stream. - - Each choice is processed with its own stateful ``response_parser`` (parsers - hold incremental tag-buffering state). Deltas are emitted as they arrive - from any generator, each tagged with its choice ``index``. After all - choices finish, a final aggregated usage chunk is emitted (when - ``include_usage``), followed by ``[DONE]``. Errors propagate to the whole - request. - """ - n = len(generators) - queue: asyncio.Queue = asyncio.Queue() - _DONE = 'done' - _DELTA = 'delta' - _ERROR = 'error' - - async def consume(index: int, gen: AsyncGenerator, parser) -> None: - streaming_tools = False - final_usage: UsageInfo | None = None - try: - async for res in gen: - logprobs = None - output_token_logprobs = None - if request.logprobs and res.logprobs: - logprobs = _create_chat_completion_logprobs( - tokenizer, res.token_ids, res.logprobs) - if request.return_logprob: - output_token_logprobs = _create_output_token_logprobs( - res.token_ids, res.logprobs) - if res.finish_reason and include_usage: - final_usage = UsageInfo.build( - prompt_tokens=res.input_token_len, - completion_tokens=res.generate_token_len, - cached_tokens=res.cached_tokens, - ) - delta_token_ids = (res.token_ids - if res.token_ids is not None else []) - stream_deltas = parser.stream_chunk(res.response, - delta_token_ids) - if not stream_deltas: - if res.finish_reason is None and not delta_token_ids: - continue - stream_deltas = [(DeltaMessage(role='assistant', - content=''), False)] - should_validate_complete = ( - res.finish_reason in ('stop', 'length') and - (request.return_token_ids - or request.return_routed_experts)) - if should_validate_complete and not parser.validate_complete(): - res.finish_reason = 'parse_error' - for delta_index, (delta_message, - tool_emitted) in enumerate(stream_deltas): - if tool_emitted: - streaming_tools = True - is_last_delta = delta_index == len(stream_deltas) - 1 - finish_reason = res.finish_reason if is_last_delta else None - chunk_logprobs = logprobs if is_last_delta else None - chunk_output_token_logprobs = (output_token_logprobs - if is_last_delta else None) - if (request.tool_choice != 'none' - and parser.tool_parser is not None): - if finish_reason == 'stop' and streaming_tools is True: - finish_reason = 'tool_calls' - routed_experts = (res.routed_experts - if finish_reason is not None else None) - stream_output_ids = delta_token_ids if ( - request.return_token_ids - and is_last_delta) else None - choice_data = ChatCompletionResponseStreamChoice( - index=index, - delta=delta_message, - finish_reason=finish_reason, - logprobs=chunk_logprobs, - output_token_logprobs=chunk_output_token_logprobs, - output_ids=stream_output_ids, - routed_experts=routed_experts, - ) - choice_data = maybe_filter_parallel_tool_calls( - choice_data, request) - response = ChatCompletionStreamResponse( - id=request_id, - created=created_time, - model=model_name, - choices=[choice_data], - usage=None, - ) - response_dict = response.model_dump(mode='json', - exclude_none=True) - if include_usage: - response_dict['usage'] = None - if res.cache_block_ids is not None and is_last_delta: - response_dict['cache_block_ids'] = res.cache_block_ids - response_dict['remote_token_ids'] = res.token_ids - await queue.put((_DELTA, response_dict)) - await queue.put((_DONE, index, final_usage)) - except Exception as e: # noqa: BLE001 - await queue.put((_ERROR, e)) - raise - - tasks = [ - asyncio.create_task(consume(idx, gen, parsers[idx])) - for idx, gen in generators - ] - pending_usages: dict[int, UsageInfo] = {} - done_count = 0 - try: - while done_count < n: - item = await queue.get() - kind = item[0] - if kind == _DELTA: - yield f'data: {json.dumps(item[1])}\n\n' - elif kind == _DONE: - done_count += 1 - _, idx, final_usage = item - if final_usage is not None: - pending_usages[idx] = final_usage - elif kind == _ERROR: - raise item[1] - if include_usage and pending_usages: - prompt_tokens = next(iter( - pending_usages.values())).prompt_tokens - completion_tokens = sum( - u.completion_tokens for u in pending_usages.values()) - cached_tokens = sum( - (u.prompt_tokens_details.cached_tokens - if u.prompt_tokens_details is not None else 0) - for u in pending_usages.values()) - agg_usage = UsageInfo.build( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - cached_tokens=cached_tokens, - ) - usage_resp = ChatCompletionStreamResponse( - id=request_id, - created=created_time, - model=model_name, - choices=[], - usage=agg_usage, - ) - yield f'data: {usage_resp.model_dump_json(exclude_none=True)}\n\n' - yield 'data: [DONE]\n\n' - finally: - for t in tasks: - if not t.done(): - t.cancel() - for t in tasks: - try: - await t - except (asyncio.CancelledError, Exception): - pass - - def register(router: APIRouter, server_context) -> None: @router.post('/v1/chat/completions', @@ -354,7 +64,7 @@ async def chat_completions_v1(request: ChatCompletionRequest, probable tokens with probabilities that add up to top_p or higher are kept for generation. - **n** (int): How many chat completion choices to generate for each input - message. **Only support one here**. + message. Accepts values from 1 to 128. - **stream**: whether to stream the results or not. Default to false. - **stream_options**: Options for streaming response. Only set this when you set stream: true. @@ -422,10 +132,23 @@ async def chat_completions_v1(request: ChatCompletionRequest, - **presence_penalty** (replaced with repetition_penalty) - **frequency_penalty** (replaced with repetition_penalty) """ - error_check_ret = validate_request(request, server_context, - check_request) + json_request = await raw_request.json() + error_check_ret = validate_request( + request, + server_context, + check_request, + json_request=json_request, + ) if error_check_ret is not None: return error_check_ret + if request.n is not None and request.n > 1: + return await fanout_chat_completions( + chat_completions_v1, + request, + raw_request, + json_request, + ) + # Resolve input: messages has priority over input_ids/image_data messages_empty = (request.messages is None or request.messages == '' or (isinstance(request.messages, list) @@ -456,7 +179,6 @@ async def chat_completions_v1(request: ChatCompletionRequest, # input_ids only — engine requires messages=None request.messages = None - json_request = await raw_request.json() migration_request = json_request.pop('migration_request', None) with_cache = json_request.pop('with_cache', False) preserve_cache = json_request.pop('preserve_cache', False) @@ -522,209 +244,32 @@ async def chat_completions_v1(request: ChatCompletionRequest, '`enable_thinking` in `chat_template_kwargs` will override the value in request.' ) + session = server_context.create_session(request.session_id) + try: + preprocessed = await server_context.async_engine.preprocess( + request.messages, + session, + gen_config=gen_config, + tools=request.tools, + reasoning_effort=request.reasoning_effort, + do_preprocess=do_preprocess, + adapter_name=adapter_name, + chat_template_kwargs=chat_template_kwargs or None, + input_ids=resolved_input_ids, + media_io_kwargs=request.media_io_kwargs, + mm_processor_kwargs=request.mm_processor_kwargs) + except RequestError as error: + return create_request_error_response(error) + try: + result_generator = server_context.async_engine.generate( + preprocessed, + stream_response=True) # always use stream to enable batching + except Exception: + server_context.session_manager.remove(session) + raise include_usage = bool(request.stream_options and request.stream_options.include_usage) - # ------------------------------------------------------------------ - # n > 1: server-side fan-out. - # A single n>1 request becomes N independent engine.generate() calls - # with distinct random seeds, collated into N choices. This is - # engine-agnostic handler-layer logic, so both pytorch and turbomind - # are covered without per-engine changes: each inner generate() still - # runs with n=1 (the engine-level n>1 fallback warning in - # async_engine.py is intentionally kept). n==1 keeps the original - # single-generator fast path below (no overhead). - # ------------------------------------------------------------------ - if request.n and request.n > 1: - n_choices = request.n - # The single session created earlier is unused on the fan-out - # path (each choice gets its own session below); release it. - server_context.session_manager.remove(session) - # Sub-sessions are internal: always auto-generate distinct ids - # (passing None) rather than reusing request.session_id. Reusing an - # explicit user session_id N times would collide in - # SessionManager.map_user_session_id on the 2nd call. Auto-gen - # avoids both the collision and any user-facing ambiguity about - # which choice owns the user-visible session id. - fanout_sessions = [ - server_context.create_session(None) - for _ in range(n_choices) - ] - fanout_parsers = [parser_cls(request) for _ in range(n_choices)] - fanout_generators: list[tuple[int, AsyncGenerator]] = [] - for i in range(n_choices): - sub_gen_config = deepcopy(gen_config) - # Per-choice seed: derive seed+i when request.seed is set, else - # leave None so the engine randomizes each choice independently - # (handled in AsyncEngine._determine_gen_config). - sub_gen_config.random_seed = ((request.seed + i) - if request.seed is not None - else None) - gen = server_context.async_engine.generate( - request.messages, - fanout_sessions[i], - gen_config=sub_gen_config, - tools=request.tools, - reasoning_effort=request.reasoning_effort, - stream_response=True, # always stream to enable batching - do_preprocess=do_preprocess, - adapter_name=adapter_name, - chat_template_kwargs=chat_template_kwargs or None, - input_ids=resolved_input_ids, - media_io_kwargs=request.media_io_kwargs, - mm_processor_kwargs=request.mm_processor_kwargs, - ) - fanout_generators.append((i, gen)) - - gen_list = [g for _, g in fanout_generators] - - # Streaming fan-out: interleave deltas from all N generators. - if request.stream: - stream_gen = _fanout_generate_stream( - fanout_generators, - fanout_parsers, - request, - request_id, - created_time, - model_name, - tokenizer, - include_usage, - ) - stream_generator = with_request_cleanup( - stream_gen, gen_list, fanout_sessions, - server_context.session_manager) - return StreamingResponse(stream_generator, - media_type='text/event-stream') - - # Non-streaming fan-out: consume all N generators concurrently. - async def _fanout_nonstream(): - # Mirror the n==1 non-streaming path: ensure engine generators - # are closed and fanout sessions removed on EVERY exit path - # (success, parse-error, disconnect, generator-error). Without - # this, every non-streaming n>1 request would leak N sessions. - try: - try: - results, usage_dict = await _fanout_generate_collect( - fanout_generators, - prompt_tokens=None, - disconnect_check=raw_request.is_disconnected, - ) - except _ClientDisconnected: - for s in fanout_sessions: - await s.async_abort() - return create_error_response( - HTTPStatus.BAD_REQUEST, 'Client disconnected') - - choices = [] - for res in results: - sub_parser = fanout_parsers[res.index] - tool_calls = None - reasoning_content = None - try: - raw_text = res.text - text, tool_calls, reasoning_content = \ - sub_parser.parse_complete( - res.text, res.token_ids) - should_validate_complete = ( - res.final_res.finish_reason in ('stop', - 'length') - and (request.return_token_ids - or request.return_routed_experts)) - if should_validate_complete and not \ - sub_parser.validate_complete(raw_text): - res.final_res.finish_reason = 'parse_error' - if isinstance(tool_calls, list) and len(tool_calls): - if res.final_res.finish_reason == 'stop': - res.final_res.finish_reason = 'tool_calls' - except Exception as e: # noqa: BLE001 - logger.error( - f'Failed to parse {res.text}. ' - f'Exception: {e}.') - return create_error_response( - HTTPStatus.BAD_REQUEST, - 'Failed to parse fc related info to ' - 'json format!') - - message = ChatMessage( - role='assistant', - content=text, - tool_calls=tool_calls, - reasoning_content=reasoning_content, - ) - choice_logprobs = None - if request.logprobs and len(res.logprobs): - choice_logprobs = _create_chat_completion_logprobs( - tokenizer, res.token_ids, res.logprobs) - choice_output_token_logprobs = None - if request.return_logprob and len(res.logprobs): - choice_output_token_logprobs = \ - _create_output_token_logprobs( - res.token_ids, res.logprobs) - choice_data = ChatCompletionResponseChoice( - index=res.index, - message=message, - logprobs=choice_logprobs, - output_token_logprobs=choice_output_token_logprobs, - finish_reason=res.final_res.finish_reason, - output_ids=(res.token_ids - if request.return_token_ids - else None), - routed_experts=(res.final_res.routed_experts - if request.return_routed_experts - else None), - ) - choice_data = maybe_filter_parallel_tool_calls( - choice_data, request) - choices.append(choice_data) - - usage = UsageInfo.build( - prompt_tokens=usage_dict['prompt_tokens'], - completion_tokens=usage_dict['completion_tokens'], - cached_tokens=usage_dict['cached_tokens'], - ) - response = ChatCompletionResponse( - id=request_id, - created=created_time, - model=model_name, - choices=choices, - usage=usage, - ).model_dump() - - # Disaggregation cache metadata (mirrors n==1 path). For - # fan-out the per-choice block-id lists are flattened: the - # first choice's first block id and the last choice's last - # remote token ids, matching the n==1 single-block shape. - if with_cache and results: - first = results[0] - last = results[-1] - response['cache_block_ids'] = ( - first.cache_block_ids[0] - if first.cache_block_ids else None) - response['remote_token_ids'] = [ - last.remote_token_ids[-1] - ] if last.remote_token_ids else [] - return response - finally: - await cleanup_result_generators( - gen_list, fanout_sessions, - server_context.session_manager) - - return await _fanout_nonstream() - - result_generator = server_context.async_engine.generate( - request.messages, - session, - gen_config=gen_config, - tools=request.tools, - reasoning_effort=request.reasoning_effort, - stream_response=True, # always use stream to enable batching - do_preprocess=do_preprocess, - adapter_name=adapter_name, - chat_template_kwargs=chat_template_kwargs or None, - input_ids=resolved_input_ids, - media_io_kwargs=request.media_io_kwargs, - mm_processor_kwargs=request.mm_processor_kwargs) - def create_stream_response_json( index: int, delta_message: DeltaMessage, @@ -855,11 +400,12 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: # Streaming response if request.stream: - stream_generator = with_request_cleanup( - completion_stream_generator(), [result_generator], [session], - server_context.session_manager) - return StreamingResponse(stream_generator, - media_type='text/event-stream') + return ManagedStreamingResponse( + completion_stream_generator(), + result_generators=[result_generator], + sessions=[session], + session_mgr=server_context.session_manager, + media_type='text/event-stream') # Non-streaming response final_logprobs = [] diff --git a/lmdeploy/serve/openai/chat_completions/streaming_response.py b/lmdeploy/serve/openai/chat_completions/streaming_response.py new file mode 100644 index 0000000000..7fcb46e8ea --- /dev/null +++ b/lmdeploy/serve/openai/chat_completions/streaming_response.py @@ -0,0 +1,114 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Streaming-response resource ownership for chat completions. + +Chat-completion fan-out creates multiple single-choice responses. Each child +allocates an engine session and result generator before Starlette starts its +body iterator. If another child fails, or the client disconnects before the +combined stream starts, cleanup placed only in an iterator ``finally`` block +is never activated and those resources can leak. + +``ManagedStreamingResponse`` makes that ownership explicit on the response +itself. Normal iteration still performs cleanup, while ``close()`` also lets +fan-out and the ASGI response lifecycle release resources whose iterators were +never started. +""" +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Iterable + +from fastapi.responses import StreamingResponse + +from lmdeploy.serve.utils.request_cleanup import cleanup_result_generators +from lmdeploy.utils import get_logger + +logger = get_logger('lmdeploy') + + +class ManagedStreamingResponse(StreamingResponse): + """A chat-completion response with explicit, cancellation-safe cleanup. + + Keeping cleanup on the response makes resource ownership independent of whether Starlette entered its body iterator. + """ + + def __init__( + self, + content, + *, + result_generators: Iterable = (), + sessions: Iterable = (), + session_mgr=None, + cleanup_callbacks: Iterable[Callable[[], Awaitable[None]]] = (), + **kwargs, + ): + self._result_generators = tuple(result_generators) + self._sessions = tuple(sessions) + self._session_mgr = session_mgr + self._cleanup_callbacks = tuple(cleanup_callbacks) + self._resource_cleanup_task: asyncio.Task | None = None + self._close_task: asyncio.Task | None = None + if self._result_generators or self._sessions: + content = self._with_resource_cleanup(content) + super().__init__(content, **kwargs) + + async def _with_resource_cleanup(self, content): + try: + async for item in content: + yield item + finally: + await self._cleanup_resources() + + async def _cleanup_resources(self) -> None: + if not self._result_generators and not self._sessions: + return + if self._resource_cleanup_task is None: + self._resource_cleanup_task = asyncio.create_task( + cleanup_result_generators( + self._result_generators, + self._sessions, + self._session_mgr, + ), + name='streaming_response_resource_cleanup') + await asyncio.shield(self._resource_cleanup_task) + + async def _close(self) -> None: + body_iterator = self.body_iterator + close_iterator = getattr(body_iterator, 'aclose', None) + if close_iterator is not None: + try: + await close_iterator() + except (asyncio.CancelledError, GeneratorExit): + pass + except Exception: + logger.exception('Close response body iterator failed.') + + await self._cleanup_resources() + for callback in self._cleanup_callbacks: + try: + await callback() + except (asyncio.CancelledError, GeneratorExit): + pass + except Exception: + logger.exception('Streaming response cleanup callback failed.') + + async def close(self) -> None: + """Close the body and its resources exactly once.""" + if (not self._cleanup_callbacks + and self._resource_cleanup_task is not None + and self._resource_cleanup_task.done()): + return + if self._close_task is None: + self._close_task = asyncio.create_task( + self._close(), name='streaming_response_cleanup') + try: + await asyncio.shield(self._close_task) + except (asyncio.CancelledError, GeneratorExit): + raise + except Exception: + logger.exception('Streaming response cleanup failed.') + + async def __call__(self, scope, receive, send) -> None: + try: + await super().__call__(scope, receive, send) + finally: + await self.close() diff --git a/lmdeploy/serve/openai/chat_completions/validation.py b/lmdeploy/serve/openai/chat_completions/validation.py index 9c6d9418ff..4554395ceb 100644 --- a/lmdeploy/serve/openai/chat_completions/validation.py +++ b/lmdeploy/serve/openai/chat_completions/validation.py @@ -9,7 +9,9 @@ _MAX_FANOUT_N = 128 -def check_request(request: ChatCompletionRequest, server_context) -> str: +def check_request(request: ChatCompletionRequest, + server_context, + json_request: dict | None = None) -> str: engine_config = server_context.engine_config session_manager = server_context.session_manager try: @@ -39,24 +41,28 @@ def check_request(request: ChatCompletionRequest, server_context) -> str: return f'The session_id {request.session_id!r} is occupied.' # check sampling settings - if request.n <= 0: + if request.n is not None and request.n <= 0: return f'The n {request.n!r} must be a positive int.' # n > 1 is implemented as server-side fan-out (N independent engine # generate() calls). Cap it to prevent unbounded resource use. - if request.n > _MAX_FANOUT_N: + if request.n is not None and request.n > _MAX_FANOUT_N: return (f'The n {request.n!r} exceeds the maximum supported ' f'choices ({_MAX_FANOUT_N}).') + if request.n is not None and request.n > 1 and request.session_id not in ( + None, -1): + return 'n > 1 cannot be used with an explicit session_id.' + if request.n is not None and request.n > 1 and json_request is not None: + if any( + json_request.get(key) + for key in ('migration_request', 'with_cache', + 'preserve_cache')): + return 'n > 1 is not supported with cache migration.' if request.top_p is not None and not (0 < request.top_p <= 1): return f'The top_p {request.top_p!r} must be in (0, 1].' if request.top_k is not None and request.top_k < 0: return f'The top_k {request.top_k!r} cannot be a negative integer.' if request.temperature is not None and not (0 <= request.temperature <= 2): return f'The temperature {request.temperature!r} must be in [0, 2]' - # seed validation: per-choice seeds are derived as `seed + i` for n > 1, - # so a negative seed could collide with engine internals; reject it. - if request.seed is not None and request.seed < 0: - return f'The seed {request.seed!r} must be a non-negative int.' - # Validate input_ids and image_data constraints. # messages has higher priority. input_ids and image_data are only used when # messages is empty (None, '', or []). image_data requires input_ids. diff --git a/lmdeploy/serve/openai/endpoints/common.py b/lmdeploy/serve/openai/endpoints/common.py index 8b222ff4f0..41b080a015 100644 --- a/lmdeploy/serve/openai/endpoints/common.py +++ b/lmdeploy/serve/openai/endpoints/common.py @@ -23,7 +23,8 @@ def build_serving_generation_config(request, server_context, ) -def validate_request(request, server_context, request_validator): +def validate_request(request, server_context, request_validator, + **validator_kwargs): """Validate the selected model and endpoint-specific request contract.""" if hasattr( request, @@ -32,7 +33,8 @@ def validate_request(request, server_context, request_validator): HTTPStatus.NOT_FOUND, f'The model {request.model!r} does not exist.') - error_message = request_validator(request, server_context) + error_message = request_validator(request, server_context, + **validator_kwargs) if error_message: return create_error_response(HTTPStatus.BAD_REQUEST, error_message) return None diff --git a/lmdeploy/serve/proxy/proxy.py b/lmdeploy/serve/proxy/proxy.py index 7afd875c2b..015925ad61 100644 --- a/lmdeploy/serve/proxy/proxy.py +++ b/lmdeploy/serve/proxy/proxy.py @@ -587,7 +587,7 @@ async def chat_completions_v1(request: ChatCompletionRequest, raw_request: Reque probable tokens with probabilities that add up to top_p or higher are kept for generation. - **n** (int): How many chat completion choices to generate for each input - message. **Only support one here**. + message. Accepts values from 1 to 128, except in DistServe mode. - **stream**: whether to stream the results or not. Default to false. - **max_completion_tokens** (int | None): output token nums. Default to None. - **max_tokens** (int | None): output token nums. Default to None. @@ -650,6 +650,11 @@ async def chat_completions_v1(request: ChatCompletionRequest, raw_request: Reque check_response = await node_manager.check_request_model(request.model) if check_response is not None: return check_response + if (node_manager.serving_strategy == ServingStrategy.DistServe + and request.n is not None and request.n > 1): + return create_error_response( + HTTPStatus.BAD_REQUEST, + 'n > 1 is not supported with the DistServe serving strategy.') if node_manager.serving_strategy == ServingStrategy.Hybrid: node_url = node_manager.get_node_url(request.model) diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py b/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py index 521a9d6a48..3d47578669 100644 --- a/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py +++ b/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py @@ -30,9 +30,13 @@ def __init__(self): self.call_count = 0 self.gen_configs = [] - def generate(self, prompt, session, **kwargs): - self.call_count += 1 + async def preprocess(self, prompt, session, **kwargs): + """Return the minimal preprocessed input consumed by the fake.""" self.gen_configs.append(kwargs.get('gen_config')) + return SimpleNamespace(prompt=prompt, session=session) + + def generate(self, preprocessed, **kwargs): + self.call_count += 1 call_index = self.call_count async def _generator(): diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py index 50b54e9de6..41b6d284f3 100644 --- a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py +++ b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py @@ -1,357 +1,245 @@ # Copyright (c) OpenMMLab. All rights reserved. -"""Unit tests for ``n > 1`` server-side fan-out in the chat completions -handler. - -The fan-out is engine-agnostic handler-layer logic: a single ``n > 1`` request -becomes N independent ``engine.generate()`` calls with distinct random seeds, -collated into N choices. These tests cover the pure aggregation helper -``_fanout_generate_collect`` using fake async generators that mimic the engine's -``GenOut`` yields. Both pytorch and turbomind engines are covered because the -fan-out lives entirely in the handler and treats the engine as a black box. - -Note: the repo uses ``asyncio.run`` (no ``pytest-asyncio`` dependency), so async -test bodies are driven through ``asyncio.run``. -""" +"""Regression tests for multiple chat completion choices.""" from __future__ import annotations import asyncio +import json from types import SimpleNamespace import pytest +from fastapi.responses import JSONResponse -from lmdeploy.serve.openai.chat_completions.serving import ( - _fanout_generate_collect, -) +from lmdeploy.serve.openai.protocol import ChatCompletionRequest -def _fake_gen(outputs): - """Build an async generator yielding ``GenOut``-like objects.""" +class _PreprocessingEngine: + """Provide the preprocessing stage expected by the serving endpoint.""" - async def _gen(): - for o in outputs: - yield o + async def preprocess(self, prompt, session, **kwargs): + self.gen_configs.append(kwargs.get('gen_config')) + return SimpleNamespace(prompt=prompt, session=session) - return _gen() - -def _genout(text, completion_tokens, *, prompt_tokens=5, finish_reason='stop'): - return SimpleNamespace( - response=text, - input_token_len=prompt_tokens, - generate_token_len=completion_tokens, - finish_reason=finish_reason, - token_ids=[], - logprobs=None, - cached_tokens=0, - routed_experts=None, - cache_block_ids=None, +def _request(**kwargs): + return ChatCompletionRequest( + model='fake-model', + messages=[{ + 'role': 'user', + 'content': 'hi' + }], + **kwargs, ) -def test_fanout_assigns_distinct_indices_and_aggregates_completion_tokens(): - """Three generators producing 2/3/1 completion tokens are collated into - three choices with distinct indices; prompt_tokens counted once, - completion_tokens summed.""" - gens = [ - (0, _fake_gen([_genout('a', 2)])), - (1, _fake_gen([_genout('b', 3)])), - (2, _fake_gen([_genout('c', 1)])), - ] - choices, usage = asyncio.run(_fanout_generate_collect(gens, prompt_tokens=5)) - assert len(choices) == 3 - assert {c.index for c in choices} == {0, 1, 2} - assert usage['prompt_tokens'] == 5 # counted once - assert usage['completion_tokens'] == 6 # 2 + 3 + 1 - - -def test_fanout_prompt_tokens_from_generator_overrides_when_unspecified(): - """When prompt_tokens is passed explicitly it is used as the single prompt- - token count (never summed across choices).""" - gens = [(0, _fake_gen([_genout('a', 2, prompt_tokens=99)]))] - choices, usage = asyncio.run(_fanout_generate_collect(gens, prompt_tokens=7)) - assert usage['prompt_tokens'] == 7 - assert usage['completion_tokens'] == 2 - - -def test_fanout_propagates_error_to_whole_request(): - """If any inner generator raises, the whole fan-out request fails.""" - - async def _boom(): - raise RuntimeError('choice 1 failed') - yield # noqa: unreachable, makes it an async generator - - with pytest.raises(RuntimeError, match='choice 1 failed'): - asyncio.run(_fanout_generate_collect([(0, _boom())], prompt_tokens=1)) - - -# --------------------------------------------------------------------------- -# Handler-level integration: exercise the n>1 branch end-to-end with a fake -# engine. Validates wiring (N sessions, N parsers, distinct seeds, aggregated -# usage, N choices) for both streaming and non-streaming. -# --------------------------------------------------------------------------- - -from lmdeploy.serve.openai.protocol import ( # noqa: E402 - ChatCompletionRequest, -) - - def _sse_payloads(text): - import json payloads = [] for line in text.splitlines(): if line.startswith('data: '): data = line.removeprefix('data: ') - if data == '[DONE]': - continue - payloads.append(json.loads(data)) + if data != '[DONE]': + payloads.append(json.loads(data)) return payloads -def test_handler_n3_nonstream_returns_three_choices_with_aggregated_usage( +async def _collect_stream(response): + chunks = [] + async for chunk in response.body_iterator: + chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk) + return ''.join(chunks) + + +def test_handler_n3_nonstream_collates_single_choice_path( chat_endpoint, fake_raw_request): - """N=3 non-streaming: 3 distinct choices, prompt counted once, - completion_tokens summed; engine called 3 times with distinct seeds.""" endpoint, context = chat_endpoint - request = ChatCompletionRequest(model='fake-model', - messages=[{'role': 'user', - 'content': 'hi'}], - n=3, - seed=42, - stream=False) - response = asyncio.run(endpoint(request, fake_raw_request)) + response = asyncio.run(endpoint(_request(n=3, seed=42), fake_raw_request)) assert response['object'] == 'chat.completion' - assert len(response['choices']) == 3 - assert {c['index'] for c in response['choices']} == {0, 1, 2} - # Each choice got distinct text from a distinct generate() call. - assert {c['message']['content'] - for c in response['choices']} == {'choice-1', 'choice-2', - 'choice-3'} - # prompt_tokens counted once (4), completion_tokens = 1 + 2 + 3 = 6. + assert [choice['index'] for choice in response['choices']] == [0, 1, 2] + assert {choice['message']['content'] + for choice in response['choices'] + } == {'choice-1', 'choice-2', 'choice-3'} assert response['usage']['prompt_tokens'] == 4 assert response['usage']['completion_tokens'] == 6 - # Engine was invoked 3 times with derived seeds 42, 43, 44. + assert [config.random_seed + for config in context.async_engine.gen_configs] == [42, 43, 44] assert context.async_engine.call_count == 3 - seeds = [gc.random_seed for gc in context.async_engine.gen_configs] - assert seeds == [42, 43, 44] - # All N fan-out sessions (plus the single pre-fan-out session) are removed - # after the request — no session leak on the non-streaming path. - assert len(context.session_manager.removed) == 3 + 1 - # And no sessions remain live in the manager. + assert len(context.session_manager.removed) == 3 assert context.session_manager.sessions == {} -def test_handler_n3_stream_interleaves_three_indices_and_aggregates_usage( +def test_handler_n3_stream_interleaves_indices_and_aggregates_usage( chat_endpoint, fake_raw_request): - """N=3 streaming: deltas carry indices 0/1/2, final usage chunk sums - completion tokens across choices.""" endpoint, context = chat_endpoint - request = ChatCompletionRequest(model='fake-model', - messages=[{'role': 'user', - 'content': 'hi'}], - n=3, - stream=True, - stream_options={'include_usage': True}) + request = _request( + n=3, + stream=True, + stream_options={'include_usage': True}, + ) response = asyncio.run(endpoint(request, fake_raw_request)) - - # StreamingResponse.body is an async iterable; collect it. - body_iterator = response.body_iterator - - async def _collect(): - chunks = [] - async for chunk in body_iterator: - chunks.append(chunk.decode() - if isinstance(chunk, bytes) else chunk) - return ''.join(chunks) - - text = asyncio.run(_collect()) + text = asyncio.run(_collect_stream(response)) payloads = _sse_payloads(text) - choice_indices = set() - for p in payloads: - for c in p.get('choices', []): - choice_indices.add(c['index']) - assert choice_indices == {0, 1, 2} - - # The final chunk carries aggregated usage (prompt once, completion sum). - usage_chunks = [p for p in payloads if p.get('usage') is not None] - assert usage_chunks, 'expected a final usage chunk' - final_usage = usage_chunks[-1]['usage'] - assert final_usage['prompt_tokens'] == 4 - assert final_usage['completion_tokens'] == 6 # 1 + 2 + 3 + indices = { + choice['index'] + for payload in payloads + for choice in payload.get('choices', []) + } + assert indices == {0, 1, 2} + usage_chunks = [ + payload for payload in payloads if payload.get('usage') is not None + ] + assert len(usage_chunks) == 1 + assert usage_chunks[0]['choices'] == [] + assert usage_chunks[0]['usage']['prompt_tokens'] == 4 + assert usage_chunks[0]['usage']['completion_tokens'] == 6 assert text.rstrip().endswith('data: [DONE]') - # Streaming fan-out must also clean up all N fan-out sessions (plus the - # pre-fan-out single session) once the stream completes. - assert len(context.session_manager.removed) == 3 + 1 + assert len(context.session_manager.removed) == 3 assert context.session_manager.sessions == {} -def test_handler_n1_keeps_single_generator_fast_path(chat_endpoint, - fake_raw_request): - """N=1 (default) must not fan out: exactly one engine.generate() call.""" +@pytest.mark.parametrize('n', [None, 1]) +def test_handler_single_choice_keeps_fast_path(n, chat_endpoint, + fake_raw_request): endpoint, context = chat_endpoint - request = ChatCompletionRequest(model='fake-model', - messages=[{'role': 'user', - 'content': 'hi'}], - stream=False) - response = asyncio.run(endpoint(request, fake_raw_request)) + response = asyncio.run(endpoint(_request(n=n), fake_raw_request)) assert len(response['choices']) == 1 assert context.async_engine.call_count == 1 -def test_handler_n3_unseeded_leaves_random_seed_none(chat_endpoint, - fake_raw_request): - """When request.seed is unset, each sub gen_config keeps random_seed=None - so the engine randomizes each choice independently.""" +def test_handler_unseeded_choices_leave_seed_resolution_to_engine( + chat_endpoint, fake_raw_request): endpoint, context = chat_endpoint - request = ChatCompletionRequest(model='fake-model', - messages=[{'role': 'user', - 'content': 'hi'}], - n=3, - stream=False) - asyncio.run(endpoint(request, fake_raw_request)) - seeds = [gc.random_seed for gc in context.async_engine.gen_configs] - assert seeds == [None, None, None] - - -def test_validation_rejects_oversized_n(): - """Fan-out resource cap: n above _MAX_FANOUT_N is rejected.""" - from types import SimpleNamespace - - from lmdeploy.serve.openai.chat_completions.validation import _MAX_FANOUT_N, check_request - - request = ChatCompletionRequest(model='fake-model', - messages=[{'role': 'user', - 'content': 'hi'}], - n=_MAX_FANOUT_N + 1) - ctx = SimpleNamespace( - engine_config=SimpleNamespace(logprobs_mode=None, adapters=[]), - session_manager=SimpleNamespace(has=lambda sid: False), - response_parser_cls=None, - ) - msg = check_request(request, ctx) - assert 'exceeds the maximum' in msg + asyncio.run(endpoint(_request(n=3), fake_raw_request)) + assert [config.random_seed for config in context.async_engine.gen_configs + ] == [None, None, None] -def test_validation_rejects_negative_seed(): - from types import SimpleNamespace +def test_handler_negative_seed_is_mapped_to_engine_seed_domain( + chat_endpoint, fake_raw_request): + endpoint, context = chat_endpoint + asyncio.run(endpoint(_request(n=2, seed=-1), fake_raw_request)) + assert [config.random_seed + for config in context.async_engine.gen_configs] == [(1 << 64) - 1, + 0] + + +@pytest.mark.parametrize('n, expected', [ + (0, 'positive int'), + (129, 'maximum supported'), +]) +def test_validation_rejects_invalid_n(n, expected, chat_endpoint, + fake_raw_request): + endpoint, context = chat_endpoint + response = asyncio.run(endpoint(_request(n=n), fake_raw_request)) + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert expected in response.body.decode() + assert context.async_engine.call_count == 0 - from lmdeploy.serve.openai.chat_completions.validation import check_request - request = ChatCompletionRequest(model='fake-model', - messages=[{'role': 'user', - 'content': 'hi'}], - seed=-7) - ctx = SimpleNamespace( - engine_config=SimpleNamespace(logprobs_mode=None, adapters=[]), - session_manager=SimpleNamespace(has=lambda sid: False), - response_parser_cls=None, - ) - msg = check_request(request, ctx) - assert 'non-negative' in msg +def test_handler_rejects_explicit_session_id_for_multiple_choices( + chat_endpoint, fake_raw_request): + endpoint, context = chat_endpoint + response = asyncio.run( + endpoint(_request(n=2, session_id=777), fake_raw_request)) + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert 'explicit session_id' in response.body.decode() + assert context.async_engine.call_count == 0 + assert context.session_manager.sessions == {} + + +def test_handler_rejects_cache_migration_for_multiple_choices( + chat_endpoint, fake_raw_request): + endpoint, context = chat_endpoint + fake_raw_request._payload = {'with_cache': True} + response = asyncio.run(endpoint(_request(n=2), fake_raw_request)) + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert 'cache migration' in response.body.decode() + assert context.async_engine.call_count == 0 -# --------------------------------------------------------------------------- -# Fix-round-1 regression tests: session-id collision, session cleanup, sibling -# cancellation, multi-chunk interleaving. -# --------------------------------------------------------------------------- +def test_distserve_proxy_rejects_multiple_choices(monkeypatch): + from lmdeploy.pytorch.disagg.config import ServingStrategy + from lmdeploy.serve.proxy import proxy + async def model_exists(model): + return None -def test_handler_n3_with_explicit_session_id_does_not_crash(chat_endpoint, - fake_raw_request): - """An explicit user session_id + n>1 must not collide in - SessionManager.map_user_session_id. + monkeypatch.setattr(proxy.node_manager, 'check_request_model', + model_exists) + monkeypatch.setattr(proxy.node_manager, 'serving_strategy', + ServingStrategy.DistServe) + response = asyncio.run(proxy.chat_completions_v1(_request(n=2))) + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert 'DistServe' in response.body.decode() + + +def test_prompt_cache_usage_is_counted_once(chat_endpoint, fake_raw_request): + + class CachedEngine(_PreprocessingEngine): + model_name = 'fake-model' + backend_config = SimpleNamespace(adapters=[], logprobs_mode=None) + + def __init__(self, original_engine): + self.session_mgr = original_engine.session_mgr + self.tokenizer = original_engine.tokenizer + self.call_count = 0 + self.gen_configs = [] + + def generate(self, preprocessed, **kwargs): + self.call_count += 1 + index = self.call_count + + async def generate(): + yield SimpleNamespace( + response=f'choice-{index}', + token_ids=[index], + input_token_len=4, + generate_token_len=1, + finish_reason='stop', + logprobs=None, + cached_tokens=index, + routed_experts=None, + cache_block_ids=None, + ) + + return generate() - Fan-out sub-sessions are auto-generated (None), so the user id is mapped at most once and N distinct internal - sessions are created. Regression for the crash + leaked-session bug. - """ endpoint, context = chat_endpoint - request = ChatCompletionRequest(model='fake-model', - messages=[{'role': 'user', - 'content': 'hi'}], - n=3, - session_id=777, - stream=False) - response = asyncio.run(endpoint(request, fake_raw_request)) + context.async_engine = CachedEngine(context.async_engine) + response = asyncio.run(endpoint(_request(n=2), fake_raw_request)) + assert response['usage']['prompt_tokens'] == 4 + assert response['usage']['completion_tokens'] == 2 + assert response['usage']['prompt_tokens_details']['cached_tokens'] == 1 - assert len(response['choices']) == 3 - assert {c['index'] for c in response['choices']} == {0, 1, 2} - # The user session_id was mapped exactly once (to the pre-fan-out single - # session, which is then removed). - assert 777 not in context.session_manager.user_session_id_map - # N distinct internal fan-out sessions were created and all cleaned up. - assert context.async_engine.call_count == 3 - assert context.session_manager.sessions == {} +def test_streaming_multiple_choices_preserves_each_inner_parser( + chat_endpoint, fake_raw_request): -def test_fanout_cancels_sibling_generators_on_error(): - """When one fan-out generator raises, the still-running siblings are - cancelled and their generators closed BEFORE the error propagates out of - _fanout_generate_collect (not only at event-loop shutdown). - - Regression for the asyncio.gather-doesn't-cancel-siblings bug. - """ - from lmdeploy.serve.openai.chat_completions.serving import _fanout_generate_collect - - sibling_closed_before_error = {'value': False} - - async def _boom(): - raise RuntimeError('choice 0 failed') - yield # noqa: unreachable - - async def _long_running(): - try: - # Pretend to produce forever; should be cancelled before done. - while True: - yield _genout('x', 1) - await asyncio.sleep(0.01) - except (asyncio.CancelledError, GeneratorExit): - sibling_closed_before_error['value'] = True - raise - - async def _run_and_record_order(): - # The sibling must be cancelled BEFORE _fanout_generate_collect raises. - # We record the closure state synchronously in the except block, while - # still inside the event loop (before asyncio.run tears it down). - with pytest.raises(RuntimeError, match='choice 0 failed'): - await _fanout_generate_collect( - [(0, _boom()), (1, _long_running())], prompt_tokens=1) - return sibling_closed_before_error['value'] - - closed_before = asyncio.run(_run_and_record_order()) - assert closed_before, \ - 'sibling generator was not cancelled/closed before the error propagated' - - -def test_handler_n2_stream_interleaves_multi_chunk_per_choice(chat_endpoint, - fake_raw_request): - """Streaming fan-out where each generator yields multiple chunks: deltas - from both choices are interleaved and each choice's index appears with its - full text content across chunks.""" - - class MultiChunkEngine: + class MultiChunkEngine(_PreprocessingEngine): model_name = 'fake-model' backend_config = SimpleNamespace(adapters=[], logprobs_mode=None) - def __init__(self): - self.session_mgr = None # wired from the existing context below - self.tokenizer = SimpleNamespace( - model=SimpleNamespace(model='fake-tokenizer')) + def __init__(self, original_engine): + self.session_mgr = original_engine.session_mgr + self.tokenizer = original_engine.tokenizer self.call_count = 0 self.gen_configs = [] - def generate(self, prompt, session, **kwargs): + def generate(self, preprocessed, **kwargs): self.call_count += 1 - self.gen_configs.append(kwargs.get('gen_config')) - idx = self.call_count + index = self.call_count - async def _gen(): - for piece in (f'{idx}-a', f'{idx}-b', f'{idx}-c'): + async def generate(): + for suffix in ('a', 'b', 'c'): yield SimpleNamespace( - response=piece, - token_ids=[len(piece)], + response=f'{index}-{suffix}', + token_ids=[index], input_token_len=3, - generate_token_len=len(piece), + generate_token_len=1, finish_reason=None, logprobs=None, cached_tokens=0, @@ -362,7 +250,7 @@ async def _gen(): response='', token_ids=[], input_token_len=3, - generate_token_len=0, + generate_token_len=3, finish_reason='stop', logprobs=None, cached_tokens=0, @@ -370,45 +258,232 @@ async def _gen(): cache_block_ids=None, ) - return _gen() + return generate() endpoint, context = chat_endpoint - # Swap in a multi-chunk engine while reusing the context's session manager. - original_engine = context.async_engine - multi_engine = MultiChunkEngine() - multi_engine.session_mgr = original_engine.session_mgr - context.async_engine = multi_engine - try: - request = ChatCompletionRequest( - model='fake-model', - messages=[{'role': 'user', 'content': 'hi'}], - n=2, - stream=True, - stream_options={'include_usage': True}) - response = asyncio.run(endpoint(request, fake_raw_request)) - - async def _collect(): - chunks = [] - async for chunk in response.body_iterator: - chunks.append(chunk.decode() - if isinstance(chunk, bytes) else chunk) - return ''.join(chunks) - - text = asyncio.run(_collect()) - finally: - context.async_engine = original_engine + context.async_engine = MultiChunkEngine(context.async_engine) + response = asyncio.run( + endpoint(_request(n=2, stream=True), fake_raw_request)) + payloads = _sse_payloads(asyncio.run(_collect_stream(response))) + content = {0: '', 1: ''} + for payload in payloads: + for choice in payload.get('choices', []): + content[choice['index']] += choice['delta'].get('content') or '' + assert content == {0: '1-a1-b1-c', 1: '2-a2-b2-c'} + + +def test_fanout_error_cancels_siblings_and_cleans_sessions( + chat_endpoint, fake_raw_request): - payloads = _sse_payloads(text) - # Both choices appear, and the concatenated content per index reconstructs - # the full multi-chunk text for that choice. - per_index = {} - for p in payloads: - for c in p.get('choices', []): - per_index.setdefault(c['index'], '') - content = c['delta'].get('content') if c.get('delta') else None - if content: - per_index[c['index']] += content - assert set(per_index) == {0, 1} - assert per_index[0] == '1-a1-b1-c' - assert per_index[1] == '2-a2-b2-c' - assert text.rstrip().endswith('data: [DONE]') + class FailingEngine(_PreprocessingEngine): + model_name = 'fake-model' + backend_config = SimpleNamespace(adapters=[], logprobs_mode=None) + + def __init__(self, original_engine): + self.session_mgr = original_engine.session_mgr + self.tokenizer = original_engine.tokenizer + self.call_count = 0 + self.gen_configs = [] + self.sibling_started = asyncio.Event() + self.sibling_closed = False + + def generate(self, preprocessed, **kwargs): + self.call_count += 1 + index = self.call_count + + async def fail(): + await self.sibling_started.wait() + raise RuntimeError('choice failed') + yield # noqa: unreachable + + async def wait_forever(): + self.sibling_started.set() + try: + await asyncio.Event().wait() + yield # noqa: unreachable + finally: + self.sibling_closed = True + + return fail() if index == 1 else wait_forever() + + endpoint, context = chat_endpoint + engine = FailingEngine(context.async_engine) + context.async_engine = engine + + with pytest.raises(RuntimeError, match='choice failed'): + asyncio.run(endpoint(_request(n=2), fake_raw_request)) + assert engine.sibling_closed + assert context.session_manager.sessions == {} + + +def test_early_stream_close_cleans_all_choice_sessions(chat_endpoint, + fake_raw_request): + endpoint, context = chat_endpoint + response = asyncio.run( + endpoint(_request(n=2, stream=True), fake_raw_request)) + + async def consume_one_chunk(): + iterator = response.body_iterator + await anext(iterator) + await iterator.aclose() + + asyncio.run(consume_one_chunk()) + assert context.session_manager.sessions == {} + + +def test_asgi_disconnect_before_stream_start_cleans_all_choice_sessions( + chat_endpoint, fake_raw_request): + endpoint, context = chat_endpoint + + async def disconnect_before_stream_start(): + response = await endpoint( + _request(n=2, stream=True), fake_raw_request) + + async def receive(): + return {'type': 'http.disconnect'} + + async def send(message): + if message['type'] == 'http.response.start': + await asyncio.Event().wait() + + scope = { + 'type': 'http', + 'asgi': { + 'version': '3.0', + 'spec_version': '2.3' + }, + } + await response(scope, receive, send) + + asyncio.run(asyncio.wait_for(disconnect_before_stream_start(), 1)) + assert context.session_manager.sessions == {} + + +def test_streaming_cancelled_choice_fails_without_hanging( + chat_endpoint, fake_raw_request): + + class CancelledChoiceEngine(_PreprocessingEngine): + model_name = 'fake-model' + backend_config = SimpleNamespace(adapters=[], logprobs_mode=None) + + def __init__(self, original_engine): + self.session_mgr = original_engine.session_mgr + self.tokenizer = original_engine.tokenizer + self.call_count = 0 + self.gen_configs = [] + self.sibling_started = asyncio.Event() + self.sibling_closed = False + + def generate(self, preprocessed, **kwargs): + self.call_count += 1 + index = self.call_count + + async def cancel(): + await self.sibling_started.wait() + raise asyncio.CancelledError + yield # noqa: unreachable + + async def wait_forever(): + self.sibling_started.set() + try: + await asyncio.Event().wait() + yield # noqa: unreachable + finally: + self.sibling_closed = True + + return cancel() if index == 1 else wait_forever() + + endpoint, context = chat_endpoint + engine = CancelledChoiceEngine(context.async_engine) + context.async_engine = engine + + async def collect(): + response = await endpoint( + _request(n=2, stream=True), fake_raw_request) + await asyncio.wait_for(_collect_stream(response), 1) + + with pytest.raises(RuntimeError, match='choice 0 was cancelled'): + asyncio.run(collect()) + assert engine.sibling_closed + assert context.session_manager.sessions == {} + + +def test_streaming_setup_failure_closes_completed_choice_responses( + chat_endpoint, fake_raw_request): + + class SetupFailingEngine(_PreprocessingEngine): + model_name = 'fake-model' + backend_config = SimpleNamespace(adapters=[], logprobs_mode=None) + + def __init__(self, original_engine): + self.session_mgr = original_engine.session_mgr + self.tokenizer = original_engine.tokenizer + self.call_count = 0 + self.gen_configs = [] + + def generate(self, preprocessed, **kwargs): + self.call_count += 1 + if self.call_count == 2: + raise RuntimeError('choice setup failed') + + async def wait_forever(): + await asyncio.Event().wait() + yield # noqa: unreachable + + return wait_forever() + + endpoint, context = chat_endpoint + context.async_engine = SetupFailingEngine(context.async_engine) + + with pytest.raises(RuntimeError, match='choice setup failed'): + asyncio.run( + endpoint(_request(n=2, stream=True), fake_raw_request)) + assert context.session_manager.sessions == {} + + +def test_stream_usage_is_omitted_when_a_choice_has_no_usage( + chat_endpoint, fake_raw_request): + + class IncompleteUsageEngine(_PreprocessingEngine): + model_name = 'fake-model' + backend_config = SimpleNamespace(adapters=[], logprobs_mode=None) + + def __init__(self, original_engine): + self.session_mgr = original_engine.session_mgr + self.tokenizer = original_engine.tokenizer + self.call_count = 0 + self.gen_configs = [] + + def generate(self, preprocessed, **kwargs): + self.call_count += 1 + index = self.call_count + + async def generate(): + yield SimpleNamespace( + response=f'choice-{index}', + token_ids=[index], + input_token_len=4, + generate_token_len=1, + finish_reason='stop' if index == 1 else None, + logprobs=None, + cached_tokens=0, + routed_experts=None, + cache_block_ids=None, + ) + + return generate() + + endpoint, context = chat_endpoint + context.async_engine = IncompleteUsageEngine(context.async_engine) + request = _request( + n=2, + stream=True, + stream_options={'include_usage': True}, + ) + response = asyncio.run(endpoint(request, fake_raw_request)) + payloads = _sse_payloads(asyncio.run(_collect_stream(response))) + + assert not [ + payload for payload in payloads if payload.get('usage') is not None + ] + assert context.session_manager.sessions == {} From 07c3b0c6f849db4ca5f572e931f33ff2dc7a058d Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Wed, 12 Aug 2026 08:51:10 +0000 Subject: [PATCH 4/7] perf(chat): batch ready fan-out stream chunks --- .../serve/openai/chat_completions/fanout.py | 55 +++++++++++++++---- .../chat_completions/test_n_completions.py | 34 ++++++++++++ 2 files changed, 79 insertions(+), 10 deletions(-) diff --git a/lmdeploy/serve/openai/chat_completions/fanout.py b/lmdeploy/serve/openai/chat_completions/fanout.py index 1caef1c0bc..e253e03415 100644 --- a/lmdeploy/serve/openai/chat_completions/fanout.py +++ b/lmdeploy/serve/openai/chat_completions/fanout.py @@ -225,6 +225,26 @@ async def _stream_choice( await _close_streaming_response(response) +def _batch_stream_payloads(payloads: list[dict]) -> list[dict]: + """Combine each choice's Nth ready delta into the Nth output batch.""" + batches: list[dict] = [] + next_batch_by_index: dict[int, int] = {} + + for payload in payloads: + choice = payload['choices'][0] + index = choice['index'] + target = next_batch_by_index.get(index, 0) + if target == len(batches): + batches.append(payload) + else: + batches[target]['choices'].append(choice) + next_batch_by_index[index] = target + 1 + + for payload in batches: + payload['choices'].sort(key=lambda choice: choice['index']) + return batches + + async def _collate_streams( responses: list[StreamingResponse], request: ChatCompletionRequest, @@ -247,16 +267,31 @@ async def _collate_streams( # consume the streaming responses of each fan-out request, and yield to the client try: while completed < len(tasks): - item = await queue.get() - if item[0] == 'data': - yield f'data: {json.dumps(item[1])}\n\n' - elif item[0] == 'usage': - # item[1]: index, iterm[2]: usage payload - usages[item[1]] = item[2] - elif item[0] == 'done': - completed += 1 - else: - raise item[1] + items = [await queue.get()] + while True: + try: + items.append(queue.get_nowait()) + except asyncio.QueueEmpty: + break + + payloads = [] + stream_error = None + for item in items: + if item[0] == 'data': + payloads.append(item[1]) + elif item[0] == 'usage': + # item[1]: index, item[2]: usage payload + usages[item[1]] = item[2] + elif item[0] == 'done': + completed += 1 + else: + stream_error = item[1] + break + + for payload in _batch_stream_payloads(payloads): + yield f'data: {json.dumps(payload)}\n\n' + if stream_error is not None: + raise stream_error if include_usage and len(usages) == len(tasks): ordered_usages = [usages[index] for index in range(len(tasks))] usage_response = { diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py index 41b6d284f3..1f7b91cdeb 100644 --- a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py +++ b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py @@ -9,6 +9,7 @@ import pytest from fastapi.responses import JSONResponse +from lmdeploy.serve.openai.chat_completions.fanout import _batch_stream_payloads from lmdeploy.serve.openai.protocol import ChatCompletionRequest @@ -48,6 +49,39 @@ async def _collect_stream(response): return ''.join(chunks) +def _stream_payload(index, content, **metadata): + return { + 'id': 'chatcmpl-test', + 'object': 'chat.completion.chunk', + 'created': 1, + 'model': 'fake-model', + 'choices': [{ + 'index': index, + 'delta': { + 'content': content + } + }], + **metadata, + } + + +def test_ready_stream_chunks_are_batched_by_choice(): + payloads = [ + _stream_payload(0, '0-a'), + _stream_payload(0, '0-b'), + _stream_payload(1, '1-a'), + ] + + batches = _batch_stream_payloads(payloads) + + assert [[choice['index'] for choice in batch['choices']] + for batch in batches] == [[0, 1], [0]] + assert [choice['delta']['content'] + for batch in batches + for choice in batch['choices'] if choice['index'] == 0 + ] == ['0-a', '0-b'] + + def test_handler_n3_nonstream_collates_single_choice_path( chat_endpoint, fake_raw_request): endpoint, context = chat_endpoint From 8110e02efc54ed20a62841d731a760468d641e0e Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Wed, 12 Aug 2026 09:12:14 +0000 Subject: [PATCH 5/7] docs(chat): document fan-out helpers --- .../serve/openai/chat_completions/fanout.py | 19 +++++++++++++++++++ .../chat_completions/streaming_response.py | 3 +++ .../openai/chat_completions/validation.py | 1 + 3 files changed, 23 insertions(+) diff --git a/lmdeploy/serve/openai/chat_completions/fanout.py b/lmdeploy/serve/openai/chat_completions/fanout.py index e253e03415..cea8e34493 100644 --- a/lmdeploy/serve/openai/chat_completions/fanout.py +++ b/lmdeploy/serve/openai/chat_completions/fanout.py @@ -23,6 +23,8 @@ @dataclass class _FanoutResponseError(Exception): + """Carry an HTTP error response out of concurrent choice invocation.""" + response: Response @@ -34,17 +36,21 @@ def __init__(self, request: Request, payload: dict): self._payload = payload def __getattr__(self, name): + """Delegate request attributes not overridden by this wrapper.""" return getattr(self._request, name) async def json(self) -> dict: + """Return an isolated copy of this choice's raw JSON payload.""" return deepcopy(self._payload) async def is_disconnected(self) -> bool: + """Report the connection state of the original client request.""" return await self._request.is_disconnected() def _choice_request(request: ChatCompletionRequest, index: int) -> ChatCompletionRequest: + """Create an independent single-choice request for one fan-out index.""" choice_request = request.model_copy(deep=True) choice_request.n = 1 choice_request.session_id = -1 @@ -54,6 +60,7 @@ def _choice_request(request: ChatCompletionRequest, async def _cancel_tasks(tasks: list[asyncio.Task]) -> list: + """Cancel unfinished tasks and collect every terminal result.""" for task in tasks: if not task.done(): task.cancel() @@ -61,6 +68,7 @@ async def _cancel_tasks(tasks: list[asyncio.Task]) -> list: async def _close_streaming_response(response: StreamingResponse) -> None: + """Close a child response through its explicit or iterator lifecycle.""" close_response = getattr(response, 'close', None) if close_response is not None: await close_response() @@ -71,6 +79,7 @@ async def _close_streaming_response(response: StreamingResponse) -> None: async def _close_responses(responses) -> None: + """Close all streaming responses in an invocation result collection.""" await asyncio.gather(*( _close_streaming_response(response) for response in responses @@ -79,11 +88,13 @@ async def _close_responses(responses) -> None: async def _cleanup_invocations(tasks: list[asyncio.Task]) -> None: + """Cancel choice invocations and close responses they already created.""" results = await _cancel_tasks(tasks) await _close_responses(results) def _consume_cleanup_result(task: asyncio.Task) -> None: + """Retrieve a detached cleanup task's result to suppress task warnings.""" try: task.result() except BaseException: # cleanup is best-effort after caller cancellation @@ -91,6 +102,7 @@ def _consume_cleanup_result(task: asyncio.Task) -> None: async def _shield_cleanup(awaitable, name: str) -> None: + """Let cleanup continue if cancellation interrupts its caller.""" cleanup_task = asyncio.create_task(awaitable, name=name) try: await asyncio.shield(cleanup_task) @@ -105,8 +117,10 @@ async def _invoke_choices( raw_request: Request, payload: dict, ) -> list[dict | StreamingResponse] | Response: + """Invoke the single-choice endpoint concurrently for every choice.""" async def invoke(index: int): + """Invoke and validate one indexed single-choice response.""" choice_request = _choice_request(request, index) choice_payload = deepcopy(payload) choice_payload.update(n=1, session_id=-1, seed=choice_request.seed) @@ -133,11 +147,13 @@ async def invoke(index: int): def _cached_tokens(usage: dict) -> int: + """Read cached prompt tokens from an OpenAI-compatible usage object.""" details = usage.get('prompt_tokens_details') or {} return details.get('cached_tokens', 0) def _aggregate_usage(usages: list[dict]) -> UsageInfo: + """Count the shared prompt once and sum completion tokens by choice.""" first_usage = usages[0] return UsageInfo.build( prompt_tokens=first_usage.get('prompt_tokens', 0), @@ -152,6 +168,7 @@ def _collate_responses( request_id: str, created_time: int, ) -> dict: + """Combine single-choice JSON responses into one multi-choice response.""" response = deepcopy(responses[0]) response['id'] = request_id response['created'] = created_time @@ -180,6 +197,7 @@ async def _stream_choice( model_name: str, stopping: asyncio.Event, ) -> None: + """Parse one child SSE stream and forward normalized events to a queue.""" buffer = '' try: async for chunk in response.body_iterator: @@ -251,6 +269,7 @@ async def _collate_streams( request_id: str, created_time: int, ) -> AsyncGenerator[str, None]: + """Interleave child streams into one multi-choice SSE response.""" queue: asyncio.Queue = asyncio.Queue(maxsize=max(1, len(responses) * 2)) stopping = asyncio.Event() # produce the streaming responses for each fan-out request diff --git a/lmdeploy/serve/openai/chat_completions/streaming_response.py b/lmdeploy/serve/openai/chat_completions/streaming_response.py index 7fcb46e8ea..6919c12055 100644 --- a/lmdeploy/serve/openai/chat_completions/streaming_response.py +++ b/lmdeploy/serve/openai/chat_completions/streaming_response.py @@ -52,6 +52,7 @@ def __init__( super().__init__(content, **kwargs) async def _with_resource_cleanup(self, content): + """Clean owned engine resources after normal body iteration.""" try: async for item in content: yield item @@ -59,6 +60,7 @@ async def _with_resource_cleanup(self, content): await self._cleanup_resources() async def _cleanup_resources(self) -> None: + """Close result generators and sessions once, despite cancellation.""" if not self._result_generators and not self._sessions: return if self._resource_cleanup_task is None: @@ -72,6 +74,7 @@ async def _cleanup_resources(self) -> None: await asyncio.shield(self._resource_cleanup_task) async def _close(self) -> None: + """Close the response body, owned resources, and child callbacks.""" body_iterator = self.body_iterator close_iterator = getattr(body_iterator, 'aclose', None) if close_iterator is not None: diff --git a/lmdeploy/serve/openai/chat_completions/validation.py b/lmdeploy/serve/openai/chat_completions/validation.py index 4554395ceb..18ef3571f0 100644 --- a/lmdeploy/serve/openai/chat_completions/validation.py +++ b/lmdeploy/serve/openai/chat_completions/validation.py @@ -12,6 +12,7 @@ def check_request(request: ChatCompletionRequest, server_context, json_request: dict | None = None) -> str: + """Validate chat-completion options and fan-out compatibility.""" engine_config = server_context.engine_config session_manager = server_context.session_manager try: From e5ade8b38f35b983de1d344450002ae6ab5f73b7 Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Thu, 13 Aug 2026 03:21:39 +0000 Subject: [PATCH 6/7] fix ut --- tests/test_lmdeploy/serve/openai/chat_completions/conftest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py b/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py index 3d47578669..64327b28d3 100644 --- a/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py +++ b/tests/test_lmdeploy/serve/openai/chat_completions/conftest.py @@ -64,6 +64,7 @@ def __init__(self, request): self.request = request self.tool_parser = None self._chunks = [] + self.reasoning_tokens = 0 def stream_chunk(self, delta_text, delta_token_ids, **kwargs): if not delta_text: From f5258484cbcfd1b1d7fd1d335b4499c60292f42e Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Fri, 14 Aug 2026 08:04:07 +0000 Subject: [PATCH 7/7] fix(chat): aggregate reasoning tokens for n choices --- lmdeploy/serve/openai/chat_completions/fanout.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/lmdeploy/serve/openai/chat_completions/fanout.py b/lmdeploy/serve/openai/chat_completions/fanout.py index cea8e34493..b9ae570d3d 100644 --- a/lmdeploy/serve/openai/chat_completions/fanout.py +++ b/lmdeploy/serve/openai/chat_completions/fanout.py @@ -153,13 +153,21 @@ def _cached_tokens(usage: dict) -> int: def _aggregate_usage(usages: list[dict]) -> UsageInfo: - """Count the shared prompt once and sum completion tokens by choice.""" + """Count shared prompt usage once and sum per-choice completion usage.""" first_usage = usages[0] + completion_details = [ + usage.get('completion_tokens_details') for usage in usages + ] + reasoning_tokens = None + if all(details is not None for details in completion_details): + reasoning_tokens = sum( + details['reasoning_tokens'] for details in completion_details) return UsageInfo.build( prompt_tokens=first_usage.get('prompt_tokens', 0), completion_tokens=sum( usage.get('completion_tokens') or 0 for usage in usages), cached_tokens=_cached_tokens(first_usage), + reasoning_tokens=reasoning_tokens, )