From ab63e4acb52399fa751c40ed296801c5fffc313c Mon Sep 17 00:00:00 2001 From: A5bhinav Date: Tue, 27 Oct 2026 14:54:08 -0700 Subject: [PATCH 1/2] fix(server): refuse n != 1 on /v1/completions instead of ignoring it --- python/freetoken/server/openai_api.py | 2 ++ tests/server/test_openai_api.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 99a4283549..15992930e3 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -396,6 +396,8 @@ async def handle_completion( unsupported = _completion_unsupported_reason(req) if unsupported is not None: return create_error_response(unsupported) + if req.n != 1: + return create_error_response("Only n=1 is supported", param="n") try: # surfaces an out-of-range max_tokens as a 400 rather than a 500 from the worker default_max_tokens = ( getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index e33018e195..3dcef25f34 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -4,6 +4,7 @@ import json from types import SimpleNamespace +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from freetoken.message import TokenizeMsg, UserReply @@ -356,6 +357,24 @@ def test_completion_rejects_token_id_prompts(): assert "token-id prompt" in body["error"]["message"] +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("n", [0, 2]) +def test_completion_refuses_n_other_than_one(n, stream): + state = FakeState([UserReply(uid=42, incremental_output="hello", finished=True)]) + response = run( + handle_completion( + CompletionRequest(model="client-model", prompt="say hi", n=n, stream=stream), + request=None, + state=state, + model_sampling={}, + ) + ) + + assert response.status_code == 400 + assert json.loads(response.body)["error"]["param"] == "n" + assert state.sent is None + + def test_completion_accepts_text_prompt(): state = FakeState( [UserReply(uid=42, incremental_output="hello", finished=True, prompt_tokens_delta=2, completion_tokens_delta=1)] From 662c266e56a6aad5dc431a7d08ccb13c29e446a6 Mon Sep 17 00:00:00 2001 From: A5bhinav Date: Tue, 27 Oct 2026 16:22:53 -0700 Subject: [PATCH 2/2] feat(server): support n > 1 on /v1/completions --- python/freetoken/server/api_server.py | 13 +- python/freetoken/server/generation.py | 105 +++++++++++++- python/freetoken/server/openai_api.py | 191 +++++++++++++++----------- tests/server/test_openai_api.py | 130 +++++++++++++++++- 4 files changed, 350 insertions(+), 89 deletions(-) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index ce04b34102..4bf034632a 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -375,16 +375,23 @@ async def stream_generate(self, uid: int): yield b"data: [DONE]\n\n" logger.debug("Finished streaming response for user %s", uid) - async def stream_with_cancellation(self, generator, request: Request, uid: int): + async def stream_with_cancellation( + self, generator, request: Request, uid: int | Callable[[], list[int]] + ): + """``uid`` may be a callable returning the uids to abort, for a stream (OpenAI ``n``) + whose requests are submitted after it starts.""" try: async for chunk in generator: # detect if the client has disconnected if await request.is_disconnected(): - logger.info("Client disconnected for user %s", uid) + logger.info( + "Client disconnected for user %s", uid if isinstance(uid, int) else "(n > 1)" + ) raise asyncio.CancelledError yield chunk except asyncio.CancelledError: - asyncio.create_task(self.abort_user(uid)) + for live_uid in uid() if callable(uid) else [uid]: + asyncio.create_task(self.abort_user(live_uid)) raise async def abort_user(self, uid: int): diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index ef90402aab..d933371cf3 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -16,7 +16,7 @@ import asyncio import json import time -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass, field from typing import Any @@ -507,6 +507,109 @@ async def with_keepalive(events: AsyncIterator[GenEvent], interval: float): task.cancel() +_CHOICE_DONE = object() + + +@dataclass +class _ChoiceFailed: + exc: Exception + + +class ChoiceFanOut: + """Runs ``n`` generations of one prompt (OpenAI ``n``) and merges their outputs as + ``(choice_index, item)`` pairs. Choice 0 is submitted first and the rest follow its first + output for which ``releases_siblings`` holds: by then its prompt is in the prefix cache, + so they skip a second prefill. + + ``submit(uid)`` sends one request; ``consume(uid)`` yields its outputs, ends at its + terminal one and raises on failure. Closing ``events()`` early (error, disconnect) + aborts every choice not yet finished, including ones never submitted.""" + + def __init__( + self, + n: int, + state: Any, + submit: Callable[[int], Awaitable[None]], + consume: Callable[[int], AsyncIterator[Any]], + releases_siblings: Callable[[Any], bool], + ) -> None: + self._n = n + self._state = state + self._submit = submit + self._consume = consume + self._releases_siblings = releases_siblings + self._uids: list[int] = [] + self._live: set[int] = set() + self._closed = False + + async def start(self) -> None: + """Admit every choice and submit choice 0. Admission is decided once for the whole + request, so it fails here, before a stream opens, rather than mid-response.""" + try: + for _ in range(self._n): + uid = self._state.new_user() + self._uids.append(uid) + self._live.add(uid) + await self._submit(self._uids[0]) + except BaseException: + self.abort_live() + raise + + def take_live_uids(self) -> list[int]: + """Stop launching choices and hand over the uids not yet finished, to abort.""" + self._closed = True + uids = list(self._live) + self._live.clear() + return uids + + def abort_live(self) -> None: + """Abort every choice not yet finished; a no-op once ``events()`` has cleaned up.""" + for uid in self.take_live_uids(): + asyncio.create_task(self._state.abort_user(uid)) + + async def events(self) -> AsyncIterator[tuple[int, Any]]: + if not self._uids: + await self.start() + queue: asyncio.Queue[tuple[int, Any]] = asyncio.Queue() + released = asyncio.Event() + + async def run(index: int) -> None: + uid = self._uids[index] + try: + if index: + await released.wait() + if self._closed: + return + await self._submit(uid) + async for item in self._consume(uid): + if self._releases_siblings(item): + released.set() + queue.put_nowait((index, item)) + self._live.discard(uid) + except Exception as exc: # noqa: BLE001 -- re-raised by the consumer below + self._closed = True + queue.put_nowait((index, _ChoiceFailed(exc))) + finally: + released.set() + queue.put_nowait((index, _CHOICE_DONE)) + + tasks = [asyncio.create_task(run(index)) for index in range(self._n)] + try: + pending = self._n + while pending: + index, item = await queue.get() + if item is _CHOICE_DONE: + pending -= 1 + elif isinstance(item, _ChoiceFailed): + raise item.exc + else: + yield index, item + finally: + for task in tasks: + task.cancel() + self.abort_live() + + def _record_generation( *, source: str | None, diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 15992930e3..dd88ed54ae 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib import json import time import uuid @@ -9,6 +10,7 @@ from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, StreamingResponse +from starlette.background import BackgroundTask from freetoken.core import SamplingParams from freetoken.message import TokenizeMsg from freetoken.tokenizer.effort import EFFORT_SCALE, KNOWN_REASONING_EFFORTS @@ -24,6 +26,7 @@ from .request_logger import log_request from .generation import ( DEFAULT_MAX_OUTPUT_TOKENS, + ChoiceFanOut, ContentDelta, GenDone, GenerationError, @@ -387,6 +390,10 @@ async def stream_chat_completion_chunks( yield b"data: [DONE]\n\n" +#: OpenAI's own ceiling for ``n``. +MAX_N = 128 + + async def handle_completion( req: CompletionRequest, request: Request | None, @@ -396,8 +403,8 @@ async def handle_completion( unsupported = _completion_unsupported_reason(req) if unsupported is not None: return create_error_response(unsupported) - if req.n != 1: - return create_error_response("Only n=1 is supported", param="n") + if not 1 <= req.n <= MAX_N: + return create_error_response(f"n must be between 1 and {MAX_N}", param="n") try: # surfaces an out-of-range max_tokens as a 400 rather than a 500 from the worker default_max_tokens = ( getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS @@ -406,50 +413,59 @@ async def handle_completion( except ValueError as exc: return create_error_response(str(exc), param="max_tokens") + def sampling() -> SamplingParams: + return _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) + prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt assert isinstance(prompts, list) if req.stream: if len(prompts) != 1: return create_error_response("Streaming completions only support a single text prompt") - uid = state.new_user() - await state.send_one( - TokenizeMsg(uid=uid, text=prompts[0], sampling_params=_resolve_sampling( - req, model_sampling, default_max_tokens=default_max_tokens - )) - ) - chunks = stream_completion_chunks(uid, req, state) + fan_out = _completion_fan_out(req.n, prompts[0], sampling, state) + await fan_out.start() + chunks = stream_completion_chunks(fan_out, req, state) if request is not None: - chunks = state.stream_with_cancellation(chunks, request, uid) - return StreamingResponse(chunks, media_type="text/event-stream") + chunks = state.stream_with_cancellation(chunks, request, fan_out.take_live_uids) + + async def abort_if_never_streamed() -> None: + # Choice 0 is already submitted; a body that never starts never reaches cleanup. + fan_out.abort_live() + + return StreamingResponse( + chunks, media_type="text/event-stream", background=BackgroundTask(abort_if_never_streamed) + ) choices: list[dict[str, Any]] = [] prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 - for index, prompt in enumerate(prompts): - uid = state.new_user() - await state.send_one( - TokenizeMsg( - uid=uid, - text=prompt, - sampling_params=_resolve_sampling( - req, model_sampling, default_max_tokens=default_max_tokens - ), - ) + for prompt_index, prompt in enumerate(prompts): + texts = [""] * req.n + finish_reasons = ["stop"] * req.n + fan_out = _completion_fan_out(req.n, prompt, sampling, state) + await fan_out.start() + try: + async with contextlib.aclosing(fan_out.events()) as events: + async for index, ack in events: + # Every choice re-reports the shared prompt; OpenAI counts it once per prompt. + if index == 0: + prompt_tokens += ack.prompt_tokens_delta + cached_tokens += ack.cached_tokens + completion_tokens += ack.completion_tokens_delta + texts[index] += ack.incremental_output + if ack.finished: + finish_reasons[index] = getattr(ack, "finish_reason", None) or "stop" + except GenerationError as exc: + return create_error_response(str(exc)) + choices.extend( + { + "index": prompt_index * req.n + index, + "text": text, + "finish_reason": finish_reason, + "logprobs": None, + } + for index, (text, finish_reason) in enumerate(zip(texts, finish_reasons)) ) - text = "" - finish_reason = "stop" - async for ack in state.wait_for_ack(uid): - if getattr(ack, "error", None): - return create_error_response(ack.error) - prompt_tokens += ack.prompt_tokens_delta - completion_tokens += ack.completion_tokens_delta - cached_tokens += ack.cached_tokens - text += ack.incremental_output - if ack.finished: - finish_reason = getattr(ack, "finish_reason", None) or "stop" - break - choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None}) return { "id": f"cmpl-{uuid.uuid4().hex}", @@ -461,62 +477,71 @@ async def handle_completion( } -async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) -> AsyncIterator[bytes]: - prompt_tokens = 0 - completion_tokens = 0 - cached_tokens = 0 - finish_reason = "stop" - async for ack in state.wait_for_ack(uid): - if getattr(ack, "error", None): - yield _sse({"error": {"message": ack.error, "type": "invalid_request_error", "code": None}}) - yield b"data: [DONE]\n\n" - return - prompt_tokens += ack.prompt_tokens_delta - completion_tokens += ack.completion_tokens_delta - cached_tokens += ack.cached_tokens - if ack.incremental_output: - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [ - { - "text": ack.incremental_output, - "index": 0, - "finish_reason": None, - "logprobs": None, - } - ], - } - ) - if ack.finished: - finish_reason = getattr(ack, "finish_reason", None) or "stop" - break +def _completion_fan_out( + n: int, prompt: str, sampling: Callable[[], SamplingParams], state: Any +) -> ChoiceFanOut: + async def submit(uid: int) -> None: + await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling())) - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [{"text": "", "index": 0, "finish_reason": finish_reason, "logprobs": None}], - } - ) - if req.stream_options and req.stream_options.include_usage: - yield _sse( + async def consume(uid: int) -> AsyncIterator[Any]: + async for ack in state.wait_for_ack(uid): + if getattr(ack, "error", None): + raise GenerationError(ack.error) + yield ack + if ack.finished: + return + + def releases_siblings(ack: Any) -> bool: + # The first reply only reports admission; the prompt is cached by the first sampled token. + return ack.finished or ack.completion_tokens_delta > 0 + + return ChoiceFanOut(n, state, submit, consume, releases_siblings) + + +async def stream_completion_chunks( + fan_out: ChoiceFanOut, req: CompletionRequest, state: Any +) -> AsyncIterator[bytes]: + completion_id = f"cmpl-{uuid.uuid4().hex}" + + def chunk(choices: list[dict[str, Any]], **extra: Any) -> bytes: + return _sse( { - "id": f"cmpl-{uid}", + "id": completion_id, "object": "text_completion.chunk", "created": int(time.time()), "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), + "choices": choices, + **extra, } ) + + prompt_tokens = 0 + completion_tokens = 0 + cached_tokens = 0 + try: + async with contextlib.aclosing(fan_out.events()) as events: + async for index, ack in events: + if index == 0: + prompt_tokens += ack.prompt_tokens_delta + cached_tokens += ack.cached_tokens + completion_tokens += ack.completion_tokens_delta + if ack.incremental_output: + yield chunk( + [{"text": ack.incremental_output, "index": index, "finish_reason": None, "logprobs": None}] + ) + if ack.finished: + finish_reason = getattr(ack, "finish_reason", None) or "stop" + yield chunk([{"text": "", "index": index, "finish_reason": finish_reason, "logprobs": None}]) + except GenerationError as exc: + yield _sse({"error": {"message": str(exc), "type": "invalid_request_error", "code": None}}) + yield b"data: [DONE]\n\n" + return + + if req.stream_options and req.stream_options.include_usage: + yield chunk( + [], + usage=_usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)), + ) yield b"data: [DONE]\n\n" diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index 3dcef25f34..1a447a2de5 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -358,8 +358,8 @@ def test_completion_rejects_token_id_prompts(): @pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("n", [0, 2]) -def test_completion_refuses_n_other_than_one(n, stream): +@pytest.mark.parametrize("n", [0, 129]) +def test_completion_refuses_n_out_of_range(n, stream): state = FakeState([UserReply(uid=42, incremental_output="hello", finished=True)]) response = run( handle_completion( @@ -411,6 +411,132 @@ def test_completion_forwards_length_finish_reason(): assert response["choices"][0]["finish_reason"] == "length" +class MultiUserState(FakeState): + """One reply list per uid, in submission order; logs sends, acks and aborts.""" + + def __init__(self, replies_by_uid: dict[int, list[UserReply]]) -> None: + super().__init__([]) + self.replies_by_uid = replies_by_uid + self.next_uid = 42 + self.log: list[tuple[str, int]] = [] + + def new_user(self) -> int: + self.next_uid += 1 + return self.next_uid - 1 + + async def send_one(self, msg): + self.log.append(("send", msg.uid)) + + async def wait_for_ack(self, uid: int): + for reply in self.replies_by_uid[uid]: + await asyncio.sleep(0) # each engine reply is a separate message, so siblings can run + self.log.append(("ack", uid)) + yield reply + await asyncio.Event().wait() # an unfinished request stays open, as in the engine + + async def abort_user(self, uid: int): + self.log.append(("abort", uid)) + + +def _reply(uid: int, text: str, finished: bool = True, **kwargs) -> UserReply: + return UserReply(uid=uid, incremental_output=text, finished=finished, **kwargs) + + +def test_completion_n_orders_choices_and_counts_prompt_once(): + usage = {"prompt_tokens_delta": 5, "completion_tokens_delta": 1} + state = MultiUserState({ + 42: [_reply(42, "a0", **usage)], + 43: [_reply(43, "a1", **usage)], + 44: [_reply(44, "b0", **usage)], + 45: [_reply(45, "b1", **usage)], + }) + response = run(handle_completion( + CompletionRequest(model="m", prompt=["a", "b"], n=2), + request=None, state=state, model_sampling={}, + )) + + assert [(c["index"], c["text"]) for c in response["choices"]] == [ + (0, "a0"), (1, "a1"), (2, "b0"), (3, "b1"), + ] + assert response["usage"] == {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14} + + +def test_completion_n_submits_siblings_after_first_sampled_token(): + # The engine's first reply is the admission notice, sent before the prompt is cached. + state = MultiUserState({ + 42: [_reply(42, "", finished=False, prompt_tokens_delta=5), + _reply(42, "x", finished=False, completion_tokens_delta=1), + _reply(42, "y", completion_tokens_delta=1)], + 43: [_reply(43, "z", completion_tokens_delta=1)], + }) + run(handle_completion( + CompletionRequest(model="m", prompt="p", n=2), + request=None, state=state, model_sampling={}, + )) + + acks_42 = [i for i, entry in enumerate(state.log) if entry == ("ack", 42)] + assert state.log.index(("send", 43)) > acks_42[1] + + +def test_completion_n_error_aborts_sibling_choices(): + state = MultiUserState({ + 42: [_reply(42, "partial", finished=False, completion_tokens_delta=1)], + 43: [_reply(43, "", error="prompt too long")], + }) + + async def go(): + response = await handle_completion( + CompletionRequest(model="m", prompt="p", n=2), + request=None, state=state, model_sampling={}, + ) + await asyncio.sleep(0) # let the scheduled aborts run + return response + + response = run(go()) + assert response.status_code == 400 + assert ("abort", 42) in state.log + + +def test_completion_n_choice_zero_error_submits_no_siblings(): + state = MultiUserState({ + 42: [_reply(42, "", error="prompt too long")], + 43: [_reply(43, "x", completion_tokens_delta=1)], + 44: [_reply(44, "y", completion_tokens_delta=1)], + }) + response = run(handle_completion( + CompletionRequest(model="m", prompt="p", n=3), + request=None, state=state, model_sampling={}, + )) + + assert response.status_code == 400 + assert [entry for entry in state.log if entry[0] == "send"] == [("send", 42)] + + +def test_stream_completion_n_tags_chunks_and_finishes_each_choice_once(): + state = MultiUserState({ + 42: [_reply(42, "a", finished=False, prompt_tokens_delta=5, completion_tokens_delta=1), + _reply(42, "b", completion_tokens_delta=1, finish_reason="length")], + 43: [_reply(43, "c", prompt_tokens_delta=5, completion_tokens_delta=1)], + }) + + async def collect(): + response = await handle_completion( + CompletionRequest(model="m", prompt="p", n=2, stream=True, + stream_options={"include_usage": True}), + request=None, state=state, model_sampling={}, + ) + return [chunk async for chunk in response.body_iterator] + + events = parse_sse(run(collect())) + assert events[-1] == "[DONE]" + choices = [c for e in events[:-1] for c in e["choices"]] + text = {i: "".join(c["text"] for c in choices if c["index"] == i) for i in (0, 1)} + finishes = [(c["index"], c["finish_reason"]) for c in choices if c["finish_reason"]] + assert text == {0: "ab", 1: "c"} + assert sorted(finishes) == [(0, "length"), (1, "stop")] + assert events[-2]["usage"] == {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8} + + def test_omitted_max_tokens_honors_server_default(): from freetoken.server.generation import DEFAULT_MAX_OUTPUT_TOKENS