diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index e7a90b4307..e78c93b6f3 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -8,6 +8,7 @@ instead of torch/transformers for model loading and generation. import json import os import threading +from contextlib import contextmanager from typing import Optional, Generator from core.inference.message_content import content_to_text from core.inference.runtime_context import runtime_context_length @@ -20,6 +21,63 @@ from loggers import get_logger logger = get_logger(__name__) +def _mlx_adapter_modules(model): + """Return bypassable adapter entries and unsupported wrapper paths.""" + adapters = [] + unsupported = [] + for path, module in model.named_modules(): + if not path or not (hasattr(module, "lora_a") and hasattr(module, "lora_b")): + continue + base = getattr(module, "linear", None) + if base is None: + base = getattr(module, "embedding", None) + if base is None: + unsupported.append(path) + else: + adapters.append((path, module, base)) + return adapters, unsupported + + +@contextmanager +def _temporary_mlx_adapter_state(model, use_adapter): + """Select base or adapter modules for one request, then restore the tree.""" + if use_adapter is None: + yield + return + if isinstance(use_adapter, str): + raise NotImplementedError( + "Unsloth MLX: named adapter selection is not supported; use True for " + "the loaded adapter or False for the base model." + ) + if use_adapter is not True and use_adapter is not False: + raise TypeError("Unsloth MLX: use_adapter must be None, True, False, or a string.") + + adapters, unsupported = _mlx_adapter_modules(model) + if use_adapter is True: + if not adapters and not unsupported: + logger.warning("MLX adapter requested, but the active model has no adapter layers") + yield + return + if unsupported: + raise RuntimeError( + "Unsloth MLX: cannot disable adapter layers without their base modules: " + + ", ".join(unsupported[:5]) + ) + if not adapters: + yield + return + + from mlx.utils import tree_unflatten + + base_modules = tree_unflatten([(path, base) for path, _, base in adapters]) + adapter_modules = tree_unflatten([(path, wrapper) for path, wrapper, _ in adapters]) + try: + model.update_modules(base_modules) + yield + finally: + model.update_modules(adapter_modules) + + def _mlx_vlm_model_config(model): """Return the loaded MLX model config and its type, preferring whichever of config / _config actually carries a model_type.""" @@ -508,6 +566,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ) -> Generator[str, None, None]: if self._model is None: raise RuntimeError("No model loaded") @@ -552,6 +611,7 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, + _adapter_state = _adapter_state, ) else: stream = self._generate_text( @@ -568,6 +628,7 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, + _adapter_state = _adapter_state, ) yield from stream @@ -587,6 +648,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ): from mlx_lm import stream_generate from mlx_lm.sample_utils import make_sampler, make_logits_processors @@ -635,10 +697,6 @@ class MLXInferenceBackend: think_prefix = detect_think_prefill( prompt, getattr(self._tokenizer, "all_special_tokens", None) ) - # Emit it before the first token so the block renders during prefill. - if think_prefix: - yield think_prefix - sampler = make_sampler( temp = temperature, top_p = top_p, @@ -680,9 +738,12 @@ class MLXInferenceBackend: type(self._model).__name__, type(self._tokenizer).__name__, ) - with self._generation_lock: + with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): final_response = None try: + # Enter request-scoped model state before yielding any response. + if think_prefix: + yield think_prefix gen_kwargs = dict( prompt = prompt, max_tokens = max_new_tokens, @@ -749,6 +810,7 @@ class MLXInferenceBackend: reasoning_effort = None, preserve_thinking = None, presence_penalty = 0.0, + _adapter_state = None, ): from mlx_vlm import stream_generate as vlm_stream @@ -852,9 +914,6 @@ class MLXInferenceBackend: # Re-emit an open prefill from the prompt (see _generate_text). cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) - # Emit it before the first token so the block renders during prefill. - if cumulative: - yield cumulative logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), @@ -891,9 +950,18 @@ class MLXInferenceBackend: def _stream_vlm_snapshots(): nonlocal cumulative - with self._generation_lock: + # Hold the generation lock AND the request-scoped adapter state for the + # whole stream so Base-vs-LoRA compare mode honors use_adapter and the + # wrapper tree is restored on completion, cancellation, or close. + with self._generation_lock, _temporary_mlx_adapter_state(self._model, _adapter_state): final_response = None try: + # Emit any prefilled block before the first token so the + # UI renders it during prefill, matching _generate_text. Done + # inside the adapter context so an unsupported request raises + # before any output escapes. + if cumulative: + yield cumulative for response in vlm_stream( self._model, self._processor, @@ -927,8 +995,11 @@ class MLXInferenceBackend: cancel_event = None, **gen_kwargs, ) -> Generator[str, None, None]: - # MLX LoRA adapter toggling not yet supported; generate normally - yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) + yield from self.generate_chat_response( + cancel_event = cancel_event, + _adapter_state = use_adapter, + **gen_kwargs, + ) def reset_generation_state(self): import mlx.core as mx diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 3afda74411..eaa474d9b8 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -1502,14 +1502,27 @@ class InferenceOrchestrator: Uses the dispatcher path (no _gen_lock) so compare-mode requests don't block each other; the subprocess serializes them via its - sequential command loop. + sequential command loop. Backend failures raise instead of becoming + assistant text. """ - yield from self._generate_dispatched( + stream = self._generate_dispatched( use_adapter = use_adapter, cancel_event = cancel_event, stats_holder = stats_holder, **gen_kwargs, ) + try: + for chunk in stream: + if isinstance(chunk, GenStreamError): + # Preserve the public/operational flag so the route can surface + # the real message (e.g. "model is being unloaded") instead of a + # generic error. Mirrors the safetensors tool loop's _single_turn. + raise GenStreamErrorRaised(str(chunk), public = chunk.public) + yield chunk + finally: + close = getattr(stream, "close", None) + if callable(close): + close() def _generate_inner( self, diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index e4628dcea8..9f301ba37e 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -513,20 +513,25 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: logger.info("Starting text generation for request_id=%s", request_id) - for cumulative_text in generator: - # cancel_event is an mp.Event — checked instantly, no queue polling. - if cancel_event.is_set(): - logger.info("Generation cancelled for request %s", request_id) - break + try: + for cumulative_text in generator: + # cancel_event is an mp.Event — checked instantly, no queue polling. + if cancel_event.is_set(): + logger.info("Generation cancelled for request %s", request_id) + break - _send_response( - resp_queue, - { - "type": "token", - "request_id": request_id, - "text": cumulative_text, - }, - ) + _send_response( + resp_queue, + { + "type": "token", + "request_id": request_id, + "text": cumulative_text, + }, + ) + finally: + close = getattr(generator, "close", None) + if callable(close): + close() _send_response( resp_queue, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 52ea1f86a3..9299e26d56 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -9443,6 +9443,13 @@ async def openai_chat_completions( backend.reset_generation_state() api_monitor.finish(monitor_id, "cancelled") raise + except GenStreamErrorRaised as exc: + # Adapter-controlled (compare-mode) backend failure. Honor the + # public flag so operational errors surface their real message. + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}}) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI streaming: {e}", exc_info = True) @@ -9591,6 +9598,13 @@ async def openai_chat_completions( except HTTPException: raise + except GenStreamErrorRaised as exc: + # Adapter-controlled (compare-mode) backend failure. Honor the public + # flag so operational errors surface their real message. + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI completion: {e}", exc_info = True) diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index fa50cd84d6..7b8aefb722 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -2,6 +2,7 @@ import sys import types +from contextlib import contextmanager from types import SimpleNamespace import pytest @@ -40,12 +41,16 @@ class _DummyModel: def _install_fake_mlx(monkeypatch): mlx_pkg = types.ModuleType("mlx") mlx_core = types.ModuleType("mlx.core") + mlx_utils = types.ModuleType("mlx.utils") mlx_core.metal = _DummyMetal() mlx_core.set_wired_limit = _DummyMX.set_wired_limit mlx_core.device_info = _DummyMX.device_info + mlx_utils.tree_unflatten = dict mlx_pkg.core = mlx_core + mlx_pkg.utils = mlx_utils monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + monkeypatch.setitem(sys.modules, "mlx.utils", mlx_utils) def _install_fake_fast_mlx(monkeypatch, calls): @@ -68,6 +73,99 @@ def _install_fake_fast_mlx(monkeypatch, calls): monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.loader", mlx_loader) +class _AdapterTree: + def __init__(self, modules): + self.modules = dict(modules) + + def named_modules(self): + return list(self.modules.items()) + + def update_modules(self, modules): + self.modules.update(modules) + + +def test_temporary_mlx_adapter_state_bypasses_and_restores_wrappers(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + base = object() + wrapper = SimpleNamespace(lora_a = object(), lora_b = object(), linear = base, m = object()) + model = _AdapterTree({"model.layers.0.proj": wrapper}) + + with pytest.raises(RuntimeError, match = "generation failed"): + with _temporary_mlx_adapter_state(model, False): + assert model.modules["model.layers.0.proj"] is base + raise RuntimeError("generation failed") + assert model.modules["model.layers.0.proj"] is wrapper + + +def test_temporary_mlx_adapter_state_validates_requests(): + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + wrapper = SimpleNamespace(lora_a = object(), lora_b = object(), embedding = object()) + model = _AdapterTree({"embed_tokens": wrapper}) + with _temporary_mlx_adapter_state(model, True): + assert model.modules["embed_tokens"] is wrapper + with pytest.raises(NotImplementedError, match = "named adapter"): + with _temporary_mlx_adapter_state(model, "other"): + pass + + base_model = _AdapterTree({"proj": object()}) + with _temporary_mlx_adapter_state(base_model, None): + pass + with _temporary_mlx_adapter_state(base_model, True): + pass + + unsupported = _AdapterTree({"proj": SimpleNamespace(lora_a = object(), lora_b = object())}) + with _temporary_mlx_adapter_state(unsupported, True): + pass + with pytest.raises(RuntimeError, match = "without their base modules"): + with _temporary_mlx_adapter_state(unsupported, False): + pass + + +def test_temporary_mlx_adapter_state_uses_real_mlx_module_tree(): + nn = pytest.importorskip("mlx.nn") + pytest.importorskip("mlx_lm") + from mlx_lm.models.switch_layers import SwitchLinear + from mlx_lm.tuner.dora import DoRALinear + from mlx_lm.tuner.lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear + + from core.inference.mlx_inference import _temporary_mlx_adapter_state + + class _Layer(nn.Module): + def __init__(self): + super().__init__() + quantized = nn.QuantizedLinear.from_linear(nn.Linear(32, 32), group_size = 32, bits = 4) + self.quantized_proj = LoRALinear.from_base(quantized) + self.dora_proj = DoRALinear.from_base(nn.Linear(4, 4)) + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.layers = [_Layer()] + self.embed_tokens = LoRAEmbedding.from_base(nn.Embedding(16, 4)) + self.experts = LoRASwitchLinear.from_base(SwitchLinear(4, 4, 2)) + + model = _Model() + wrappers = { + path: module + for path, module in model.named_modules() + if hasattr(module, "lora_a") and hasattr(module, "lora_b") + } + bases = { + path: getattr(module, "linear", getattr(module, "embedding", None)) + for path, module in wrappers.items() + } + + with _temporary_mlx_adapter_state(model, False): + live = dict(model.named_modules()) + assert all(live[path] is base for path, base in bases.items()) + + restored = dict(model.named_modules()) + assert all(restored[path] is wrapper for path, wrapper in wrappers.items()) + + def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): _install_fake_mlx(monkeypatch) calls = [] @@ -333,10 +431,87 @@ def test_mlx_generate_chat_response_accepts_template_kwargs(): ), f"{name!r} must default to None so existing callers stay valid" +def test_mlx_vlm_reemits_think_prefill_inside_adapter_context(monkeypatch): + """A prefilled block must be re-emitted as the first VLM snapshot, + inside the adapter context (so unsupported requests still raise first), so + the UI renders the thinking block during prefill and a pre-first-token + cancel does not drop it. Mirrors _generate_text.""" + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend + + order = [] + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + order.append("adapter_enter") + try: + yield + finally: + order.append("adapter_exit") + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) + monkeypatch.setattr( + "core.inference.chat_template_helpers.detect_think_prefill", + lambda *_a, **_k: "\n", + ) + + prompt_utils = SimpleNamespace( + MODEL_CONFIG = {"deepseek_vl_v2": object()}, + apply_chat_template = lambda *_a, **_k: " model-aware", + ) + mlx_vlm = types.ModuleType("mlx_vlm") + mlx_vlm.prompt_utils = prompt_utils + + def _vlm_stream(*_a, **_k): + # The prefill must have been emitted before any generated token. + assert order[-1] == "adapter_enter" + yield SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1) + + mlx_vlm.stream_generate = _vlm_stream + monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda _t, _m, **_k: " model-aware", + ) + + backend = MLXInferenceBackend() + backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"}) + backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) + args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) + + gen = backend._generate_vlm(*args, _adapter_state = False) + # First snapshot is the prefill alone, emitted after entering the adapter context. + assert next(gen) == "\n" + assert order == ["adapter_enter"] + # Subsequent snapshots are cumulative (prefill + generated text). + assert next(gen) == "\nok" + gen.close() + assert order == ["adapter_enter", "adapter_exit"] + + def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend calls = {"generic": [], "model": [], "stream": []} + adapter_events = [] + adapter_active = {"value": False} + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + adapter_events.append(("enter", state)) + adapter_active["value"] = True + try: + yield + finally: + adapter_active["value"] = False + adapter_events.append(("exit", state)) + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) state = {"generic": "serialized", "model": " model-aware"} prompt_utils = SimpleNamespace( MODEL_CONFIG = {"deepseek_vl_v2": object()}, @@ -346,10 +521,13 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): ) mlx_vlm = types.ModuleType("mlx_vlm") mlx_vlm.prompt_utils = prompt_utils - mlx_vlm.stream_generate = lambda *_args, **kwargs: ( - calls["stream"].append((_args, kwargs)) - or iter([SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)]) - ) + + def _vlm_stream(*args, **kwargs): + assert adapter_active["value"] + calls["stream"].append((args, kwargs)) + yield SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1) + + mlx_vlm.stream_generate = _vlm_stream monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) def generic(_target, _messages, **kwargs): @@ -369,7 +547,11 @@ def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) tools = [{"function": {"name": "search"}}] - assert list(backend._generate_vlm(*args)) == ["ok"] + generator = backend._generate_vlm(*args, _adapter_state = False) + assert next(generator) == "ok" + assert adapter_active["value"] and backend._generation_lock.locked() + generator.close() + assert adapter_events == [("enter", False), ("exit", False)] assert calls["model"][0]["num_images"] == 1 assert calls["stream"][0][0][2] == " model-aware" with pytest.raises(RuntimeError, match = "dropping requested tools"): @@ -449,7 +631,10 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): """Mac text path must route through apply_chat_template_for_generation so reasoning / tool kwargs reach the tokenizer.""" _install_fake_mlx(monkeypatch) - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference import mlx_inference + + MLXInferenceBackend = mlx_inference.MLXInferenceBackend + real_adapter_state = mlx_inference._temporary_mlx_adapter_state # The text path renders once with tools, then the native-template fallback makes a second no- # tools probe call (tools=None) to detect whether the template dropped the schema. @@ -474,11 +659,31 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): mlx_lm_sample.make_sampler = lambda **_kw: object() mlx_lm_sample.make_logits_processors = lambda **_kw: None + adapter_events = [] + adapter_active = {"value": False} + stream_state = {"fail": False} + + @contextmanager + def _adapter_state(_model, state): + assert backend._generation_lock.locked() + adapter_events.append(("enter", state)) + adapter_active["value"] = True + try: + yield + finally: + adapter_active["value"] = False + adapter_events.append(("exit", state)) + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", _adapter_state) + class _Resp: def __init__(self, tok): self.token = tok def _stream_generate(_model, _tokenizer, **_kw): + assert adapter_active["value"] + if stream_state["fail"]: + raise RuntimeError("generation failed") yield _Resp(1) mlx_lm_pkg.stream_generate = _stream_generate @@ -500,17 +705,45 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): backend._tokenizer = _Tok() backend._is_vlm = False - out = list( - backend.generate_chat_response( - messages = [{"role": "user", "content": "ping"}], - tools = [{"function": {"name": "web_search"}}], - enable_thinking = True, - reasoning_effort = "medium", - preserve_thinking = True, - max_new_tokens = 1, - ) + generator = backend.generate_with_adapter_control( + use_adapter = False, + messages = [{"role": "user", "content": "ping"}], + tools = [{"function": {"name": "web_search"}}], + enable_thinking = True, + reasoning_effort = "medium", + preserve_thinking = True, + max_new_tokens = 1, ) - assert out == ["hi"] + assert next(generator) == "hi" + assert adapter_active["value"] and backend._generation_lock.locked() + generator.close() + assert adapter_events == [("enter", False), ("exit", False)] + stream_state["fail"] = True + with pytest.raises(RuntimeError, match = "generation failed"): + list( + backend.generate_with_adapter_control( + use_adapter = False, + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 1, + ) + ) + assert adapter_events[-2:] == [("enter", False), ("exit", False)] + assert not backend._generation_lock.locked() + + monkeypatch.setattr(mlx_inference, "_temporary_mlx_adapter_state", real_adapter_state) + monkeypatch.setattr( + "core.inference.chat_template_helpers.detect_think_prefill", + lambda *_args, **_kwargs: "", + ) + stream_state["fail"] = False + named = backend.generate_with_adapter_control( + use_adapter = "named", + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 1, + ) + with pytest.raises(NotImplementedError, match = "named adapter"): + next(named) + assert not adapter_active["value"] and not backend._generation_lock.locked() # The toggled kwargs must reach the chat-template helper on the real render # (one of the calls carries the tools; the fallback probe passes tools=None). tool_renders = [ diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index fb80b6d061..3a36500aee 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -34,6 +34,70 @@ def _bare_orchestrator(): return o +def test_adapter_control_raises_stream_errors(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr( + o, + "_generate_dispatched", + lambda **_kwargs: iter([orch_mod.GenStreamError("Error: adapter failed")]), + ) + + with pytest.raises(RuntimeError, match = "adapter failed"): + list(o.generate_with_adapter_control(use_adapter = False)) + + closed = [] + + def _stream(**_kwargs): + try: + yield "token" + yield "late token" + finally: + closed.append(True) + + monkeypatch.setattr(o, "_generate_dispatched", _stream) + generator = o.generate_with_adapter_control(use_adapter = False) + assert next(generator) == "token" + generator.close() + assert closed == [True] + + +def test_worker_closes_cancelled_generator_before_gen_done(): + from core.inference.worker import _handle_generate + + events = [] + + class _Backend: + last_generation_stats = None + + def generate_with_adapter_control(self, **_kwargs): + try: + yield "token" + yield "late token" + finally: + events.append("closed") + + class _Responses: + def __init__(self): + self.items = [] + + def put(self, item): + if item["type"] == "gen_done": + assert events == ["closed"] + self.items.append(item) + + responses = _Responses() + cancel = threading.Event() + cancel.set() + _handle_generate( + _Backend(), + {"request_id": "r1", "messages": [], "use_adapter": False}, + responses, + cancel, + ) + + assert [item["type"] for item in responses.items] == ["gen_done"] + + def test_unload_cancels_inflight_generation_then_unloads(monkeypatch): o = _bare_orchestrator() monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)