fix: serialize generation with _gen_lock to prevent concurrent queue readers
Two overlapping /chat/completions requests could both read from the shared resp_queue, consuming and dropping each other's token events. Replace the request_id filtering (which silently dropped non-matching messages) with a threading.Lock that serializes generation — correct for single-GPU inference. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
9470957bb9
commit
609e3168a1
1 changed files with 44 additions and 9 deletions
|
|
@ -44,6 +44,7 @@ class InferenceOrchestrator:
|
|||
self._resp_queue: Any = None
|
||||
self._cancel_event: Any = None # mp.Event — set to cancel generation instantly
|
||||
self._lock = threading.Lock()
|
||||
self._gen_lock = threading.Lock() # Serializes generation — one request at a time
|
||||
|
||||
# Local state mirrors (updated from subprocess responses)
|
||||
self.active_model_name: Optional[str] = None
|
||||
|
|
@ -393,7 +394,12 @@ class InferenceOrchestrator:
|
|||
cancel_event=None,
|
||||
use_adapter=None,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Inner generation logic — sends command to subprocess, yields tokens."""
|
||||
"""Inner generation logic — sends command to subprocess, yields tokens.
|
||||
|
||||
Serialized by _gen_lock: only one generation runs at a time.
|
||||
This prevents concurrent readers from consuming each other's
|
||||
tokens off the shared resp_queue.
|
||||
"""
|
||||
if not self._ensure_subprocess_alive():
|
||||
yield "Error: Inference subprocess is not running"
|
||||
return
|
||||
|
|
@ -402,6 +408,39 @@ class InferenceOrchestrator:
|
|||
yield "Error: No active model"
|
||||
return
|
||||
|
||||
# Serialize generation — single GPU, one generation at a time.
|
||||
# Without this lock, two concurrent readers on the same resp_queue
|
||||
# can consume and drop each other's token events.
|
||||
with self._gen_lock:
|
||||
yield from self._generate_locked(
|
||||
messages=messages,
|
||||
system_prompt=system_prompt,
|
||||
image=image,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
top_k=top_k,
|
||||
min_p=min_p,
|
||||
max_new_tokens=max_new_tokens,
|
||||
repetition_penalty=repetition_penalty,
|
||||
cancel_event=cancel_event,
|
||||
use_adapter=use_adapter,
|
||||
)
|
||||
|
||||
def _generate_locked(
|
||||
self,
|
||||
messages: list = None,
|
||||
system_prompt: str = "",
|
||||
image=None,
|
||||
temperature: float = 0.7,
|
||||
top_p: float = 0.9,
|
||||
top_k: int = 40,
|
||||
min_p: float = 0.0,
|
||||
max_new_tokens: int = 256,
|
||||
repetition_penalty: float = 1.1,
|
||||
cancel_event=None,
|
||||
use_adapter=None,
|
||||
) -> Generator[str, None, None]:
|
||||
"""Actual generation logic — must be called under _gen_lock."""
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
# Convert PIL Image to base64 if needed
|
||||
|
|
@ -432,7 +471,8 @@ class InferenceOrchestrator:
|
|||
yield f"Error: {exc}"
|
||||
return
|
||||
|
||||
# Yield tokens from response queue
|
||||
# Yield tokens from response queue — we are the only reader
|
||||
# because _gen_lock is held.
|
||||
while True:
|
||||
resp = self._read_resp(timeout=30.0)
|
||||
|
||||
|
|
@ -443,23 +483,18 @@ class InferenceOrchestrator:
|
|||
return
|
||||
continue
|
||||
|
||||
# Skip responses for other requests
|
||||
resp_rid = resp.get("request_id")
|
||||
rtype = resp.get("type", "")
|
||||
|
||||
# Status messages don't have request_id
|
||||
# Status messages — skip
|
||||
if rtype == "status":
|
||||
continue
|
||||
|
||||
# Error without request_id = subprocess-level error
|
||||
resp_rid = resp.get("request_id")
|
||||
if rtype == "error" and not resp_rid:
|
||||
yield f"Error: {resp.get('error', 'Unknown error')}"
|
||||
return
|
||||
|
||||
# Skip responses for other requests
|
||||
if resp_rid and resp_rid != request_id:
|
||||
continue
|
||||
|
||||
if rtype == "token":
|
||||
# Check cancel from route (e.g. SSE connection closed)
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue