From 19f873ece43c076f06fa9ccb63b1f8372cf5b6fa Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Sun, 9 Aug 2026 08:27:22 +0000 Subject: [PATCH 1/2] 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 | 429 +++++++++++++++++- .../openai/chat_completions/validation.py | 13 + .../serve/openai/chat_completions/conftest.py | 142 ++++++ .../chat_completions/test_n_completions.py | 241 ++++++++++ 4 files changed, 822 insertions(+), 3 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 2d9c0728ff..afd4ae3dea 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 @@ -36,6 +39,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', @@ -228,6 +498,161 @@ async def chat_completions_v1(request: ChatCompletionRequest, '`enable_thinking` in `chat_template_kwargs` will override the value in request.' ) + 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, @@ -241,8 +666,6 @@ async def chat_completions_v1(request: ChatCompletionRequest, input_ids=resolved_input_ids, media_io_kwargs=request.media_io_kwargs, mm_processor_kwargs=request.mm_processor_kwargs) - include_usage = bool(request.stream_options - and request.stream_options.include_usage) def create_stream_response_json( index: int, 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 3d7e19f6267ea7489c79553d696844511551d469 Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Sun, 9 Aug 2026 09:04:38 +0000 Subject: [PATCH 2/2] 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 afd4ae3dea..f2598ae58e 100644 --- a/lmdeploy/serve/openai/chat_completions/serving.py +++ b/lmdeploy/serve/openai/chat_completions/serving.py @@ -28,7 +28,7 @@ UsageInfo, ) from lmdeploy.serve.openai.utils import create_error_response, 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 @@ -67,7 +67,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 @@ -104,19 +105,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') @@ -130,8 +131,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) @@ -516,8 +537,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)] @@ -568,88 +595,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]')