vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-10 00:08:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Studio: serialize the compare-mode dispatcher lifecycle to fix a start race (#6922)

Browse source 
* Studio: serialize the compare-mode dispatcher lifecycle to fix a start race

_generate_dispatched (compare mode) bypasses _gen_lock so two concurrent
compare requests can both reach _start_dispatcher. The check-then-spawn there
had no lock, so both could observe no live dispatcher and each spawn one. The
extra dispatcher is orphaned (self._dispatcher_thread tracks only the last) and
during a later unload it can consume the 'unloaded' reply off _resp_queue before
unload_model's _wait_response, hanging the unload on its timeout.

Add _dispatcher_lifecycle_lock and take it around the whole body of both
_start_dispatcher and _stop_dispatcher, so start/stop cannot interleave and the
second concurrent starter sees the dispatcher alive and returns. _start_dispatcher
now returns whether it actually spawned the thread, and _generate_dispatched
derives dispatcher_preexisting from that atomic result instead of a separate
unlocked is_alive() read.

No call site holds _mailbox_lock when calling start/stop, so joining the
dispatcher (which takes _mailbox_lock) under the new lock cannot deadlock; the
lock order is always _gen_lock then _dispatcher_lifecycle_lock and is never
inverted.

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

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

* Studio: refuse dispatcher start queued behind an unload's stop

A compare request could pass the early _unload_pending check, then block in _start_dispatcher on _dispatcher_lifecycle_lock behind an unload's _stop_dispatcher. When the unload released the lock the start spawned a fresh dispatcher, which became the resp_queue reader and consumed the worker's unroutable 'unloaded' reply before unload_model's _wait_response saw it, hanging the unload for 300s.

Gate _start_dispatcher on _unload_pending under the lifecycle lock, and set _unload_pending under the same lock ahead of the stop, so any start queued behind the stop observes the unload and refuses. Ordering stays _gen_lock -> _dispatcher_lifecycle_lock. Adds a regression test forcing the queued-behind-stop interleaving.

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-06 22:09:41 -07:00 • committed by GitHub
parent 3506371677
commit 9674e882c2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 257 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -79,6 +79,12 @@ class InferenceOrchestrator:
self._mailbox_lock = threading.Lock()
self._dispatcher_thread: Optional[threading.Thread] = None
self._dispatcher_stop = threading.Event()
# Serializes dispatcher start/stop. _generate_dispatched (compare mode) bypasses
# _gen_lock, so two concurrent compare requests can both reach _start_dispatcher;
# without this lock both could observe no live dispatcher and each spawn one,
# orphaning the extra thread (self._dispatcher_thread tracks only the last). The
# orphan later steals the "unloaded" reply off resp_queue and hangs unload_model.
self._dispatcher_lifecycle_lock = threading.Lock()
# Local state mirrors (updated from subprocess responses)
self.active_model_name: Optional[str] = None
@ -514,33 +520,56 @@ class InferenceOrchestrator:
# Dispatcher — per-request mailbox routing for compare mode
# ------------------------------------------------------------------
def _start_dispatcher(self) -> None:
def _start_dispatcher(self) -> bool:
"""Start the dispatcher thread if not already running.
The dispatcher reads the shared resp_queue and routes responses to
per-request mailbox queues, letting multiple adapter-controlled
(compare) requests 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")
The whole check-then-spawn runs under _dispatcher_lifecycle_lock so
concurrent compare requests (which bypass _gen_lock) can't both observe
no live dispatcher and each spawn one. Returns True only for the caller
that actually started a new thread; False if one was already alive.
"""
with self._dispatcher_lifecycle_lock:
# Refuse to start while an unload is in progress. unload_model sets
# _unload_pending under this same lock before it stops the idle
# dispatcher, so a start queued behind that stop observes the unload
# here and bails. Without this a fresh dispatcher would be spawned
# after the stop, become the resp_queue reader, and consume the
# worker's "unloaded" reply (unroutable, so dropped) before
# unload_model's _wait_response sees it -- hanging the unload 300s.
if self._unload_pending:
return False
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
return False
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")
return True
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")
"""Signal the dispatcher to stop and wait for it.
Runs under _dispatcher_lifecycle_lock (paired with _start_dispatcher) so
a stop can't interleave with a concurrent start. Callers must NOT hold
_mailbox_lock here: this joins the dispatcher, and the dispatcher loop
takes _mailbox_lock, so holding it would deadlock the join.
"""
with self._dispatcher_lifecycle_lock:
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."""
@ -628,13 +657,14 @@ class InferenceOrchestrator:
yield "Error: model is being unloaded"
return
# Ensure the dispatcher runs. Track whether it was already running: if this call
# starts it and then bails on a racing unload, it must stop it again (see the
# unloading bail below).
dispatcher_preexisting = (
self._dispatcher_thread is not None and self._dispatcher_thread.is_alive()
)
self._start_dispatcher()
# Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under
# _dispatcher_lifecycle_lock and returns True only for the caller that actually spawned
# the thread, so at most one dispatcher ever exists even when two compare requests race
# here. Derive dispatcher_preexisting from that atomic result (not a separate unlocked
# is_alive() read): if THIS call started the dispatcher and then bails on a racing
# unload, it must stop it again (see the unloading bail below).
started = self._start_dispatcher()
dispatcher_preexisting = not started
request_id = str(uuid.uuid4())
@ -1041,7 +1071,15 @@ class InferenceOrchestrator:
# The subprocess runs commands sequentially, so a bare unload queues behind a
# running generate (a 2-3 min hang). Cancel first (via the mp.Event the worker
# polls each token), then take _gen_lock as sole resp_queue reader (like GGUF).
self._unload_pending = True
#
# Set _unload_pending under _dispatcher_lifecycle_lock so it is ordered ahead of
# the dispatcher stop that _wait_dispatcher_idle runs under the same lock: a
# compare request's _start_dispatcher queued behind that stop then observes the
# unload and refuses to spawn a fresh dispatcher that would eat the "unloaded"
# reply off resp_queue. This is a standalone acquisition (no _gen_lock held yet),
# so it keeps the _gen_lock -> _dispatcher_lifecycle_lock order and can't deadlock.
with self._dispatcher_lifecycle_lock:
self._unload_pending = True
# Cancelling only the running generation isn't enough: the worker clears
# cancel_event at each generate start, so a queued one would clear it and run the
# outgoing model to completion. drain_event, never cleared, makes any generate

192
studio/backend/tests/test_orchestrator_unload_cancel.py
View file

@ -25,6 +25,8 @@ def _bare_orchestrator():
o._cmd_queue = object()
o._resp_queue = object()
o._dispatcher_thread = None
o._dispatcher_stop = threading.Event()
o._dispatcher_lifecycle_lock = threading.Lock()
o._unload_pending = False
o.active_model_name = "m"
o.models = {"m": {}}
@ -1253,6 +1255,7 @@ def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch):
def fake_start():
started["v"] = True
o._dispatcher_thread = _AliveDispatcher()
return True # _start_dispatcher returns True for the caller that spawned it
def fake_stop():
stopped["v"] = True
@ -1428,3 +1431,192 @@ def test_load_model_aborts_publish_when_cancelled_after_wait_response(monkeypatc
assert o.active_model_name is None, "must not publish a cancelled model's active name"
assert o.models == {}, "must not publish a cancelled model's mirror"
assert "m" not in o.loading_models
# ----------------------------------------------------------------------------
# Concurrent compare-mode requests must not each spawn a dispatcher. Compare mode
# (_generate_dispatched) deliberately bypasses _gen_lock, so two requests can reach
# _start_dispatcher at once. Without _dispatcher_lifecycle_lock the check-then-spawn
# races: both observe no live dispatcher and each start one. The extra dispatcher is
# orphaned (self._dispatcher_thread tracks only the last) and later consumes the
# "unloaded" reply off the shared resp_queue before unload_model's _wait_response,
# hanging the unload on its 300s timeout. The lifecycle lock must serialize the
# check-then-spawn so exactly one dispatcher thread is ever created.
# ----------------------------------------------------------------------------
def test_concurrent_start_dispatcher_spawns_exactly_one():
import queue as _queue
o = _bare_orchestrator()
o._resp_queue = _queue.Queue() # real queue so the dispatcher loop blocks and stays alive
o._mailbox_lock = threading.Lock()
o._mailboxes = {}
o._dispatcher_thread = None
o._dispatcher_stop = threading.Event()
o._dispatcher_lifecycle_lock = threading.Lock()
n = 32
# A barrier aligns every thread on the check-then-spawn window: without the lifecycle
# lock several would clear the "is a dispatcher alive?" check together and each spawn one.
barrier = threading.Barrier(n)
results: list = []
results_lock = threading.Lock()
def racer():
barrier.wait()
started = o._start_dispatcher()
with results_lock:
results.append(started)
threads = [threading.Thread(target = racer, name = f"racer-{i}") for i in range(n)]
for t in threads:
t.start()
for t in threads:
t.join(timeout = 5)
try:
# _start_dispatcher returns True only for the caller that actually spawned a thread.
# Exactly one caller may win; every other must observe the dispatcher alive and bail.
assert results.count(True) == 1, f"expected exactly one spawn, got {results.count(True)}"
assert results.count(False) == n - 1
# And exactly one live dispatcher thread exists -- no orphan racing resp_queue.
live = [
t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()
]
assert len(live) == 1, f"expected one live dispatcher, found {len(live)}"
assert o._dispatcher_thread is live[0]
finally:
o._stop_dispatcher()
# Stop joins and clears it; no dispatcher thread must survive.
assert o._dispatcher_thread is None
remaining = [
t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()
]
assert remaining == [], "dispatcher must be stopped and joined"
# ----------------------------------------------------------------------------
# A compare request whose _start_dispatcher is queued behind an unload's
# _stop_dispatcher must NOT spawn a fresh dispatcher. The idle-dispatcher stop
# and the queued start both serialize on _dispatcher_lifecycle_lock; if the
# queued start spawned a new dispatcher after the stop, it would become the
# resp_queue reader and consume unload_model's "unloaded" reply (unroutable, so
# dropped) before _wait_response saw it -- hanging the unload on its 300s
# timeout. unload_model sets _unload_pending under the SAME lifecycle lock ahead
# of the stop, so _start_dispatcher observes it and refuses.
# ----------------------------------------------------------------------------
def test_start_dispatcher_refuses_while_unload_pending():
# Direct unit guard: with an unload in progress (_unload_pending set under the
# lifecycle lock by unload_model), _start_dispatcher must refuse and spawn nothing,
# even though no dispatcher is currently running.
import queue as _queue
o = _bare_orchestrator()
o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive
o._dispatcher_thread = None
o._dispatcher_stop = threading.Event()
o._dispatcher_lifecycle_lock = threading.Lock()
o._unload_pending = True
started = o._start_dispatcher()
assert started is False, "must not start a dispatcher while an unload is pending"
assert o._dispatcher_thread is None, "no dispatcher thread may be created"
live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()]
assert live == [], "no dispatcher may exist to consume the unloaded reply"
def test_start_dispatcher_resumes_after_unload_clears():
# Guard the other direction: once the unload finishes and clears _unload_pending, a
# later compare request must be able to start the dispatcher again (the gate must not
# wedge). Proves the refusal above is scoped to the unload, not permanent.
import queue as _queue
o = _bare_orchestrator()
o._resp_queue = _queue.Queue()
o._dispatcher_thread = None
o._dispatcher_stop = threading.Event()
o._dispatcher_lifecycle_lock = threading.Lock()
o._unload_pending = False
try:
assert (
o._start_dispatcher() is True
), "a fresh dispatcher must start once no unload is pending"
assert o._dispatcher_thread is not None and o._dispatcher_thread.is_alive()
finally:
o._stop_dispatcher()
assert o._dispatcher_thread is None
def test_queued_start_behind_unload_stop_spawns_no_dispatcher():
# Codex's exact ordering, forced deterministically: an unload holds
# _dispatcher_lifecycle_lock across its _stop_dispatcher (the idle dispatcher's join
# is gated by an event), while a compare request's _start_dispatcher is queued behind
# it on the same lock. When the stop releases the lock the queued start must observe
# _unload_pending (set under the lock ahead of the stop) and refuse: no fresh
# dispatcher may be left running to steal the "unloaded" reply.
import queue as _queue
o = _bare_orchestrator()
o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive
o._mailbox_lock = threading.Lock()
o._mailboxes = {}
o._dispatcher_stop = threading.Event()
o._dispatcher_lifecycle_lock = threading.Lock()
o._unload_pending = False
start_queued = threading.Event() # release the stop's join once the start is queued behind it
join_may_finish = threading.Event()
class _IdleDispatcher:
# Stand-in for the idle compare-mode dispatcher the unload stops. Its join blocks
# until we confirm the compare _start_dispatcher is queued behind the stop, so the
# stop provably holds _dispatcher_lifecycle_lock across that window.
def is_alive(self):
return True
def join(self, timeout = None):
assert start_queued.wait(timeout = 5), "compare start must queue behind the stop"
assert join_may_finish.wait(timeout = 5)
o._dispatcher_thread = _IdleDispatcher()
def unload_side():
# unload_model's sequence: set _unload_pending under the lifecycle lock, then stop
# the idle dispatcher (also under the lock, via _wait_dispatcher_idle).
with o._dispatcher_lifecycle_lock:
o._unload_pending = True
o._stop_dispatcher()
started_result = {}
def compare_side():
started_result["v"] = o._start_dispatcher()
u = threading.Thread(target = unload_side, name = "unload-side")
u.start()
# Let the unload set _unload_pending, enter _stop_dispatcher, and block in the gated join
# while holding the lifecycle lock.
time.sleep(0.2)
c = threading.Thread(target = compare_side, name = "compare-side")
c.start()
# Let the compare _start_dispatcher block on the lifecycle lock (queued behind the stop).
time.sleep(0.2)
start_queued.set() # the start is now queued behind the stop
join_may_finish.set() # let the stop's join complete and release the lock
u.join(timeout = 5)
c.join(timeout = 5)
assert started_result.get("v") is False, "the queued start must refuse while unloading"
assert o._dispatcher_thread is None, "the stop cleared it and the queued start spawned nothing"
live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()]
assert live == [], "no fresh dispatcher may be left to consume the unloaded reply"