Keep interactive streams live during inference

Exact source: ttracx/thoxroute@f56b047f0c269ed6559d728ad42e06cd8bb54201

Dockerfile

@@ -25,6 +25,8 @@ ENV HF_HOME=/home/mambauser/.cache/huggingface \
     THOX_FAST_MODEL_FILE=qwen2.5-0.5b-instruct-q4_k_m.gguf \
     THOX_FAST_POOL_SIZE=1 \
     THOX_FAST_QUEUE_TIMEOUT_S=6 \
+    THOX_FAST_MAX_OUTPUT_TOKENS=16 \
+    THOX_STREAM_HEARTBEAT_S=2 \
     THOX_MAX_OUTPUT_TOKENS=128
 
 # Bake the immutable public interactive model into the image. A RUNNING Space

README.md

@@ -19,7 +19,8 @@ This Space exposes one OpenAI-compatible endpoint:
 `GET /healthz` is ready only after the interactive context is loaded and has
 completed a one-token startup warmup. Streaming sends a role chunk before model
 evaluation so upstream time-to-first-byte watchdogs do not abandon healthy CPU
-work. First-party 15-second requests are capped at 16 output tokens by
-ThoxRoute.
+work; two-second SSE comments keep the stream live during synchronous prompt
+evaluation. The interactive provider itself caps every request at 16 output
+tokens, so malformed or older callers cannot leave minutes of abandoned work.
 The interactive model is pinned to immutable model revision
 `9217f5db79a29953eb74d5343926648285ec7e67` and baked into the image.

app.py

@@ -50,6 +50,12 @@ FAST_QUEUE_TIMEOUT_S = max(
 MAX_OUTPUT_TOKENS = max(
     1, min(int(os.environ.get("THOX_MAX_OUTPUT_TOKENS", "128")), 512)
 )
+FAST_MAX_OUTPUT_TOKENS = max(
+    1, min(int(os.environ.get("THOX_FAST_MAX_OUTPUT_TOKENS", "16")), 32)
+)
+STREAM_HEARTBEAT_S = max(
+    0.25, min(float(os.environ.get("THOX_STREAM_HEARTBEAT_S", "2")), 5.0)
+)
 MAX_MESSAGES = 64
 MAX_MESSAGE_CHARS = 65_536
 MAX_REQUEST_CHARS = 131_072
@@ -219,7 +225,8 @@ def _messages(req: ChatRequest) -> list[dict[str, str]]:
 
 
 def _max_tokens(req: ChatRequest) -> int:
-    return min(req.max_tokens, MAX_OUTPUT_TOKENS)
+    ceiling = FAST_MAX_OUTPUT_TOKENS if req.model == FAST_MODEL_ID else MAX_OUTPUT_TOKENS
+    return min(req.max_tokens, ceiling)
 
 
 def _runtime() -> Runtime:
@@ -240,6 +247,8 @@ def healthz() -> dict[str, Any]:
         "interactive_pool_size": FAST_POOL_SIZE,
         "interactive_available": active.interactive_available,
         "specialist_model": CODER_MODEL_ID,
+        "interactive_max_output_tokens": FAST_MAX_OUTPUT_TOKENS,
+        "stream_heartbeat_s": STREAM_HEARTBEAT_S,
         "max_output_tokens": MAX_OUTPUT_TOKENS,
         "n_ctx": N_CTX,
         "threads": _usable_cpus(),
@@ -253,31 +262,61 @@ def _completion_id() -> str:
 def _stream_completion(req: ChatRequest, lease: Lease) -> Iterator[bytes]:
     completion_id = _completion_id()
     created = int(time.time())
-    try:
-        # Send a valid role chunk before llama prompt evaluation. This proves
-        # the stream is alive inside ThoxRoute's short time-to-first-byte budget
-        # and prevents a healthy CPU inference from being abandoned at 12 s.
-        initial = {
-            "id": completion_id,
-            "object": "chat.completion.chunk",
-            "created": created,
-            "model": req.model,
-            "choices": [
-                {
-                    "index": 0,
-                    "delta": {"role": "assistant", "content": ""},
-                    "finish_reason": None,
-                }
-            ],
-        }
-        yield f"data: {json.dumps(initial, separators=(',', ':'))}\n\n".encode()
-        chunks = lease.model.create_chat_completion(
-            messages=_messages(req),
-            max_tokens=_max_tokens(req),
-            temperature=req.temperature,
-            stream=True,
-        )
-        for chunk in chunks:
+    items: queue.Queue[dict[str, Any] | object] = queue.Queue()
+    complete = object()
+    failed = object()
+
+    def generate() -> None:
+        try:
+            chunks = lease.model.create_chat_completion(
+                messages=_messages(req),
+                max_tokens=_max_tokens(req),
+                temperature=req.temperature,
+                stream=True,
+            )
+            for chunk in chunks:
+                items.put(chunk)
+            items.put(complete)
+        except Exception:
+            # Provider-controlled exception text must never cross the API boundary.
+            items.put(failed)
+        finally:
+            # A disconnected client closes the response generator, but llama.cpp
+            # can still be using the context. The worker therefore owns release.
+            lease.release()
+
+    threading.Thread(target=generate, daemon=True, name="thox-fast-generation").start()
+
+    # Prove liveness before prompt evaluation, then emit SSE comments while the
+    # synchronous llama context is busy. Upstream read timers reset on every
+    # heartbeat and comments are ignored by OpenAI-compatible parsers.
+    initial = {
+        "id": completion_id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": req.model,
+        "choices": [
+            {
+                "index": 0,
+                "delta": {"role": "assistant", "content": ""},
+                "finish_reason": None,
+            }
+        ],
+    }
+    yield f"data: {json.dumps(initial, separators=(',', ':'))}\n\n".encode()
+    while True:
+        try:
+            item = items.get(timeout=STREAM_HEARTBEAT_S)
+        except queue.Empty:
+            yield b": thox-fast-heartbeat\n\n"
+            continue
+        if item is complete:
+            yield b"data: [DONE]\n\n"
+            return
+        if item is failed:
+            raise RuntimeError("interactive generation failed")
+        if isinstance(item, dict):
+            chunk = item
             choice = chunk.get("choices", [{}])[0]
             payload = {
                 "id": completion_id,
@@ -293,9 +332,6 @@ def _stream_completion(req: ChatRequest, lease: Lease) -> Iterator[bytes]:
                 ],
             }
             yield f"data: {json.dumps(payload, separators=(',', ':'))}\n\n".encode()
-        yield b"data: [DONE]\n\n"
-    finally:
-        lease.release()
 
 
 @api.post("/v1/chat/completions")