diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 675bd9f3ea..e0e6cef6c9 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -45,6 +45,10 @@ _DISPATCH_STOP_TIMEOUT = 5.0 _DISPATCH_IDLE_TIMEOUT = 30.0 _DISPATCH_DRAIN_TIMEOUT = 5.0 +# Max wait for a cancelled generation to release _gen_lock before unload_model +# tears the subprocess down. Only bounds a wedged worker. +_UNLOAD_GEN_LOCK_TIMEOUT = 15.0 + class InferenceOrchestrator: """ @@ -60,7 +64,13 @@ class InferenceOrchestrator: self._cmd_queue: Any = None self._resp_queue: Any = None self._cancel_event: Any = None # mp.Event — set to cancel generation + # Set for the whole unload; the worker never clears it (unlike _cancel_event), + # so a generate queued behind the cancelled one is skipped, not run. + self._drain_event: Any = None self._gen_lock = threading.Lock() # Serializes generation + # Set during a switch so a generation winning the _gen_lock handoff bails + # instead of starting on the outgoing model. + self._unload_pending = False # Dispatcher state for compare mode (adapter-controlled requests): # bypass _gen_lock, send commands directly, read from per-request @@ -159,6 +169,7 @@ class InferenceOrchestrator: self._cmd_queue = _CTX.Queue() self._resp_queue = _CTX.Queue() self._cancel_event = _CTX.Event() + self._drain_event = _CTX.Event() self._proc = _CTX.Process( target = run_without_native_path_secret, @@ -167,6 +178,7 @@ class InferenceOrchestrator: "cmd_queue": self._cmd_queue, "resp_queue": self._resp_queue, "cancel_event": self._cancel_event, + "drain_event": self._drain_event, "config": config, }, daemon = True, @@ -228,6 +240,7 @@ class InferenceOrchestrator: self._cmd_queue = None self._resp_queue = None self._cancel_event = None + self._drain_event = None logger.info("Inference subprocess shut down") def _cleanup(self): @@ -456,7 +469,15 @@ class InferenceOrchestrator: cancel ack from that same source so stale events don't leak into the next request. """ + # Latch this stream's subprocess/queue: if a wedged worker is torn down and a + # later load spawns a fresh one, bail rather than re-block on the new queue + # under _gen_lock (deadlock). + initial_proc = self._proc + initial_resp_queue = self._resp_queue while True: + if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: + yield f"Error: {self._subprocess_crash_message(crash_context)}" + return resp = read_one(read_timeout) if resp is None: # Check subprocess health @@ -595,8 +616,24 @@ class InferenceOrchestrator: if not self.active_model_name: yield "Error: No active model" return + # Latch the target model so the recheck below can detect a switch that completed + # between _start_dispatcher and mailbox registration (mirrors the locked path's + # expected_model check). + expected_model = self.active_model_name - # Ensure dispatcher is running + # Switch in flight (unload waiting on _gen_lock). This path bypasses the lock, + # so without this early-out a compare request would enqueue a generate on the + # outgoing model and delay the switch. + if self._unload_pending: + 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() request_id = str(uuid.uuid4()) @@ -624,10 +661,42 @@ class InferenceOrchestrator: preserve_thinking = preserve_thinking, ) - # Create mailbox BEFORE sending command + # Create the mailbox BEFORE sending, rechecking _unload_pending under + # _mailbox_lock: an unload sets _unload_pending before _wait_dispatcher_idle + # reads _mailboxes under the same lock, so either the idle check sees this + # mailbox (and tears the dispatcher down) or we see the unload and bail. + # Registering after would orphan the mailbox and hang the compare stream forever. mailbox: queue.Queue = queue.Queue() with self._mailbox_lock: - self._mailboxes[request_id] = mailbox + # _unload_pending alone is not enough: an unload that ran fully since + # _start_dispatcher clears it in its finally and stops the dispatcher, so it + # reads False here though the dispatcher is gone and the model swapped. Also + # bail when the active model changed or the dispatcher died: a mailbox with no + # dispatcher to route gen_done/gen_error hangs the compare stream. + dispatcher_alive = ( + self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() + ) + unloading = ( + self._unload_pending + or self.active_model_name != expected_model + or not dispatcher_alive + ) + if not unloading: + self._mailboxes[request_id] = mailbox + # When bailing without a mailbox, note whether any OTHER compare request still + # routes through the dispatcher; if none and this call started it, stop it below. + orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes + if unloading: + # A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was + # stopped, then set _unload_pending. The one we just started would otherwise + # linger with no mailboxes, race unload_model's _wait_response for the "unloaded" + # reply off resp_queue, and drop it as unroutable -- hanging the unload 300s. Stop + # it here so the unload stays the sole resp_queue reader. Outside _mailbox_lock: + # _stop_dispatcher joins the dispatcher, which itself takes that lock. + if orphaned_dispatcher: + self._stop_dispatcher() + yield "Error: model is being unloaded" + return try: self._send_cmd(cmd) @@ -676,14 +745,18 @@ class InferenceOrchestrator: return logger.warning("Timed out draining mailbox after cancel") - def _wait_dispatcher_idle(self) -> None: + def _wait_dispatcher_idle(self) -> bool: """Wait for all dispatched requests to complete, then stop dispatcher. - Called by _generate_inner before the _gen_lock path so the dispatcher - thread isn't competing for resp_queue reads. + Returns True if the dispatcher was stopped (all mailboxes drained, or no + dispatcher was running), and False if it was left running because compare + requests were still active after _DISPATCH_IDLE_TIMEOUT. + + Called before the _gen_lock path so the dispatcher thread isn't competing + for resp_queue reads. """ if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive(): - return + return True # Wait for all mailboxes to be emptied (dispatched requests complete) deadline = time.monotonic() + _DISPATCH_IDLE_TIMEOUT @@ -704,8 +777,9 @@ class InferenceOrchestrator: "leaving dispatcher running for compare requests", len(self._mailboxes), ) - else: - self._stop_dispatcher() + return False + self._stop_dispatcher() + return True # ------------------------------------------------------------------ # Public API — same interface as InferenceBackend @@ -772,6 +846,19 @@ class InferenceOrchestrator: ) for attempt in range(2): + # Stop-loading (/unload -> cancel_load) aborts a load by discarding this + # model's loading marker. cancel_load only kills a live child; if the cancel + # lands before any child exists (GPU placement, or between retries) there is + # nothing to kill, and without this check the loop would spawn a worker and + # load the model after /unload reported it unloaded. Observe removal and stop. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled before spawn; not starting a worker", + model_name, + ) + self.active_model_name = None + self.models.clear() + return False logger.info( "Spawning fresh inference subprocess for '%s' " "(transformers %s.x, attempt %d/2%s)", @@ -783,6 +870,22 @@ class InferenceOrchestrator: sub_config["disable_xet"] = disable_xet self._spawn_subprocess(sub_config) + # A cancel can land after the pre-spawn recheck but while _spawn_subprocess + # is still creating the queues/process. cancel_load runs off the lifecycle + # gate, so its _shutdown_subprocess can see _proc still None and no-op, + # orphaning this fresh worker; the load would then wait for "loaded" and + # publish a model /unload reported unloaded, over a live subprocess nothing + # reaps. Recheck now the child exists and tear it down before publishing. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled during spawn; tearing the worker down", + model_name, + ) + self._shutdown_subprocess(timeout = 5) + self.active_model_name = None + self.models.clear() + return False + try: resp = self._wait_response("loaded") except DownloadStallError: @@ -803,8 +906,31 @@ class InferenceOrchestrator: ) if resp.get("success"): + # A cancel can land while we were parked in _wait_response above. + # cancel_load (off the lifecycle gate) discards this model's loading + # marker BEFORE its teardown, so a Stop-loading that fired after the + # worker queued "loaded" (which we can still consume during cancel_load's + # shutdown window) shows up here only as the marker's removal. Without + # this recheck we would publish active_model_name/models for a model + # /unload reported cancelled, over a subprocess cancel_load just killed; + # its post-teardown re-clear cannot undo a publish that lands after it + # returns. Observe the removal and abort; cancel_load owns teardown. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled while waiting for 'loaded'; " + "not publishing the cancelled model", + model_name, + ) + self.active_model_name = None + self.models.clear() + return False model_info = resp.get("model_info", {}) self.active_model_name = model_info.get("identifier", model_name) + # A load always spawns a fresh subprocess holding only this model, so + # mirror that. A lingering stale name would pass unload_model's "not in + # self.models" guard, and the worker's absent-name fallback would unload + # its *active* model, not the already-gone one. + self.models = {} self.models[self.active_model_name] = { "is_vision": model_info.get("is_vision", False), "is_lora": model_info.get("is_lora", False), @@ -837,17 +963,65 @@ class InferenceOrchestrator: self.models.clear() raise - def unload_model(self, model_name: str) -> bool: - """Unload a model from the subprocess.""" - if model_name in self.loading_models: - logger.info( - "Cancelling in-flight load for model '%s' by terminating subprocess", + def cancel_load(self, model_name: str) -> bool: + """Abort an in-flight load by terminating its subprocess. + + Returns True if a load for ``model_name`` (matched case-insensitively) was + cancelled, False if nothing was loading under that name. This only tears the + loading subprocess down -- it sends no command to a worker -- so, unlike the + rest of ``unload_model``, it is safe to run WITHOUT the inference lifecycle + gate. ``/unload`` calls it off-gate so the "stop loading" button can interrupt + a safetensors load that holds the gate for its whole (multi-minute) duration; + a gated cancel could never preempt that load. + """ + target = model_name + if target not in self.loading_models: + target = next( + (m for m in self.loading_models if m.lower() == model_name.lower()), model_name, ) - self._shutdown_subprocess(timeout = 0.5) - self.loading_models.discard(model_name) - self.active_model_name = None - self.models.clear() + if target not in self.loading_models: + return False + logger.info( + "Cancelling in-flight load for model '%s' by terminating subprocess", + target, + ) + # Discard the loading marker (and clear local state) BEFORE the teardown, not + # after. cancel_load runs off the lifecycle gate, alongside a load_model that + # rechecks this marker before each spawn. But _shutdown_subprocess can block (~1s + # tearing a live child down and joining the dispatcher), so clearing only after + # leaves a window where load_model reads the marker still set, passes its pre-spawn + # recheck, and loads the model after /unload reported it cancelled. Clear first. + self.loading_models.discard(target) + self.active_model_name = None + self.models.clear() + self._shutdown_subprocess(timeout = 0.5) + # Clear the local mirrors again AFTER the teardown. A racing off-gate load_model + # may still be parked in _wait_response("loaded"): its worker already queued a + # "loaded" reply, so during the shutdown window above (the 0.5s settle before the + # response queue is drained and nulled) that thread can consume it and repopulate + # active_model_name/models, undoing the pre-teardown clear. _shutdown_subprocess + # nulls the queue but not the mirrors, so without this second clear /unload reports + # success while the backend still advertises a killed model. The nulled queue lets + # no further "loaded" through, so re-clearing here wipes any repopulation. + self.active_model_name = None + self.models.clear() + return True + + def unload_model(self, model_name: str) -> bool: + """Unload a model from the subprocess.""" + # active_model_name can differ in case from the client's raw /unload name (the + # load path canonicalizes casing). Match case-insensitively and use the canonical + # spelling so the guard, unload command, and cleanup below hit the loaded model. + if ( + self.active_model_name is not None + and model_name != self.active_model_name + and model_name.lower() == self.active_model_name.lower() + ): + model_name = self.active_model_name + # In-flight load: tear its subprocess down (shared loading-cancel logic; no + # worker command sent). + if self.cancel_load(model_name): return True if not self._ensure_subprocess_alive(): @@ -857,30 +1031,85 @@ class InferenceOrchestrator: self.active_model_name = None return True - try: - self._send_cmd( - { - "type": "unload", - "model_name": model_name, - } - ) - resp = self._wait_response("unloaded") - - # Update local state + # Nothing loaded under this name: don't unload a stale model. The worker falls + # back to unloading its *active* model when the name is absent, so a stale unload + # (lost a race to a concurrent load) would hit the wrong one. + if model_name != self.active_model_name and model_name not in self.models: self.models.pop(model_name, None) - if self.active_model_name == model_name: - self.active_model_name = None - - logger.info("Model '%s' unloaded from subprocess", model_name) return True - except Exception as exc: - logger.error("Error unloading model '%s': %s", model_name, exc) - # Clear local state anyway - self.models.pop(model_name, None) - if self.active_model_name == model_name: - self.active_model_name = None - return False + # 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 + # 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 + # dequeued during the unload skip. + if self._drain_event is not None: + self._drain_event.set() + try: + self._cancel_generation() + acquired = self._gen_lock.acquire(timeout = _UNLOAD_GEN_LOCK_TIMEOUT) + if not acquired: + # Wedged worker: tear the subprocess down to free the GPU (next load respawns). + logger.warning( + "Unload: generation did not yield %.1fs after cancel; " + "shutting the inference subprocess down to free the model", + _UNLOAD_GEN_LOCK_TIMEOUT, + ) + self._shutdown_subprocess(timeout = 5) + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + + try: + # Stop the compare-mode dispatcher so it can't consume the "unloaded" reply + # off resp_queue before we do. A dispatched generation bypasses _gen_lock, so + # a wedged one slips past the acquire above; if the dispatcher is still active + # it owns resp_queue and the queued unload hangs _wait_response behind the + # stuck generate. Mirror the wedged locked path: tear the subprocess down. + if not self._wait_dispatcher_idle(): + logger.warning( + "Unload: compare-mode dispatcher still active after idle " + "wait; shutting the inference subprocess down to free the model" + ) + self._shutdown_subprocess(timeout = 5) + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + # Drop stale tokens so they can't be read as the unload reply. + self._drain_queue() + self._send_cmd( + { + "type": "unload", + "model_name": model_name, + } + ) + self._wait_response("unloaded") + + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + + logger.info("Model '%s' unloaded from subprocess", model_name) + return True + + except Exception as exc: + logger.error("Error unloading model '%s': %s", model_name, exc) + # Clear local state anyway + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return False + finally: + self._gen_lock.release() + finally: + self._unload_pending = False + if self._drain_event is not None: + self._drain_event.clear() def generate_chat_response( self, @@ -1068,6 +1297,7 @@ class InferenceOrchestrator: if not self.active_model_name: yield "Error: No active model" return + expected_model = self.active_model_name # Drain any prior compare-mode dispatcher so we can read resp_queue. self._wait_dispatcher_idle() @@ -1076,6 +1306,14 @@ class InferenceOrchestrator: # consume and drop each other's token events. Hold _gen_lock across the # cmd build + send + whole stream so we stay the sole resp_queue reader. with self._gen_lock: + # Recheck under the lock: an unload we raced may have cleared/swapped the model. + # _unload_pending resets after the lock releases, so it can read False by now; + # the active-model check catches that handoff and a reload that swapped models, + # so we never generate on the wrong one. + if self._unload_pending or self.active_model_name != expected_model: + # Won the lock handoff during a switch; don't start on the outgoing model. + yield "Error: model is being unloaded" + return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None cmd = self._build_generate_cmd( @@ -1143,53 +1381,62 @@ class InferenceOrchestrator: raise RuntimeError("Inference subprocess is not running") if not self.active_model_name: raise RuntimeError("No active model") + expected_model = self.active_model_name - request_id = str(uuid.uuid4()) + # Serialize under _gen_lock (sole resp_queue reader) and refuse to start on the + # outgoing model once an unload is pending, like the text and audio-input paths. + # Without this a concurrent /audio/generate could run TTS on a model being switched. + with self._gen_lock: + # Recheck under the lock (see _generate_inner): a raced unload/switch may have + # cleared or swapped the model while we waited. + if self._unload_pending or self.active_model_name != expected_model: + raise RuntimeError("model is being unloaded") - cmd = { - "type": "generate_audio", - "request_id": request_id, - "text": text, - "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 + request_id = str(uuid.uuid4()) - self._send_cmd(cmd) + cmd = { + "type": "generate_audio", + "request_id": request_id, + "text": text, + "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 - # Wait for audio_done or audio_error - deadline = time.monotonic() + 120.0 - while time.monotonic() < deadline: - remaining = max(0.1, deadline - time.monotonic()) - resp = self._read_resp(timeout = min(remaining, 1.0)) + self._send_cmd(cmd) - if resp is None: - if not self._ensure_subprocess_alive(): - raise RuntimeError(self._subprocess_crash_message("audio generation")) - continue + deadline = time.monotonic() + 120.0 + while time.monotonic() < deadline: + remaining = max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) - rtype = resp.get("type", "") + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("audio generation")) + continue - if rtype == "audio_done": - wav_bytes = base64.b64decode(resp["wav_base64"]) - sample_rate = resp["sample_rate"] - return wav_bytes, sample_rate + rtype = resp.get("type", "") - if rtype == "audio_error": - raise RuntimeError(resp.get("error", "Audio generation failed")) + if rtype == "audio_done": + wav_bytes = base64.b64decode(resp["wav_base64"]) + sample_rate = resp["sample_rate"] + return wav_bytes, sample_rate - if rtype == "error": - raise RuntimeError(resp.get("error", "Unknown error")) + if rtype == "audio_error": + raise RuntimeError(resp.get("error", "Audio generation failed")) - if rtype == "status": - continue + if rtype == "error": + raise RuntimeError(resp.get("error", "Unknown error")) - raise RuntimeError("Timeout waiting for audio generation (120s)") + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for audio generation (120s)") def generate_whisper_response( self, @@ -1254,8 +1501,15 @@ class InferenceOrchestrator: if not self.active_model_name: yield "Error: No active model" return + expected_model = self.active_model_name with self._gen_lock: + # Recheck under the lock (see _generate_inner): a raced unload/switch may have + # cleared or swapped the model while we waited. + if self._unload_pending or self.active_model_name != expected_model: + # Won the lock handoff during a switch; don't start on the outgoing model. + yield "Error: model is being unloaded" + return request_id = str(uuid.uuid4()) # numpy array -> list for mp.Queue serialization diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 4e27183d88..615c5c5a0d 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -406,6 +406,32 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: ) +def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool: + """Skip a generate queued behind a cancelled one during an unload. + + The parent sets ``drain_event`` for the whole unload. Because the parent's + per-token ``cancel_event`` is cleared at the start of every generate, a cancel + set while this generate was still queued would otherwise be lost when it is + dequeued. If the drain is in effect, emit an immediate (empty) ``gen_done`` so + the parent's stream/mailbox drains fast and the switch stays fast, and report + the generate was skipped so the caller does not clear the cancel or run it. + """ + if drain_event is None or not drain_event.is_set(): + return False + request_id = cmd.get("request_id", "") + logger.info("Skipping generate for request %s: unload draining", request_id) + _send_response( + resp_queue, + { + "type": "gen_done", + "request_id": request_id, + "cancelled": True, + "stats": None, + }, + ) + return True + + def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: """Handle a generate command: stream tokens back via resp_queue. @@ -632,7 +658,14 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None: ) -def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None: +def run_inference_process( + *, + cmd_queue: Any, + resp_queue: Any, + cancel_event, + config: dict, + drain_event = None, +) -> None: """Subprocess entrypoint. Persistent — runs the command loop until shutdown. Args: @@ -640,6 +673,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf resp_queue: mp.Queue for sending responses to parent. cancel_event: mp.Event the parent sets to cancel generation. config: Initial configuration dict with model info. + drain_event: mp.Event the parent sets for the duration of an unload. Unlike + cancel_event (cleared at the start of every generate), it is never cleared + here, so a generate still queued behind a cancelled one is skipped rather + than run — the cancel survives the queue handoff. """ os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports @@ -715,7 +752,16 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf cmd_type = cmd.get("type", "") try: if cmd_type == "generate": + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue cancel_event.clear() + # Re-check the drain after clearing: the parent sets drain_event + # then cancel_event for an unload, so if that pair landed between + # the check above and this clear, the clear just erased the unload's + # cancel. Skip here so the outgoing model is not run to completion, + # which would stall the switch until the dispatcher idle-timeout. + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue _handle_generate(backend, cmd, resp_queue, cancel_event) elif cmd_type == "load": if backend.active_model_name: @@ -918,7 +964,16 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf try: if cmd_type == "generate": + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue cancel_event.clear() + # Re-check the drain after clearing: the parent sets drain_event then + # cancel_event for an unload, so if that pair landed between the check + # above and this clear, the clear just erased the unload's cancel. Skip + # here so the outgoing model is not run to completion, which would stall + # the switch until the dispatcher idle-timeout tears the subprocess down. + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue _handle_generate(backend, cmd, resp_queue, cancel_event) elif cmd_type == "load": diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 032a6e874a..e31c03f7f8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3394,12 +3394,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() - # Unload any active Unsloth model to free VRAM + # Unload any active Unsloth model to free VRAM (off the event loop: + # unload takes _gen_lock and can wait on an in-flight stream). if unsloth_backend.active_model_name: logger.info( f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" ) - unsloth_backend.unload_model(unsloth_backend.active_model_name) + await asyncio.to_thread( + unsloth_backend.unload_model, unsloth_backend.active_model_name + ) # Inherit llama_extra_args from the previous load when the request # omits the field (the chat-settings Apply path doesn't round-trip @@ -4063,28 +4066,76 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge # A deliberate unload means "stay unloaded": drop any idle reload stash so the # next /v1 request can't resurrect this model. The idle loop unloads via the # backend directly (not this route), so clearing here never fights keep-warm. - from core.inference.llama_keepwarm import note_model_unloaded + from core.inference.llama_keepwarm import inference_lifecycle_gate, note_model_unloaded try: - # Check if the GGUF backend has this model loaded or is loading it. - llama_backend = get_llama_cpp_backend() - if llama_backend.is_active and ( - llama_backend.model_identifier == request.model_path - or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) - or not llama_backend.is_loaded + # "Stop loading" (frontend cancelLoading -> /unload) must abort a still-loading + # model promptly. /load holds the lifecycle gate for the whole (multi-minute) load, + # so gating first would make the cancel wait it out. cancel_load only tears the + # loading subprocess down (no unload command), so it is safe off-gate. + backend = get_inference_backend() + loading = getattr(backend, "get_loading_model", lambda: None)() + if ( + loading is not None + and hasattr(backend, "cancel_load") + and (request.model_path == loading or request.model_path.lower() == loading.lower()) ): - # A manual unload is a deliberate user action: tear down now even if a - # request is mid-stream (only the automatic idle loop defers to it). - llama_backend.unload_model() + if await asyncio.to_thread(backend.cancel_load, request.model_path): + note_model_unloaded() + logger.info(f"Cancelled in-flight load: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) + + # Same "stop loading" fast path for a still-loading GGUF (llama-server spawned, + # health check not yet passed). A gated unload would wait out the multi-minute + # load; unload_model() sets the cancel_event load_model polls off its own lock and + # kills the child, sending no worker command, so it is safe off-gate like + # cancel_load. The gated GGUF branch below handles the already-loaded case. Gate on + # the loading model (identifier or native label): the single llama-server loads one + # GGUF at a time, so an unload for a different model must not cancel this load. + llama_backend = get_llama_cpp_backend() + if ( + llama_backend.is_active + and not llama_backend.is_loaded + and ( + llama_backend.model_identifier == request.model_path + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) + ) + ): + await asyncio.to_thread(llama_backend.unload_model) note_model_unloaded() - logger.info(f"Unloaded GGUF model: {request.model_path}") + logger.info(f"Cancelled in-flight GGUF load: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) - # Otherwise, unload from Unsloth backend - backend = get_inference_backend() - backend.unload_model(request.model_path) - note_model_unloaded() - logger.info(f"Unloaded model: {request.model_path}") - return UnloadResponse(status = "unloaded", model = request.model_path) + # Serialize with /load under the same lifecycle gate: the Unsloth unload now runs + # off the event loop (asyncio.to_thread), so without this a concurrent /load could + # swap in a fresh subprocess mid-unload and the unload command would land on the + # new worker. The gate makes load and unload exclusive. + async with inference_lifecycle_gate(): + # Check if the GGUF backend has this model loaded or is loading it. + llama_backend = get_llama_cpp_backend() + if llama_backend.is_active and ( + llama_backend.model_identifier == request.model_path + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) + or not llama_backend.is_loaded + ): + # A manual unload is a deliberate user action: tear down now even if a + # request is mid-stream (only the automatic idle loop defers to it). + llama_backend.unload_model() + note_model_unloaded() + logger.info(f"Unloaded GGUF model: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) + + # Unload from Unsloth backend off the event loop: unload takes _gen_lock, which + # a slow SSE stream paused between tokens still holds, so a sync call would block + # the loop that drives the stream's next token and the lock release. + backend = get_inference_backend() + await asyncio.to_thread(backend.unload_model, request.model_path) + note_model_unloaded() + logger.info(f"Unloaded model: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) except Exception as e: logger.error(f"Error unloading model: {e}", exc_info = True) diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py new file mode 100644 index 0000000000..e9d0f36fe2 --- /dev/null +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -0,0 +1,1430 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""unload_model cancels an in-flight generation instead of waiting it out. + +The sequential subprocess used to queue ``unload`` behind a running ``generate``, +hanging the UI. ``unload_model`` now cancels first (the mp.Event the worker checks +each token) and takes ``_gen_lock`` before the unload round-trip. +""" + +import threading +import time + +import pytest + +from core.inference import orchestrator as orch_mod +from core.inference.orchestrator import InferenceOrchestrator + + +def _bare_orchestrator(): + """An orchestrator without the real __init__ subprocess/network.""" + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._gen_lock = threading.Lock() + o._cancel_event = threading.Event() # stands in for the mp.Event + o._drain_event = threading.Event() # stands in for the unload-drain mp.Event + o._proc = object() # truthy so _ensure_subprocess_alive reports alive + o._cmd_queue = object() + o._resp_queue = object() + o._dispatcher_thread = None + o._unload_pending = False + o.active_model_name = "m" + o.models = {"m": {}} + o.loading_models = set() + return o + + +def test_unload_cancels_inflight_generation_then_unloads(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + # A generation holds _gen_lock and releases it only once cancelled. + o._gen_lock.acquire() + + def releaser(): + o._cancel_event.wait(timeout = 5) # released only after the cancel fires + o._gen_lock.release() + + t = threading.Thread(target = releaser) + t.start() + + start = time.monotonic() + ok = o.unload_model("m") + elapsed = time.monotonic() - start + t.join(timeout = 5) + + assert ok is True + assert o._cancel_event.is_set(), "generation must be cancelled before the unload" + assert {"type": "unload", "model_name": "m"} in sent + assert o.active_model_name is None + assert "m" not in o.models + # Waited on the released-after-cancel lock, not a full generation. + assert elapsed < 2.0 + + +def test_unload_no_active_generation_unloads_normally(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + ok = o.unload_model("m") + + assert ok is True + assert {"type": "unload", "model_name": "m"} in sent + assert o.active_model_name is None + # Lock released for the next caller. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_unload_falls_back_to_shutdown_when_generation_wont_yield(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send unload when wedged")) + + # A wedged worker never releases _gen_lock, even after the cancel. + o._gen_lock.acquire() + + ok = o.unload_model("m") + + assert ok is True + assert shutdown, "should tear the subprocess down to free the GPU" + assert o.active_model_name is None + + +def test_unload_tears_down_when_compare_dispatcher_wedged(monkeypatch): + # A wedged compare-mode generation bypasses _gen_lock, so the acquire guard + # misses it and _send_cmd/_wait_response would hang on resp_queue. Unload must + # instead tear the subprocess down, like the wedged locked-generation path. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_DISPATCH_IDLE_TIMEOUT", 0.2) + + # A live dispatcher whose mailbox never drains == a wedged compare-mode gen. + o._mailbox_lock = threading.Lock() + o._mailboxes = {"req-1": object()} + + class _AliveThread: + def is_alive(self): + return True + + o._dispatcher_thread = _AliveThread() + + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send unload with a wedged dispatcher") + ) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: pytest.fail( + "must not wait on resp_queue with a wedged dispatcher" + ), + ) + + # _gen_lock is free (compare mode never took it), so the acquire guard passes. + ok = o.unload_model("m") + + assert ok is True + assert shutdown, "should tear the subprocess down to free the GPU" + assert o.active_model_name is None + assert "m" not in o.models + + +def test_consume_token_stream_bails_when_subprocess_swapped(monkeypatch): + # After a wedged-worker teardown a fresh load swaps _proc/_resp_queue; the + # still-live generation thread must detect the swap and bail, not re-block on + # the new queue while holding _gen_lock. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_subprocess_crash_message", lambda ctx: "inference subprocess restarted" + ) + + def read_one(timeout): + o._proc = object() # simulate the reload swapping the subprocess + return None + + gen = o._consume_token_stream(read_one, lambda: None, crash_context = "generation") + msg = next(gen) + + assert "restarted" in msg + with pytest.raises(StopIteration): + next(gen) + + +def test_unload_pending_clears_after_unload(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: None) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + o.unload_model("m") + + # The flag must not leak past the unload, else every later generation bails. + assert o._unload_pending is False + + +def test_generation_bails_when_unload_pending(monkeypatch): + # Winning the _gen_lock handoff mid-switch must not start on the outgoing model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + o._unload_pending = True + + out = list(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + # It released (or never held) the lock, so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_dispatched_generation_bails_when_unload_pending(monkeypatch): + # Compare-mode bypasses _gen_lock, so it must early-out on a pending switch or + # it enqueues a generate on the outgoing model and delays the unload. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: pytest.fail("must not start a generation mid-switch") + ) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") + ) + o._unload_pending = True + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + + +def test_audio_input_generation_bails_when_unload_pending(monkeypatch): + # The audio path takes _gen_lock but must also skip the outgoing model mid-switch. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") + ) + o._unload_pending = True + + out = list(o._generate_audio_input_inner(audio_array = [0.0, 0.1])) + + assert any("unloaded" in chunk.lower() for chunk in out) + # Lock released so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_audio_response_bails_when_unload_pending(monkeypatch): + # TTS (generate_audio_response) is blocking, so it RAISES rather than starting on the + # outgoing model mid-switch; it takes _gen_lock and must release it either way. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send audio generate mid-switch") + ) + o._unload_pending = True + + with pytest.raises(RuntimeError, match = "unload"): + o.generate_audio_response("hello") + + # Lock released so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +# ---------------------------------------------------------------------------- +# Preserve unload cancels across the queue handoff (drain_event) — items #1/#4. +# ---------------------------------------------------------------------------- + + +def test_worker_drain_skip_emits_cancelled_gen_done_when_draining(): + # The worker clears cancel_event at the start of every generate, so a cancel set + # while a generate is still queued would be lost when it is dequeued. drain_event + # is the durable signal: while it is set the worker skips the generate (emitting an + # immediate gen_done so the stream/mailbox drains) instead of running it. + import queue as _queue + + from core.inference.worker import _drain_skip_generate + + drain = threading.Event() + rq: _queue.Queue = _queue.Queue() + cmd = {"type": "generate", "request_id": "r1"} + + # Not draining -> run normally (do not skip, emit nothing). + assert _drain_skip_generate(cmd, rq, drain) is False + assert rq.empty() + # Missing event (older worker) -> also runs normally. + assert _drain_skip_generate(cmd, rq, None) is False + assert rq.empty() + + # Draining -> skip and emit a cancelled gen_done for this request_id. + drain.set() + assert _drain_skip_generate(cmd, rq, drain) is True + resp = rq.get_nowait() + assert resp["type"] == "gen_done" + assert resp["request_id"] == "r1" + assert resp["cancelled"] is True + + +def test_worker_generate_branches_check_drain_before_clearing_cancel(): + # Both worker command loops (MLX fast-path + GPU) must consult the drain skip + # before clearing cancel_event and running, so a queued generate can't clear an + # unload-initiated cancel and run the outgoing model to completion. Each loop + # checks the drain twice -- once before the clear and once after -- so a + # drain+cancel pair that lands in the window between them is still caught. + import inspect + + from core.inference import worker + + src = inspect.getsource(worker.run_inference_process) + assert src.count("_drain_skip_generate(cmd, resp_queue, drain_event)") == 4 + + +def test_worker_generate_rechecks_drain_after_clearing_cancel(): + # The exact interleaving item #3 describes: the drain check reads unset, then the + # parent sets drain+cancel for an unload, then the worker clears cancel_event + # (erasing that cancel). A second drain check *after* the clear catches it and + # skips the generate instead of running the outgoing model to completion. + import queue as _queue + + from core.inference.worker import _drain_skip_generate + + drain = threading.Event() + cancel = threading.Event() + rq: _queue.Queue = _queue.Queue() + cmd = {"type": "generate", "request_id": "r1"} + + # 1. Pre-clear drain check: not draining yet -> run (no skip, no emit). + assert _drain_skip_generate(cmd, rq, drain) is False + assert rq.empty() + + # 2. Parent starts an unload: sets drain, then cancel (orchestrator order). + drain.set() + cancel.set() + + # 3. Worker clears cancel at the start of the generate -- erasing the cancel. + cancel.clear() + assert not cancel.is_set() + + # 4. Post-clear drain re-check catches the erased cancel and skips. + assert _drain_skip_generate(cmd, rq, drain) is True + resp = rq.get_nowait() + assert resp["type"] == "gen_done" and resp["cancelled"] is True + + +def test_unload_sets_drain_event_during_switch_and_clears_after(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + + seen = {} + + def record_send(cmd): + # drain_event must be set for the whole unload round-trip so any generate the + # worker dequeues in this window is skipped, not run. + seen["drain_set"] = o._drain_event.is_set() + + monkeypatch.setattr(o, "_send_cmd", record_send) + + assert o.unload_model("m") is True + assert seen.get("drain_set") is True + # Cleared on exit so a later generation (e.g. unloading a non-active model, or a + # reused subprocess) is not wrongly skipped. + assert o._drain_event.is_set() is False + + +def test_unload_clears_drain_event_even_on_wedged_teardown(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send when wedged")) + + # A wedged worker never releases _gen_lock; unload tears the subprocess down. The + # real teardown nulls _drain_event, so emulate that so the finally exercises its guard. + def fake_shutdown(timeout = 5): + o._drain_event = None + + monkeypatch.setattr(o, "_shutdown_subprocess", fake_shutdown) + o._gen_lock.acquire() + + assert o.unload_model("m") is True # must not raise in the drain_event clear + + +# ---------------------------------------------------------------------------- +# Recheck the active model after the lock wait — items #2/#3. +# ---------------------------------------------------------------------------- + + +def test_generation_rechecks_model_after_lock_wait(monkeypatch): + # A request passes the pre-lock active-model check, then blocks on _gen_lock while + # an unload clears/swaps the model. Even if _unload_pending was already reset (the + # unload's finally runs after the lock release), the under-lock active-model recheck + # must make it bail instead of sending a generate to the wrong/unloaded backend. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate on a swapped/unloaded model") + ) + + reached_lock = threading.Event() + # _wait_dispatcher_idle runs after the pre-lock check and before acquiring the lock; + # signalling here means the generator captured the model and is about to block. + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + + o.active_model_name = "m" + o._unload_pending = False + o._gen_lock.acquire() # stand in for an in-flight unload holding the lock + + out: list = [] + + def run(): + out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + + t = threading.Thread(target = run) + t.start() + assert reached_lock.wait(timeout = 5) + # Unload finished: model swapped, pending already cleared. Release the lock. + o.active_model_name = "other" + o._gen_lock.release() + t.join(timeout = 5) + + assert out and any("unloaded" in chunk.lower() for chunk in out) + + +def test_generation_rechecks_model_when_unloaded_to_none(monkeypatch): + # Same race, but the unload left no active model (a plain unload, not a switch). + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate after the model was unloaded") + ) + reached_lock = threading.Event() + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + + o.active_model_name = "m" + o._unload_pending = False + o._gen_lock.acquire() + + out: list = [] + t = threading.Thread( + target = lambda: out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + ) + t.start() + assert reached_lock.wait(timeout = 5) + o.active_model_name = None + o._gen_lock.release() + t.join(timeout = 5) + + assert out and any("unloaded" in chunk.lower() for chunk in out) + + +# ---------------------------------------------------------------------------- +# Don't unload a stale model name (worker's active-model fallback) — item #5. +# ---------------------------------------------------------------------------- + + +def test_unload_of_stale_name_does_not_touch_active_model(monkeypatch): + # If the named model isn't loaded (e.g. a concurrent load already swapped in a + # different one), unload must not send a command the worker would satisfy by + # unloading its *active* model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + ) + o.active_model_name = "current" + o.models = {"current": {}} + + assert o.unload_model("stale") is True + # The active model is left intact. + assert o.active_model_name == "current" + assert "current" in o.models + + +def test_unload_matches_active_model_case_insensitively(monkeypatch): + # active_model_name can differ in case from the raw model_path a client sends + # to /unload (the load path canonicalizes casing). The stale-name guard must + # match case-insensitively too; otherwise it no-ops the unload and leaves the + # model resident while reporting success. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + o.active_model_name = "unsloth/Qwen3-4B" + o.models = {"unsloth/Qwen3-4B": {}} + + # Client unloads with the casing it originally typed, before canonicalization. + assert o.unload_model("unsloth/qwen3-4b") is True + # The guard did not no-op: an unload for the canonical active model reached + # the worker (not the raw lowercase name, so the worker matches it directly). + assert {"type": "unload", "model_name": "unsloth/Qwen3-4B"} in sent + # Local state is cleared for the canonical name, not left stale. + assert o.active_model_name is None + assert o.models == {} + + +def test_unload_of_stale_name_still_no_ops_after_case_insensitive_match(monkeypatch): + # The case-insensitive match must only rescue the active model; a genuinely + # different model name (case-insensitively too) must still no-op so the + # worker's absent-name fallback can't tear down the active model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + ) + o.active_model_name = "unsloth/Qwen3-4B" + o.models = {"unsloth/Qwen3-4B": {}} + + assert o.unload_model("unsloth/Llama-3.1-8B") is True + assert o.active_model_name == "unsloth/Qwen3-4B" + assert "unsloth/Qwen3-4B" in o.models + + +def test_load_does_not_accumulate_stale_models_defeating_the_unload_guard(monkeypatch): + # A load always spawns a fresh subprocess holding only the new model, so + # self.models must mirror that instead of accumulating the previous model's name. + # Otherwise switching A -> B leaves 'A' in self.models, so a later unload('A') + # passes the "not in self.models" guard and the worker's absent-name fallback + # unloads the *active* model B. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + monkeypatch.setattr(orch_mod.time, "sleep", lambda *_a, **_k: None) + + def _load(name): + monkeypatch.setattr( + o, + "_wait_response", + lambda expected, timeout = 300.0: { + "type": "loaded", + "success": True, + "model_info": {"identifier": name, "display_name": name}, + }, + ) + assert o.load_model(types.SimpleNamespace(identifier = name, gguf_variant = None)) is True + + _load("modelA") + _load("modelB") # switch to B without unloading A first + + # self.models mirrors the single live model; the swapped-out name is gone. + assert o.active_model_name == "modelB" + assert set(o.models) == {"modelB"} + + # A stale unload of the swapped-out model must not reach the worker (whose + # absent-name fallback would unload the active model B). + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("stale unload reached the worker")) + assert o.unload_model("modelA") is True + assert o.active_model_name == "modelB" + assert "modelB" in o.models + + +def test_unload_route_serializes_with_loads_via_lifecycle_gate(monkeypatch): + # Item #5: /unload must hold the same lifecycle gate as /load so a concurrent load + # can't swap the backend subprocess/queues mid-unload. + import asyncio + + import routes.inference as inference_route + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + class _Llama: + is_active = False + is_loaded = False + model_identifier = None + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + + unloaded: list = [] + + class _Backend: + active_model_name = "m" + models = {"m": {}} + + def unload_model(self, name): + unloaded.append(name) + return True + + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: _Backend()) + + async def scenario(): + # Hold the real gate, exactly as an in-flight /load would. + assert kw._lifecycle_lock.acquire(blocking = False) + try: + task = asyncio.ensure_future( + inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + ) + # Yield to the loop repeatedly: the route must stay blocked on the gate. + for _ in range(10): + await asyncio.sleep(0.01) + assert unloaded == [], "unload ran while the lifecycle gate was held" + assert not task.done() + finally: + kw._lifecycle_lock.release() + resp = await task + assert resp.status == "unloaded" + assert unloaded == ["m"] + + asyncio.run(scenario()) + + +# ---------------------------------------------------------------------------- +# Cancel an in-flight load OFF the lifecycle gate (Stop-loading regression). +# /load holds the gate for the whole load, so a gated /unload could never +# interrupt it; cancel_load only tears the loading subprocess down. +# ---------------------------------------------------------------------------- + + +def test_cancel_load_terminates_loading_subprocess_and_sends_no_command(monkeypatch): + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + ) + + assert o.cancel_load("m") is True + assert shutdown, "must tear the loading subprocess down" + assert "m" not in o.loading_models + assert o.active_model_name is None + # A name that is not loading -> no-op, returns False so the caller takes the gate. + assert o.cancel_load("other") is False + + +def test_cancel_load_matches_loading_model_case_insensitively(monkeypatch): + o = _bare_orchestrator() + o.loading_models = {"unsloth/Qwen3-4B"} + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) + + assert o.cancel_load("unsloth/qwen3-4b") is True + assert o.loading_models == set() + + +def test_unload_model_cancels_a_loading_model_via_cancel_load(monkeypatch): + # unload_model still cancels an in-flight load (shared logic with cancel_load). + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send a command to cancel a load") + ) + + assert o.unload_model("m") is True + assert shutdown + assert "m" not in o.loading_models + + +def test_unload_route_cancels_in_flight_load_without_waiting_on_gate(monkeypatch): + # The regression: /unload wrapped its whole body in the lifecycle gate, so the + # Stop-loading button (cancelLoading -> /unload) could not interrupt a safetensors + # load that holds the gate for its full duration. The cancel must run off-gate. + import asyncio + + import routes.inference as inference_route + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + class _Llama: + is_active = False + is_loaded = False + model_identifier = None + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + + cancelled: list = [] + + class _Backend: + active_model_name = None + models: dict = {} + + def get_loading_model(self): + return "m" + + def cancel_load(self, name): + cancelled.append(name) + return True + + def unload_model(self, name): + pytest.fail("must not take the gated unload path for a still-loading model") + + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: _Backend()) + + async def scenario(): + # Hold the real gate, exactly as an in-flight /load would. + assert kw._lifecycle_lock.acquire(blocking = False) + try: + # Even with the gate held, the loading-cancel must go through. + resp = await inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + assert resp.status == "unloaded" + assert cancelled == ["m"] + finally: + kw._lifecycle_lock.release() + + asyncio.run(scenario()) + + +# ---------------------------------------------------------------------------- +# A dispatched (compare-mode) request that races an unload must not orphan its +# mailbox after _wait_dispatcher_idle stops the dispatcher. +# ---------------------------------------------------------------------------- + + +def test_dispatched_bails_when_unload_flips_before_mailbox_registration(monkeypatch): + # The request passes the pre-work _unload_pending check, then an unload sets + # _unload_pending and _wait_dispatcher_idle stops the dispatcher (mailboxes empty) + # before this request registers its mailbox. The recheck under _mailbox_lock must + # make it bail, or the worker's skipped-generate reply has nothing to route it and + # the compare stream hangs on an orphaned mailbox. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + # Flip the unload flag after the pre-work check (626) but before mailbox + # registration -- exactly the window _wait_dispatcher_idle exploits. + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +# ---------------------------------------------------------------------------- +# Dispatched path: bail when a cleared-pending unload swapped the model or +# tore the dispatcher down during the pre-registration window -- item #2. +# ---------------------------------------------------------------------------- + + +class _AliveDispatcher: + """Stand-in dispatcher thread that reports itself alive.""" + + def is_alive(self): + return True + + +def test_dispatched_bails_when_model_swapped_before_mailbox_registration(monkeypatch): + # The request passes the pre-work checks, then a full unload+reload completes + # (clearing _unload_pending) before this request registers its mailbox. The + # under-lock recheck must notice active_model_name changed and bail, instead of + # sending a generate that lands on the swapped-in model. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + # Swap the active model after the pre-work check but before registration, + # with _unload_pending already back to False (the unload finally ran). + def swap(*a, **k): + o.active_model_name = "other" + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", swap) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate on the swapped-in model") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration(monkeypatch): + # Same window, but the unload was a same-model reload so active_model_name is + # unchanged; the give-away is that the dispatcher was stopped. Registering a + # mailbox with no dispatcher to route the reply would hang the compare stream. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + def stop_dispatcher(*a, **k): + o._dispatcher_thread = None # unload's _stop_dispatcher cleared it + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", stop_dispatcher) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate with the dispatcher stopped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +def test_dispatched_happy_path_registers_and_sends(monkeypatch): + # Guard against a false bail: with the model unchanged and the dispatcher alive, + # the recheck must let the generate through (register a mailbox and send). + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + monkeypatch.setattr( + o, "_build_generate_cmd", lambda *a, **k: {"type": "generate", "request_id": "r1"} + ) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + + # Feed one gen_done so the consumer returns promptly. + def fake_consume(read_mailbox, drainer, **k): + mbox = o._mailboxes.get("r1") + if mbox is not None: + mbox.put({"type": "gen_done", "request_id": "r1"}) + yield "" + + monkeypatch.setattr(o, "_consume_token_stream", fake_consume) + + list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert sent, "happy path must send the generate command" + assert o._mailboxes == {}, "mailbox popped in finally" + + +# ---------------------------------------------------------------------------- +# load_model observes a cancel that discarded its loading marker -- item #4. +# ---------------------------------------------------------------------------- + + +def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch): + # Stop-loading during GPU placement discards the loading marker (cancel_load) with + # no child yet to kill. load_model must observe the removal and not spawn a worker + # that loads the model after /unload already reported it unloaded. + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = set() + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + monkeypatch.setattr( + o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn a worker after a cancel") + ) + + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + + # cancel_load discards the marker while we resolve GPU placement. + def cancel_during_gpu(gpu_ids, **k): + o.loading_models.discard("m") + return ([0], "sel") + + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", cancel_during_gpu) + + class _Cfg: + identifier = "m" + + ok = o.load_model(_Cfg()) + + assert ok is False + assert o.active_model_name is None + assert o.models == {} + + +def test_load_model_proceeds_when_not_cancelled(monkeypatch): + # Guard against a false abort: an uncancelled load keeps its marker and spawns. + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = set() + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + + spawned = [] + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: spawned.append(cfg)) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: {"success": True, "model_info": {"identifier": "m"}}, + ) + + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + + class _Cfg: + identifier = "m" + + ok = o.load_model(_Cfg()) + + assert ok is True + assert spawned, "uncancelled load must spawn a worker" + assert o.active_model_name == "m" + + +def test_load_model_aborts_when_cancelled_during_spawn(monkeypatch): + # Stop-loading can land AFTER the pre-spawn marker recheck but while + # _spawn_subprocess is still creating the queues/process, so cancel_load's + # _shutdown_subprocess finds _proc not yet alive and no-ops. load_model must + # recheck the marker once the child exists and tear the orphaned worker down, + # instead of waiting for "loaded" and publishing a model /unload already + # reported as unloaded (a live subprocess nothing later reaps). + import types + + from utils import transformers_version as tv + + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = {"m"} + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + + # The cancel lands during the spawn window: cancel_load already discarded the + # marker, but its teardown no-oped because _proc was not alive yet. + def spawn_then_cancel(cfg): + o.loading_models.discard("m") + + monkeypatch.setattr(o, "_spawn_subprocess", spawn_then_cancel) + + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: pytest.fail( + "must not wait for 'loaded' after a cancel during spawn" + ), + ) + + ok = o.load_model(types.SimpleNamespace(identifier = "m", gguf_variant = None)) + + assert ok is False + assert shutdown, "must tear the orphaned worker down" + assert o.active_model_name is None + assert o.models == {} + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# /unload cancels a still-loading GGUF off the lifecycle gate -- item #1. +# ---------------------------------------------------------------------------- + + +def test_unload_cancels_loading_gguf_off_gate(monkeypatch): + # A still-loading GGUF (is_active, not is_loaded) must be cancelled off the gate: + # /load holds the lifecycle gate for the whole load, so a gated unload would wait + # it out. Assert the gate is never entered and unload_model() runs. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = False + model_identifier = "gguf-model" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None # no Unsloth load in flight -> Unsloth fast path skipped + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-model") + resp = _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert getattr(resp, "status", None) == "unloaded" + assert llama.unloaded is True, "must cancel the loading GGUF via unload_model()" + assert gate_entered["v"] is False, "must handle the loading GGUF off the lifecycle gate" + + +def test_unload_loaded_gguf_still_uses_gate(monkeypatch): + # Guard: an already-loaded GGUF (is_loaded True) is NOT caught by the off-gate + # fast path; it goes through the gate as before. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = True + model_identifier = "gguf-model" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(ri, "is_registered_native_path_label", lambda a, b: False) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-model") + resp = _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert getattr(resp, "status", None) == "unloaded" + assert llama.unloaded is True + assert gate_entered["v"] is True, "loaded GGUF unload must still take the gate" + + +def test_unload_of_mismatched_loading_gguf_skips_off_gate_fast_path(monkeypatch): + # A still-loading GGUF X (is_active, not is_loaded) must NOT be torn down by the + # off-gate fast path when /unload names a DIFFERENT model Y. The single llama-server + # can only load one GGUF at a time, so this fast path is "stop loading THIS model"; + # without a target check it fires for any in-flight GGUF and would abort an unrelated + # load (e.g. a second tab unloading Y kills the load of X). A mismatched target must + # fall through to the lifecycle gate (where, in production, it waits out X's /load and + # then no-ops) instead of taking the off-gate teardown. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = False + model_identifier = "gguf-X" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None # no Unsloth load in flight -> Unsloth fast path skipped + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(ri, "is_registered_native_path_label", lambda a, b: False) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-Y") # different from the loading model X + _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert gate_entered["v"] is True, ( + "a mismatched-target unload must not use the off-gate GGUF fast path; " + "it would cancel the wrong in-flight load" + ) + + +# ---------------------------------------------------------------------------- +# cancel_load clears its loading marker BEFORE tearing the subprocess down, so a +# racing off-gate load_model observes the cancel during the shutdown window. +# ---------------------------------------------------------------------------- + + +def test_cancel_load_clears_marker_before_shutdown(monkeypatch): + # cancel_load runs off the lifecycle gate, concurrently with a load_model that + # rechecks the loading marker before each spawn to observe the cancel. + # _shutdown_subprocess can block (tearing a live child down / joining the compare + # dispatcher), so discarding the marker only AFTER it leaves a long window in which + # that load_model reads the marker still set, passes its pre-spawn recheck, and + # spawns + loads the model after /unload already reported it cancelled. The marker + # (and local state) must be cleared before the teardown. + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = "m" + o.models = {"m": {}} + + at_shutdown = {} + + def record_shutdown(timeout = 5): + at_shutdown["marker_present"] = "m" in o.loading_models + at_shutdown["active"] = o.active_model_name + at_shutdown["models"] = dict(o.models) + + monkeypatch.setattr(o, "_shutdown_subprocess", record_shutdown) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + ) + + assert o.cancel_load("m") is True + assert at_shutdown.get("marker_present") is False, ( + "the loading marker must be cleared before _shutdown_subprocess so a concurrent " + "load_model pre-spawn recheck observes the cancel during the shutdown window" + ) + assert at_shutdown.get("active") is None + assert at_shutdown.get("models") == {} + assert "m" not in o.loading_models + assert o.active_model_name is None + assert o.models == {} + + +def test_cancel_load_reclears_state_when_racing_load_repopulates_during_teardown(monkeypatch): + # cancel_load (off the lifecycle gate) can race a load_model whose worker already + # queued its successful "loaded" reply. cancel_load discards the loading marker and + # clears the local mirrors, then tears the subprocess down; but the still-running + # load_model thread can consume that "loaded" DURING the teardown window and repopulate + # active_model_name/models. _shutdown_subprocess nulls the queues but never touches those + # mirrors, so without a second clear /unload reports success while the backend keeps + # advertising a model whose worker was just killed. cancel_load must re-clear after the + # teardown so no phantom loaded model survives. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + o._proc = None # no prior subprocess -> load_model goes straight to the spawn loop + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + + parked = threading.Event() # load_model is parked in _wait_response("loaded") + release_loaded = threading.Event() # cancel_load lets the load consume "loaded" + load_done = threading.Event() + + def blocking_wait_response(expected, timeout = 300.0): + parked.set() + assert release_loaded.wait(timeout = 5) + return { + "type": "loaded", + "success": True, + "model_info": {"identifier": "m", "display_name": "m"}, + } + + monkeypatch.setattr(o, "_wait_response", blocking_wait_response) + + load_result: dict = {} + + def run_load(): + try: + load_result["ok"] = o.load_model( + types.SimpleNamespace(identifier = "m", gguf_variant = None) + ) + except Exception as exc: # noqa: BLE001 + load_result["exc"] = exc + finally: + load_done.set() + + loader = threading.Thread(target = run_load) + loader.start() + assert parked.wait(timeout = 5), "load_model must reach _wait_response" + + # The teardown IS the window in which the racing load repopulates the mirrors: the + # marker is already discarded here, so release the load and wait for it to finish + # repopulating, mirroring the 0.5s cancel-settle inside the real _shutdown_subprocess. + def racing_shutdown(timeout = 0.5): + release_loaded.set() + assert load_done.wait(timeout = 5), "the racing load must repopulate during teardown" + + monkeypatch.setattr(o, "_shutdown_subprocess", racing_shutdown) + + assert o.cancel_load("m") is True + loader.join(timeout = 5) + + # Fail-without: load_model set active_model_name/models during racing_shutdown and + # cancel_load left them set, so the backend advertises a model whose worker was killed. + assert o.active_model_name is None, "cancel_load must not leave a repopulated active model" + assert o.models == {}, "cancel_load must not leave a repopulated models mirror" + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# A dispatched (compare-mode) request that starts the dispatcher and then bails on +# a racing unload must stop the dispatcher it started, or that orphaned dispatcher +# steals the worker's "unloaded" reply and hangs unload_model on its 300s timeout. +# ---------------------------------------------------------------------------- + + +def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch): + # The request passes the pre-work _unload_pending check and starts the dispatcher + # (none was running), then an unload sets _unload_pending so the under-lock recheck + # bails. The just-started dispatcher, left running with no mailboxes, competes with + # unload_model()'s _wait_response for the worker's "unloaded" reply off the shared + # resp_queue and drops it as unroutable, hanging the unload until its 300s timeout. + # The bail must stop the dispatcher it started. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = None # none running -> this call starts it + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + + started = {"v": False} + stopped = {"v": False} + + def fake_start(): + started["v"] = True + o._dispatcher_thread = _AliveDispatcher() + + def fake_stop(): + stopped["v"] = True + o._dispatcher_thread = None + + monkeypatch.setattr(o, "_start_dispatcher", fake_start) + monkeypatch.setattr(o, "_stop_dispatcher", fake_stop) + + # An unload flips _unload_pending after the pre-work check but before registration. + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert started["v"], "this call started the dispatcher" + assert stopped["v"], "the bail must stop the dispatcher it started (no other mailboxes)" + assert o._mailboxes == {} + + +def test_dispatched_bail_keeps_dispatcher_with_other_active_mailbox(monkeypatch): + # Guard against over-stopping: if another compare request registered a mailbox on the + # dispatcher this call started, the bail must NOT stop it, or that request's token + # routing dies mid-stream. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: setattr(o, "_dispatcher_thread", _AliveDispatcher()) + ) + monkeypatch.setattr( + o, + "_stop_dispatcher", + lambda: pytest.fail("must not stop a dispatcher another compare request is using"), + ) + + # A concurrent compare request registers its mailbox, then an unload flips the flag. + def flip(*a, **k): + o._mailboxes["other"] = object() + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert set(o._mailboxes) == {"other"}, "the other request's mailbox is untouched" + + +def test_dispatched_bail_keeps_preexisting_dispatcher(monkeypatch): + # Guard: if the dispatcher was already running before this request (an earlier compare + # request started it), a bail must not stop it even with no mailboxes now -- this + # request did not start it and another may re-use it. Only the call that starts an + # otherwise-idle dispatcher during the race is responsible for stopping it. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() # already running + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + monkeypatch.setattr( + o, "_stop_dispatcher", lambda: pytest.fail("must not stop a pre-existing dispatcher") + ) + + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + + +# ---------------------------------------------------------------------------- +# load_model rechecks the loading marker AFTER _wait_response("loaded") and +# BEFORE publishing -- item #6. cancel_load's post-teardown re-clear only wipes a +# repopulation that lands during its shutdown; a publish that lands after +# cancel_load returns survives it, so the recheck must abort the publish itself. +# ---------------------------------------------------------------------------- + + +def test_load_model_aborts_publish_when_cancelled_after_wait_response(monkeypatch): + # cancel_load (off the lifecycle gate) discards the loading marker BEFORE its teardown + # and re-clears the mirrors AFTER it. A racing load_model can consume its worker's + # already-queued "loaded" reply and reach the publish block only AFTER cancel_load has + # fully returned -- so cancel_load's post-teardown re-clear cannot undo that publish. + # Without a marker recheck between _wait_response("loaded") and the publish, load_model + # advertises active_model_name/models for a model /unload already reported cancelled, + # over a subprocess cancel_load just killed. The recheck must observe the discarded + # marker and abort the publish. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + o._proc = None # no prior subprocess -> load_model goes straight to the spawn loop + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + # cancel_load tears the worker down; a no-op keeps the test off real subprocesses. + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) + + parked = threading.Event() # load_model reached _wait_response("loaded") + cancel_done = threading.Event() # cancel_load fully returned (marker discarded + re-clear) + load_done = threading.Event() + + def blocking_wait_response(expected, timeout = 300.0): + parked.set() + # Do not consume "loaded" until cancel_load has fully returned, so the publish + # would land AFTER cancel_load's post-teardown re-clear -- the window the + # re-clear alone cannot cover. + assert cancel_done.wait(timeout = 5) + return { + "type": "loaded", + "success": True, + "model_info": {"identifier": "m", "display_name": "m"}, + } + + monkeypatch.setattr(o, "_wait_response", blocking_wait_response) + + load_result: dict = {} + + def run_load(): + try: + load_result["ok"] = o.load_model( + types.SimpleNamespace(identifier = "m", gguf_variant = None) + ) + except Exception as exc: # noqa: BLE001 + load_result["exc"] = exc + finally: + load_done.set() + + loader = threading.Thread(target = run_load) + loader.start() + assert parked.wait(timeout = 5), "load_model must reach _wait_response" + + # cancel_load runs to completion while the load is parked: it discards the marker and + # re-clears the mirrors (post-teardown), then returns. Only then let the load consume + # "loaded" and attempt to publish. + assert o.cancel_load("m") is True + cancel_done.set() + + loader.join(timeout = 5) + assert load_done.is_set() + + # Fail-without: load_model published active_model_name/models for 'm' AFTER cancel_load + # returned, advertising a cancelled model over a killed subprocess. + assert load_result.get("ok") is False, "the cancelled load must not report success" + 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