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

fix(studio): honor MLX adapter state in compare mode (#7196)

Browse source 
* fix(studio): add MLX adapter state control

* fix(studio): honor MLX adapter comparison state

* fix(studio): keep enabled MLX adapters permissive

* Studio: preserve public error message on MLX compare-mode adapter failures

generate_with_adapter_control raised a plain RuntimeError, which the compare
route handled with the generic handler that drops the operational message.
Raise GenStreamErrorRaised(public=chunk.public) instead and catch it in the
streaming and non-streaming consumers, matching the safetensors tool loop, so
errors like 'model is being unloaded' surface their real message.

* Studio: re-emit VLM think prefill inside the adapter context

The compare-mode merge dropped _generate_vlm's upfront yield of the prefilled
<think> block. Restore it as the first snapshot inside the lock+adapter context
(matching _generate_text) so the UI renders the thinking block during prefill
and a cancel/error before the first token does not drop it. Adds a regression
test asserting the prefill is emitted first, after entering the adapter context.

---------

Co-authored-by: danielhanchen <unslothshared@gmail.com>
This commit is contained in:
Long Yixing 2026-07-19 15:27:55 +08:00 • committed by GitHub
commit 4e4af72b9c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 442 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

93
studio/backend/core/inference/mlx_inference.py
View file

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

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

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

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

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

14
studio/backend/routes/inference.py
View file

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

267
studio/backend/tests/test_mlx_inference_backend.py
View file

@ -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 <think> 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: "<think>\n",
)
prompt_utils = SimpleNamespace(
MODEL_CONFIG = {"deepseek_vl_v2": object()},
apply_chat_template = lambda *_a, **_k: "<image> 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: "<image> 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) == "<think>\n"
assert order == ["adapter_enter"]
# Subsequent snapshots are cumulative (prefill + generated text).
assert next(gen) == "<think>\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": "<image> 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] == "<image> 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: "<think>",
)
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 = [

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

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