ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: use mp.Event for instant cross-process generation cancel

Browse source 
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
This commit is contained in:
Roland Tannous 2026-03-05 18:54:17 +00:00
commit 7fc563731a
2 changed files with 54 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

54
studio/backend/core/inference/orchestrator.py
View file

@ -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:

39
studio/backend/core/inference/worker.py
View file

@ -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()