From 7fc563731a737cdad9e6caff85a8268941cacf21 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Thu, 5 Mar 2026 18:54:17 +0000 Subject: [PATCH] fix: use mp.Event for instant cross-process generation cancel Replaces cmd_queue-based cancel polling with a shared mp.Event. Fixes two issues: - Loading a new model while generating no longer hangs (cancel is instant) - Subprocess shuts down cleanly after explicit stop generation --- studio/backend/core/inference/orchestrator.py | 54 +++++++++++++------ studio/backend/core/inference/worker.py | 39 ++++++-------- 2 files changed, 54 insertions(+), 39 deletions(-) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index a3637cfb84..22ca93ddeb 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -42,6 +42,7 @@ class InferenceOrchestrator: self._proc: Optional[mp.Process] = None self._cmd_queue: Any = None self._resp_queue: Any = None + self._cancel_event: Any = None # mp.Event — set to cancel generation instantly self._lock = threading.Lock() # Local state mirrors (updated from subprocess responses) @@ -74,12 +75,14 @@ class InferenceOrchestrator: self._cmd_queue = _CTX.Queue() self._resp_queue = _CTX.Queue() + self._cancel_event = _CTX.Event() self._proc = _CTX.Process( target=run_inference_process, kwargs={ "cmd_queue": self._cmd_queue, "resp_queue": self._resp_queue, + "cancel_event": self._cancel_event, "config": config, }, daemon=True, @@ -87,26 +90,37 @@ class InferenceOrchestrator: self._proc.start() logger.info("Inference subprocess started (pid=%s)", self._proc.pid) + def _cancel_generation(self) -> None: + """Cancel any ongoing generation in the subprocess (instant).""" + if self._cancel_event is not None: + self._cancel_event.set() + def _shutdown_subprocess(self, timeout: float = 10.0) -> None: """Gracefully shut down the inference subprocess.""" - with self._lock: - if self._proc is None or not self._proc.is_alive(): - self._proc = None - return + if self._proc is None or not self._proc.is_alive(): + self._proc = None + return - # Send shutdown command - try: - self._cmd_queue.put({"type": "shutdown"}) - except (OSError, ValueError): - pass + # 1. Cancel any ongoing generation first (instant via mp.Event) + self._cancel_generation() + time.sleep(0.5) # Brief wait for generation to stop - # Wait for graceful shutdown + # 2. Drain stale responses from queue + self._drain_queue() + + # 3. Send shutdown command + try: + self._cmd_queue.put({"type": "shutdown"}) + except (OSError, ValueError): + pass + + # 4. Wait for graceful shutdown try: self._proc.join(timeout=timeout) except Exception: pass - # Force kill if still alive + # 5. Force kill if still alive if self._proc is not None and self._proc.is_alive(): logger.warning("Inference subprocess did not exit gracefully, terminating") try: @@ -125,6 +139,7 @@ class InferenceOrchestrator: self._proc = None self._cmd_queue = None self._resp_queue = None + self._cancel_event = None logger.info("Inference subprocess shut down") def _cleanup(self): @@ -250,6 +265,12 @@ class InferenceOrchestrator: } if self._ensure_subprocess_alive(): + # Cancel any ongoing generation first (user may be loading + # a new model while the current one is still generating) + self._cancel_generation() + time.sleep(0.3) + self._drain_queue() + if needed_major == self._current_transformers_major: # Reuse existing subprocess — send load command logger.info( @@ -457,12 +478,10 @@ class InferenceOrchestrator: continue if rtype == "token": - # Check cancel from route + # Check cancel from route (e.g. SSE connection closed) if cancel_event is not None and cancel_event.is_set(): - try: - self._send_cmd({"type": "cancel", "request_id": request_id}) - except RuntimeError: - pass + # Set the subprocess mp.Event — instant, no queue needed + self._cancel_generation() return yield resp.get("text", "") @@ -474,7 +493,8 @@ class InferenceOrchestrator: return def reset_generation_state(self): - """Send cancel/reset to subprocess.""" + """Cancel any ongoing generation and reset state.""" + self._cancel_generation() if not self._ensure_subprocess_alive(): return try: diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 6fcc4546be..3e98fb0e3c 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -17,7 +17,6 @@ import logging import os import queue as _queue import sys -import threading import time import traceback from io import BytesIO @@ -93,19 +92,6 @@ def _send_response(resp_queue: Any, response: dict) -> None: logger.error("Failed to send response: %s", exc) -def _check_cancel(cmd_queue: Any, cancel_event: threading.Event) -> None: - """Non-blocking poll of cmd_queue for cancel commands.""" - try: - cmd = cmd_queue.get_nowait() - if cmd and cmd.get("type") == "cancel": - cancel_event.set() - logger.info("Cancel signal received for request %s", cmd.get("request_id", "*")) - except _queue.Empty: - pass - except (EOFError, OSError): - pass - - def _build_model_config(config: dict): """Build a ModelConfig from the config dict.""" from utils.models import ModelConfig @@ -201,10 +187,15 @@ def _handle_generate( backend, cmd: dict, resp_queue: Any, - cmd_queue: Any, - cancel_event: threading.Event, + cancel_event, ) -> None: - """Handle a generate command: stream tokens back via resp_queue.""" + """Handle a generate command: stream tokens back via resp_queue. + + cancel_event is an mp.Event shared with the parent process. + The parent can set it at any time (e.g. user stops generation, + or user loads a new model while generating) and generation + stops within 1-2 tokens. + """ request_id = cmd.get("request_id", "") try: @@ -240,9 +231,9 @@ def _handle_generate( generator = backend.generate_chat_response(**gen_kwargs) for cumulative_text in generator: - # Check for cancel commands (non-blocking) - _check_cancel(cmd_queue, cancel_event) + # cancel_event is an mp.Event — checked instantly, no queue polling if cancel_event.is_set(): + logger.info("Generation cancelled for request %s", request_id) break _send_response(resp_queue, { @@ -297,6 +288,7 @@ def run_inference_process( *, cmd_queue: Any, resp_queue: Any, + cancel_event, config: dict, ) -> None: """Subprocess entrypoint. Persistent — runs command loop until shutdown. @@ -304,6 +296,7 @@ def run_inference_process( Args: cmd_queue: mp.Queue for receiving commands from parent. resp_queue: mp.Queue for sending responses to parent. + cancel_event: mp.Event shared with parent — set by parent to cancel generation. config: Initial configuration dict with model info and project_root. """ os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -371,7 +364,8 @@ def run_inference_process( return # ── 4. Command loop — process commands until shutdown ── - cancel_event = threading.Event() + # cancel_event is an mp.Event shared with parent — parent can set it + # at any time to cancel generation instantly (no queue polling needed). logger.info("Inference subprocess ready, entering command loop") while True: @@ -392,7 +386,7 @@ def run_inference_process( try: if cmd_type == "generate": cancel_event.clear() - _handle_generate(backend, cmd, resp_queue, cmd_queue, cancel_event) + _handle_generate(backend, cmd, resp_queue, cancel_event) elif cmd_type == "load": # Load a new model (reusing this subprocess) @@ -405,8 +399,9 @@ def run_inference_process( _handle_unload(backend, cmd, resp_queue) elif cmd_type == "cancel": + # Redundant with mp.Event but handle gracefully cancel_event.set() - logger.info("Cancel signal received") + logger.info("Cancel command received") elif cmd_type == "reset": cancel_event.set()