From 9dd6e985e673b863d1440a55cbc7ab760517e258 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:46:53 -0700 Subject: [PATCH 1/7] api: sentence splitter for long text-to-speech inputs --- mstar/api_server/openai/speech_chunking.py | 100 +++++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 mstar/api_server/openai/speech_chunking.py diff --git a/mstar/api_server/openai/speech_chunking.py b/mstar/api_server/openai/speech_chunking.py new file mode 100644 index 0000000000..e21cadf931 --- /dev/null +++ b/mstar/api_server/openai/speech_chunking.py @@ -0,0 +1,100 @@ +"""Sentence chunking for ``/v1/audio/speech``. + +Long inputs are split into sentence groups that are synthesized as separate +requests and played back in order. This keeps every autoregressive +text-to-speech request inside the length its model was trained for, lets the +engine batch the pieces of one long input like independent requests, and +starts the second piece while the first is still streaming (see +``serving_speech``). Splitting is purely textual and model-agnostic. +""" + +from __future__ import annotations + +import re + +# One sentence: a lazy body up to a Latin terminator (with any closing quotes +# or brackets) that precedes whitespace or the end, a CJK terminator, a +# paragraph break, or the end of the text. +_SENTENCE = re.compile( + r""".+?(?: + [.!?;]["'”’)\]]*(?=\s|\Z) + | [。！？；] + | (?=\n[ \t]*\n) + | \Z + )""", + re.VERBOSE | re.DOTALL, +) +# Soft break points inside an over-long sentence, most preferred first. +_SOFT_BREAKS = ( + re.compile(r"(?<=[,;:，；：])\s*"), + re.compile(r"\s+"), +) + + +def _hard_wrap(sentence: str, max_chars: int) -> list[str]: + """Split one over-long sentence at clause boundaries, then at spaces.""" + pieces = [sentence] + for pattern in _SOFT_BREAKS: + wrapped: list[str] = [] + for piece in pieces: + if len(piece) <= max_chars: + wrapped.append(piece) + continue + current = "" + for part in (p.strip() for p in pattern.split(piece)): + if not part: + continue + candidate = f"{current} {part}" if current else part + if current and len(candidate) > max_chars: + wrapped.append(current) + current = part + else: + current = candidate + if current: + wrapped.append(current) + pieces = wrapped + return pieces + + +def split_sentences(text: str, max_chars: int = 400, min_chars: int = 24) -> list[str]: + """Group ``text`` into sentence chunks of at most ``max_chars`` characters. + + Sentences are never cut unless one alone exceeds ``max_chars`` (then it is + wrapped at clause boundaries, or spaces as a last resort). A trailing + fragment shorter than ``min_chars`` is merged into its predecessor so the + model is not asked to voice a lone "Okay." Returns ``[text]`` when nothing + needs splitting and ``[]`` for blank input. + """ + if max_chars <= 0: + raise ValueError("max_chars must be positive") + text = text.strip() + if not text: + return [] + if len(text) <= max_chars: + return [text] + + sentences: list[str] = [] + for match in _SENTENCE.finditer(text): + piece = " ".join(match.group(0).split()) + if not piece: + continue + sentences.extend(_hard_wrap(piece, max_chars) if len(piece) > max_chars else [piece]) + + chunks: list[str] = [] + current = "" + for sentence in sentences: + if current and len(current) + 1 + len(sentence) > max_chars: + chunks.append(current) + current = sentence + else: + current = f"{current} {sentence}".strip() + if current: + chunks.append(current) + + if ( + len(chunks) > 1 + and len(chunks[-1]) < min_chars + and len(chunks[-2]) + 1 + len(chunks[-1]) <= max_chars * 5 // 4 + ): + chunks[-2:] = [f"{chunks[-2]} {chunks[-1]}"] + return chunks From 6d1e734f9290d72c65b094fa8ce0abf397d9d3ad Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:46:53 -0700 Subject: [PATCH 2/7] api: per-adapter sentence-chunking thresholds for /v1/audio/speech --- mstar/api_server/openai/adapters.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/mstar/api_server/openai/adapters.py b/mstar/api_server/openai/adapters.py index 5a715e2735..06bd3fd498 100644 --- a/mstar/api_server/openai/adapters.py +++ b/mstar/api_server/openai/adapters.py @@ -185,6 +185,16 @@ class OpenAIAdapter: supports_videos: bool = False # POST /v1/videos/generations supports_realtime: bool = False # /v1/realtime (bidirectional speech WebSocket) + # ``/v1/audio/speech`` sentence chunking (``serving_speech``): inputs of at + # least ``speech_chunk_min_chars`` characters are split into sentence + # groups of about ``speech_chunk_max_chars`` and synthesized as ordered + # sub-requests. ``None`` keeps the whole text in one request unless the + # client sends ``sentence_chunking: true``. + speech_chunk_min_chars: int | None = None + speech_chunk_max_chars: int = 400 + # sub-requests kept in flight ahead of the one being streamed + speech_chunk_lookahead: int = 2 + def chat_to_request(self, req: ChatCompletionRequest, upload_dir: Path) -> SubmitArgs: # noqa: ARG002 # Output modalities vary by model: e.g. Qwen3-Omni speech output also # emits text, whereas BAGEL chat is text-only. From 474ef233be03f396d76fc716d7a0abffe04a78a6 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:46:53 -0700 Subject: [PATCH 3/7] api: synthesize long speech inputs as ordered sentence chunks --- mstar/api_server/openai/serving_speech.py | 82 ++++++++++++++++++----- 1 file changed, 65 insertions(+), 17 deletions(-) diff --git a/mstar/api_server/openai/serving_speech.py b/mstar/api_server/openai/serving_speech.py index 335f491f6a..8833186756 100644 --- a/mstar/api_server/openai/serving_speech.py +++ b/mstar/api_server/openai/serving_speech.py @@ -3,14 +3,48 @@ Non-streaming returns the full audio as a container blob (WAV by default). Streaming returns a single open-ended WAV response (header + PCM16 frames) as the audio is produced. + +Long inputs can be synthesized as ordered sentence chunks (one engine request +per chunk, ``speech_chunking.split_sentences``): the adapter's +``speech_chunk_min_chars`` turns this on for long texts, and a client can force +or suppress it per request with ``sentence_chunking: true|false``. The next +chunks are submitted while the current one streams, so the engine batches them +and playback never waits for a prefill. """ from __future__ import annotations +from collections.abc import Callable + from fastapi.responses import Response, StreamingResponse from mstar.api_server import media_io from mstar.api_server.openai._util import rid +from mstar.api_server.openai.speech_chunking import split_sentences + + +def _plan_chunks(req, adapter, text: str) -> list[str]: + """The texts to synthesize, in playback order (``[text]`` when unchunked).""" + # Extra request fields live in pydantic's ``model_extra``; plain objects + # (tests, other frontends) may carry the attribute directly. + requested = (getattr(req, "model_extra", None) or {}).get("sentence_chunking") + if requested is None: + requested = getattr(req, "sentence_chunking", None) + min_chars = getattr(adapter, "speech_chunk_min_chars", None) + if requested is False or (requested is None and (min_chars is None or len(text) < min_chars)): + return [text] + chunks = split_sentences(text, max_chars=getattr(adapter, "speech_chunk_max_chars", 400)) + return chunks or [text] + + +def _chunk_kwargs(model_kwargs: dict, index: int) -> dict: + """Per-chunk model kwargs: identical, except a client seed advances per chunk.""" + kwargs = dict(model_kwargs) + kwargs.pop("sentence_chunking", None) + seed = kwargs.get("seed") + if isinstance(seed, int) and not isinstance(seed, bool): + kwargs["seed"] = seed + index + return kwargs async def create_speech(api, model_name, adapter, req, raw_request=None): # noqa: ARG001 @@ -18,32 +52,46 @@ async def create_speech(api, model_name, adapter, req, raw_request=None): # noq request_id = rid("speech") sample_rate = api.model.get_output_sample_rate("audio") if api.model is not None else 24000 fmt = (req.response_format or "wav").lower() + chunks = _plan_chunks(req, adapter, args.text or "") - api.submit_request( - text=args.text, - file_paths=args.file_paths, - input_modalities=args.input_modalities, - output_modalities=args.output_modalities, - model_kwargs=args.model_kwargs, - streaming=bool(req.stream), - request_id=request_id, - ) + def submit(index: int) -> str: + chunk_id = request_id if len(chunks) == 1 else f"{request_id}-{index}" + return api.submit_request( + text=chunks[index], + file_paths=args.file_paths, + input_modalities=args.input_modalities, + output_modalities=args.output_modalities, + model_kwargs=_chunk_kwargs(args.model_kwargs, index), + streaming=bool(req.stream), + request_id=chunk_id, + ) + lookahead = max(1, int(getattr(adapter, "speech_chunk_lookahead", 2))) if req.stream: return StreamingResponse( - _stream_wav(api, request_id, sample_rate), + _stream_wav(api, submit, len(chunks), lookahead, sample_rate), media_type="audio/wav", headers={"Cache-Control": "no-cache"}, ) - chunks = await api.collect_results(request_id, raw_request) - pcm = b"".join(c.data for c in chunks if c.modality == "audio") - audio_bytes, mime = media_io.pcm16_to_container(pcm, sample_rate, fmt) + pcm_parts: list[bytes] = [] + pending: list[str] = [submit(i) for i in range(min(lookahead, len(chunks)))] + for index in range(len(chunks)): + if len(pending) < len(chunks): + pending.append(submit(len(pending))) + results = await api.collect_results(pending[index], raw_request) + pcm_parts.append(b"".join(c.data for c in results if c.modality == "audio")) + audio_bytes, mime = media_io.pcm16_to_container(b"".join(pcm_parts), sample_rate, fmt) return Response(content=audio_bytes, media_type=mime) -async def _stream_wav(api, request_id, sample_rate): +async def _stream_wav(api, submit: Callable[[int], str], num_chunks: int, lookahead: int, sample_rate: int): yield media_io.wav_stream_header(sample_rate) - async for c in api.iter_result_chunks(request_id): - if c.modality == "audio" and c.data: - yield c.data + pending: list[str] = [submit(i) for i in range(min(lookahead, num_chunks))] + for index in range(num_chunks): + if len(pending) < num_chunks: + # Keep the next chunk generating while this one plays. + pending.append(submit(len(pending))) + async for c in api.iter_result_chunks(pending[index]): + if c.modality == "audio" and c.data: + yield c.data From d99c95a77fe014ed53f06deab0ade82c0940d429 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 02:46:53 -0700 Subject: [PATCH 4/7] test: sentence chunking on /v1/audio/speech --- test/modular/test_speech_chunking.py | 209 +++++++++++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 test/modular/test_speech_chunking.py diff --git a/test/modular/test_speech_chunking.py b/test/modular/test_speech_chunking.py new file mode 100644 index 0000000000..37a80e329a --- /dev/null +++ b/test/modular/test_speech_chunking.py @@ -0,0 +1,209 @@ +"""Sentence chunking on ``/v1/audio/speech``: the splitter and the ordered sub-requests. + +The router is mounted on a FastAPI app with a stubbed APIServer (as in +``test_openai_router.py``); every sub-request the handler submits is recorded +so ordering, ids, seeds and playback concatenation can be checked without an +engine. +""" + +import sys +import tempfile +import types +from pathlib import Path + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("pydantic") +pytest.importorskip("httpx") +np = pytest.importorskip("numpy") + +from fastapi import FastAPI # noqa: E402 +from fastapi.testclient import TestClient # noqa: E402 + +from mstar.api_server.openai import adapters # noqa: E402 +from mstar.api_server.openai.speech_chunking import split_sentences # noqa: E402 + +LONG_TEXT = ( + "The train to the coast leaves at seven tomorrow morning, so please pack your bag tonight. " + "She opened the window and let the cool evening air drift into the kitchen. " + "Our meeting has been moved to Thursday afternoon because the room is being repainted! " + "A gentle rain fell over the harbor while the fishing boats returned one by one. " + "Remember to water the tomatoes twice a week during the hottest part of the summer? " + "The museum's new exhibit traces the history of printing from wooden blocks to modern presses." +) + + +# --------------------------------------------------------------------------- +# splitter +# --------------------------------------------------------------------------- + + +def test_short_text_is_one_chunk(): + assert split_sentences("Hello there. How are you?", max_chars=400) == ["Hello there. How are you?"] + assert split_sentences(" ", max_chars=400) == [] + + +def test_sentences_are_grouped_up_to_max_chars_and_never_cut(): + chunks = split_sentences(LONG_TEXT, max_chars=200) + assert len(chunks) == 3 + assert all(len(c) <= 200 for c in chunks) + assert " ".join(chunks) == " ".join(LONG_TEXT.split()) + # every chunk ends where a sentence ends + assert all(c[-1] in ".!?" for c in chunks) + + +def test_cjk_terminators_and_paragraph_breaks_split(): + text = "今天天气很好。我们去公园散步吧！你觉得怎么样？\n\nSecond paragraph here." + chunks = split_sentences(text, max_chars=8, min_chars=1) + assert chunks[:3] == ["今天天气很好。", "我们去公园散步吧！", "你觉得怎么样？"] + assert " ".join(chunks[3:]) == "Second paragraph here." + + paragraphs = split_sentences("First paragraph no period\n\nsecond paragraph no period", max_chars=30, min_chars=1) + assert paragraphs == ["First paragraph no period", "second paragraph no period"] + + +def test_quotes_after_terminators_stay_with_their_sentence(): + text = 'He said "Wait!" Then she left. "Really?" she asked. Yes.' + chunks = split_sentences(text, max_chars=20, min_chars=1) + assert chunks[0] == 'He said "Wait!"' + assert chunks[1] == "Then she left." + assert chunks[2] == '"Really?" she asked.' + + +def test_overlong_sentence_is_wrapped_at_clauses_then_spaces(): + sentence = "alpha beta gamma, delta epsilon zeta, eta theta iota, kappa lambda mu nu xi omicron" + chunks = split_sentences(sentence, max_chars=30, min_chars=1) + assert all(len(c) <= 30 for c in chunks) + assert " ".join(chunks).replace(" ,", ",") == sentence + assert chunks[0] == "alpha beta gamma," + + +def test_tiny_trailing_fragment_merges_into_previous_chunk(): + text = ("A sentence that is fairly long and goes on for a while to fill the chunk nicely. " + "Yes.") + chunks = split_sentences(text, max_chars=90, min_chars=24) + assert chunks == [" ".join(text.split())] + + +def test_max_chars_must_be_positive(): + with pytest.raises(ValueError): + split_sentences("a. b.", max_chars=0) + + +# --------------------------------------------------------------------------- +# handler +# --------------------------------------------------------------------------- + + +class _Chunk: + def __init__(self, modality, data, metadata=None): + self.modality = modality + self.data = data + self.metadata = metadata or {} + + +class _StubModel: + def get_output_sample_rate(self, modality="audio"): + return 24000 + + +def _pcm(*vals): + return np.array(vals, dtype=" bytes: + return content[44:] + + +def test_long_input_is_synthesized_as_ordered_sentence_chunks(client_and_stub): + client, stub = client_and_stub + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT, "voice": "tara", "seed": 10}) + assert r.status_code == 200 and r.headers["content-type"] == "audio/wav" + texts = [s["text"] for s in stub.submits] + assert texts == split_sentences(LONG_TEXT, max_chars=200) + # One id per chunk, derived from the request id; seeds advance per chunk; + # the model kwargs are otherwise identical and carry no chunking flag. + ids = [s["request_id"] for s in stub.submits] + assert [i.rsplit("-", 1)[1] for i in ids] == ["0", "1", "2"] and len({i.rsplit("-", 1)[0] for i in ids}) == 1 + assert [s["model_kwargs"]["seed"] for s in stub.submits] == [10, 11, 12] + for submit in stub.submits: + assert submit["model_kwargs"]["voice"] == "tara" + assert "sentence_chunking" not in submit["model_kwargs"] + # Playback order == submission order. + assert _wav_pcm(r.content) == _pcm(0, 0) + _pcm(1, 1) + _pcm(2, 2) + + +def test_streaming_chunks_are_concatenated_in_order(client_and_stub): + client, stub = client_and_stub + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT, "stream": True}) + assert r.status_code == 200 and r.content[:4] == b"RIFF" + assert len(stub.submits) == 3 and all(s["streaming"] is True for s in stub.submits) + assert _wav_pcm(r.content) == _pcm(0, 0) + _pcm(1, 1) + _pcm(2, 2) + + +def test_short_input_and_opt_out_keep_a_single_request(client_and_stub): + client, stub = client_and_stub + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": "Hi there. All good?"}) + assert r.status_code == 200 and len(stub.submits) == 1 + assert stub.submits[0]["request_id"].startswith("speech-") and "-" not in stub.submits[0]["request_id"][7:] + + stub.submits.clear() + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT, "sentence_chunking": False}) + assert r.status_code == 200 and len(stub.submits) == 1 + assert stub.submits[0]["text"] == LONG_TEXT and "sentence_chunking" not in stub.submits[0]["model_kwargs"] + + +def test_client_can_force_chunking_below_the_threshold(client_and_stub, monkeypatch): + client, stub = client_and_stub + monkeypatch.setattr(adapters.OrpheusAdapter, "speech_chunk_min_chars", None) + monkeypatch.setattr(adapters.OrpheusAdapter, "speech_chunk_max_chars", 60) + text = LONG_TEXT[:150] + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": text, "sentence_chunking": True}) + assert r.status_code == 200 and len(stub.submits) >= 2 + assert " ".join(s["text"] for s in stub.submits) == text + assert all(len(s["text"]) <= 60 for s in stub.submits) From 45eaef1c157a27e4fbb06272ff1a66de41f1bc02 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:38:54 -0700 Subject: [PATCH 5/7] api: streaming speech surfaces an up-front engine error as its HTTP status --- mstar/api_server/openai/serving_speech.py | 38 ++++++++++++++++++++--- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/mstar/api_server/openai/serving_speech.py b/mstar/api_server/openai/serving_speech.py index 8833186756..c56471130d 100644 --- a/mstar/api_server/openai/serving_speech.py +++ b/mstar/api_server/openai/serving_speech.py @@ -10,12 +10,17 @@ or suppress it per request with ``sentence_chunking: true|false``. The next chunks are submitted while the current one streams, so the engine batches them and playback never waits for a prefill. + +A streaming request only commits to HTTP 200 once its first result chunk has +arrived and is not an error; an error chunk before that becomes the HTTP error +it carries (the non-streaming path gets the same from ``collect_results``). """ from __future__ import annotations from collections.abc import Callable +from fastapi import HTTPException from fastapi.responses import Response, StreamingResponse from mstar.api_server import media_io @@ -68,8 +73,15 @@ def submit(index: int) -> str: lookahead = max(1, int(getattr(adapter, "speech_chunk_lookahead", 2))) if req.stream: + pending = [submit(i) for i in range(min(lookahead, len(chunks)))] + # Look at the first result before committing to a 200: a request the + # engine rejects (bad voice, dead worker, ...) must surface as an HTTP + # error, not as an empty WAV. + first_iter = api.iter_result_chunks(pending[0]) + first = await anext(first_iter, None) + _raise_if_error(first) return StreamingResponse( - _stream_wav(api, submit, len(chunks), lookahead, sample_rate), + _stream_wav(api, submit, len(chunks), pending, first_iter, first, sample_rate), media_type="audio/wav", headers={"Cache-Control": "no-cache"}, ) @@ -85,13 +97,31 @@ def submit(index: int) -> str: return Response(content=audio_bytes, media_type=mime) -async def _stream_wav(api, submit: Callable[[int], str], num_chunks: int, lookahead: int, sample_rate: int): +def _raise_if_error(chunk) -> None: + """A data-worker failure arrives as an ``error`` chunk; turn it into the HTTP error it carries.""" + if chunk is not None and chunk.modality == "error": + raise HTTPException( + status_code=int((chunk.metadata or {}).get("status", 500)), + detail=chunk.data.decode("utf-8", "replace") if isinstance(chunk.data, bytes) else str(chunk.data), + ) + + +async def _stream_wav(api, submit: Callable[[int], str], num_chunks: int, pending: list[str], + first_iter, first, sample_rate: int): yield media_io.wav_stream_header(sample_rate) - pending: list[str] = [submit(i) for i in range(min(lookahead, num_chunks))] for index in range(num_chunks): if len(pending) < num_chunks: # Keep the next chunk generating while this one plays. pending.append(submit(len(pending))) - async for c in api.iter_result_chunks(pending[index]): + if index == 0: + iterator, head = first_iter, first + else: + iterator, head = api.iter_result_chunks(pending[index]), None + if head is not None and head.modality == "audio" and head.data: + yield head.data + async for c in iterator: + # Mid-stream the status is already sent; closing the stream is the + # only honest signal left, so raise rather than end quietly. + _raise_if_error(c) if c.modality == "audio" and c.data: yield c.data From 20d5de17f36ab7e6248b8b525044cf96188ecd59 Mon Sep 17 00:00:00 2001 From: merceod Date: Fri, 18 Sep 2026 03:38:54 -0700 Subject: [PATCH 6/7] test: streaming speech returns the engine error status --- test/modular/test_speech_chunking.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/test/modular/test_speech_chunking.py b/test/modular/test_speech_chunking.py index 37a80e329a..cc30b17609 100644 --- a/test/modular/test_speech_chunking.py +++ b/test/modular/test_speech_chunking.py @@ -207,3 +207,30 @@ def test_client_can_force_chunking_below_the_threshold(client_and_stub, monkeypa assert r.status_code == 200 and len(stub.submits) >= 2 assert " ".join(s["text"] for s in stub.submits) == text assert all(len(s["text"]) <= 60 for s in stub.submits) + + +def test_streaming_request_that_fails_up_front_returns_the_error_status(client_and_stub): + client, stub = client_and_stub + + def failing_submit(**kw): + stub.submits.append(kw) + stub._chunks[kw["request_id"]] = [ + _Chunk("error", b"Unsupported Qwen3-TTS speaker 'nobody'", {"status": 400}), + ] + return kw["request_id"] + + stub.submit_request = failing_submit + payload = {"model": "orpheus", "input": "hi there", "voice": "nobody", "stream": True} + r = client.post("/v1/audio/speech", json=payload) + assert r.status_code == 400 + assert "nobody" in r.json()["error"]["message"] + # and the non-streaming path keeps returning the error too + stub._chunks.clear() + + async def collect(request_id, raw_request=None): + from fastapi import HTTPException + raise HTTPException(status_code=400, detail="Unsupported Qwen3-TTS speaker 'nobody'") + + stub.collect_results = collect + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": "hi there", "voice": "nobody"}) + assert r.status_code == 400 From 0260eb5373f0a1567e3364b13a9ca14582110ea9 Mon Sep 17 00:00:00 2001 From: NSagan271 Date: Thu, 1 Oct 2026 17:03:05 +0000 Subject: [PATCH 7/7] respond to PR comments --- mstar/api_server/entrypoint.py | 11 ++ mstar/api_server/openai/adapters.py | 2 + mstar/api_server/openai/serving_speech.py | 120 ++++++++++++------ test/modular/test_speech_chunking.py | 145 ++++++++++++++++++++++ 4 files changed, 243 insertions(+), 35 deletions(-) diff --git a/mstar/api_server/entrypoint.py b/mstar/api_server/entrypoint.py index 4f0efcb572..2c01c11f10 100644 --- a/mstar/api_server/entrypoint.py +++ b/mstar/api_server/entrypoint.py @@ -766,6 +766,17 @@ def abort_request(self, request_id: str) -> None: logger.info("Client cancelled request %s; releasing resources", request_id) self.preprocess_worker.abort_request(request_id) + def release_request(self, request_id: str) -> None: + """Drop a request whose results will never be read: abort it if it is + still running, else just forget it (an ended request needs no abort).""" + with self.request_lock: + req = self.pending_requests.get(request_id) + finished = req is not None and req.event.is_set() + if finished: + self.pending_requests.pop(request_id) + if not finished: + self.abort_request(request_id) + # ---------------------------------------------------------- # Cleanup # ---------------------------------------------------------- diff --git a/mstar/api_server/openai/adapters.py b/mstar/api_server/openai/adapters.py index 61904b7496..a6d67e3a65 100644 --- a/mstar/api_server/openai/adapters.py +++ b/mstar/api_server/openai/adapters.py @@ -194,6 +194,8 @@ class OpenAIAdapter: speech_chunk_max_chars: int = 400 # sub-requests kept in flight ahead of the one being streamed speech_chunk_lookahead: int = 2 + # most sub-requests one input may split into; longer inputs get a 400 + speech_chunk_max_pieces: int = 32 def chat_to_request(self, req: ChatCompletionRequest, upload_dir: Path) -> SubmitArgs: # noqa: ARG002 # Output modalities vary by model: e.g. Qwen3-Omni speech output also diff --git a/mstar/api_server/openai/serving_speech.py b/mstar/api_server/openai/serving_speech.py index a10fdae481..b7e55e2643 100644 --- a/mstar/api_server/openai/serving_speech.py +++ b/mstar/api_server/openai/serving_speech.py @@ -11,7 +11,8 @@ ``speech_chunk_min_chars`` turns this on for long texts, and a client can force or suppress it per request with ``sentence_chunking: true|false``. The next chunks are submitted while the current one streams, so the engine batches them -and playback never waits for a prefill. +and playback never waits for a prefill. An input needing more than the +adapter's ``speech_chunk_max_pieces`` chunks is rejected. A streaming request only commits to HTTP 200 once its first result chunk has arrived and is not an error; an error chunk before that becomes the HTTP error @@ -42,6 +43,12 @@ def _plan_chunks(req, adapter, text: str) -> list[str]: if requested is False or (requested is None and (min_chars is None or len(text) < min_chars)): return [text] chunks = split_sentences(text, max_chars=getattr(adapter, "speech_chunk_max_chars", 400)) + max_pieces = getattr(adapter, "speech_chunk_max_pieces", 32) + if len(chunks) > max_pieces: + raise HTTPException( + status_code=400, + detail=f"input splits into {len(chunks)} sentence chunks; at most {max_pieces} are allowed", + ) return chunks or [text] @@ -87,15 +94,23 @@ def submit(index: int) -> str: ) lookahead = max(1, int(getattr(adapter, "speech_chunk_lookahead", 2))) + pending = [submit(i) for i in range(min(lookahead, len(chunks)))] if req.stream: raw_pcm = fmt == "pcm" - pending = [submit(i) for i in range(min(lookahead, len(chunks)))] # Look at the first result before committing to a 200: a request the # engine rejects (bad voice, dead worker, ...) must surface as an HTTP # error, not as an empty WAV. first_iter = api.iter_result_chunks(pending[0]) - first = await anext(first_iter, None) - _raise_if_error(first) + try: + first = await anext(first_iter, None) + if _is_error(first): + # Read it to the end so the finished request isn't aborted. + await _drain(first_iter) + raise _http_error(first) + except BaseException: + # The first piece released itself; release the ones submitted ahead. + _release_unread(api, pending, 1) + raise return StreamingResponse( _stream_pcm(api, submit, len(chunks), pending, first_iter, first, sample_rate, with_wav_header=not raw_pcm), @@ -104,42 +119,77 @@ def submit(index: int) -> str: ) pcm_parts: list[bytes] = [] - pending: list[str] = [submit(i) for i in range(min(lookahead, len(chunks)))] - for index in range(len(chunks)): - if len(pending) < len(chunks): - pending.append(submit(len(pending))) - results = await api.collect_results(pending[index], raw_request) - pcm_parts.append(b"".join(c.data for c in results if c.modality == "audio")) + read = 0 # pieces collect_results has released; the rest are released on early exit + try: + for index in range(len(chunks)): + if len(pending) < len(chunks): + pending.append(submit(len(pending))) + try: + results = await api.collect_results(pending[index], raw_request) + except HTTPException: + read = index + 1 + raise + read = index + 1 + if not results and raw_request is not None and await raw_request.is_disconnected(): + break + pcm_parts.append(b"".join(c.data for c in results if c.modality == "audio")) + finally: + _release_unread(api, pending, read) audio_bytes, mime = media_io.pcm16_to_container(b"".join(pcm_parts), sample_rate, fmt) return Response(content=audio_bytes, media_type=mime) -def _raise_if_error(chunk) -> None: - """A data-worker failure arrives as an ``error`` chunk; turn it into the HTTP error it carries.""" - if chunk is not None and chunk.modality == "error": - raise HTTPException( - status_code=int((chunk.metadata or {}).get("status", 500)), - detail=chunk.data.decode("utf-8", "replace") if isinstance(chunk.data, bytes) else str(chunk.data), - ) +def _is_error(chunk) -> bool: + return chunk is not None and chunk.modality == "error" + + +def _http_error(chunk) -> HTTPException: + """A data-worker failure arrives as an ``error`` chunk; the HTTP error it carries.""" + return HTTPException( + status_code=int((chunk.metadata or {}).get("status", 500)), + detail=chunk.data.decode("utf-8", "replace") if isinstance(chunk.data, bytes) else str(chunk.data), + ) + + +async def _drain(iterator) -> None: + async for _ in iterator: + pass + + +def _release_unread(api, pending: list[str], read: int) -> None: + """Release every submitted piece from ``pending[read]`` on (not yet fully read).""" + for chunk_id in pending[read:]: + api.release_request(chunk_id) async def _stream_pcm(api, submit: Callable[[int], str], num_chunks: int, pending: list[str], first_iter, first, sample_rate: int, with_wav_header: bool = True): - if with_wav_header: - yield media_io.wav_stream_header(sample_rate) - for index in range(num_chunks): - if len(pending) < num_chunks: - # Keep the next chunk generating while this one plays. - pending.append(submit(len(pending))) - if index == 0: - iterator, head = first_iter, first - else: - iterator, head = api.iter_result_chunks(pending[index]), None - if head is not None and head.modality == "audio" and head.data: - yield head.data - async for c in iterator: - # Mid-stream the status is already sent; closing the stream is the - # only honest signal left, so raise rather than end quietly. - _raise_if_error(c) - if c.modality == "audio" and c.data: - yield c.data + read = 0 # pieces read to the end; the rest are released if the stream stops early + try: + if with_wav_header: + yield media_io.wav_stream_header(sample_rate) + for index in range(num_chunks): + if len(pending) < num_chunks: + # Keep the next chunk generating while this one plays. + pending.append(submit(len(pending))) + if index == 0: + iterator, head = first_iter, first + else: + iterator, head = api.iter_result_chunks(pending[index]), None + if head is not None and head.modality == "audio" and head.data: + yield head.data + error = None + async for c in iterator: + # Read past an error to the end: leaving early would abort a + # request that has already finished. + if error is None and _is_error(c): + error = c + elif error is None and c.modality == "audio" and c.data: + yield c.data + read = index + 1 + if error is not None: + # Mid-stream the status is already sent; closing the stream is + # the only honest signal left, so raise rather than end quietly. + raise _http_error(error) + finally: + _release_unread(api, pending, read) diff --git a/test/modular/test_speech_chunking.py b/test/modular/test_speech_chunking.py index cc30b17609..828b0e35b3 100644 --- a/test/modular/test_speech_chunking.py +++ b/test/modular/test_speech_chunking.py @@ -6,6 +6,7 @@ engine. """ +import asyncio import sys import tempfile import types @@ -121,6 +122,8 @@ def __init__(self): self.upload_dir = Path(tempfile.mkdtemp()) self.submits: list[dict] = [] self._chunks: dict[str, list] = {} + self.released: list[str] = [] + self.read_to_end: set[str] = set() def submit_request(self, **kw): index = len(self.submits) @@ -134,6 +137,22 @@ async def collect_results(self, request_id, raw_request=None): async def iter_result_chunks(self, request_id): for c in self._chunks[request_id]: yield c + self.read_to_end.add(request_id) + + def release_request(self, request_id): + self.released.append(request_id) + + def fail(self, index, status=400): + """Make the ``index``-th submitted piece end with an error chunk.""" + submit = self.submit_request + + def failing_submit(**kw): + request_id = submit(**kw) + if len(self.submits) - 1 == index: + self._chunks[request_id] = [_Chunk("error", b"piece failed", {"status": status})] + return request_id + + self.submit_request = failing_submit @pytest.fixture @@ -234,3 +253,129 @@ async def collect(request_id, raw_request=None): stub.collect_results = collect r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": "hi there", "voice": "nobody"}) assert r.status_code == 400 + + +def _ids(stub): + return [s["request_id"] for s in stub.submits] + + +def test_chunk_count_is_capped(client_and_stub, monkeypatch): + client, stub = client_and_stub + monkeypatch.setattr(adapters.OrpheusAdapter, "speech_chunk_max_pieces", 2) + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT, "sentence_chunking": True}) + assert r.status_code == 400 and "at most 2" in r.json()["error"]["message"] + assert stub.submits == [] + + +def test_streaming_error_up_front_releases_only_the_pieces_ahead(client_and_stub): + client, stub = client_and_stub + stub.fail(0) + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT, "stream": True}) + assert r.status_code == 400 + # the failed piece was read to its end (so it isn't aborted); the lookahead is released + ids = _ids(stub) + assert ids[0] in stub.read_to_end and stub.released == ids[1:] + + +def test_non_streaming_error_releases_only_the_pieces_ahead(client_and_stub): + client, stub = client_and_stub + from fastapi import HTTPException + + async def collect(request_id, raw_request=None): + if request_id == _ids(stub)[1]: + raise HTTPException(status_code=400, detail="piece failed") + return stub._chunks[request_id] + + stub.collect_results = collect + r = client.post("/v1/audio/speech", json={"model": "orpheus", "input": LONG_TEXT}) + assert r.status_code == 400 + assert stub.released == _ids(stub)[2:] and len(stub.submits) == 3 + + +def _stream(stub, num_chunks): + ids = [f"speech-x-{i}" for i in range(num_chunks)] + + def submit(index): + return stub.submit_request(request_id=ids[index], text="") + + async def start(): + from mstar.api_server.openai import serving_speech + + pending = [submit(0), submit(1)] + first_iter = stub.iter_result_chunks(pending[0]) + first = await anext(first_iter) + return serving_speech._stream_pcm(stub, submit, num_chunks, pending, first_iter, first, 24000) + + return ids, start + + +def test_closing_a_chunked_stream_releases_every_unread_piece(): + stub = _RecordingAPI() + ids, start = _stream(stub, 4) + + async def run(): + gen = await start() + assert (await anext(gen))[:4] == b"RIFF" + assert await anext(gen) == _pcm(0, 0) # piece 0 playing, piece 2 submitted ahead + await gen.aclose() # the client went away + + asyncio.run(run()) + assert stub.released == ids[:3] and len(stub.submits) == 3 + + +def test_mid_stream_error_closes_the_stream_and_releases_the_pieces_ahead(): + from fastapi import HTTPException + + stub = _RecordingAPI() + stub.fail(1) + ids, start = _stream(stub, 4) + + async def run(): + gen = await start() + with pytest.raises(HTTPException): + async for _ in gen: + pass + + asyncio.run(run()) + assert ids[1] in stub.read_to_end and stub.released == ids[2:] + + +def test_non_streaming_disconnect_stops_submitting_and_releases_the_pieces_ahead(client_and_stub): + _, stub = client_and_stub + from mstar.api_server.openai.protocol import SpeechRequest + from mstar.api_server.openai.serving_speech import create_speech + + collected = [] + + async def collect(request_id, raw_request=None): + collected.append(request_id) + return [] # what collect_results returns once it sees the disconnect + + class _Gone: + async def is_disconnected(self): + return True + + stub.collect_results = collect + req = SpeechRequest(model="orpheus", input=LONG_TEXT) + asyncio.run(create_speech(stub, "orpheus", adapters.OrpheusAdapter(), req, _Gone())) + ids = _ids(stub) + assert collected == ids[:1] and len(ids) == 3 and stub.released == ids[1:] + + +def test_release_request_aborts_only_requests_still_running(): + import threading + + from mstar.api_server.entrypoint import APIServer + + aborted = [] + done, running = threading.Event(), threading.Event() + done.set() + server = types.SimpleNamespace( + request_lock=threading.Lock(), + pending_requests={"done": types.SimpleNamespace(event=done), + "running": types.SimpleNamespace(event=running)}, + abort_request=aborted.append, + ) + APIServer.release_request(server, "done") + APIServer.release_request(server, "running") + assert aborted == ["running"] and "done" not in server.pending_requests