Upload app.py

app.py

@@ -10,6 +10,7 @@ import os
 import time
 import uuid
 import json
+import queue
 import asyncio
 import logging
 import threading
@@ -49,6 +50,12 @@ _llm_lock   = threading.Lock()
 _llm_ready  = threading.Event()   # set once the model is loaded
 _llm_error: Optional[str] = None  # set if loading failed
 
+# llama.cpp contexts are NOT safe to call concurrently -- there's a single
+# KV cache / sampling state shared by every call into the same Llama
+# instance. This lock serializes all generation calls (streaming and
+# non-streaming) so two requests can never run inference at the same time.
+_inference_lock = threading.Lock()
+
 
 def _download_model() -> None:
     """Download the GGUF file from MODEL_URL if MODEL_PATH doesn't exist."""
@@ -196,27 +203,46 @@ def _make_chunk(delta_content: str, finish_reason: Optional[str], request_id: st
     return f"data: {json.dumps(chunk)}\n\n"
 
 
-async def _stream_response(request: ChatCompletionRequest, request_id: str) -> AsyncIterator[str]:
-    llm = _get_llm()
-    messages = [{"role": m.role, "content": m.content} for m in request.messages]
+async def _stream_response(llm: Llama, messages: list, request: ChatCompletionRequest, request_id: str) -> AsyncIterator[str]:
     loop = asyncio.get_event_loop()
-
-    def _run():
-        return llm.create_chat_completion(
-            messages=messages,
-            max_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            stop=request.stop or [],
-            stream=True,
-        )
-
-    gen = await loop.run_in_executor(None, _run)
+    q: "queue.Queue" = queue.Queue()
+    _SENTINEL = object()
+
+    def _produce():
+        # Holds the lock for the *entire* generation, not just creation --
+        # this is the only place token generation actually happens, and it
+        # must never overlap with another request's call into the same
+        # Llama instance.
+        with _inference_lock:
+            try:
+                gen = llm.create_chat_completion(
+                    messages=messages,
+                    max_tokens=request.max_tokens,
+                    temperature=request.temperature,
+                    top_p=request.top_p,
+                    stop=request.stop or [],
+                    stream=True,
+                )
+                for chunk in gen:
+                    q.put(chunk)
+            except Exception as exc:  # surfaced to the consumer below
+                q.put(exc)
+            finally:
+                q.put(_SENTINEL)
+
+    # Fire-and-forget: runs on a worker thread, the consumer below just
+    # drains the queue without ever blocking the asyncio event loop.
+    loop.run_in_executor(None, _produce)
 
     yield _make_chunk("", None, request_id)  # opening delta
 
-    for chunk in gen:
-        choice  = chunk["choices"][0]
+    while True:
+        item = await loop.run_in_executor(None, q.get)
+        if item is _SENTINEL:
+            break
+        if isinstance(item, Exception):
+            raise item
+        choice  = item["choices"][0]
         delta   = choice.get("delta", {})
         content = delta.get("content", "")
         finish  = choice.get("finish_reason")
@@ -265,14 +291,15 @@ async def list_models():
 
 @app.post("/v1/chat/completions")
 async def chat_completions(request: ChatCompletionRequest):
-    llm = _get_llm()  # raises 503 if not ready
+    llm = _get_llm()  # raises 503/500 *before* we commit to a response,
+                       # including for the streaming branch below
 
     messages = [{"role": m.role, "content": m.content} for m in request.messages]
 
     if request.stream:
         request_id = f"chatcmpl-{uuid.uuid4().hex}"
         return StreamingResponse(
-            _stream_response(request, request_id),
+            _stream_response(llm, messages, request, request_id),
             media_type="text/event-stream",
             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
         )
@@ -281,14 +308,15 @@ async def chat_completions(request: ChatCompletionRequest):
     loop = asyncio.get_event_loop()
 
     def _run():
-        return llm.create_chat_completion(
-            messages=messages,
-            max_tokens=request.max_tokens,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            stop=request.stop or [],
-            stream=False,
-        )
+        with _inference_lock:  # never overlap with another generation call
+            return llm.create_chat_completion(
+                messages=messages,
+                max_tokens=request.max_tokens,
+                temperature=request.temperature,
+                top_p=request.top_p,
+                stop=request.stop or [],
+                stream=False,
+            )
 
     result = await loop.run_in_executor(None, _run)
     choice = result["choices"][0]