fix: warn before using a likely too-big chunk size for sglang-omni models

3 files changed, 167 insertions, 18 deletions

diff --git a/app/backends/servers.py b/app/backends/servers.py
index d352717..fd28866 100644
--- a/app/backends/servers.py
+++ b/app/backends/servers.py
@@ -23,6 +23,7 @@ boot screen, which renders the same events. Pid/log files live under
 """
 
 import os
+import re
 import signal
 import subprocess
 import sys
@@ -62,6 +63,12 @@ _BOOT_HINTS = (
      "re-install the model via Configure Backends (checking the model "
      "installs its companion packages), or pip-install it into the "
      "backend's venv manually"),
+    # A launcher that could not take its configured port — uvicorn's bind
+    # failure dies outright, sglang-omni's silently moves to a random one
+    # (the live detection for that is _PORT_FALLBACK_RE below).
+    ("already in use",
+     "the configured port is held by another process — stop that process "
+     "or move this server to a free port, then start it again"),
 )
 
 # Progress callback: called with an event dict. KIND is one of:
@@ -70,6 +77,10 @@ _BOOT_HINTS = (
 #   "ready"     {name, url}                       server is up and answering
 #   "exited"    {name, returncode, log_tail, hint}  process died while booting
 #   "timeout"   {name, seconds, log_tail, hint}   readiness deadline elapsed
+#   "port_taken" {name, taken, moved, message, log_tail}
+#                                                 launcher moved the boot
+#                                                 to a random port (the
+#                                                 configured one was taken)
 #   "running"   {name, url}                       already up (no spawn)
 #   "cancelled" {name}                            boot aborted via cancel
 #   "error"     {message}                         could not spawn the executable
@@ -104,6 +115,9 @@ def _console_progress(event: dict) -> None:
         _print_tail(event.get("log_tail"))
         if event.get("hint"):
             print(f"[WARNING] hint: {event['hint']}")
+    elif kind == "port_taken":
+        print(f"[ERROR] {event['message']}")
+        _print_tail(event.get("log_tail"))
     elif kind == "error":
         print(f"[ERROR] {event['message']}")
 
@@ -154,6 +168,62 @@ def _print_tail(tail: List[str]) -> None:
         print("---")
 
 
+# sglang-omni's launcher, finding the requested port taken, silently binds
+# a random one instead — the boot then stays healthy but unreachable at
+# the configured URL (every client keeps polling the taken port until the
+# start timeout). Scanned from the boot log so the boot fails in seconds
+# with both port numbers named instead.
+_PORT_FALLBACK_RE = re.compile(
+    r"Port (?P<taken>\d+) is already in use.*?"
+    r"Using port (?P<moved>\d+) instead", re.DOTALL)
+
+# How much of the previous log read is rescanned with the next chunk, so a
+# fallback message split across two reads still matches (the launcher
+# prints its two lines back to back; half a kilobyte is generous).
+_LOG_CARRY_BYTES = 512
+
+
+def _read_new_log(path: Path, offset: int) -> tuple:
+    """Read the bytes appended to the server log since OFFSET (best effort).
+
+    Returns ``(new_offset, text)`` — where to resume and what was read
+    (undecodable bytes replaced). An unreadable or unchanged file yields
+    the offset unchanged and ""; a shrunken file (truncated or rotated)
+    restarts from zero so nothing new is skipped.
+    """
+    try:
+        size = path.stat().st_size
+    except OSError:
+        return offset, ""
+    if size < offset:
+        offset = 0
+    if size == offset:
+        return offset, ""
+    try:
+        with path.open("rb") as fh:
+            fh.seek(offset)
+            raw = fh.read()
+    except OSError:
+        return offset, ""
+    return offset + len(raw), raw.decode("utf-8", errors="replace")
+
+
+def _port_taken_message(spec) -> str:
+    """The refusal message for spawning while SPEC's URL already answers."""
+    return (f"another process is listening at {spec.url} but it is not a "
+            f"usable {spec.name} server — stop that process (or move this "
+            f"server to a free port) and start again")
+
+
+def _port_fallback_message(spec, taken: str, moved: str) -> str:
+    """The boot-failure message when the launcher moved ports on us."""
+    return (f"the {spec.name} launcher moved the server from port "
+            f"{taken} (already in use by another process) to port {moved}; "
+            f"clients poll {taken}, so this boot cannot become ready — "
+            f"stop whatever holds port {taken} (or move this server to a "
+            f"free port) and start again")
+
+
 def _pid_alive(pid: int) -> bool:
     """True when a process with PID is still running (POSIX signal-0 probe)."""
     if sys.platform == "win32":
@@ -286,15 +356,19 @@ def _kill_pid(pid: int) -> bool:
     return True
 
 
-def _server_ready(spec) -> bool:
+def _server_ready(spec, listening: Optional[bool] = None) -> bool:
     """True when the server described by SPEC is usable, not just listening.
 
     Without an IDENTITY this is the plain TCP-connect check. With one, the
     server must also answer HTTP as that backend (``probe.identify_server``);
     for the faster backend (whose model loads after the port opens) the
-    ``/health`` model_loaded flag must additionally be true.
+    ``/health`` model_loaded flag must additionally be true. LISTENING, when
+    given, is a caller's fresh ``common.server_running`` result — reused so
+    one start pass probes the port only once.
     """
-    if not common.server_running(spec.url):
+    if listening is None:
+        listening = common.server_running(spec.url)
+    if not listening:
         return False
     identity = getattr(spec, "identity", None)
     if identity is None:
@@ -320,6 +394,13 @@ def start(spec, progress: ProgressCallback = None,
     exit reports the log tail and returns False. A no-op (True) when the
     server is already running.
 
+    Two port-conflict guards fail fast instead of letting a doomed boot
+    run out the clock: a foreign process already listening on the spec's
+    URL (but not answering as the backend) refuses the spawn outright, and
+    a launcher that logs a port fallback mid-boot ("Port N is already in
+    use ... Using port M instead" — sglang-omni's) aborts the boot with a
+    "port_taken" event and kills the misdirected server.
+
     PROGRESS, when given, receives each boot event (see ProgressCallback);
     the default ``_console_progress`` prints them, preserving the old console
     output. CANCEL (a threading.Event) aborts the boot: the spawned process
@@ -333,7 +414,11 @@ def start(spec, progress: ProgressCallback = None,
                 "message": f"server executable not found: {exe}. Run 'Set "
                            "up a backend' first."})
         return False
-    if _server_ready(spec):
+    # One TCP probe feeds both checks: an already-usable server takes the
+    # "running" path, and a listener that is NOT usable (identity probe
+    # failed) is exactly the foreign-holder conflict refused below.
+    listening = common.server_running(spec.url)
+    if _server_ready(spec, listening):
         report({"kind": "running", "name": spec.name, "url": spec.url})
         return True
 
@@ -350,6 +435,15 @@ def start(spec, progress: ProgressCallback = None,
                            f"{pid_for(spec.name)}) is already starting or "
                            "running; stop it first"})
         return False
+    # Refuse to spawn onto a port a foreign process already holds: TCP-up
+    # but probe-down means the listener is not a usable instance of this
+    # server. A fresh spawn would then either die on the bind or (launchers
+    # that fall back silently, like sglang-omni) move to a random port and
+    # leave every client polling the taken one — the boot watchdog below
+    # catches that late, so name the conflict here.
+    if listening:
+        report({"kind": "error", "message": _port_taken_message(spec)})
+        return False
     pid_file = _pid_path(spec.name)
     if pid_file.exists():
         try:
@@ -419,6 +513,13 @@ def start(spec, progress: ProgressCallback = None,
     started = time.time()
     next_heartbeat = started + 15
     deadline = started + start_timeout
+    # Boot-log watchdog state: scan only what this boot appends (the log
+    # file accumulates across boots, so a previous boot's fallback lines
+    # must not re-fire here), carrying a tail between reads so a fallback
+    # message split across two polls still matches.
+    log_path = _log_path(spec.name)
+    log_offset = log_path.stat().st_size if log_path.exists() else 0
+    log_carry = ""
     while time.time() < deadline:
         if cancel is not None and cancel.is_set():
             # User cancelled while booting: kill what we spawned (the
@@ -441,6 +542,31 @@ def start(spec, progress: ProgressCallback = None,
             except OSError:
                 pass
             return False
+        # Watchdog: a launcher that silently moved to another port (the
+        # configured one was taken) keeps booting healthily where no
+        # client will ever call it — abort now instead of polling the
+        # taken port until the timeout.
+        log_offset, new_text = _read_new_log(log_path, log_offset)
+        if new_text:
+            scan = log_carry + new_text
+            match = _PORT_FALLBACK_RE.search(scan)
+            if match is not None:
+                # Kill the misdirected server: it would serve on a port
+                # no client will call while holding GPU memory.
+                _kill_pid(proc.pid)
+                try:
+                    pid_file.unlink()
+                except OSError:
+                    pass
+                report({"kind": "port_taken", "name": spec.name,
+                        "taken": int(match.group("taken")),
+                        "moved": int(match.group("moved")),
+                        "message": _port_fallback_message(
+                            spec, match.group("taken"),
+                            match.group("moved")),
+                        "log_tail": _read_log_tail(spec.name)})
+                return False
+            log_carry = scan[-_LOG_CARRY_BYTES:]
         if _server_ready(spec):
             report({"kind": "ready", "name": spec.name, "url": spec.url})
             return True
diff --git a/app/backends/sglomni/catalog.py b/app/backends/sglomni/catalog.py
index 1e4ca2d..160eb8b 100644
--- a/app/backends/sglomni/catalog.py
+++ b/app/backends/sglomni/catalog.py
@@ -69,11 +69,20 @@ class ModelEntry:
     # 86.13 fps, ~12 s) and Higgs's to 2048 frames (75 fps, ~27 s, and
     # per-request values are clamped to the engine cap, so its vendored
     # config raises the cap too) — and silently truncate longer text.
-    # 12288 frames covers a full 250-word CHUNK_SIZE sub-chunk on both and
-    # stays under the KV-pool admission check on every GPU that can host
-    # the model (the scheduler rejects prompt + max_new_tokens above it,
-    # and these requests are not auto-clamped).
+    # 12288 frames covers a full 250-word CHUNK_SIZE sub-chunk on ZONOS2
+    # (whose KV pools are >= 28083 tokens on every GPU that can host it)
+    # but NOT on Higgs: its thinker engine pins context_length at 4096,
+    # and the scheduler rejects any request whose prompt tokens plus
+    # max_new_tokens exceed that window (kv_capacity=4095, on every GPU —
+    # the pool side is never the binding constraint there).
     max_new_tokens: Optional[int] = None
+    # The largest sub-request (words) the model can narrate within its
+    # admission window, for models whose engine cannot cover a full
+    # CHUNK_SIZE sub-chunk (None = no cap; config.CHUNK_SIZE stands).
+    # Higgs: 3000 frames ≈ 40 s at 75 fps ≈ 80 words of narration, and
+    # the run pre-flight offers to clamp CHUNK_SIZE for the run (the
+    # client's adaptive retry still rescues requests the window rejects).
+    chunk_words: Optional[int] = None
 
 
 # The Qwen3-TTS CustomVoice speaker table — the same built-in speakers the
@@ -177,9 +186,18 @@ ENTRIES: Tuple[ModelEntry, ...] = (
         requires_reference=False,
         # The engine's 2048-frame default is ~27 s of speech at the codec's
         # 75 fps; requests are clamped to the engine cap server-side, so the
-        # yaml raises the cap and every request carries 12288 frames
-        # (~164 s) — enough for a full CHUNK_SIZE sub-chunk.
-        max_new_tokens=12288,
+        # yaml raises the cap and every request carries 3000 frames (~40 s)
+        # — the most the admission window allows: upstream pins the thinker
+        # engine's context_length at 4096, and the scheduler rejects any
+        # request whose prompt (including the reference-audio tokens) plus
+        # max_new_tokens exceeds it. 3000 frames leaves ~1095 tokens of
+        # prompt headroom (an 80-word chunk with a 20.5 s reference
+        # measured 684). Sub-requests cap at 80 words so the text fits the
+        # window too — the pre-flight offers to clamp CHUNK_SIZE for the
+        # run, and the client refits rejected requests to whatever the
+        # server reports as its capacity.
+        max_new_tokens=3000,
+        chunk_words=80,
         notes="zero-shot narration, cloning from a reference clip",
     ),
     ModelEntry(
diff --git a/app/backends/sglomni/configs/higgs_audio_v3_tts.yaml b/app/backends/sglomni/configs/higgs_audio_v3_tts.yaml
index 74736fa..8a42771 100644
--- a/app/backends/sglomni/configs/higgs_audio_v3_tts.yaml
+++ b/app/backends/sglomni/configs/higgs_audio_v3_tts.yaml
@@ -9,22 +9,27 @@
 # "CUDA out of memory. Tried to allocate 14.00 MiB".
 #
 # 0.80 trims the engine's static pool by ~1.2 GB per 24 GB of VRAM while
-# leaving a KV cache pool (~10 GB on a 24 GB card) far larger than any
-# narration request needs. Cards with heavy other-GPU-process usage can go
-# lower (e.g. 0.75).
+# leaving a KV cache pool far larger than any narration request needs.
+# Cards with heavy other-GPU-process usage can go lower (e.g. 0.75).
 #
 # The tts_engine factory also caps every request at max_new_tokens=2048
 # audio frames, and per-request values are clamped to that cap server-side
 # (make_higgs_scheduler_adapters) — the Higgs codec runs 75 frames per
 # second (24 kHz / 320 downsample), so the default is ~27 s of speech, which
 # silently truncates this tool's full 250-word sub-chunks (~100 s). Raising
-# the factory cap is the only way past it; the catalog also sends
-# max_new_tokens=12288 per request (the same value ZONOS2 uses) so a request
-# may use the room: 12288 frames ≈ 164 s.
+# the factory cap is the only way past it, but the ceiling is hard: upstream
+# pins the thinker engine's context_length at 4096 (HiggsTtsEngineBuilder —
+# not overridable), and the scheduler rejects any request whose prompt
+# tokens (including the reference-audio tokens) plus max_new_tokens exceed
+# that window ("Request requires more tokens than the thinker KV cache can
+# hold", kv_capacity=4095, on every GPU). The cap therefore lands at 3000
+# frames ≈ 40 s — the most the window allows with prompt headroom (an
+# 80-word chunk with a 20.5 s reference measured 684 prompt tokens) — and
+# the catalog caps sub-requests at 80 words to match (chunk_words).
 config_cls: HiggsTtsPipelineConfig
 model_path: bosonai/higgs-audio-v3-tts-4b
 stages:
   tts_engine:
     gpu_memory_fraction: 0.80
     factory:
-      max_new_tokens: 12288
+      max_new_tokens: 3000