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

[Lyxot]

Summary

• make Studio's MLX Base-vs-LoRA comparison honor use_adapter=False and use_adapter=True
• switch adapter state request-by-request for both text and VLM generation while holding the existing generation lock
• restore the loaded adapter after normal completion, cancellation, explicit generator close, or generation failure
• close worker-owned generators before emitting terminal responses and propagate adapter-controlled backend failures as errors instead of assistant text

Problem

Studio already sends use_adapter=False for the base side of a comparison and use_adapter=True for the LoRA side. The CUDA backend applies that state through PEFT, but the MLX backend accepted the argument and ignored it, so both comparison panes could generate with the loaded adapter. Sampling randomness could make the outputs appear different and hide the invalid comparison.

The fix also has to preserve the base model exactly. Setting each LoRA scale to zero is insufficient for trained DoRA adapters because their learned magnitude parameters can still change the output. Adapter selection must therefore bypass the complete wrapper rather than numerically suppressing only its low-rank delta.

During review of the request lifecycle, two related failure modes were identified. A cancelled subprocess generation could report gen_done while its backend generator remained suspended inside the adapter context, and GenStreamError values from adapter-controlled generation could be treated as successful model text because the type subclasses str.

Implementation

The MLX backend now discovers live adapter wrappers through their lora_a and lora_b parameters and retains each wrapper's underlying linear or embedding base module. For use_adapter=False, it temporarily replaces wrappers through MLX's tree_unflatten and update_modules APIs. The original wrappers are restored in finally. use_adapter=True keeps the loaded adapter active, None preserves existing behavior, and named multi-adapter selection fails explicitly because MLX currently loads a single adapter repository.

Text and VLM generation pass a private adapter state into their streaming implementations. The existing generation lock is acquired before changing the module tree, and the adapter context remains active for the complete stream. Reasoning-prefill output is emitted only after state validation, preventing partial successful output for an unsupported request.

The inference worker now closes the generator it owns before sending gen_done, which guarantees adapter restoration before terminal state is visible to the parent process. The orchestrator converts GenStreamError from adapter-controlled generation into an exception while preserving close propagation to its inner dispatched stream, allowing streaming and non-streaming routes to use their normal error handling.

Behavior and compatibility

• ordinary MLX requests with use_adapter=None are unchanged
• base or merged models remain permissive when use_adapter=True is supplied without adapter wrappers
• CUDA inference behavior is unchanged
• text and VLM adapters share the same state mechanism
• quantized linear, DoRA linear, embedding, and SwitchLinear wrappers are covered by a real MLX module-tree test
• named MLX adapter selection remains out of scope and now fails clearly instead of being ignored

Validation

python -m pytest studio/backend/tests/test_mlx_inference_backend.py studio/backend/tests/test_orchestrator_unload_cancel.py studio/backend/tests/test_presence_penalty.py -q

Result: 83 passed.

The repository's Ruff and keyword-spacing pre-commit hooks also pass, and git diff --check origin/main...HEAD is clean.

[gemini-code-assist]

Code Review

This pull request introduces dynamic MLX LoRA adapter toggling on a per-request basis. It adds a context manager _temporary_mlx_adapter_state to temporarily swap adapter modules with their base modules when an adapter is disabled, and updates the text and VLM generation pipelines to support this state. Additionally, it ensures proper generator cleanup upon cancellation or completion in both the orchestrator and worker. The feedback recommends refactoring _mlx_adapter_modules to defer raising a RuntimeError for unsupported adapter layers, ensuring that requests where use_adapter is True do not fail unnecessarily when unsupported layers are present.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Copilot]

Pull request overview

This PR fixes Studio’s MLX “Base vs LoRA” compare mode by making the MLX backend actually honor use_adapter=True/False request-by-request, ensuring the base pane can bypass LoRA wrappers entirely (important for DoRA correctness). It also tightens the generation lifecycle so adapter state is always restored before terminal responses are emitted and adapter-controlled backend failures propagate as real errors instead of being treated as assistant text.

Changes:

• Add MLX request-scoped adapter switching via a _temporary_mlx_adapter_state(...) context that swaps LoRA wrappers with their base modules under the existing generation lock.
• Ensure generator/stream cleanup happens before gen_done (worker) and convert GenStreamError chunks into raised exceptions (orchestrator), while still propagating .close().
• Add targeted tests covering adapter swapping/restoration, stream error raising, and generator closure behavior on cancellation.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
studio/backend/core/inference/mlx_inference.py	Implements request-scoped MLX adapter bypass/restore and threads adapter state through text/VLM streaming under the generation lock.
studio/backend/core/inference/orchestrator.py	Raises on GenStreamError during adapter-controlled streaming and ensures dispatched streams are closed in finally.
studio/backend/core/inference/worker.py	Closes the backend-owned generator in finally before emitting gen_done, ensuring adapter restoration is complete first.
studio/backend/tests/test_mlx_inference_backend.py	Adds unit and integration-style tests validating adapter swapping/restoration and adapter-state lifetime for text/VLM generation.
studio/backend/tests/test_orchestrator_unload_cancel.py	Adds regression tests for raising streamed adapter errors and for closing cancelled generators before gen_done.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[danielhanchen]

Thank you for the PR! The goal here is to make Studio's MLX Base-vs-LoRA compare mode actually honor each pane's requested adapter state. As a summary, this PR makes the MLX backend temporarily swap each live LoRA/DoRA wrapper for its underlying base linear/embedding module when use_adapter=False (restoring it in a finally on normal exit, cancellation, error, and generator close), keeps the adapter active for use_adapter=True, leaves use_adapter=None unchanged, closes the worker's text generator before gen_done so the adapter is restored before the terminal message is visible, and raises adapter-path GenStreamError values instead of streaming them as assistant text.

I reviewed and tested this end to end. The core fix is correct and safe to merge.

Before vs after

• Before: the MLX backend accepted use_adapter and ignored it, so both compare panes generated through the loaded adapter and sampling randomness masked the invalid comparison. Zeroing the LoRA scale would not have been enough for a trained DoRA adapter, whose learned magnitude still perturbs the base output, so bypassing the whole wrapper is the right call.
• After: use_adapter=False swaps in the exact base module, so the base pane is genuinely adapter free while the LoRA pane keeps the adapter.

Validation

• Real MLX module-tree behavior confirmed on Apple Silicon and on the Linux CPU MLX wheel: the focused suite is 83 passed with test_temporary_mlx_adapter_state_uses_real_mlx_module_tree executed (not skipped).
• Numeric check on real LoRALinear, quantized LoRALinear, LoRAEmbedding, and DoRALinear with non-zero adapter weights: the base pane output equals the true base module, the LoRA pane output equals the adapter, and the two differ. Restore is exact after normal exit, an exception inside the context, and generator close.
• Backwards compatible with Studio's minimum MLX (mlx / mlx-lm 0.22.0): update_modules, tree_unflatten, and the .linear / .embedding base attributes on the LoRA/DoRA wrappers all exist and are unchanged from 0.22.0 through 0.32.0, and the approach matches mlx-lm's own linear_to_lora_layers pattern.
• No regression to existing paths: the CUDA backend and use_adapter=None are untouched, and the rest of the Studio backend suite passes. Cross-checked on Linux, macOS, and Windows (where MLX is absent and the three MLX tests skip cleanly).

Non-blocking follow-ups (not required for this PR)

Reviewers	Severity	Finding
7/12	Low	The close-before-gen_done generator cleanup is applied to the text path (worker.py _handle_generate) but not the audio-input / Whisper path (_handle_generate_audio_input). This predates the PR and does not involve an adapter, but the ordering is the same, so it is worth mirroring in a follow-up so a cancelled ASR generator's finally (thread join and lock release) runs before gen_done.
4/12	Nit	test_temporary_mlx_adapter_state_uses_real_mlx_module_tree asserts only module identity; if wrapper discovery ever returns an empty set both all(...) checks pass vacuously. Consider assert len(wrappers) == 4.
3/12	Nit	Consider a small numeric test with non-zero adapter weights so the base vs adapter outputs are proven to differ, not just the module identities.

Note: the failing Repo tests (CPU) check is test_sidebar_account_block_uses_leading_tight, a frontend sidebar-CSS assertion unrelated to this PR. It also fails on main at the base commit, so it is pre-existing.

Nice work, this is a clean and correctly scoped fix.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 619035edce

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Re-emit VLM think prefill before streaming tokens

For VLM prompts that end with an open <think> prefill, this path now only seeds cumulative and then waits for vlm_stream to produce the first response. Unlike _generate_text, the prefilled reasoning marker is no longer yielded after entering the request-scoped adapter state, so the UI won't render the thinking block during prefill and a cancel/error before the first generated token drops the prefilled reasoning entirely. Yield the non-empty cumulative inside this context before iterating vlm_stream.

Useful? React with 👍 / 👎.

[danielhanchen]

Good catch. Fixed in 65e6880: _generate_vlm now re-emits the prefilled block as the first snapshot inside the lock plus adapter context, matching _generate_text, so it renders during prefill and a cancel or error before the first token no longer drops it. Added a regression test asserting the prefill is emitted first, after entering the adapter context.

[chatgpt-codex-connector]

Detect legacy MLX-VLM LoRA wrappers

When a VLM adapter was saved in the legacy MLX-VLM format (adapter_config.json with rank/alpha), the runtime wrapper exposes A/B and original_layer rather than lora_a/lora_b and linear. This discovery loop skips those layers completely, so use_adapter=False takes the not adapters path and the base side of compare mode still generates with the adapter active. Include that wrapper shape (or fail loudly) so affected MLX-VLM adapters do not silently produce invalid comparisons.

Useful? React with 👍 / 👎.

[danielhanchen]

I checked the mlx-vlm source and this does not apply to the supported versions. The mlx-vlm LoRA and DoRA wrappers (mlx_vlm/trainer/lora_layers.py and dora_layers.py) expose lora_a / lora_b with the base module on .linear or .embedding, identical to mlx-lm, on both main and the pinned floor v0.4.4. There is no A / B / original_layer wrapper anywhere in mlx-vlm, so _mlx_adapter_modules discovers these layers and use_adapter=False correctly bypasses them. Studio also pins mlx-vlm>=0.4.4, so an older format is not loaded.

[danielhanchen]

Pushed two fixes on top of the main merge:

• 65e6880: re-emit the VLM <think> prefill inside the lock plus adapter context (addresses the Codex P2), with a regression test.
• e5aad81: preserve the public error message on compare-mode adapter failures by raising GenStreamErrorRaised(public=chunk.public) and catching it in the streaming and non-streaming routes, matching the safetensors tool loop.

The failing Repo tests (CPU) check is test_sidebar_account_block_uses_leading_tight, a frontend sidebar CSS assertion unrelated to this backend change; it also fails on main at the current head, so it is pre-existing. The focused MLX backend suite is green (87 passed) on Linux CPU MLX and macOS.

[danielhanchen]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Nice work!

Reviewed commit: 65e6880544

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".