Fix: Compare Mode Deadlock, Cancel Event Poisoning & IPC Optimization (#4303)

* fix: resolve compare mode deadlock, cancel_event poisoning, and add dispatcher-based IPC optimization

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* revert to 2048 tokens

* refactor: extract dispatcher timeout values into named constants

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: guard dispatcher shutdown against active compare mailboxes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

studio/backend/core/inference/inference.py

@@ -886,6 +886,7 @@ class InferenceBackend:
             output = ""
             from queue import Empty
 
+            generation_complete = False
             try:
                 while True:
                     if cancel_event is not None and cancel_event.is_set():
@@ -893,9 +894,11 @@ class InferenceBackend:
                     try:
                         new_token = next(streamer)
                     except StopIteration:
+                        generation_complete = True
                         break
                     except Empty:
                         if not thread.is_alive():
+                            generation_complete = True
                             break
                         continue
                     if new_token:
@@ -903,7 +906,7 @@ class InferenceBackend:
                         cleaned = self._clean_generated_text(output)
                         yield cleaned
             finally:
-                if cancel_event is not None:
+                if cancel_event is not None and not generation_complete:
                     cancel_event.set()
                 thread.join(timeout = 10)
                 if thread.is_alive():
@@ -1167,6 +1170,7 @@ class InferenceBackend:
             output = ""
             from queue import Empty
 
+            generation_complete = False
             try:
                 while True:
                     if cancel_event is not None and cancel_event.is_set():
@@ -1174,9 +1178,11 @@ class InferenceBackend:
                     try:
                         new_token = next(streamer)
                     except StopIteration:
+                        generation_complete = True
                         break
                     except Empty:
                         if not thread.is_alive():
+                            generation_complete = True
                             break
                         continue
                     if new_token:
@@ -1184,7 +1190,12 @@ class InferenceBackend:
                         cleaned = self._clean_generated_text(output)
                         yield cleaned
             finally:
-                if cancel_event is not None:
+                # Only set cancel_event when we exited early (user cancel),
+                # NOT on normal completion.  cancel_event is a shared mp.Event
+                # — setting it unconditionally would leave a stale cancel
+                # signal that could interfere with the next serialized
+                # generation request (e.g. in compare mode).
+                if cancel_event is not None and not generation_complete:
                     cancel_event.set()
                 thread.join(timeout = 10)
                 if thread.is_alive():

studio/backend/core/inference/orchestrator.py

@@ -32,6 +32,13 @@ logger = get_logger(__name__)
 
 _CTX = mp.get_context("spawn")
 
+# Dispatcher timeout constants (seconds)
+_DISPATCH_READ_TIMEOUT = 30.0
+_DISPATCH_POLL_INTERVAL = 0.5
+_DISPATCH_STOP_TIMEOUT = 5.0
+_DISPATCH_IDLE_TIMEOUT = 30.0
+_DISPATCH_DRAIN_TIMEOUT = 5.0
+
 
 class InferenceOrchestrator:
     """
@@ -53,6 +60,15 @@ class InferenceOrchestrator:
             threading.Lock()
         )  # Serializes generation — one request at a time
 
+        # Dispatcher state — for compare mode (adapter-controlled requests).
+        # Instead of serializing via _gen_lock, adapter-controlled requests
+        # send commands directly to the subprocess and read from per-request
+        # mailboxes. A dispatcher thread routes resp_queue events by request_id.
+        self._mailboxes: dict[str, queue.Queue] = {}
+        self._mailbox_lock = threading.Lock()  # Protects _mailboxes dict
+        self._dispatcher_thread: Optional[threading.Thread] = None
+        self._dispatcher_stop = threading.Event()
+
         # Local state mirrors (updated from subprocess responses)
         self.active_model_name: Optional[str] = None
         self.models: dict = {}
@@ -105,6 +121,7 @@ class InferenceOrchestrator:
 
     def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
         """Gracefully shut down the inference subprocess."""
+        self._stop_dispatcher()  # Stop dispatcher before killing subprocess
         if self._proc is None or not self._proc.is_alive():
             self._proc = None
             return
@@ -256,6 +273,229 @@ class InferenceOrchestrator:
                 return
         logger.warning("Timed out waiting for gen_done after cancel")
 
+    # ------------------------------------------------------------------
+    # Dispatcher — per-request mailbox routing for compare mode
+    # ------------------------------------------------------------------
+
+    def _start_dispatcher(self) -> None:
+        """Start the dispatcher thread if not already running.
+
+        The dispatcher reads from the shared resp_queue and routes
+        responses to per-request mailbox queues. This allows multiple
+        adapter-controlled (compare) requests to be in-flight without
+        holding _gen_lock.
+        """
+        if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
+            return
+
+        self._dispatcher_stop.clear()
+        self._dispatcher_thread = threading.Thread(
+            target = self._dispatcher_loop,
+            daemon = True,
+            name = "inference-dispatcher",
+        )
+        self._dispatcher_thread.start()
+        logger.debug("Dispatcher thread started")
+
+    def _stop_dispatcher(self) -> None:
+        """Signal the dispatcher to stop and wait for it."""
+        if self._dispatcher_thread is None:
+            return
+        self._dispatcher_stop.set()
+        self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT)
+        self._dispatcher_thread = None
+        logger.debug("Dispatcher thread stopped")
+
+    def _dispatcher_loop(self) -> None:
+        """Background loop: read resp_queue → route to mailboxes by request_id."""
+        while not self._dispatcher_stop.is_set():
+            if self._resp_queue is None:
+                break
+
+            try:
+                resp = self._resp_queue.get(timeout = _DISPATCH_POLL_INTERVAL)
+            except queue.Empty:
+                continue
+            except (EOFError, OSError, ValueError):
+                break
+
+            rid = resp.get("request_id")
+            rtype = resp.get("type", "")
+
+            # Status messages — log and skip
+            if rtype == "status":
+                logger.info("Subprocess status: %s", resp.get("message", ""))
+                continue
+
+            # Route to mailbox if a matching request_id exists
+            if rid:
+                with self._mailbox_lock:
+                    mbox = self._mailboxes.get(rid)
+                if mbox is not None:
+                    mbox.put(resp)
+                    continue
+
+            # No matching mailbox — might be for a _gen_lock reader or orphaned
+            # Push it back so _read_resp can pick it up. But we can't un-get
+            # from mp.Queue, so log a warning.
+            if rtype not in ("status",):
+                logger.debug(
+                    "Dispatcher: no mailbox for request_id=%s type=%s, dropping",
+                    rid,
+                    rtype,
+                )
+
+    def _generate_dispatched(
+        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]:
+        """Dispatched generation — sends command without holding _gen_lock.
+
+        Uses a per-request mailbox to receive tokens. This allows two
+        compare-mode requests to be queued in the subprocess simultaneously,
+        eliminating the inter-generation round-trip overhead.
+
+        The subprocess processes commands sequentially from its cmd_queue,
+        so generation is still serialized at the GPU level — we just avoid
+        the orchestrator-level lock contention.
+        """
+        if not self._ensure_subprocess_alive():
+            yield "Error: Inference subprocess is not running"
+            return
+
+        if not self.active_model_name:
+            yield "Error: No active model"
+            return
+
+        # Ensure dispatcher is running
+        self._start_dispatcher()
+
+        request_id = str(uuid.uuid4())
+
+        # Convert PIL Image to base64 if needed
+        image_b64 = None
+        if image is not None:
+            image_b64 = self._pil_to_base64(image)
+
+        cmd = {
+            "type": "generate",
+            "request_id": request_id,
+            "messages": messages or [],
+            "system_prompt": system_prompt,
+            "image_base64": image_b64,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "min_p": min_p,
+            "max_new_tokens": max_new_tokens,
+            "repetition_penalty": repetition_penalty,
+        }
+
+        if use_adapter is not None:
+            cmd["use_adapter"] = use_adapter
+
+        # Create mailbox BEFORE sending command
+        mailbox: queue.Queue = queue.Queue()
+        with self._mailbox_lock:
+            self._mailboxes[request_id] = mailbox
+
+        try:
+            self._send_cmd(cmd)
+        except RuntimeError as exc:
+            with self._mailbox_lock:
+                self._mailboxes.pop(request_id, None)
+            yield f"Error: {exc}"
+            return
+
+        # Read tokens from our private mailbox
+        try:
+            while True:
+                try:
+                    resp = mailbox.get(timeout = _DISPATCH_READ_TIMEOUT)
+                except queue.Empty:
+                    # Timeout — check subprocess health
+                    if not self._ensure_subprocess_alive():
+                        yield "Error: Inference subprocess crashed during generation"
+                        return
+                    continue
+
+                rtype = resp.get("type", "")
+
+                if rtype == "token":
+                    # Check cancel from route (e.g. SSE connection closed)
+                    if cancel_event is not None and cancel_event.is_set():
+                        self._cancel_generation()
+                        # Drain remaining events for this request
+                        self._drain_mailbox(mailbox, timeout = 5.0)
+                        return
+                    yield resp.get("text", "")
+
+                elif rtype == "gen_done":
+                    return
+
+                elif rtype == "gen_error":
+                    yield f"Error: {resp.get('error', 'Unknown error')}"
+                    return
+        finally:
+            with self._mailbox_lock:
+                self._mailboxes.pop(request_id, None)
+
+    def _drain_mailbox(self, mailbox: queue.Queue, timeout: float = 5.0) -> None:
+        """Drain a mailbox until gen_done/gen_error, discarding tokens."""
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            try:
+                resp = mailbox.get(
+                    timeout = min(_DISPATCH_POLL_INTERVAL, deadline - time.monotonic())
+                )
+            except queue.Empty:
+                continue
+            rtype = resp.get("type", "")
+            if rtype in ("gen_done", "gen_error"):
+                return
+        logger.warning("Timed out draining mailbox after cancel")
+
+    def _wait_dispatcher_idle(self) -> None:
+        """Wait for all dispatched requests to complete, then stop dispatcher.
+
+        Called by _generate_inner before using the _gen_lock path, to ensure
+        the dispatcher thread isn't competing for resp_queue reads.
+        """
+        if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive():
+            return
+
+        # Wait for all mailboxes to be emptied (dispatched requests complete)
+        deadline = time.monotonic() + _DISPATCH_IDLE_TIMEOUT
+        while time.monotonic() < deadline:
+            with self._mailbox_lock:
+                if not self._mailboxes:
+                    break
+            time.sleep(0.1)
+
+        # Only stop dispatcher if all mailboxes drained.  If compare
+        # requests are still active, leave the dispatcher running so
+        # their token routing isn't killed mid-stream.
+        with self._mailbox_lock:
+            still_active = bool(self._mailboxes)
+        if still_active:
+            logger.warning(
+                "Dispatcher still has %d active mailbox(es); "
+                "leaving dispatcher running for compare requests",
+                len(self._mailboxes),
+            )
+        else:
+            self._stop_dispatcher()
+
     # ------------------------------------------------------------------
     # Public API — same interface as InferenceBackend
     # ------------------------------------------------------------------
@@ -410,8 +650,13 @@ class InferenceOrchestrator:
         cancel_event = None,
         **gen_kwargs,
     ) -> Generator[str, None, None]:
-        """Generate with adapter control, streaming tokens from subprocess."""
-        yield from self._generate_inner(
+        """Generate with adapter control, streaming tokens from subprocess.
+
+        Uses the dispatcher path (no _gen_lock) so that compare-mode
+        requests don't block each other. The subprocess naturally
+        serializes them via its sequential command loop.
+        """
+        yield from self._generate_dispatched(
             use_adapter = use_adapter,
             cancel_event = cancel_event,
             **gen_kwargs,
@@ -445,6 +690,11 @@ class InferenceOrchestrator:
             yield "Error: No active model"
             return
 
+        # If the dispatcher is running (from a previous compare-mode request),
+        # wait for all dispatched requests to finish, then stop the dispatcher
+        # so we can safely read from resp_queue directly.
+        self._wait_dispatcher_idle()
+
         # 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.

studio/backend/routes/inference.py

@@ -1051,7 +1051,22 @@ async def openai_chat_completions(
                 yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n"
 
                 prev_text = ""
-                for cumulative in generate():
+                # Run sync generator in thread pool to avoid blocking
+                # the event loop. Critical for compare mode: two SSE
+                # requests arrive concurrently but the orchestrator
+                # serializes them via _gen_lock. Without run_in_executor
+                # the second request's blocking lock acquisition would
+                # freeze the entire event loop, stalling both streams.
+                _DONE = object()  # sentinel for generator exhaustion
+                loop = asyncio.get_event_loop()
+                gen = generate()
+                while True:
+                    # next(gen, _DONE) returns _DONE instead of raising
+                    # StopIteration — StopIteration cannot propagate
+                    # through asyncio futures (Python limitation).
+                    cumulative = await loop.run_in_executor(None, next, gen, _DONE)
+                    if cumulative is _DONE:
+                        break
                     if await request.is_disconnected():
                         cancel_event.set()
                         backend.reset_generation_state()