From c5adb69a107913e7e2b2bc5f432d40dd3d70d0ec Mon Sep 17 00:00:00 2001 From: Filip Trajkovic Date: Thu, 2 Jul 2026 05:01:15 +0200 Subject: [PATCH 01/23] Fix GRPO logit scaling when model is wrapped by DDP (#5955) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- tests/python/test_grpo_ddp_model_config.py | 31 ++++++++++++++++++++++ unsloth/models/rl_replacements.py | 28 ++++++++++++++----- 2 files changed, 52 insertions(+), 7 deletions(-) create mode 100644 tests/python/test_grpo_ddp_model_config.py diff --git a/tests/python/test_grpo_ddp_model_config.py b/tests/python/test_grpo_ddp_model_config.py new file mode 100644 index 0000000000..5af31f65b8 --- /dev/null +++ b/tests/python/test_grpo_ddp_model_config.py @@ -0,0 +1,31 @@ +"""GRPO logit-scaling helpers must read config through DDP wrappers.""" + +from __future__ import annotations + +import os + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)) +SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py") + + +def _read_source() -> str: + with open(SOURCE_PATH, "r") as fh: + return fh.read() + + +def test_grpo_logit_scaling_uses_model_config_helper(): + src = _read_source() + # Helper exists and unwraps DDP/Accelerate wrappers via `.module`. + assert "def _unsloth_get_model_config(model):" in src + assert 'getattr(model.module, "config", None)' in src + # Softcapping takes the model and tolerates a missing config. + assert "logit_softcapping = _unsloth_get_final_logit_softcapping(model)" in src + assert "if config is None:" in src.split("def _unsloth_get_final_logit_softcapping")[1] + # Logit scale/divide read through the unwrapped config, not bare model.config. + assert 'getattr(model_config, "logit_scale", 0)' in src + assert 'getattr(model_config, "logits_scaling", 0)' in src + assert src.count("model_config = _unsloth_get_model_config(model)") >= 2 + # Helper source is injected into the compiled GRPO trainer. + assert "inspect.getsource(_unsloth_get_model_config)" in src + # No direct model.config access remains in the RL logit path. + assert "model.config" not in src diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 61a07b686d..d3ada23cf9 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -1337,11 +1337,12 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) temperature = self.temperature - logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) - logit_scale_multiply = getattr(model.config, "logit_scale", 0) + model_config = _unsloth_get_model_config(model) + logit_softcapping = _unsloth_get_final_logit_softcapping(model) + logit_scale_multiply = getattr(model_config, "logit_scale", 0) if logit_scale_multiply is None: logit_scale_multiply = 0 - logit_scale_divide = getattr(model.config, "logits_scaling", 0) + logit_scale_divide = getattr(model_config, "logits_scaling", 0) if logit_scale_divide is None: logit_scale_divide = 0 @@ -1471,7 +1472,15 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies) -def _unsloth_get_final_logit_softcapping(config): +def _unsloth_get_model_config(model): + """Return HuggingFace model config, unwrapping DDP/Accelerate wrappers.""" + config = getattr(model, "config", None) + if config is None and hasattr(model, "module"): + config = getattr(model.module, "config", None) + return config + + +def _unsloth_get_final_logit_softcapping(model): """Return final_logit_softcapping for a model config, falling back to the nested text sub-config for composite models. Handles both: - Gemma-4-style configs where the attribute lives on ``config.text_config`` @@ -1479,6 +1488,9 @@ def _unsloth_get_final_logit_softcapping(config): reachable via ``config.get_text_config()`` Returns 0 if unset, matching the previous behaviour. """ + config = _unsloth_get_model_config(model) + if config is None: + return 0 softcap = getattr(config, "final_logit_softcapping", None) if softcap is None: text_cfg = getattr(config, "text_config", None) @@ -1499,6 +1511,7 @@ grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"] UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"] grpo_accumulated_loss = RL_REPLACEMENTS["grpo_accumulated_loss"] grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"] +RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_model_config)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_final_logit_softcapping)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_mm_token_id)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_fix_mm_token_type_ids)) @@ -1616,11 +1629,12 @@ def grpo_trainer_compute_loss(function_name, function): input_ids = input_ids[:, -logits_to_keep:] # Get logit softcapping and logit scale - logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) # Gemma - logit_scale_multiply = getattr(model.config, "logit_scale", 0) # Cohere + model_config = _unsloth_get_model_config(model) + logit_softcapping = _unsloth_get_final_logit_softcapping(model) # Gemma + logit_scale_multiply = getattr(model_config, "logit_scale", 0) # Cohere if logit_scale_multiply is None: logit_scale_multiply = 0 - logit_scale_divide = getattr(model.config, "logits_scaling", 0) # Granite + logit_scale_divide = getattr(model_config, "logits_scaling", 0) # Granite if logit_scale_divide is None: logit_scale_divide = 0 From d91183d03feca2539a946354ca4fdbb4928e273c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 22:39:00 -0700 Subject: [PATCH 02/23] Fix gpt-oss offload_embedding and generate() kwargs, and guard offload_embedding on tied/vLLM models (#6774) * Fix gpt-oss offload_embedding and generate() logits_to_keep on fused models offload_embedding=True moved embed_tokens to CPU but left the input/output device-shuffling forward hooks commented out ('[TODO] Doesn't seem to work!'), so an eager forward/generate with CUDA input_ids hit the CPU embedding and raised a device-mismatch RuntimeError. Re-implement them in a testable helper _install_offload_embedding_hooks that saves the origin device on the module (the pre-hook returns a new tensor, so a device stashed on the original input is lost) and runs the lookup on the embedding weight's CURRENT device. Reading the weight device at call time (not a hard-coded cpu) also handles a non-quantized (bf16) embedding that a later model.to(...) pulls back onto the GPU, which the hard-coded version broke in the opposite direction. unsloth_base_fast_generate injected logits_to_keep/num_logits_to_keep whenever an inner submodule forward accepted it, but transformers validates generate kwargs against the top-level prepare_inputs_for_generation (plus forward when it takes kwargs). On fused/PEFT-wrapped gpt-oss this raised 'model_kwargs are not used by the model: [logits_to_keep]'. Only inject when the top level would accept it, mirroring transformers _validate_model_kwargs. Behavior is unchanged for every model that works today. Adds tests/test_offload_embedding_hooks.py and tests/test_generate_kwarg_gate.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload hooks: store origin device on the tensor, not the shared module The pre-hook stashed the input device on embed_tokens itself, which races when concurrent forwards share the module (serving). Ride it on the moved tensor and read it from the post-hook args instead: stateless and thread-safe. * Also strip mm_token_type_ids that generate() rejects (Qwen3-VL vision GRPO) The vision processor (Transformers 5.x path) emits mm_token_type_ids, which Qwen3-VL's generate() then rejects in _validate_model_kwargs on transformers 4.x, so vision GRPO fails at the first rollout: ValueError: The following `model_kwargs` are not used by the model: ['mm_token_type_ids'] Unlike logits_to_keep this is an incoming kwarg rather than one we inject, so drop it in unsloth_base_fast_generate when the top level generate does not accept it, reusing the same _unsloth_generate_accepts_kwarg gate. Extends the GPU-free gate test with the accept/reject mm_token_type_ids cases (7/7 pass). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim mm_token_type_ids comment * Trim comments in gpt-oss offload/logits fix (comment-only) * gpt-oss offload: return embedding output to the decoder device, not the input's When offload_embedding moves the embedding to CPU, model.device can become CPU and inputs then arrive on CPU, so returning the output to the input device left it on CPU and the CUDA decoder hit a device mismatch. Capture the decoder device before offload and always return there. This also drops the per-request tensor state (stateless, so concurrent forwards stay correct). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload: refuse offload_embedding for tied word embeddings Tied models share embed_tokens.weight with lm_head, so offloading the weight to CPU strands the output projection there (device mismatch at generate) and saves no VRAM since lm_head still needs it on GPU. Detect the shared weight via get_output_embeddings and raise NotImplementedError instead of loading into a crash. Untied models (gpt-oss, Llama-3.1-8B) offload unchanged. Adds tests/test_offload_tied_guard.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload: skip embedding offload on fast_inference (vLLM) vLLM manages its own weights, so offload_embedding cannot apply on the fast_inference path (previously it was silently ignored). Disable it with a notice, mirroring the WSL and Windows skips. * Trim offload embedding comments (comment-only) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload: track decoder device live so it survives model.to() The post-hook returned the embedding output to a device captured at load time. If a model is loaded on CPU then moved with model.to(cuda), that device is stale and the output lands on the wrong device. Read the decoder device live from the (untied) output embeddings, keeping the captured device as a fallback. Adds a stale-fallback regression test. * Make generate-kwarg-gate cases pytest-collectable Cases lived in run(), which pytest does not collect, so CI never exercised the gate. Expose them as test_generate_kwarg_gate; still runnable via __main__. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gpt-oss offload: skip a meta (disk-offloaded) lm_head as the return device A device_map that disk-offloads an untied lm_head leaves its weight on the meta device until that module's own hook runs, so reading it as the decoder device would move real hidden states to meta. Skip meta (and a missing weight) and fall back to the captured device. Adds a regression test. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_generate_kwarg_gate.py | 137 ++++++++++++++++++++++++++ tests/test_offload_embedding_hooks.py | 129 ++++++++++++++++++++++++ tests/test_offload_tied_guard.py | 62 ++++++++++++ unsloth/models/vision.py | 105 +++++++++++++++++--- 4 files changed, 422 insertions(+), 11 deletions(-) create mode 100644 tests/test_generate_kwarg_gate.py create mode 100644 tests/test_offload_embedding_hooks.py create mode 100644 tests/test_offload_tied_guard.py diff --git a/tests/test_generate_kwarg_gate.py b/tests/test_generate_kwarg_gate.py new file mode 100644 index 0000000000..6d1379d3a9 --- /dev/null +++ b/tests/test_generate_kwarg_gate.py @@ -0,0 +1,137 @@ +"""GPU-free test for the generate-kwarg gate in vision.py +(_unsloth_generate_accepts_kwarg), covering both logits_to_keep injection and mm_token_type_ids +stripping, AST-extracted so no unsloth/CUDA import is needed.""" + +import ast, inspect, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_helper(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_unsloth_generate_accepts_kwarg": + ns = {"inspect": inspect} + exec(ast.get_source_segment(src, node), ns) + return ns["_unsloth_generate_accepts_kwarg"] + raise AssertionError("_unsloth_generate_accepts_kwarg not found in vision.py") + + +accepts = _load_helper() + + +class PrepHasKwargs_ForwardHasKey: + # **kwargs on prepare unions forward params; key in forward -> ACCEPTED. + def prepare_inputs_for_generation(self, input_ids, **kwargs): ... + def forward( + self, + input_ids, + logits_to_keep = 0, + **kwargs, + ): ... + + +class PrepNoKwargs_ForwardHasKey: + # no **kwargs -> forward not unioned; key only in forward -> REJECTED (gpt-oss shape). + def prepare_inputs_for_generation( + self, + input_ids, + attention_mask = None, + ): ... + def forward( + self, + input_ids, + logits_to_keep = 0, + ): ... + + +class PrepHasKeyDirectly: + # key directly on prepare -> ACCEPTED. + def prepare_inputs_for_generation( + self, + input_ids, + logits_to_keep = 0, + ): ... + def forward(self, input_ids): ... + + +class NoPrepare: + # no prepare -> empty args, no union -> REJECTED. + def forward( + self, + input_ids, + logits_to_keep = 0, + **kwargs, + ): ... + + +class VisionRejectsMM: + # Qwen3-VL shape: neither prepare nor forward names mm_token_type_ids -> REJECTED (stripped). + def prepare_inputs_for_generation( + self, + input_ids, + attention_mask = None, + ): ... + def forward( + self, + input_ids, + pixel_values = None, + ): ... + + +class VisionAcceptsMM: + # forward names mm_token_type_ids and prepare unions it via **kwargs -> ACCEPTED (kept). + def prepare_inputs_for_generation(self, input_ids, **kwargs): ... + def forward( + self, + input_ids, + mm_token_type_ids = None, + **kwargs, + ): ... + + +# (model, key, expected) per gate case. +CASES = [ + ( + "prep(**kwargs)+forward(key) -> accept", + PrepHasKwargs_ForwardHasKey(), + "logits_to_keep", + True, + ), + ( + "prep(no kwargs)+forward(key) -> reject", + PrepNoKwargs_ForwardHasKey(), + "logits_to_keep", + False, + ), + ("prep(key) direct -> accept", PrepHasKeyDirectly(), "logits_to_keep", True), + ("no prepare_inputs_for_gen -> reject", NoPrepare(), "logits_to_keep", False), + ( + "num_logits_to_keep variant -> reject", + PrepNoKwargs_ForwardHasKey(), + "num_logits_to_keep", + False, + ), + ( + "mm_token_type_ids not accepted -> reject (strip)", + VisionRejectsMM(), + "mm_token_type_ids", + False, + ), + ("mm_token_type_ids accepted -> keep", VisionAcceptsMM(), "mm_token_type_ids", True), +] + + +def test_generate_kwarg_gate(): + for name, model, key, expected in CASES: + got = accepts(model, key) + assert got is expected, f"{name}: got {got}, expected {expected}" + + +if __name__ == "__main__": + test_generate_kwarg_gate() + for name, _, _, _ in CASES: + print(f" [PASS] {name}") + print("OK: generate-kwarg gate behaves like transformers _validate_model_kwargs") diff --git a/tests/test_offload_embedding_hooks.py b/tests/test_offload_embedding_hooks.py new file mode 100644 index 0000000000..b8be603b2a --- /dev/null +++ b/tests/test_offload_embedding_hooks.py @@ -0,0 +1,129 @@ +"""Tests _install_offload_embedding_hooks in vision.py: the offloaded lookup must work and +its output must land on the decoder device, read live from the output embeddings (lm_head) +so it tracks model.to() moves. CUDA cases skip without a GPU.""" + +import ast, os +import torch +import torch.nn as nn + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_installer(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_install_offload_embedding_hooks": + ns = {"torch": torch} + exec(ast.get_source_segment(src, node), ns) + return ns["_install_offload_embedding_hooks"] + raise AssertionError("_install_offload_embedding_hooks not found in vision.py") + + +install = _load_installer() +CPU = torch.device("cpu") + + +def _emb(): + return nn.Embedding(32, 8) + + +def _lm_head(device): + # Stand-in decoder reference (untied lm_head) whose weight device is the target. + return nn.Linear(8, 32, bias = False).to(device) + + +def test_install_and_idempotent(): + emb = _emb() + lm = _lm_head(CPU) + assert install(emb, lm, CPU) is True + assert emb._unsloth_offload_hooks_installed is True + n_pre = len(emb._forward_pre_hooks) + n_post = len(emb._forward_hooks) + assert install(emb, lm, CPU) is True + assert len(emb._forward_pre_hooks) == n_pre and len(emb._forward_hooks) == n_post + assert install(None, lm, CPU) is False + + +def test_cpu_noop_forward(): + # cpu weight + cpu decoder + cpu input -> output stays cpu. + emb = _emb() + install(emb, _lm_head(CPU), CPU) + out = emb(torch.randint(0, 32, (2, 5))) + assert out.shape == (2, 5, 8) + assert out.device.type == "cpu" + + +def test_cuda_input_roundtrip(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # CPU weight, CUDA decoder + input -> lookup on cpu, output back on cuda. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +def test_cpu_input_still_returns_to_decoder(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # P1: offload makes the input arrive on cpu; the output must still reach the cuda decoder. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cpu")) + assert out.device.type == "cuda", out.device + + +def test_live_decoder_over_stale_fallback(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # P2: fallback captured as cpu (model loaded on cpu), but the decoder later lives on cuda. + # The output must follow the live lm_head device, not the stale cpu fallback. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), CPU) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +def test_meta_lm_head_falls_back(): + # A disk-offloaded (meta) lm_head must not be used as the return device: moving hidden + # states to meta is unrecoverable, so fall back to the captured device. No GPU needed. + emb = _emb().to("cpu") + lm = _lm_head(CPU) + lm.weight = nn.Parameter(lm.weight.to("meta")) + install(emb, lm, CPU) + out = emb(torch.randint(0, 32, (2, 5))) + assert out.device.type == "cpu", out.device + + +def test_cuda_weight_pulled_back_to_gpu(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # bf16 weight later pulled back to gpu + cuda input -> no-op, stays on cuda. + emb = _emb().to("cuda") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +if __name__ == "__main__": + test_install_and_idempotent() + print("[PASS] install + idempotent") + test_cpu_noop_forward() + print("[PASS] cpu no-op forward") + test_cuda_input_roundtrip() + print("[PASS] cuda input roundtrip") + test_cpu_input_still_returns_to_decoder() + print("[PASS] cpu input still returns to cuda decoder (P1)") + test_live_decoder_over_stale_fallback() + print("[PASS] live decoder device beats stale fallback (P2)") + test_meta_lm_head_falls_back() + print("[PASS] meta lm_head falls back to captured device (P2)") + test_cuda_weight_pulled_back_to_gpu() + print("[PASS] cuda weight-on-gpu no-op") + print("OK: offloaded embedding output always lands on the live decoder device") diff --git a/tests/test_offload_tied_guard.py b/tests/test_offload_tied_guard.py new file mode 100644 index 0000000000..096fba116d --- /dev/null +++ b/tests/test_offload_tied_guard.py @@ -0,0 +1,62 @@ +"""Tests _embeddings_are_tied in vision.py: offload_embedding must detect a shared +embed_tokens/lm_head weight so the loader can refuse to offload tied embeddings +(offloading would strand the output projection on CPU). No GPU needed.""" + +import ast, os +import torch +import torch.nn as nn + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_fn(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_embeddings_are_tied": + ns = {"torch": torch} + exec(ast.get_source_segment(src, node), ns) + return ns["_embeddings_are_tied"] + raise AssertionError("_embeddings_are_tied not found in vision.py") + + +tied = _load_fn() + + +def test_untied_separate_weights(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + assert tied(emb, lm) is False + + +def test_tied_shared_parameter(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + lm.weight = emb.weight # transformers-style weight tying + assert tied(emb, lm) is True + + +def test_tied_by_storage_even_if_distinct_parameter(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + lm.weight = nn.Parameter(emb.weight.detach()) # distinct Parameter, shared storage + assert tied(emb, lm) is True + + +def test_none_output_is_untied(): + emb = nn.Embedding(32, 8) + assert tied(emb, None) is False + assert tied(None, nn.Linear(8, 32)) is False + + +if __name__ == "__main__": + test_untied_separate_weights() + print("[PASS] untied separate weights -> False") + test_tied_shared_parameter() + print("[PASS] tied shared parameter -> True") + test_tied_by_storage_even_if_distinct_parameter() + print("[PASS] tied by storage -> True") + test_none_output_is_untied() + print("[PASS] missing lm_head -> untied (safe to offload)") + print("OK: tied embeddings are detected so offload_embedding can refuse them") diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index bdc2bd9ef6..5ab55152db 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -238,6 +238,71 @@ def _attach_bnb_multidevice_hooks( global NUM_LOGITS_TO_KEEP NUM_LOGITS_TO_KEEP = dict() + +def _unsloth_generate_accepts_kwarg(model, key): + # True if the top level accepts this generate kwarg (some models expose it on an inner forward only). + try: + model_args = set(inspect.signature(model.prepare_inputs_for_generation).parameters) + except (TypeError, ValueError, AttributeError): + model_args = set() + if "kwargs" in model_args or "model_kwargs" in model_args: + try: + model_args |= set(inspect.signature(model.forward).parameters) + except (TypeError, ValueError, AttributeError): + pass + return key in model_args + + +def _install_offload_embedding_hooks(embed_tokens, output_embeddings, return_device): + # Lookup runs on the weight's current device (CPU when offloaded); the output returns to the + # decoder device read live from output_embeddings (lm_head, untied here) so it tracks + # model.to() moves. A meta (disk-offloaded) or missing lm_head falls back to return_device. + if embed_tokens is None: + return False + if getattr(embed_tokens, "_unsloth_offload_hooks_installed", False): + return True + + def _decoder_device(): + weight = getattr(output_embeddings, "weight", None) + if weight is not None and weight.device.type != "meta": + return weight.device + return return_device + + def _unsloth_offload_pre_hook(module, args): + if not args: + return args + inp = args[0] + if not hasattr(inp, "device"): + return args + weight = getattr(module, "weight", None) + target = weight.device if weight is not None else _decoder_device() + if target is None or inp.device == target: + return args + return (inp.to(target),) + tuple(args[1:]) + + def _unsloth_offload_post_hook(module, args, output): + target = _decoder_device() + if target is not None and hasattr(output, "device") and output.device != target: + return output.to(target) + return output + + embed_tokens.register_forward_pre_hook(_unsloth_offload_pre_hook, prepend = True) + embed_tokens.register_forward_hook(_unsloth_offload_post_hook, prepend = True) + embed_tokens._unsloth_offload_hooks_installed = True + return True + + +def _embeddings_are_tied(input_embeddings, output_embeddings): + # Tied lm_head reuses this weight; offloading to CPU would strand the output projection. + if input_embeddings is None or output_embeddings is None: + return False + w_in = getattr(input_embeddings, "weight", None) + w_out = getattr(output_embeddings, "weight", None) + if w_in is None or w_out is None: + return False + return w_in is w_out or w_in.data_ptr() == w_out.data_ptr() + + VLLM_SUPPORTED_VLM = [ "qwen2_5_vl", "gemma3", @@ -321,6 +386,13 @@ def unsloth_base_fast_generate(self, *args, **kwargs): kwargs.pop("token_type_ids", None) # kwargs.pop("token_type_ids", None) + # Vision processors emit mm_token_type_ids that generate() rejects (Qwen3-VL); unlike + # logits_to_keep it is an incoming kwarg, so drop it when generate does not accept it. + if "mm_token_type_ids" in kwargs and not _unsloth_generate_accepts_kwarg( + self, "mm_token_type_ids" + ): + kwargs.pop("mm_token_type_ids", None) + # VLMs do not allow logits_to_keep global NUM_LOGITS_TO_KEEP if arch not in NUM_LOGITS_TO_KEEP: @@ -339,7 +411,7 @@ def unsloth_base_fast_generate(self, *args, **kwargs): if arch not in NUM_LOGITS_TO_KEEP: NUM_LOGITS_TO_KEEP[arch] = None key = NUM_LOGITS_TO_KEEP[arch] - if key is not None and key not in kwargs: + if key is not None and key not in kwargs and _unsloth_generate_accepts_kwarg(self, key): kwargs[key] = 1 model_eos_token_id = getattr(self.config, "eos_token_id", None) @@ -1024,6 +1096,12 @@ class FastBaseModel: raise_handler = RaiseUninitialized() try: + if offload_embedding and fast_inference: + # vLLM manages its own weights; embedding offload does not apply. + print( + "Unsloth: Not offloading embeddings; incompatible with fast_inference (vLLM)." + ) + offload_embedding = False if not fast_inference: # Prevent load_in_fp8 from being forwarded into HF internal model loading load_in_fp8 = kwargs.pop("load_in_fp8", None) @@ -1070,21 +1148,26 @@ class FastBaseModel: pass else: embed_tokens = model.get_input_embeddings() + out_embed = ( + model.get_output_embeddings() + if hasattr(model, "get_output_embeddings") + else None + ) + if _embeddings_are_tied(embed_tokens, out_embed): + raise NotImplementedError( + "offload_embedding = True is not supported for models with tied word " + "embeddings (embed_tokens shares its weight with lm_head). Offloading " + "would strand the output projection on CPU and saves no VRAM. Set " + "offload_embedding = False for this model." + ) nbytes = embed_tokens.weight.numel() * embed_tokens.weight.itemsize ngb = round(nbytes / 1024 / 1024 / 1024, 2) print(f"Unsloth: Offloading embeddings to RAM to save {ngb} GB.") + _embed_device = embed_tokens.weight.device # decoder device, before offload embed_tokens.to("cpu") - # Add hooks to move inputs to CPU and back to CUDA - # [TODO] Doesn't seem to work! - # def pre_hook(module, args): - # args[0]._old_device = args[0].device - # return (args[0].to("cpu", non_blocking = True)) - # def post_hook(module, args, output): - # old_device = getattr(args[0], "_old_device", "cuda") - # return output.to(old_device, non_blocking = True) - # embed_tokens.register_forward_pre_hook(pre_hook, prepend = True) - # embed_tokens.register_forward_hook (post_hook, prepend = True) + # Device-safe embedding offload. + _install_offload_embedding_hooks(embed_tokens, out_embed, _embed_device) # Must free GPU memory otherwise will not free! torch.cuda.empty_cache() gc.collect() From ac6ba96f9e21fb91f4d9a93bc31c0beb0feceafe Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Thu, 2 Jul 2026 02:18:20 -0700 Subject: [PATCH 03/23] Add a fits-on-device filter to the model selects (#6802) * Add a shared fits-on-device filter to the model selects The chat model selector gains an Only show models that fit on this device tick under its filter row, and the Hub page gains a matching Fits device pill next to the sort menu. Both read one persisted preference (unsloth_models_fit_on_device_only), so toggling either applies to both. The filter reuses the Recommended sort's existing fit math, extracted into hfModelFitsDevice: size from safetensors metadata, GGUF param count, or the repo name, against the 0.7 GPU + 0.7 RAM budget, with unsizable models hidden. In the chat selector it extends the fit filtering to the Trending and Recent sorts and to search results; downloaded models stay visible regardless. An unknown device budget keeps everything. The preference is cleared by Reset all local preferences like the other picker toggles. * Move the device-fit toggle into the sort dropdowns * Tighten sort menu footer spacing and shorten the label * Align the footer checkbox with the option text * Make the footer checkbox circular with a smaller tick * Clear menu highlight when the pointer leaves the options * Address review: fit filter coverage and sizing Exempt on-disk models from the Hub fit filter, apply it to the feed trending rows and curated search results, size safetensors and MLX rows by the quantized load estimate instead of checkpoint bytes, and replace the native title hint with the app Tooltip. * Make the whole device-fit row toggle the filter --- .../assistant-ui/model-selector/pickers.tsx | 124 ++++++++++++------ .../model-selector/recommended-fit.ts | 32 +++++ .../chat/stores/chat-runtime-store.ts | 11 ++ .../features/hub/catalog/hub-option-menu.tsx | 36 ++++- .../features/hub/catalog/models-toolbar.tsx | 33 +++++ studio/frontend/src/features/hub/hub-page.tsx | 32 ++++- .../features/settings/tabs/general-tab.tsx | 1 + 7 files changed, 222 insertions(+), 47 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index e11a08f8ab..20000c82ee 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1,6 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -102,6 +103,7 @@ import { type FormatFilter, estimateQuantBytes, fitsDevice, + hfModelFitsDevice, isMlxId, isMobileVariant, isRecommendableFormat, @@ -1340,6 +1342,9 @@ export function HubModelPicker({ }, []); // When on, On Device GGUF repos show their quantizations without a click. const expandQuantizations = useChatRuntimeStore((s) => s.expandQuantizations); + // Shared with the Hub page: list only models sized within the device budget. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); // Repos the user clicked to collapse while expand-by-default is on. Kept in // memory only, so it resets on reload (and when the setting is toggled). const [collapsedGguf, setCollapsedGguf] = useState>( @@ -1717,34 +1722,19 @@ export function HubModelPicker({ formatFilter === "all" ? rows.filter((r) => isRecommendableFormat(r.id, r.isGguf, isMac)) : rows.filter((r) => matchesFormatFilter(r.id, r.isGguf, formatFilter)); - if (recommendedSort !== "recommended") return rows; + // The "recommended" sort always applies the device-fit filter; the shared + // "Fits on device" tick extends it to the other sorts too. + if (recommendedSort !== "recommended" && !fitOnDeviceOnly) return rows; return rows.filter((r) => { // Downloaded models always show, regardless of device fit. if (downloadedSet.has(r.id.toLowerCase())) return true; - // Unified-memory hosts (Mac / no discrete GPU) still report system RAM, - // so fall back to that budget instead of skipping the fit check entirely. - const hasDeviceBudget = - gpu.memoryTotalGb > 0 || gpu.systemRamAvailableGb > 0; - if (!hasDeviceBudget) return true; - // GGUF/MLX repos rarely expose safetensors metadata, so fall back to the - // GGUF param count, then the repo name, for a size estimate. Anything we - // still cannot size is hidden (requireKnown) so over-budget models like a - // 1T GGUF don't slip into Recommended. - const params = r.totalParams ?? paramsFromId(r.id); - const sizeBytes = - r.estimatedSizeBytes ?? - (params ? estimateQuantBytes(params) : undefined); - return fitsDevice({ - sizeBytes, - gpuGb: gpu.memoryTotalGb, - systemRamGb: gpu.systemRamAvailableGb, - requireKnown: true, - }); + return hfModelFitsDevice(r, gpu); }); }, [ recommendedSearch.results, downloadedSet, recommendedSort, + fitOnDeviceOnly, formatFilter, isMac, gpu, @@ -1976,23 +1966,6 @@ export function HubModelPicker({ [visibleCachedModelRows], ); - // Recommended models that match the current search query - const filteredRecommendedIds = useMemo(() => { - if (!showHfSection) return []; - const q = normalizeForSearch(debouncedQuery.trim()); - return recommendedIds - .filter((id) => normalizeForSearch(id).includes(q)) - .filter((id) => - matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), - ); - }, [ - showHfSection, - debouncedQuery, - recommendedIds, - formatFilter, - isKnownGgufRepo, - ]); - // Param counts come straight off the unsloth listings the picker already // loaded, so no extra per-id fetch is needed for the VRAM badges. const recommendedParamCountById = useMemo(() => { @@ -2003,6 +1976,42 @@ export function HubModelPicker({ return map; }, [results, recommendedSearch.results]); + // Recommended models that match the current search query + const filteredRecommendedIds = useMemo(() => { + if (!showHfSection) return []; + const q = normalizeForSearch(debouncedQuery.trim()); + return recommendedIds + .filter((id) => normalizeForSearch(id).includes(q)) + .filter((id) => + matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), + ) + // Curated defaults obey the fit toggle like the live HF rows, else large + // defaults resurface in search results with the filter on. + .filter( + (id) => + !fitOnDeviceOnly || + downloadedSet.has(id.toLowerCase()) || + hfModelFitsDevice( + { + id, + totalParams: recommendedParamCountById.get(id), + isGguf: isKnownGgufRepo(id), + }, + gpu, + ), + ); + }, [ + showHfSection, + debouncedQuery, + recommendedIds, + formatFilter, + isKnownGgufRepo, + fitOnDeviceOnly, + downloadedSet, + recommendedParamCountById, + gpu, + ]); + const recommendedSet = useMemo( () => new Set(filteredRecommendedIds), [filteredRecommendedIds], @@ -2013,6 +2022,12 @@ export function HubModelPicker({ if (!showHfSection || section !== "recommended") return []; return results .filter(isChatSupported) + .filter( + (r) => + !fitOnDeviceOnly || + downloadedSet.has(r.id.toLowerCase()) || + hfModelFitsDevice(r, gpu), + ) .map((result) => result.id) .filter((id) => !isHiddenModelId(id)) .filter((id) => id.toLowerCase().startsWith("unsloth/")) @@ -2035,6 +2050,9 @@ export function HubModelPicker({ isKnownGgufRepo, isChatSupported, formatFilter, + fitOnDeviceOnly, + downloadedSet, + gpu, isMac, ]); @@ -2323,6 +2341,35 @@ export function HubModelPicker({ // selected-item checkmark never overlaps the label. const sortMenuContentClassName = "!p-1 !rounded-[14px] [&_[role=option]]:!pl-2 [&_[role=option]]:!py-1.5 [&_[role=option]]:!text-xs [&_[role=option]]:!rounded-[10px]"; + // Device-fit toggle lives inside the sort menu (shared with the Hub page). + // The whole row is the click target (a button): a Checkbox renders as a + // + + + Hides models larger than this device's memory budget. Downloaded models + stay visible. + + + ); const sectionSortDropdown = section === "recommended" ? ( ) : section === "downloaded" ? ( ) : ( ); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts index 24f0edc784..7c2ed266c0 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts @@ -114,3 +114,35 @@ export function fitsDevice(opts: { } return requireKnown ? false : true; } + +/** Fit predicate for one Hub listing row, shared by the chat model selector + * and the Hub page "Fits on device" filter. GGUF repos: metadata size (actual + * weights) or the smallest-quant estimate from the param count. Safetensors / + * MLX repos: always the params-based smallest-quant estimate, matching the + * VRAM badge's quantized-load assumption; their estimatedSizeBytes is the + * full-precision checkpoint and would wrongly hide models the quantized load + * path can run. Anything unsizable is hidden (requireKnown) so over-budget + * models with no metadata don't slip through. An unknown device budget keeps + * everything. */ +export function hfModelFitsDevice( + model: { + id: string; + totalParams?: number; + estimatedSizeBytes?: number; + isGguf?: boolean; + }, + gpu: { memoryTotalGb: number; systemRamAvailableGb: number }, +): boolean { + if (gpu.memoryTotalGb <= 0 && gpu.systemRamAvailableGb <= 0) return true; + const params = model.totalParams ?? paramsFromId(model.id); + const quantBytes = params ? estimateQuantBytes(params) : undefined; + const sizeBytes = isGgufId(model.id, model.isGguf) + ? (model.estimatedSizeBytes ?? quantBytes) + : (quantBytes ?? model.estimatedSizeBytes); + return fitsDevice({ + sizeBytes, + gpuGb: gpu.memoryTotalGb, + systemRamGb: gpu.systemRamAvailableGb, + requireKnown: true, + }); +} diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index ca4bb7afde..7c9685d6d0 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -42,6 +42,8 @@ export const CHAT_EXPAND_QUANTIZATIONS_KEY = "unsloth_chat_expand_quantizations"; export const CHAT_SHOW_ALL_QUANTIZATIONS_KEY = "unsloth_chat_show_all_quantizations"; +export const MODELS_FIT_ON_DEVICE_ONLY_KEY = + "unsloth_models_fit_on_device_only"; export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions"; export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = "unsloth_chat_web_fetch_tools_enabled"; @@ -671,6 +673,9 @@ type ChatRuntimeStore = { expandQuantizations: boolean; /** Persisted: show non-downloaded quantizations too, not just downloaded. */ showAllQuantizations: boolean; + /** Persisted, shared by the chat model selector and the Hub page: list only + * models whose size fits this device's memory budget. */ + fitOnDeviceOnly: boolean; /** A local model picked while `loadOnSelection` is off: staged, not loaded. * The settings sheet shows its load knobs and a Load button. */ pendingSelection: PendingModelSelection | null; @@ -793,6 +798,7 @@ type ChatRuntimeStore = { setLoadOnSelection: (value: boolean) => void; setExpandQuantizations: (value: boolean) => void; setShowAllQuantizations: (value: boolean) => void; + setFitOnDeviceOnly: (value: boolean) => void; setPendingSelection: (selection: PendingModelSelection | null) => void; /** Stage a pick for a deferred load: revert knobs to the loaded baseline, * record the selection, and open the settings sheet. */ @@ -1111,6 +1117,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true), expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false), showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true), + fitOnDeviceOnly: loadBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, false), pendingSelection: null, loadedIsMultimodal: false, loadedIsDiffusion: false, @@ -1582,6 +1589,10 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, showAllQuantizations); set({ showAllQuantizations }); }, + setFitOnDeviceOnly: (fitOnDeviceOnly) => { + saveBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, fitOnDeviceOnly); + set({ fitOnDeviceOnly }); + }, setPendingSelection: (pendingSelection) => set({ pendingSelection }), stageModel: (selection) => { // Refuse staging mid-load: post-load cleanup would silently drop the queued diff --git a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx index 7895a89254..38464d36e4 100644 --- a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx +++ b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx @@ -38,6 +38,7 @@ export function HubOptionMenu({ showChevron = true, title, triggerContent, + footer, }: { value: T; options: readonly HubOption[]; @@ -49,9 +50,12 @@ export function HubOptionMenu({ showChevron?: boolean; title?: string; triggerContent?: ReactNode; + /** Rendered under the options behind a separator; clicks keep the menu open. */ + footer?: ReactNode; }) { const [open, setOpen] = useState(false); - const [activeIndex, setActiveIndex] = useState(0); + // -1 = nothing highlighted (no hover, no keyboard nav yet). + const [activeIndex, setActiveIndex] = useState(-1); const triggerRef = useRef(null); const listboxRef = useRef(null); const idBase = useId(); @@ -63,9 +67,9 @@ export function HubOptionMenu({ }, [options, value]); const selected = options[selectedIndex]; const resolvedActiveIndex = - options.length === 0 + options.length === 0 || activeIndex < 0 ? -1 - : Math.min(Math.max(activeIndex, 0), options.length - 1); + : Math.min(activeIndex, options.length - 1); const activeOptionId = resolvedActiveIndex >= 0 ? `${idBase}-option-${resolvedActiveIndex}` : undefined; @@ -92,11 +96,13 @@ export function HubOptionMenu({ (nextOpen: boolean) => { setOpen(nextOpen); if (nextOpen) { - activateIndex(selectedIndex); + // Nothing highlighted until the user hovers or uses the keyboard; + // keyboard nav anchors on the selected option (handleContentKeyDown). + activateIndex(-1); requestAnimationFrame(() => listboxRef.current?.focus()); } }, - [activateIndex, selectedIndex], + [activateIndex], ); const handleContentKeyDown = useCallback( @@ -112,12 +118,21 @@ export function HubOptionMenu({ } if (event.key === "ArrowDown") { event.preventDefault(); - setActiveIndex((currentIndex + 1) % options.length); + // First arrow press highlights the selected option, then steps. + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex + 1) % options.length, + ); return; } if (event.key === "ArrowUp") { event.preventDefault(); - setActiveIndex((currentIndex - 1 + options.length) % options.length); + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex - 1 + options.length) % options.length, + ); return; } if (event.key === "Home") { @@ -197,6 +212,7 @@ export function HubOptionMenu({ aria-activedescendant={activeOptionId} tabIndex={0} onKeyDown={handleContentKeyDown} + onPointerLeave={() => activateIndex(-1)} className="outline-none" > {options.map((option, index) => { @@ -235,6 +251,12 @@ export function HubOptionMenu({ ); })} + {footer && ( + // -mt-3 cancels the surface's 16px flex gap down to 4px. No side + // padding: the footer label carries the same padding as the options + // so its checkbox lines up with the option text. +
{footer}
+ )} ); diff --git a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx index 48f7fcffaa..7c08c9c486 100644 --- a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx +++ b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx @@ -1,6 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -68,6 +69,8 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange, capabilityFilter, onCapabilityFilterChange, + fitOnDeviceOnly, + onFitOnDeviceOnlyChange, onManageLocalFolders, onOpenFineTune, }: { @@ -84,6 +87,9 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange: (value: ModelFormatFilter) => void; capabilityFilter: CapabilityFilter; onCapabilityFilterChange: (value: CapabilityFilter) => void; + /** Shared with the chat model selector: hide models over the device budget. */ + fitOnDeviceOnly: boolean; + onFitOnDeviceOnlyChange: (value: boolean) => void; onManageLocalFolders: () => void; /** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a * format-dropdown option rather than a standalone feed section. */ @@ -350,6 +356,33 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onValueChange={onSortChange} ariaLabel="Sort models" className={cn(triggerBase, "w-[128px]")} + footer={ + isDataset ? undefined : ( + + + + + + Hides models larger than this device's memory budget. + Downloaded models stay visible. + + + ) + } /> )} diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index c696862283..56aa07335d 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -5,6 +5,7 @@ import { loadRememberedLoadSettings, rememberedLoadSettingsKey, } from "@/components/assistant-ui/model-selector/remembered-load-settings"; +import { hfModelFitsDevice } from "@/components/assistant-ui/model-selector/recommended-fit"; import { useHubInventory } from "@/features/hub/inventory"; import { useDebouncedValue } from "@/hooks/use-debounced-value"; import { useGpuInfo } from "@/hooks/use-gpu-info"; @@ -327,6 +328,9 @@ export function ModelsPage() { const activeCheckpoint = checkpoint && !isExternalModelId(checkpoint) ? checkpoint : null; const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + // Shared with the chat model selector: list only models sized for this device. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); useEffect(() => { let cancelled = false; @@ -697,7 +701,12 @@ export function ModelsPage() { !isHiddenModelId(row.id) && matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) && matchesCapability(row.capabilities, deferredCapabilityFilter) && - (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)), + (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)) && + // Models already on disk stay visible regardless of device fit, + // matching the chat model selector. + (!fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu)), ); }, [ discoverRows, @@ -705,6 +714,8 @@ export function ModelsPage() { effectiveDiscoverFormat, deferredCapabilityFilter, activeChannel, + fitOnDeviceOnly, + gpu, ]); const listRows = filteredDiscoverRows; @@ -724,8 +735,21 @@ export function ModelsPage() { effectiveLocalRows, ) .filter((row) => !isHiddenModelId(row.id)) - .filter((row) => matchesFormat(row.result.isGguf, "gguf")), - [hubFeed.trending.results, modelDiscoveryInventorySignature], + .filter((row) => matchesFormat(row.result.isGguf, "gguf")) + // Same fit filter as the main Discover list, so the feed carousel + // honors the toggle too. + .filter( + (row) => + !fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu), + ), + [ + hubFeed.trending.results, + modelDiscoveryInventorySignature, + fitOnDeviceOnly, + gpu, + ], ); const feedRows = useMemo(() => { if (!isFeedMode) return []; @@ -1448,6 +1472,8 @@ export function ModelsPage() { onFormatFilterChange={setFormatFilter} capabilityFilter={capabilityFilter} onCapabilityFilterChange={setCapabilityFilter} + fitOnDeviceOnly={fitOnDeviceOnly} + onFitOnDeviceOnlyChange={setFitOnDeviceOnly} onManageLocalFolders={handleManageLocalFolders} onOpenFineTune={() => handleOpenList("finetune")} /> diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 247d5040fb..4e02f7f14f 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -81,6 +81,7 @@ const PREFS_KEYS: string[] = [ "unsloth_chat_load_on_selection", "unsloth_chat_expand_quantizations", "unsloth_chat_show_all_quantizations", + "unsloth_models_fit_on_device_only", // Chat presets "unsloth_chat_custom_presets", "unsloth_chat_active_preset", From 62e96442665892dc8d06084731dc700573af0f28 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 03:38:48 -0700 Subject: [PATCH 04/23] Studio RAG: fix RTL/Indic PDF corruption and dropped DOCX tables (#6780) * Studio RAG: fix RTL/Indic PDF corruption and dropped DOCX tables The RAG parser prefers pymupdf4llm.to_markdown for PDFs, but that rebuilds text from positioned glyphs and mangles complex-shaping scripts (RTL Arabic/Hebrew come back as shaped Presentation Forms, Indic matras drop to U+FFFD) and can silently drop most of a heavy-RTL page. _pdf now compares the Markdown against PyMuPDF's logical-order get_text() per page and falls back to it when the Markdown looks corrupted (shaped Presentation Forms or U+FFFD above a small floor/ratio) or holds far fewer letters than the raw layer. Latin PDFs are unaffected and keep their Markdown tables/headings. _docx walked document.paragraphs, which excludes table cells, so DOCX tables were dropped entirely. It now walks body content in document order via iter_inner_content, emitting each table row as pipe-joined cells (deduped across merged cells); the preview locator already anchors on pipes. Adds parser tests for the corruption and incompleteness fallbacks and for DOCX table extraction. These mirror the chat document-extractor guard raised in the unslothai/ unsloth#5351 review; the RAG parser is a separate module and needed its own fix. * RAG DOCX: keep empty table cells and collapse in-cell newlines Skipping empty cells shifted later cells left and broke column alignment across rows; a cell with internal paragraphs (newlines) also broke the pipe-joined row. Keep every cell (dropping the row only when all are empty) and normalize each cell with " ".join(split()) so multi-paragraph cells stay on one row. Adds a test for both. * RAG DOCX: dedup merged table cells on the element directly Store the shared lxml element in the seen set instead of its id(); it is hashable and compares by the underlying node, so it dedups spanned/merged cells the same way without relying on id(). Adds a merged-cell test. * RAG DOCX: align merged cells, pad skipped grid columns, flatten nested tables * RAG DOCX: walk cells in document order so nested tables keep in-cell position * RAG DOCX: dedup vertically merged cells so a spanning label is indexed once --------- Co-authored-by: danielhanchen Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/rag/parsers.py | 103 ++++++++++- studio/backend/tests/test_rag_parsing.py | 209 +++++++++++++++++++++++ 2 files changed, 307 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index ba248cf9a6..9afddf1d9e 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -12,6 +12,7 @@ from __future__ import annotations import logging import os +import re from dataclasses import dataclass from html.parser import HTMLParser @@ -69,6 +70,39 @@ def _html(raw: str) -> list[Page]: return [_page("\n".join(parser.out), 1)] +# pymupdf4llm rebuilds text from positioned glyphs, which mangles complex-shaping +# scripts (RTL Arabic/Hebrew emerge as shaped Presentation Forms, Indic matras drop to +# U+FFFD) and can silently drop most of a heavy-RTL page. When Markdown trips these +# signals we fall back to PyMuPDF's logical-order get_text(). Thresholds mirror the chat +# extractor guard (unslothai/unsloth#5351 review). +_SHAPED_PRESENTATION_FORMS = re.compile("[\ufb1d-\ufdff\ufe70-\ufefc]") +_PDF_FALLBACK_MIN_BAD_GLYPHS = 5 +_PDF_FALLBACK_BAD_GLYPH_RATIO = 0.0005 +_PDF_INCOMPLETE_RATIO = 0.75 +_PDF_INCOMPLETE_MIN_LETTERS = 200 + + +def _markdown_corrupted(text: str) -> bool: + """True when pymupdf4llm's glyph reconstruction mangled the text: shaped RTL + Presentation Forms or U+FFFD replacements above a small floor/ratio (so a lone + legitimate shaped glyph does not force the fallback).""" + if not text: + return False + threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text)) + shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text)) + return shaped > threshold or text.count("\ufffd") > threshold + + +def _markdown_incomplete(markdown: str, plain: str) -> bool: + """True when ``markdown`` holds far fewer letters than the raw ``get_text`` layer -- a + coarse guard for heavy-RTL pages pymupdf4llm silently drops without shaped glyphs.""" + plain_letters = sum(1 for c in plain if c.isalnum()) + if plain_letters < _PDF_INCOMPLETE_MIN_LETTERS: + return False + markdown_letters = sum(1 for c in markdown if c.isalnum()) + return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters + + def _pdf_markdown(doc) -> list[str] | None: """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index i maps to page i+1. Returns None when the lib is missing, extraction fails, or the @@ -100,9 +134,19 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: try: md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None for i, page in enumerate(doc): - # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval); - # fall back to plain text when Markdown is off, unavailable, or empty here. - text = (md[i] if md else "") or page.get_text("text") or "" + plain = page.get_text("text") or "" + candidate = md[i] if md else "" + # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval), + # but drop to PyMuPDF's logical-order text when Markdown is off/empty or when + # pymupdf4llm mangled it (RTL/Indic) or dropped most of the page. + if ( + candidate + and not _markdown_corrupted(candidate) + and not _markdown_incomplete(candidate, plain) + ): + text = candidate + else: + text = plain pages.append(_page(text, i + 1)) if want_images: for img in page.get_images(full = True): @@ -308,12 +352,61 @@ def render_pdf_pages( doc.close() +def _docx_table_rows(table) -> list[str]: + """Each row as pipe-joined cell text (the locator splits anchors on pipes). + Columns stay aligned to the layout grid (merged cells fill their spanned slots, + skipped leading/trailing grid columns become empty fields). Cells are walked in + document order so a nested table, and any text after it, flattens in place.""" + from docx.table import Table + from docx.text.paragraph import Paragraph + + rows: list[str] = [] + seen: set = set() # already emitted; dedups merges spanning columns or rows + for row in table.rows: + cells: list[str] = [""] * getattr(row, "grid_cols_before", 0) + trailing: list[str] = [] # nested rows + any post-nested text, kept in order + for cell in row.cells: + # A merged cell shares one across the columns and rows it spans: + # emit its text once, then placeholders, so columns and rows stay aligned. + if cell._tc in seen: + cells.append("") + continue + seen.add(cell._tc) + # Paragraph text before the first nested table is the aligned field; the + # nested table and anything after it flatten below the row, in order. + field: list[str] = [] + after_table = False + for item in cell.iter_inner_content(): + if isinstance(item, Table): + after_table = True + trailing.extend(_docx_table_rows(item)) + elif isinstance(item, Paragraph): + text = " ".join(item.text.split()) # collapse in-cell newlines + if text: + (trailing if after_table else field).append(text) + cells.append(" ".join(field)) # empty cells kept so columns line up + cells.extend([""] * getattr(row, "grid_cols_after", 0)) + if any(c.strip() for c in cells): + rows.append(" | ".join(cells)) + rows.extend(trailing) + return rows + + def _docx(path: str) -> list[Page]: import docx + from docx.table import Table + from docx.text.paragraph import Paragraph document = docx.Document(path) - text = "\n".join(p.text for p in document.paragraphs) - return [_page(text, None)] + lines: list[str] = [] + # Walk body content in document order: paragraphs alone drop tables entirely. + for block in document.iter_inner_content(): + if isinstance(block, Paragraph): + if block.text.strip(): + lines.append(block.text) + elif isinstance(block, Table): + lines.extend(_docx_table_rows(block)) + return [_page("\n".join(lines), None)] def parse(path: str, *, want_images: bool = False): diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py index 4c46f49495..14ab0efe2e 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -86,3 +86,212 @@ def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch): _table_pdf(pdf) pages = parsers.parse(str(pdf)) assert pages and "Quarter" in pages[0].text + + +def _long_text_pdf(path): + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + body = "The quick brown fox jumps over the lazy dog. " * 12 # >200 letters + page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def test_pdf_markdown_corruption_falls_back_to_plain(tmp_path, monkeypatch): + # pymupdf4llm can emit shaped RTL Presentation Forms for Arabic/Hebrew; the parser + # detects that and uses PyMuPDF's logical-order text instead of the mangled Markdown. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + shaped = "".join(chr(c) for c in range(0xFE8D, 0xFEA0)) * 20 # heavy shaped forms + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: [shaped] * doc.page_count) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Quarter" in text # real logical-order text recovered + assert not parsers._markdown_corrupted(text) # shaped garbage not carried through + + +def test_pdf_markdown_incomplete_falls_back_to_plain(tmp_path, monkeypatch): + # If pymupdf4llm silently drops most of a page, the parser prefers the fuller raw layer. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: ["x"] * doc.page_count) + pdf = tmp_path / "long.pdf" + _long_text_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "quick brown fox" in text # fuller raw layer used, not the near-empty Markdown + + +def _docx_with_table(path): + import docx + + document = docx.Document() + document.add_paragraph("Intro before table.") + table = document.add_table(rows = 2, cols = 2) + table.cell(0, 0).text = "NAME" + table.cell(0, 1).text = "SCORE" + table.cell(1, 0).text = "Alice" + table.cell(1, 1).text = "97pts" + document.add_paragraph("Outro after table.") + document.save(str(path)) + + +def test_docx_extracts_table_cells(tmp_path): + # document.paragraphs alone drops tables; the parser walks body content in order so + # table cells survive (pipe-joined, which the preview locator anchors on). + pytest.importorskip("docx") + from core.rag import parsers + + docx_path = tmp_path / "t.docx" + _docx_with_table(docx_path) + text = "\n".join(p.text for p in parsers.parse(str(docx_path))) + assert all(v in text for v in ("NAME", "SCORE", "Alice", "97pts")) # cells kept + assert "Alice | 97pts" in text # row cells joined + assert text.index("Intro") < text.index("NAME") < text.index("Outro") # order kept + + +def test_docx_table_keeps_columns_and_collapses_cell_newlines(tmp_path): + # Empty cells are kept (so columns stay aligned across rows) and a cell's internal + # newlines are collapsed to spaces (so a multi-paragraph cell can't break the row). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "A" + table.cell(0, 1).text = "" # empty middle cell + table.cell(0, 2).text = "C" + multiline = table.cell(1, 0) + multiline.text = "line1" + multiline.add_paragraph("line2") # cell now holds an internal newline + table.cell(1, 1).text = "mid" + table.cell(1, 2).text = "end" + path = tmp_path / "aligned.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "A | | C" in text # empty cell preserved -> columns line up + assert "line1 line2 | mid | end" in text # internal newline collapsed to a space + + +def test_docx_table_merged_cell_keeps_grid_alignment(tmp_path): + # A horizontally merged cell repeats across the spanned columns: emit its text once + # then a placeholder, so the row keeps as many fields as its siblings (columns stay + # aligned) without duplicating the merged text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "WIDE" + table.cell(0, 2).text = "END" + table.cell(0, 0).merge(table.cell(0, 1)) # span the first two columns + table.cell(1, 0).text = "a" + table.cell(1, 1).text = "b" + table.cell(1, 2).text = "c" + path = tmp_path / "merged.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("WIDE") == 1 # merged cell not duplicated across spanned columns + assert "WIDE | | END" in text # placeholder keeps 3 fields, aligned with "a | b | c" + assert "a | b | c" in text + + +def test_docx_table_pads_omitted_grid_columns(tmp_path): + # A row that skips leading grid columns exposes the gap via grid_cols_before; pad it + # with empty fields so the value stays under the right header instead of shifting left. + pytest.importorskip("docx") + import docx + from docx.oxml.ns import qn + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "H1" + table.cell(0, 1).text = "H2" + table.cell(0, 2).text = "H3" + tr = table.rows[1]._tr # drop the first cell and mark it skipped via + tr.remove(tr.tc_lst[0]) + trPr = tr.get_or_add_trPr() + trPr.insert(0, trPr.makeelement(qn("w:gridBefore"), {qn("w:val"): "1"})) + table.rows[1].cells[0].text = "X" # sits in column 2 + path = tmp_path / "gap.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert " | X | " in text # leading gap padded so X lines up under H2, not H1 + + +def test_docx_flattens_nested_table(tmp_path): + # cell.text ignores tables nested inside a cell; walk cell.tables so nested rows are + # not silently dropped from the indexed text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + outer = document.add_table(rows = 1, cols = 1).cell(0, 0) + outer.text = "outer" + nested = outer.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + path = tmp_path / "nested.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "NESTED-A | NESTED-B" in text # nested table flattened, not dropped + + +def test_docx_nested_table_keeps_in_cell_order(tmp_path): + # A cell holding paragraph, nested table, paragraph must serialize in that order + # (cell.text alone would emit both paragraphs before the nested rows). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + cell = document.add_table(rows = 1, cols = 1).cell(0, 0) + cell.text = "before" + nested = cell.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + cell.add_paragraph("after") + path = tmp_path / "nested_order.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.index("before") < text.index("NESTED-A") < text.index("after") + + +def test_docx_table_vertical_merge_emitted_once(tmp_path): + # A vertically merged cell maps every continuation row back to the origin ; + # emit it once and leave placeholders below so a row-spanning label isn't repeated. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 3, cols = 2) + table.cell(0, 0).merge(table.cell(1, 0)).merge(table.cell(2, 0)).text = "SECTION" + table.cell(0, 1).text = "r0" + table.cell(1, 1).text = "r1" + table.cell(2, 1).text = "r2" + path = tmp_path / "vmerge.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("SECTION") == 1 # not repeated on each spanned row + assert "SECTION | r0" in text and " | r1" in text and " | r2" in text From 4f24b12cc933db37457a9630d4db4fd799af5a13 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Thu, 2 Jul 2026 05:26:33 -0700 Subject: [PATCH 05/23] Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) * Add customizable RAG embedding model setting and reorganize settings tabs Chat with files, project sources, and knowledge bases previously always embedded with unsloth/bge-small-en-v1.5. This adds a Settings option to pick any Hugging Face embedding model (or local path), with HF search autocomplete, server-side verification that the repo is actually an embedding model, and a save anyway escape hatch for offline or local models. The setting persists in app_settings and applies at runtime to both the sentence-transformers and llama-server GGUF embedder backends without a restart. Also reorganizes the General settings tab: Documents & RAG sits above Uploads, Helper LLM moved above the danger zone, and Model auto-switch (OpenAI API) moved to the bottom of the API tab. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Support local model paths on the GGUF embedder and normalize default saves Found by simulation testing of the embedding model setting: Local paths saved as the embedding model now work on the llama-server GGUF backend (the default backend on macOS and CPU). A path to a .gguf file is used directly and a directory is scanned for a variant-matching non-mmproj .gguf, with a clear error when none exists. Previously a local path was sent to the HF hub API and failed with a repo lookup error. Saving the default model explicitly no longer stores an override, so is_custom stays false and the UI does not show a reset button for the default value. * Address review: stale-vector handling, GGUF derivation, save-time guards Review follow-ups, each verified by new tests: Re-uploading a document after an embedding model change now re-indexes instead of deduping by content hash. Documents record the embedder that produced their vectors (lazy embedding_model column, NULL legacy rows keep deduping) and a mismatch replaces the old document. A vector width change no longer bricks the dense index. ensure_vec drops and recreates chunks_vec when the dim changes (old vectors are in a foreign space and only block inserts) and search_dense returns empty on a width mismatch instead of surfacing a vec0 error, so lexical search keeps working until documents are re-uploaded. Saving a local sentence-transformers folder with no .gguf now returns 409 with a clear message when the install embeds via llama-server, instead of failing at first index. force still saves. A custom RAG_EMBEDDING_MODEL env without RAG_EMBED_GGUF_REPO now derives the -GGUF companion repo instead of silently keeping the bge GGUF on CPU and macOS installs. The resolved GGUF path is tagged with the repo captured at entry, so a setting change during a download cannot mark the old model as current. GGUF repo detection matches gguf as a whole name segment rather than a substring, hf_token is trimmed before verification, and the settings combobox drops a redundant state mirror of its controlled value. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Shrink embedding model font to 11px in the input and dropdown The combobox wrapper applies className to the outer input group, so the size utility must target the inner input element; the previous text-xs never reached it and the field rendered at the browser default. * Show curated unsloth embedding models when the search field is empty The empty-query listing was the global top-downloads page, which holds no unsloth mirrors for the unsloth-first float to reorder, so the dropdown opened on third-party models. Match the model picker: curated unsloth listing when empty, whole-Hub search once a query is typed. * Address review: settings resilience and index consistency Keep the last known embedding model on settings store errors, remove the re-entrant dim lock in the llama-server backend, accept local GGUF saves and verify GGUF availability for HF repos on that backend, match local path embedders exactly in model list filters, drop same-width stale vectors from dense search, pin the embedder per ingestion job, and only replace completed documents after the re-index succeeds. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Consolidate the GGUF repo derivation tests * Trim to a single core embedding-model test * Address review: GGUF repo saves and cache race Accept a GGUF-named HF repo on the llama-server backend by verifying GGUF availability instead of the sentence-transformers metadata gate, and guard the settings cache with a generation counter so a read overlapping a save cannot repopulate it with the pre-save value. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/rag/config.py | 41 +++- studio/backend/core/rag/embed_llama_server.py | 107 +++++++-- studio/backend/core/rag/embeddings.py | 2 +- studio/backend/core/rag/ingestion.py | 45 +++- studio/backend/core/rag/retrieval.py | 8 +- studio/backend/core/rag/store.py | 40 +++- studio/backend/routes/models.py | 39 ++- studio/backend/routes/rag.py | 2 +- studio/backend/routes/settings.py | 191 ++++++++++++++- studio/backend/storage/rag_db.py | 33 ++- .../tests/test_embedding_model_settings.py | 55 +++++ .../tests/test_rag_embed_llama_server.py | 2 + .../backend/utils/embedding_model_settings.py | 124 ++++++++++ .../features/settings/api/embedding-model.ts | 82 +++++++ .../components/embedding-model-combobox.tsx | 134 +++++++++++ .../features/settings/tabs/api-keys-tab.tsx | 3 + .../features/settings/tabs/general-tab.tsx | 225 +++++++++++++++--- studio/frontend/src/i18n/locales/en.ts | 14 ++ 18 files changed, 1072 insertions(+), 75 deletions(-) create mode 100644 studio/backend/tests/test_embedding_model_settings.py create mode 100644 studio/backend/utils/embedding_model_settings.py create mode 100644 studio/frontend/src/features/settings/api/embedding-model.ts create mode 100644 studio/frontend/src/features/settings/components/embedding-model-combobox.tsx diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 54a224d081..2de32a68e4 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -6,8 +6,10 @@ from __future__ import annotations import os +import re -EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "unsloth/bge-small-en-v1.5") +DEFAULT_EMBEDDING_MODEL = "unsloth/bge-small-en-v1.5" +EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL) # Under bge's 512 limit, leaving headroom for the 2 special tokens (else overflow: # llama-server 500s, ST truncates). Keep <= embedder_max - ~12. CHUNK_TOKENS = int(os.environ.get("RAG_CHUNK_TOKENS", "500")) @@ -66,6 +68,43 @@ OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048")) # wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes # the vectors, so the index must be rebuilt. EMBED_BACKEND = os.environ.get("RAG_EMBED_BACKEND", "auto") + + +def effective_embedding_model() -> str: + """The embedding model actually in use: the persisted Settings override when + one is stored, else ``EMBEDDING_MODEL`` (env/default). Read at call time so a + Settings change applies without a restart.""" + try: + from utils.embedding_model_settings import get_rag_embedding_model + return get_rag_embedding_model() + except Exception: # noqa: BLE001 - settings store unavailable (tests, early boot) + return EMBEDDING_MODEL + + +def _names_gguf(model: str) -> bool: + """True when "gguf" appears as a whole name segment, so plain substrings + like "bigguf" don't count.""" + return "gguf" in re.split(r"[^a-z0-9]+", model.lower()) + + +def effective_gguf_repo() -> str: + """GGUF repo for the llama-server backend, tracking the effective model. + + An explicit ``RAG_EMBED_GGUF_REPO`` env always wins. Otherwise any custom + model (saved in Settings or via ``RAG_EMBEDDING_MODEL``) maps to its + ``-GGUF`` companion repo (the unsloth convention the default pair follows), + or is used as-is when it already names a GGUF repo. + """ + if "RAG_EMBED_GGUF_REPO" in os.environ: + return EMBED_GGUF_REPO + model = effective_embedding_model() + if model == DEFAULT_EMBEDDING_MODEL: + return EMBED_GGUF_REPO + if _names_gguf(model): + return model + return f"{model}-GGUF" + + # llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this # tiny model) and exact vs fp32, for ~30MB more on disk. EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF") diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index f53478463c..46a282c939 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -55,9 +55,15 @@ class LlamaServerBackend: self._port: int | None = None self._stdout_lines: list[str] = [] self._stdout_thread: threading.Thread | None = None + # No lock: probes are idempotent (a duplicate 1-text encode is benign) + # and dim() -> encode() -> _ensure_ready() -> _resolve_model_path() can + # re-enter on a mid-probe model change, which would self-deadlock a + # non-reentrant lock held across the probe. self._dim: int | None = None - self._dim_lock = threading.Lock() self._model_path: str | None = None + # Effective GGUF repo the cached path/dim belong to; a Settings change + # makes it stale, forcing a re-resolve + respawn (see _ensure_ready). + self._model_repo: str | None = None self._binary: str | None = None # Sticky after an auto GPU start fails: later spawns stay on CPU. self._force_cpu = False @@ -114,24 +120,77 @@ class LlamaServerBackend: "RAG_EMBED_BACKEND=llama-server requires an embeddings-capable build" ) + @staticmethod + def _resolve_local_gguf(model: str) -> str | None: + """A custom model may be a local .gguf file or a directory holding one; + resolve it without the hub. None when the value is not a local path.""" + p = Path(model).expanduser() + if p.is_file() and p.suffix.lower() == ".gguf": + return str(p) + if p.is_dir(): + files = [ + f + for f in p.iterdir() + if f.suffix.lower() == ".gguf" and "mmproj" not in f.name.lower() + ] + if not files: + raise RuntimeError(f"no .gguf file found in local model dir {model!r}") + variant = config.EMBED_GGUF_VARIANT.lower() + match = [f for f in files if variant in f.name.lower()] or files + return str(sorted(match, key = lambda f: len(f.name))[0]) + return None + def _resolve_model_path(self) -> str: """Download (or cache-hit) the variant-matching, non-mmproj GGUF embedder, - returning its local path.""" - if self._model_path is not None: + returning its local path. Re-resolves when the effective repo changed (a + custom model was saved in Settings).""" + # Captured once: if the setting changes mid-download, the path must stay + # tagged with the repo it was resolved FOR, so _current() sees the new + # setting as stale and respawns instead of serving the old model. + desired = config.effective_gguf_repo() + if self._model_path is not None and self._model_repo == desired: + return self._model_path + local = self._resolve_local_gguf(config.effective_embedding_model()) + if local is not None: + self._model_path = local + self._model_repo = desired + self._dim = None return self._model_path from huggingface_hub import hf_hub_download, list_repo_files - repo = config.EMBED_GGUF_REPO token = os.environ.get("HF_TOKEN") or None - files = [f for f in list_repo_files(repo, token = token) if f.lower().endswith(".gguf")] - files = [f for f in files if "mmproj" not in f.lower()] + # A custom model derives its "-GGUF" companion repo; when that guess does + # not exist, the model repo itself may host the .gguf files. + repo = desired + candidates = [repo] + model = config.effective_embedding_model() + if model != repo: + candidates.append(model) + files: list[str] = [] + errors: list[str] = [] + for candidate in candidates: + try: + files = [ + f + for f in list_repo_files(candidate, token = token) + if f.lower().endswith(".gguf") and "mmproj" not in f.lower() + ] + except Exception as e: # noqa: BLE001 - missing/gated repo -> next candidate + errors.append(f"{candidate!r}: {e}") + continue + if files: + repo = candidate + break + errors.append(f"{candidate!r}: no .gguf files") if not files: - raise RuntimeError(f"no .gguf file found in embedder repo {repo!r}") + raise RuntimeError("no .gguf embedder found; tried " + "; ".join(errors)) variant = config.EMBED_GGUF_VARIANT.lower() match = [f for f in files if variant in f.lower()] or files filename = sorted(match, key = len)[0] logger.info("resolving GGUF embedder %s/%s", repo, filename) self._model_path = hf_hub_download(repo_id = repo, filename = filename, token = token) + self._model_repo = desired + self._dim = None return self._model_path # Min free VRAM (MiB) for the embedder; below this, auto stays on CPU. @@ -316,13 +375,19 @@ class LlamaServerBackend: def _process_alive(self) -> bool: return self._process is not None and self._process.poll() is None + def _current(self) -> bool: + """Alive AND serving the effective repo (a Settings model change makes a + live server stale).""" + return self._process_alive() and self._model_repo == config.effective_gguf_repo() + def _ensure_ready(self) -> None: - """Guarantee a live server, (re)spawning if needed. Double-checked so the - alive path takes no lock; self-heals after the chat reaper kills us.""" - if self._process_alive(): + """Guarantee a live server on the effective model, (re)spawning if needed. + Double-checked so the current path takes no lock; self-heals after the + chat reaper kills us and re-resolves after a Settings model change.""" + if self._current(): return with self._lifecycle_lock: - if self._process_alive(): + if self._current(): return self._kill_process() self._spawn() @@ -424,14 +489,18 @@ class LlamaServerBackend: return arr def dim(self, *, model_name = None) -> int: - """Embedding width, probed once via a 1-text encode and cached.""" - if self._dim is not None: - return self._dim - with self._dim_lock: - if self._dim is None: - vec = self.encode(["x"], normalize = False) - self._dim = int(vec.shape[1]) - return self._dim + """Embedding width, probed via a 1-text encode and cached per model + (_resolve_model_path clears it when the effective repo changes). + Unlocked: concurrent probes are benign, and locking would deadlock when + the probe's encode respawns onto a changed model (see __init__).""" + self._ensure_ready() + cached = self._dim + if cached is not None: + return cached + vec = self.encode(["x"], normalize = False) + width = int(vec.shape[1]) + self._dim = width + return width def warm(self, *, model_name = None) -> None: """Start the server and probe dim off the request path.""" diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 4c8d690302..345b4dd853 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -67,7 +67,7 @@ def _get(model_name: str | None = None): """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 for a ~1.5x speedup at negligible accuracy loss.""" global _model, _name - name = model_name or config.EMBEDDING_MODEL + name = model_name or config.effective_embedding_model() with _lock: if _model is None or _name != name: _install_torchao_stub_once() diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 04365ab76b..cba076f1be 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -151,6 +151,19 @@ def _ocr_scanned_pages( return out, ocred +def _replace_old_document(conn, replaces: tuple[str, str | None] | None, keep_path: str) -> None: + """Drop the document this ingestion replaced (stale embedder / empty prior + ingest), called only after the replacement completed successfully.""" + if replaces is None: + return + old_id, old_path = replaces + try: + store.delete_document(conn, old_id) + _remove_upload(old_path, keep_path = keep_path) + except Exception: # noqa: BLE001 - the new document is already live + logger.warning("failed to remove replaced document %s", old_id, exc_info = True) + + def _run( job_id: str, document_id: str, @@ -159,6 +172,7 @@ def _run( model_name: str | None, ocr: bool | None = None, caption: bool | None = None, + replaces: tuple[str, str | None] | None = None, ) -> None: conn = rag_db.get_connection() try: @@ -213,6 +227,7 @@ def _run( ) if not chunks: store.set_document_status(conn, document_id, "completed", num_chunks = 0) + _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": 0}) return @@ -233,6 +248,7 @@ def _run( _progress(conn, job_id, "storing", 0.9) store.add_chunks(conn, scope, document_id, chunks, vectors, regions) store.set_document_status(conn, document_id, "completed", num_chunks = len(chunks)) + _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": len(chunks)}) @@ -274,17 +290,32 @@ def start_ingestion( sha = _sha256_file(stored_path) conn = rag_db.get_connection() try: + effective_model = model_name or config.effective_embedding_model() + # (old_document_id, old_stored_path) replaced by this upload; deleted by + # the worker only after the replacement completes, so a failed re-index + # never destroys the still-searchable original. + replaces: tuple[str, str | None] | None = None existing = store.document_by_hash(conn, scope, sha) if existing is not None: doc = store.get_document(conn, existing) empty_completed = ( doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks") ) - if empty_completed: + # Vectors from a different embedder are stale; re-uploading must + # re-index, not dedupe. NULL (legacy rows) is assumed current. Only + # completed rows are replaceable: a pending/running duplicate has a + # live worker whose writes must not land on a deleted document. + stale_model = ( + doc is not None + and doc.get("status") == "completed" + and doc.get("embedding_model") is not None + and doc.get("embedding_model") != effective_model + ) + if empty_completed or stale_model: # A prior ingest of identical bytes yielded zero chunks (e.g. a scanned - # PDF uploaded before a vision model loaded). Re-ingest, don't dedupe. - store.delete_document(conn, existing) - _remove_upload(doc.get("stored_path"), keep_path = stored_path) + # PDF uploaded before a vision model loaded), or was embedded with a + # different model. Re-ingest, don't dedupe. + replaces = (existing, doc.get("stored_path")) else: job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) _remove_upload(stored_path) @@ -310,6 +341,7 @@ def start_ingestion( project_id = project_id, status = "pending", stored_path = stored_path, + embedding_model = effective_model, ) job_id = _new_job(conn, document_id, scope) finally: @@ -319,7 +351,10 @@ def start_ingestion( _jobs[job_id] = queue.Queue() threading.Thread( target = _run, - args = (job_id, document_id, scope, stored_path, model_name, ocr, caption), + # effective_model (not the raw model_name) pins the embedder for the + # whole job: a Settings change mid-ingestion must not switch tokenizer + # or embedder between batches of one document. + args = (job_id, document_id, scope, stored_path, effective_model, ocr, caption, replaces), daemon = True, ).start() return document_id, job_id diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index fe6a033a52..6f933e089e 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -39,8 +39,12 @@ def retrieve_dense( model_name: str | None = None, ) -> list[Hit]: k = k or config.TOP_K_DENSE - vec = embeddings.encode([query], model_name = model_name, normalize = True)[0] - return [Hit(cid, s, dense_score = s) for cid, s in store.search_dense(conn, scope, vec, k)] + effective = model_name or config.effective_embedding_model() + vec = embeddings.encode([query], model_name = effective, normalize = True)[0] + return [ + Hit(cid, s, dense_score = s) + for cid, s in store.search_dense(conn, scope, vec, k, embedding_model = effective) + ] def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]: diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index 8e59c5fbf6..f9128d1715 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -109,11 +109,12 @@ def create_document( status: str = "pending", stored_path: str | None = None, document_id: str | None = None, + embedding_model: str | None = None, ) -> str: document_id = document_id or str(uuid.uuid4()) conn.execute( "INSERT INTO documents(id, scope, kb_id, thread_id, project_id, filename, sha256, " - "status, stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?,?)", + "status, stored_path, created_at, embedding_model) VALUES(?,?,?,?,?,?,?,?,?,?,?)", ( document_id, scope, @@ -125,6 +126,7 @@ def create_document( status, stored_path, _now(), + embedding_model, ), ) conn.commit() @@ -261,20 +263,50 @@ def search_lexical(conn: sqlite3.Connection, scope, query: str, k: int): return [(r["chunk_id"], -r["s"]) for r in rows] -def search_dense(conn: sqlite3.Connection, scope, vector, k: int): +def search_dense( + conn: sqlite3.Connection, + scope, + vector, + k: int, + *, + embedding_model: str | None = None, +): """Cosine KNN over vec0 for one scope or several. Returns [(chunk_id, 1 - distance)]. vec0 KNN constrains its partition key by - equality, so multi-scope runs one query per scope and merges by score.""" + equality, so multi-scope runs one query per scope and merges by score. + ``embedding_model`` drops hits from documents indexed under a different + (same-width) model, whose vectors live in another space; NULL-model legacy + documents are assumed current, matching the ingestion dedupe rule.""" if not rag_db.vec_table_exists(conn): return [] + dim = rag_db.vec_table_dim(conn) + if dim is not None and dim != len(vector): + # Embedding model switched widths and nothing re-indexed yet; the stale + # table cannot answer new-model queries (vec0 errors on the MATCH). + return [] + # Over-fetch when filtering so stale-model hits don't starve the top-k. + fetch = k * 3 if embedding_model else k out: list[tuple[str, float]] = [] for s in _scopes(scope): rows = conn.execute( "SELECT chunk_id, distance FROM chunks_vec " "WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?", - (s, _f32(vector), k), + (s, _f32(vector), fetch), ).fetchall() out.extend((r["chunk_id"], 1.0 - r["distance"]) for r in rows) + if embedding_model and out: + ids = [cid for cid, _ in out] + placeholders = ",".join("?" * len(ids)) + valid = { + r["id"] + for r in conn.execute( + f"SELECT c.id FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.id IN ({placeholders}) " + f"AND (d.embedding_model IS NULL OR d.embedding_model=?)", + (*ids, embedding_model), + ).fetchall() + } + out = [t for t in out if t[0] in valid] out.sort(key = lambda t: t[1], reverse = True) return out[:k] diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 7c75e85227..1501868860 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -7,6 +7,7 @@ import asyncio import hashlib import json import os +import re import shutil import sys import uuid @@ -58,25 +59,51 @@ def _safe_is_dir(path) -> bool: return False +# Hub repo id shape ("owner/name", no leading separator); anything else is +# treated as a local filesystem path. +_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$") + + def _is_hidden_model(*values: str | None) -> bool: """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or EMBED_GGUF_REPO basename) or the llama.cpp install validation probe (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF). None are usable chat models; the probe can be cached as a side effect of installing the prebuilt llama-server and otherwise sorts smallest, so it - would be auto-selected.""" + would be auto-selected. A local-path embedder is matched by exact resolved + path only: a generic basename like "model" must not substring-hide + unrelated chat models.""" from core.rag import config as rag_config - needles = ( - rag_config.EMBEDDING_MODEL.split("/")[-1].lower(), - rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(), + needles = [ # The validation probe's repo (matches the cached repo id) and its exact # filename (matches the on-disk path). The filename carries the .gguf so # it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``. "ggml-org/models", "stories260k.gguf", - ) - return any(v and any(n in v.lower() for n in needles) for v in values) + ] + exact_paths: list[str] = [] + for model in ( + rag_config.effective_embedding_model(), + rag_config.effective_gguf_repo(), + ): + if _HF_REPO_ID_RE.match(model): + needles.append(model.split("/")[-1].lower()) + else: + resolved = _safe_resolve(Path(model).expanduser()) + if resolved: + exact_paths.append(resolved.lower()) + for v in values: + if not v: + continue + low = v.lower() + if any(n in low for n in needles): + return True + if exact_paths: + resolved = _safe_resolve(Path(v).expanduser()) + if resolved and resolved.lower() in exact_paths: + return True + return False def _safe_resolve(path: Path) -> Optional[str]: diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 4e35fce3c2..e20fea74a3 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -167,7 +167,7 @@ def create_knowledge_base( conn, name = payload.name.strip(), description = (payload.description or None), - embedding_model = config.EMBEDDING_MODEL, + embedding_model = config.effective_embedding_model(), ) return {"id": kb_id, "name": payload.name.strip()} finally: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 0694ae31e0..862bce8be8 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -4,7 +4,7 @@ from typing import Literal, Optional from urllib.parse import unquote, urlsplit -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, ConfigDict, Field, field_validator from auth.authentication import get_current_subject @@ -47,6 +47,15 @@ from utils.preview_sharing_settings import ( get_preview_sharing_enabled, set_preview_sharing_enabled, ) +from utils.embedding_model_settings import ( + MAX_EMBEDDING_MODEL_LENGTH, + default_embedding_model, + get_rag_embedding_model, + get_stored_embedding_model, + reset_rag_embedding_model, + set_rag_embedding_model, + validate_embedding_model, +) router = APIRouter() @@ -229,6 +238,186 @@ def update_openai_auto_switch_override( return ModelOverridesResponse(overrides = get_model_overrides()) +class EmbeddingModelPayload(BaseModel): + embedding_model: str = Field(..., min_length = 1, max_length = MAX_EMBEDDING_MODEL_LENGTH) + # Token for gated/private repos during verification (not stored). + hf_token: Optional[str] = Field(default = None, max_length = 512) + # Skip HF verification (offline installs, local paths HF can't see). + force: bool = False + + +class EmbeddingModelResponse(BaseModel): + embedding_model: str + default_embedding_model: str + is_custom: bool + + +def _embedding_model_response() -> EmbeddingModelResponse: + return EmbeddingModelResponse( + embedding_model = get_rag_embedding_model(), + default_embedding_model = default_embedding_model(), + is_custom = get_stored_embedding_model() is not None, + ) + + +def _llama_backend_active() -> bool: + """True when this install embeds via the llama-server (GGUF) backend.""" + from core.rag import config as rag_config + from core.rag import embeddings + + try: + raw = (rag_config.EMBED_BACKEND or "auto").strip().lower() + key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw + except Exception: # noqa: BLE001 - backend probe must never block saving + return False + return key in embeddings._LLAMA_ALIASES + + +def _resolves_as_local_gguf(model: str) -> bool: + """True when ``model`` is a local .gguf file or a directory holding one, so + a save on the llama-server backend needs no HF verification (the artifact + itself is the proof).""" + from core.rag.embed_llama_server import LlamaServerBackend + try: + return LlamaServerBackend._resolve_local_gguf(model) is not None + except Exception: # noqa: BLE001 - dir without .gguf, filesystem oddity + return False + + +def _local_gguf_backend_error(model: str) -> str | None: + """409 detail when ``model`` is a local dir without a .gguf but this install + embeds via llama-server (macOS/CPU default), which needs one. A + sentence-transformers-only folder would verify fine yet fail at first index. + None when not applicable. ``force`` skips this check like HF verification.""" + from pathlib import Path + + if not Path(model).expanduser().is_dir(): + return None + from core.rag.embed_llama_server import LlamaServerBackend + + if not _llama_backend_active(): + return None + try: + LlamaServerBackend._resolve_local_gguf(model) + return None + except RuntimeError: + return ( + f"{model!r} contains no .gguf file, but this install embeds with the " + "llama-server backend which requires one. Add a GGUF file to the " + "folder or use a Hugging Face repo." + ) + except Exception: # noqa: BLE001 - filesystem oddity: don't block saving + return None + + +def _hf_gguf_backend_error(model: str, hf_token: Optional[str]) -> str | None: + """409 detail when the llama-server backend would find no .gguf for an HF + repo: neither the derived companion repo nor the repo itself has one. Saves + that verify as embedding models would otherwise fail at first index. + None when not applicable; ``force`` skips this like HF verification.""" + from pathlib import Path + + if Path(model).expanduser().exists(): + return None # local paths are handled by the local checks + if not _llama_backend_active(): + return None + from core.rag import config as rag_config + + candidates = [model] if rag_config._names_gguf(model) else [f"{model}-GGUF", model] + try: + from huggingface_hub import list_repo_files + except Exception: # noqa: BLE001 - hub client unavailable: don't block saving + return None + for candidate in candidates: + try: + files = list_repo_files(candidate, token = hf_token) + except Exception: # noqa: BLE001 - missing/gated repo: try next candidate + continue + if any(f.lower().endswith(".gguf") and "mmproj" not in f.lower() for f in files): + return None + checked = " or ".join(repr(c) for c in candidates) + return ( + f"No GGUF weights found in {checked}, but this install embeds with the " + "llama-server backend which requires them. Pick a model with a GGUF " + "companion repo or GGUF files in the repo itself." + ) + + +@router.get("/embedding-model", response_model = EmbeddingModelResponse) +def get_embedding_model( + current_subject: str = Depends(get_current_subject), +) -> EmbeddingModelResponse: + return _embedding_model_response() + + +@router.put("/embedding-model", response_model = EmbeddingModelResponse) +def update_embedding_model( + payload: EmbeddingModelPayload, current_subject: str = Depends(get_current_subject) +) -> EmbeddingModelResponse: + """Set the RAG embedding model. Unless ``force`` is set, the repo is verified + to be an embedding model via HF metadata; an unverifiable model (wrong type, + typo, gated repo, or no network) returns 409 so the UI can offer "save anyway". + Documents indexed under the previous model must be re-uploaded.""" + from utils.models import is_embedding_model + + try: + model = validate_embedding_model(payload.embedding_model) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid embedding model."), + event = "settings.update_embedding_model_failed", + log = logger, + ) from exc + # The env/default model needs no verification; saving it is a no-op override. + # A local GGUF on the llama-server backend is accepted as-is: it is exactly + # what the backend loads, and HF metadata cannot verify a local path. + if ( + model != default_embedding_model() + and not payload.force + and not (_llama_backend_active() and _resolves_as_local_gguf(model)) + ): + hf_token = (payload.hf_token or "").strip() or None + from core.rag import config as rag_config + + # A GGUF-named repo on the llama-server backend is loaded from its .gguf + # files, which rarely carry sentence-transformers metadata; verify the + # GGUF is available (below) rather than the ST embedding-metadata gate, + # which would wrongly 409 a valid online GGUF embedder. + gguf_named = _llama_backend_active() and rag_config._names_gguf(model) + if not gguf_named and not is_embedding_model(model, hf_token = hf_token): + raise HTTPException( + status_code = 409, + detail = ( + f"Could not verify {model!r} as an embedding model on " + "Hugging Face (it may be the wrong model type, gated, or " + "you may be offline)." + ), + ) + gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) + if gguf_error: + raise HTTPException(status_code = 409, detail = gguf_error) + set_rag_embedding_model(model) + logger.info( + "settings.embedding_model_updated subject=%s model=%s forced=%s", + current_subject, + model, + payload.force, + ) + return _embedding_model_response() + + +@router.delete("/embedding-model", response_model = EmbeddingModelResponse) +def reset_embedding_model( + current_subject: str = Depends(get_current_subject), +) -> EmbeddingModelResponse: + """Clear the override, returning to the env/default model.""" + reset_rag_embedding_model() + logger.info("settings.embedding_model_reset subject=%s", current_subject) + return _embedding_model_response() + + class PreviewLinkRotateResponse(BaseModel): rotated: bool = True diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index ce27326562..cbd6ceb617 100644 --- a/studio/backend/storage/rag_db.py +++ b/studio/backend/storage/rag_db.py @@ -15,6 +15,7 @@ column type). """ import logging +import re import sqlite3 import threading @@ -64,7 +65,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: error TEXT, num_chunks INTEGER NOT NULL DEFAULT 0, stored_path TEXT, - created_at TEXT NOT NULL + created_at TEXT NOT NULL, + embedding_model TEXT ); CREATE INDEX IF NOT EXISTS idx_documents_scope ON documents(scope); CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(scope, sha256); @@ -107,6 +109,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)").fetchall()} if "project_id" not in cols: conn.execute("ALTER TABLE documents ADD COLUMN project_id TEXT") + # Lazy upgrade: which embedder produced a document's vectors (NULL = legacy, + # assumed current). Dedupe re-ingests when it no longer matches. + if "embedding_model" not in cols: + conn.execute("ALTER TABLE documents ADD COLUMN embedding_model TEXT") def get_connection() -> sqlite3.Connection: @@ -143,9 +149,32 @@ def get_connection() -> sqlite3.Connection: return conn +def vec_table_dim(conn: sqlite3.Connection) -> int | None: + """Embedding width baked into ``chunks_vec``, or None when absent.""" + row = conn.execute( + "SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'" + ).fetchone() + if row is None or not row["sql"]: + return None + m = re.search(r"float\[(\d+)\]", row["sql"]) + return int(m.group(1)) if m else None + + def ensure_vec(conn: sqlite3.Connection, dim: int) -> None: """Create the dense ``chunks_vec`` table once the embedding dim is known - (vec0 bakes it into the column type). Idempotent; dim fixed per db.""" + (vec0 bakes it into the column type). A width change (embedding model + switched in Settings) drops the table: the old vectors live in a foreign + space and would only block inserts, while lexical search keeps serving old + chunks until they are re-uploaded.""" + existing = vec_table_dim(conn) + if existing is not None and existing != int(dim): + logger.warning( + "chunks_vec dim changed %d -> %d (embedding model switched); dropping " + "stale dense index. Re-upload documents to restore dense search.", + existing, + int(dim), + ) + conn.execute("DROP TABLE chunks_vec") conn.execute( f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(" f"scope TEXT partition key, " diff --git a/studio/backend/tests/test_embedding_model_settings.py b/studio/backend/tests/test_embedding_model_settings.py new file mode 100644 index 0000000000..3be4af0e32 --- /dev/null +++ b/studio/backend/tests/test_embedding_model_settings.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Test for the customizable RAG embedding model: a saved override becomes the +effective model and derives its GGUF companion for the llama-server backend.""" + +from pathlib import Path +import sys +import types as _types + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest + +import utils.embedding_model_settings as ems +from core.rag import config as rag_config + + +@pytest.fixture +def settings_store(monkeypatch): + """In-memory app_settings store patched under the module's lazy imports.""" + import storage.studio_db as studio_db + + store: dict = {} + monkeypatch.setattr( + studio_db, "get_app_setting", lambda key, fallback = None: store.get(key, fallback) + ) + monkeypatch.setattr( + studio_db, "upsert_app_settings", lambda settings: store.update(settings) or store + ) + ems._invalidate_cache() + yield store + ems._invalidate_cache() + + +def test_custom_model_overrides_default_and_derives_gguf(settings_store, monkeypatch): + """The core contract: with nothing stored the default is in effect; a saved + custom model becomes the effective embedding model and derives its -GGUF + companion (what the llama-server backend loads); reset clears the override.""" + monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False) + assert ems.get_rag_embedding_model() == rag_config.EMBEDDING_MODEL + assert rag_config.effective_gguf_repo() == rag_config.EMBED_GGUF_REPO + + assert ems.set_rag_embedding_model(" org/my-embedder ") == "org/my-embedder" + assert rag_config.effective_embedding_model() == "org/my-embedder" + assert rag_config.effective_gguf_repo() == "org/my-embedder-GGUF" + + assert ems.reset_rag_embedding_model() == rag_config.EMBEDDING_MODEL + assert ems.get_stored_embedding_model() is None diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py index 8321068afd..0e1f74cefe 100644 --- a/studio/backend/tests/test_rag_embed_llama_server.py +++ b/studio/backend/tests/test_rag_embed_llama_server.py @@ -373,6 +373,8 @@ def test_ensure_ready_respawns_dead_process(monkeypatch): def fake_spawn(): spawned["n"] += 1 b._process = _FakeProc(alive = True) + # _current() now also checks the served repo, so mark it current. + b._model_repo = config.effective_gguf_repo() monkeypatch.setattr(b, "_spawn", fake_spawn) b._ensure_ready() diff --git a/studio/backend/utils/embedding_model_settings.py b/studio/backend/utils/embedding_model_settings.py new file mode 100644 index 0000000000..798ae6d364 --- /dev/null +++ b/studio/backend/utils/embedding_model_settings.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted RAG embedding-model override (Settings -> General). + +The stored value takes precedence over the ``RAG_EMBEDDING_MODEL`` env default in +``core.rag.config``. Vectors from different models live in different spaces, so +documents already indexed under the old model must be re-uploaded after a change +(the UI warns about this). +""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +EMBEDDING_MODEL_SETTING_KEY = "rag_embedding_model" +MAX_EMBEDDING_MODEL_LENGTH = 512 + +# The effective model is consulted on the embedder hot path (once per embed / +# tokenize call during ingestion), so the stored value is cached briefly instead +# of hitting sqlite each time. Writes invalidate immediately in-process; other +# readers converge within the TTL. +_CACHE_TTL_S = 2.0 +_cached: tuple[float, str | None] | None = None +# Bumped on every write/invalidate. A reader captures it before the DB read and +# only fills the cache if it is unchanged afterward, so a read that overlapped a +# save cannot repopulate the cache with the pre-save value for the whole TTL. +_generation = 0 +_lock = threading.Lock() + + +def _invalidate_cache() -> None: + global _cached, _generation + with _lock: + _cached = None + _generation += 1 + + +def default_embedding_model() -> str: + """The env/default model from rag config (``RAG_EMBEDDING_MODEL`` or bge).""" + from core.rag import config + return config.EMBEDDING_MODEL + + +def _coerce_embedding_model(value: Any) -> str | None: + if not isinstance(value, str): + return None + cleaned = value.strip() + if not cleaned or len(cleaned) > MAX_EMBEDDING_MODEL_LENGTH: + return None + # Newlines/control chars are never valid in a repo id or path. + if any(ord(ch) < 32 for ch in cleaned): + return None + return cleaned + + +def validate_embedding_model(value: Any) -> str: + cleaned = _coerce_embedding_model(value) + if cleaned is None: + raise ValueError( + "Embedding model must be a Hugging Face repo id (e.g. " + "'unsloth/bge-small-en-v1.5') or a local model path, up to " + f"{MAX_EMBEDDING_MODEL_LENGTH} characters." + ) + return cleaned + + +def get_stored_embedding_model() -> str | None: + """The persisted override, or None when unset/invalid.""" + global _cached + now = time.monotonic() + with _lock: + cached = _cached + if cached is not None and now - cached[0] < _CACHE_TTL_S: + return cached[1] + gen = _generation + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(EMBEDDING_MODEL_SETTING_KEY, None) + except Exception: + # Transient store failure: keep the last known value instead of + # silently reverting the embed/search hot path to the default model, + # which would mix vector spaces mid-ingestion. + with _lock: + if _cached is not None: + _cached = (time.monotonic(), _cached[1]) + return _cached[1] + return None + value = _coerce_embedding_model(stored) + with _lock: + # Only cache when no save landed while we were reading; otherwise this + # value may be pre-save, and caching it would mask the new one for the + # TTL. The next reader re-reads the committed value. + if _generation == gen: + _cached = (time.monotonic(), value) + return value + + +def get_rag_embedding_model() -> str: + """Effective embedding model: persisted override, else env/default.""" + return get_stored_embedding_model() or default_embedding_model() + + +def set_rag_embedding_model(value: Any) -> str: + parsed = validate_embedding_model(value) + from storage.studio_db import upsert_app_settings + + # Saving the default is not an override; keeps is_custom (and the UI's + # reset affordance) honest. + stored = parsed if parsed != default_embedding_model() else None + upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: stored}) + _invalidate_cache() + return parsed + + +def reset_rag_embedding_model() -> str: + """Clear the override; returns the (env/default) model now in effect.""" + from storage.studio_db import upsert_app_settings + + upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: None}) + _invalidate_cache() + return default_embedding_model() diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts new file mode 100644 index 0000000000..8b6bc7ee7f --- /dev/null +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type EmbeddingModelSettings = { + embeddingModel: string; + defaultEmbeddingModel: string; + isCustom: boolean; +}; + +type ApiEmbeddingModelSettings = { + // biome-ignore lint/style/useNamingConvention: API schema + embedding_model: string; + // biome-ignore lint/style/useNamingConvention: API schema + default_embedding_model: string; + // biome-ignore lint/style/useNamingConvention: API schema + is_custom: boolean; +}; + +/** 409 from the backend: the model could not be verified as an embedding model + * (wrong type, gated repo, or offline). Retry with force to save anyway. */ +export class EmbeddingModelVerificationError extends Error {} + +function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { + return { + embeddingModel: settings.embedding_model, + defaultEmbeddingModel: settings.default_embedding_model, + isCustom: settings.is_custom, + }; +} + +export async function loadEmbeddingModelSettings(): Promise { + const res = await authFetch("/api/settings/embedding-model"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load embedding model setting"), + ); + } + return fromApi(await res.json()); +} + +export async function updateEmbeddingModelSettings( + embeddingModel: string, + options?: { hfToken?: string; force?: boolean }, +): Promise { + const res = await authFetch("/api/settings/embedding-model", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + // biome-ignore lint/style/useNamingConvention: API schema + embedding_model: embeddingModel, + // biome-ignore lint/style/useNamingConvention: API schema + hf_token: options?.hfToken || null, + force: options?.force ?? false, + }), + }); + if (res.status === 409) { + throw new EmbeddingModelVerificationError( + await readFastApiError(res, "Could not verify the embedding model"), + ); + } + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to save embedding model"), + ); + } + return fromApi(await res.json()); +} + +export async function resetEmbeddingModelSettings(): Promise { + const res = await authFetch("/api/settings/embedding-model", { + method: "DELETE", + }); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to reset embedding model"), + ); + } + return fromApi(await res.json()); +} diff --git a/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx b/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx new file mode 100644 index 0000000000..b0e9a41b7d --- /dev/null +++ b/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; +import { Spinner } from "@/components/ui/spinner"; +import type { PipelineType } from "@huggingface/hub"; +import { useHubModelSearch } from "@/features/hub/hooks/use-hub-model-search"; +import { useDebouncedValue } from "@/hooks"; +import { type ReactElement, useMemo, useRef } from "react"; + +// HF pipeline filter for embedding models; matches the backend's +// is_embedding_model signals (sentence-similarity / feature-extraction). +const EMBEDDING_TASKS: readonly PipelineType[] = [ + "sentence-similarity", + "feature-extraction", +]; + +type EmbeddingModelComboboxProps = { + value: string; + /** Fires on typing, selection, and Enter with the current text. */ + onChange: (value: string) => void; + accessToken?: string; + disabled?: boolean; + placeholder?: string; + ariaLabel?: string; + className?: string; +}; + +export function EmbeddingModelCombobox({ + value, + onChange, + accessToken, + disabled, + placeholder, + ariaLabel, + className, +}: EmbeddingModelComboboxProps): ReactElement { + const selectingRef = useRef(false); + const anchorRef = useRef(null); + // Fully controlled: the parent updates value on every keystroke, so the + // prop itself is the search query. + const debouncedQuery = useDebouncedValue(value); + + const { results, isLoading } = useHubModelSearch(debouncedQuery, { + task: EMBEDDING_TASKS, + accessToken, + excludeGguf: true, + enabled: !disabled, + // Curated unsloth listing when empty (the global top-downloads page holds + // no unsloth mirrors to float); a typed query searches the whole Hub. + ownerScope: debouncedQuery.trim() ? "all" : "unsloth", + }); + + const items = useMemo(() => { + const ids = results.map((item) => item.id); + const selected = value.trim(); + if (selected && !ids.includes(selected)) { + ids.push(selected); + } + return ids; + }, [results, value]); + + return ( +
{ + if (event.key !== "Enter") return; + if (!(event.target instanceof HTMLInputElement)) return; + event.preventDefault(); + const typed = event.target.value.trim(); + if (typed) { + onChange(typed); + } else if (items.length > 0) { + onChange(items[0]); + } + }} + > + onChange(next ?? "")} + onInputValueChange={(next) => { + if (selectingRef.current) { + selectingRef.current = false; + return; + } + onChange(next); + }} + itemToStringValue={(item) => item} + autoHighlight={true} + > + + + {isLoading ? ( +
+ + Searching... +
+ ) : ( + No embedding models found + )} + + {(id: string) => ( + { + selectingRef.current = true; + }} + > + {id} + + )} + + + +
+ ); +} diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx index e83ecebfc8..f1e503f7a3 100644 --- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx @@ -17,6 +17,7 @@ import { fetchApiKeys, revokeApiKey, type ApiKey } from "../api/api-keys"; import { ApiMonitorConsole } from "../components/api-monitor-console"; import { ApiKeyRow } from "../components/api-key-row"; import { CreateKeyForm } from "../components/create-key-form"; +import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; import { KeyRevealCard } from "../components/key-reveal-card"; import { UsageExamples } from "../components/usage-examples"; @@ -171,6 +172,8 @@ export function ApiKeysTab() { + + !o && setRevokeTarget(null)}> diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 4e02f7f14f..ce69f3d910 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -41,6 +41,13 @@ import { rotatePreviewLinks, updatePreviewSharing, } from "../api/preview-sharing"; +import { + type EmbeddingModelSettings, + EmbeddingModelVerificationError, + loadEmbeddingModelSettings, + resetEmbeddingModelSettings, + updateEmbeddingModelSettings, +} from "../api/embedding-model"; import { DEFAULT_UPLOAD_LIMIT_MB, type UploadLimitSettings, @@ -48,7 +55,7 @@ import { updateUploadLimitSettings, } from "../api/upload-limit"; import { ChangePasswordDialog } from "../components/change-password-dialog"; -import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; +import { EmbeddingModelCombobox } from "../components/embedding-model-combobox"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; import { StudioVersionSection } from "../components/studio-version-section"; @@ -164,6 +171,16 @@ export function GeneralTab() { const [revokePreviewOpen, setRevokePreviewOpen] = useState(false); const [isRevokingPreview, setIsRevokingPreview] = useState(false); const [modelsFolder, setModelsFolder] = useState(null); + const [embeddingModel, setEmbeddingModel] = + useState(null); + const [draftEmbeddingModel, setDraftEmbeddingModel] = useState(""); + const [embeddingModelError, setEmbeddingModelError] = useState( + null, + ); + // Set after a 409 (unverifiable model); offers "Save anyway". + const [embeddingModelNeedsForce, setEmbeddingModelNeedsForce] = + useState(false); + const [isSavingEmbeddingModel, setIsSavingEmbeddingModel] = useState(false); const draftRef = useRef(draftToken); useEffect(() => { @@ -258,6 +275,27 @@ export function GeneralTab() { }; }, [t]); + useEffect(() => { + let cancelled = false; + void loadEmbeddingModelSettings() + .then((settings) => { + if (cancelled) return; + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + }) + .catch((error) => { + if (cancelled) return; + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + useEffect(() => { let cancelled = false; void loadModelsFolder() @@ -350,6 +388,58 @@ export function GeneralTab() { } }; + const saveEmbeddingModel = async (force: boolean) => { + const trimmed = draftEmbeddingModel.trim(); + if (!trimmed) { + setEmbeddingModelError(t("settings.general.rag.emptyError")); + return; + } + setIsSavingEmbeddingModel(true); + setEmbeddingModelError(null); + try { + const settings = await updateEmbeddingModelSettings(trimmed, { + hfToken: hfToken || undefined, + force, + }); + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + setEmbeddingModelNeedsForce(false); + toast.success(t("settings.general.rag.saved"), { + description: t("settings.general.rag.reindexWarning"), + }); + } catch (error) { + if (error instanceof EmbeddingModelVerificationError) { + setEmbeddingModelNeedsForce(true); + } + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.saveError"), + ); + } finally { + setIsSavingEmbeddingModel(false); + } + }; + + const resetEmbeddingModel = async () => { + setIsSavingEmbeddingModel(true); + setEmbeddingModelError(null); + setEmbeddingModelNeedsForce(false); + try { + const settings = await resetEmbeddingModelSettings(); + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + } catch (error) { + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.saveError"), + ); + } finally { + setIsSavingEmbeddingModel(false); + } + }; + const saveUploadLimit = async () => { const parsed = Number(draftUploadLimit); if (!Number.isInteger(parsed)) { @@ -500,38 +590,6 @@ export function GeneralTab() { - - -
- void saveHelperPrecache(enabled)} - /> - {helperPrecache?.disabledByEnv ? ( - - {t("settings.general.helperLlm.disabledByEnv")} - - ) : helperPrecacheError ? ( - - {helperPrecacheError} - - ) : null} -
- - - - - @@ -568,6 +626,77 @@ export function GeneralTab() { + + +
+
+ { + setDraftEmbeddingModel(next); + setEmbeddingModelNeedsForce(false); + setEmbeddingModelError(null); + }} + accessToken={hfToken || undefined} + disabled={!embeddingModel} + placeholder={embeddingModel?.defaultEmbeddingModel ?? ""} + ariaLabel={t("settings.general.rag.embeddingModel")} + className="w-[220px]" + /> + +
+ {embeddingModelError ? ( + + {embeddingModelError} + + ) : null} +
+ {embeddingModelNeedsForce ? ( + + ) : null} + {embeddingModel?.isCustom ? ( + + ) : null} +
+ + {t("settings.general.rag.reindexWarning")} + +
+ + + )} + + +
+ void saveHelperPrecache(enabled)} + /> + {helperPrecache?.disabledByEnv ? ( + + {t("settings.general.helperLlm.disabledByEnv")} + + ) : helperPrecacheError ? ( + + {helperPrecacheError} + + ) : null} +
+ + + diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 27dda5193a..136e8523ba 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -193,6 +193,20 @@ export const en = { maxUploadSize: "Training dataset upload cap", maxUploadSizeDescription: "Default is {defaultSize} MB.", }, + rag: { + sectionTitle: "Documents & RAG", + embeddingModel: "Embedding model", + embeddingModelDescription: + "Hugging Face model or local path used to index and search your documents. Default is {defaultModel}.", + reindexWarning: + "Only affects newly indexed documents. Re-upload existing ones after changing the model.", + emptyError: "Enter a Hugging Face model id or local path.", + loadError: "Failed to load the embedding model setting.", + saveError: "Failed to save the embedding model.", + saved: "Embedding model saved.", + saveAnyway: "Save anyway", + resetAction: "Reset to default", + }, storage: { sectionTitle: "Storage", modelsFolder: "Models folder", From 91f4ec7ba71469e24b4f412bab774e363ba996e3 Mon Sep 17 00:00:00 2001 From: Abdul Moiz Date: Thu, 2 Jul 2026 18:49:40 +0500 Subject: [PATCH 06/23] Studio: self-heal a pre-#6483-fix anyio>=4.14 stuck in existing installs (#6805) * Studio: self-heal a pre-#6483-fix anyio>=4.14 stuck in existing installs The <4.14 cap in constraints.txt/no-torch-runtime.txt only constrains new anyio resolutions. An install made before that cap existed can already be sitting on anyio 4.14+, and since it already satisfies mcp/fastmcp's anyio>=4.5 floor, every later constrained install skips it as already-satisfied -- so affected installs never recover and keep hitting the cancel-scope RuntimeError on every request (#6797, a recurrence of #6483). Force-reinstall anyio<4.14 whenever a stuck 4.14+ is detected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: also repair anyio on the update fast path setup.sh's _SKIP_PYTHON_DEPS and setup.ps1's $SkipPythonDeps skip install_python_stack.py entirely once the installed package version already matches PyPI latest, so an install stuck on anyio>=4.14 with an otherwise up-to-date package never reaches the repair added in install_python_stack.py. Probe anyio on that fast path too and fall through to the full dependency pass when it's still >=4.14, mirroring the existing ROCm/CPU-torch override right below it. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/install_python_stack.py | 41 +++++++++++++++++++++++++++++++++- studio/setup.ps1 | 22 ++++++++++++++++++ studio/setup.sh | 17 ++++++++++++++ 3 files changed, 79 insertions(+), 1 deletion(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 9060b57542..37805e2a57 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -181,6 +181,41 @@ def _probe_installed_torch_version() -> str | None: return lines[-1] if lines else None +# constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install +# from before the cap existed can already be stuck at 4.14+, which later +# constrained installs won't touch since it already satisfies mcp/fastmcp. +_ANYIO_BAD_FLOOR = (4, 14) + + +def _installed_anyio_version() -> tuple[int, int] | None: + try: + from importlib.metadata import version as _pkg_version + raw = _pkg_version("anyio") + except Exception: + return None + parts = raw.split(".") + try: + major = int(parts[0]) + minor = int(re.sub(r"[^0-9].*", "", parts[1])) if len(parts) > 1 else 0 + except (IndexError, ValueError): + return None + return (major, minor) + + +def _repair_bad_anyio() -> None: + installed = _installed_anyio_version() + if installed is None or installed < _ANYIO_BAD_FLOOR: + return + _safe_print(_dim(f" anyio {installed[0]}.{installed[1]} found -- reinstalling anyio<4.14...")) + pip_install( + "Repairing anyio version", + "--no-cache-dir", + "--force-reinstall", + "anyio<4.14.0", + constrain = False, + ) + + # AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/). # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs. _ROCM_WINDOWS_INDEX_BASE = ( @@ -1970,7 +2005,7 @@ def install_python_stack() -> int: package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth") # --local overlays a local repo checkout after updating deps. local_repo = os.environ.get("STUDIO_LOCAL_REPO", "") - base_total = 10 if IS_WINDOWS else 11 + base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b) if IS_MACOS: base_total -= 1 # triton step is skipped on macOS if not IS_MACOS and not NO_TORCH: @@ -2284,6 +2319,10 @@ def install_python_stack() -> int: req = REQ_ROOT / "studio.txt", ) + # 8b. anyio repair (#6483) + _progress("anyio check") + _repair_bad_anyio() + # 9. Data-designer dependencies _progress("data designer deps") pip_install( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 38398eb5f5..ae4e8464ec 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2645,6 +2645,28 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") { if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) { step "python" "$_PkgName $InstalledVer is up to date" $SkipPythonDeps = $true + # A pre-#6483-fix install can be stuck on anyio>=4.14 even though + # $_PkgName itself is current; the fast path above would otherwise + # never reach install_python_stack's anyio repair (#6797). + $_anyioBad = $false + try { + & python -c " +import re, sys +from importlib.metadata import version, PackageNotFoundError +try: + parts = version('anyio').split('.') + major = int(parts[0]) + minor = int(re.sub(r'[^0-9].*', '', parts[1])) if len(parts) > 1 else 0 +except (PackageNotFoundError, ValueError, IndexError): + sys.exit(1) +sys.exit(0 if (major, minor) >= (4, 14) else 1) +" 2>$null + if ($LASTEXITCODE -eq 0) { $_anyioBad = $true } + } catch {} + if ($_anyioBad) { + substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." "Cyan" + $SkipPythonDeps = $false + } # ...but not if an AMD GPU is present and installed PyTorch is CPU-only # (host predates ROCm-wheel support, or GPU added later): the fast "up to # date" path would leave the user on CPU torch with Train/Export disabled. diff --git a/studio/setup.sh b/studio/setup.sh index 43ec9fe7b9..22a922355d 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -947,6 +947,23 @@ print(version(sys.argv[1])) if [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ] && [ "$INSTALLED_VER" = "$LATEST_VER" ]; then step "python" "$_PKG_NAME $INSTALLED_VER is up to date" _SKIP_PYTHON_DEPS=true + # A pre-#6483-fix install can be stuck on anyio>=4.14 even though + # $_PKG_NAME itself is current; the fast path above would otherwise + # never reach install_python_stack's anyio repair (#6797). + if "$VENV_DIR/bin/python" -c " +import re, sys +from importlib.metadata import version, PackageNotFoundError +try: + parts = version('anyio').split('.') + major = int(parts[0]) + minor = int(re.sub(r'[^0-9].*', '', parts[1])) if len(parts) > 1 else 0 +except (PackageNotFoundError, ValueError, IndexError): + sys.exit(1) +sys.exit(0 if (major, minor) >= (4, 14) else 1) +" 2>/dev/null; then + substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." + _SKIP_PYTHON_DEPS=false + fi elif [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ]; then substep "$_PKG_NAME $INSTALLED_VER -> $LATEST_VER available, updating..." elif [ -z "$LATEST_VER" ]; then From 22cd26f75da5e43257b4d4b385a335c5bae0256e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gabriel=20Pereira=20G=C3=B3es?= <82218878+Dspofu@users.noreply.github.com> Date: Thu, 2 Jul 2026 13:06:39 -0300 Subject: [PATCH 07/23] feat: Implementation of the Portuguese (Brazil) language and VRAM/RAM monitor (#6509) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: Implementation of the Portuguese (Brazil) language and VRAM/RAM monitor. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/frontend/src/hooks/use-gpu-utilization.ts Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/backend/main.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/backend/main.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/backend/utils/hardware/hardware.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/backend/utils/hardware/hardware.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/frontend/src/features/settings/components/usage-examples.tsx Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/backend/main.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/backend/utils/hardware/hardware.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/frontend/src/features/studio/sections/progress-section.tsx Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix: resolve automated review feedback on API shape * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix review issues for PR #6509: Cpu icon, VRAM percent, system polling - model-inspector: use the exported CpuIcon (Cpu is not a Hugeicons export) - app-sidebar: guard the VRAM percent on totalVram to avoid Infinity, and reset the system poll cache only after each request settles so a slow probe is reused instead of stacking overlapping requests - use-gpu-info: populate CPU/RAM on hosts without a GPU - progress-section: label GPUs by visible_ordinal instead of array index - hub-page: base the RAM label on systemRamTotalGb - usage-examples: emit JS sampling and tool options at the top level instead of nesting them under extra_body (the JS SDK does not unwrap extra_body) - main: read torch and transformers versions from package metadata instead of importing the libraries on every system poll, and guard the VRAM math against null values - hardware: translate a leftover comment to English * Harden /api/system: guard psutil.boot_time for PR #6509 Simulating restricted containers and some VMs (where psutil.boot_time can raise) showed the /api/system endpoint would 500 on the unguarded boot_time call, the same failure class already handled for cpu_freq, disk_usage, and Process. Wrap boot_time and return uptime_seconds as null when it is unavailable so the sidebar monitor degrades gracefully instead of breaking. Widen the uptime_seconds type to number | null to match. * Studio: make the sidebar hardware monitor a toggle (default on) for PR #6509 Adds a "Show hardware monitor" switch under Settings > Appearance > Layout, backed by a localStorage preference (default on), mirroring the existing useSidebarPin pattern. When turned off, the sidebar hides the VRAM/RAM meters and useSystemInfo stops the 3s /api/system poll entirely, so no nvidia-smi / SMI probes run while the monitor is disabled. Adds the en and pt-BR strings. * Studio: default the sidebar hardware monitor to off (opt-in) for PR #6509 * Studio pt-BR: fix three small translation defects for PR #6509 - learningRateDescription: "5e-5 for CPT" -> "5e-5 para CPT" (leftover English) - exportScopeRecents: "Recents" -> "Recentes" (untranslated) - relativeMonthsAgo/relativeYearsAgo: add the missing space ("há {count} meses"/ "há {count} anos") so they no longer render as "há 3meses" * Studio pt-BR: translate the last 10 fallback keys for PR #6509 Adds the settings.general.storage block (Armazenamento) and the settings.chat.modelDisclaimer pair, so pt-BR now covers all en keys (679/679) with no English fallbacks. * Studio: hide sidebar VRAM row on CPU-only hosts for PR #6509 * Studio: tighten and trim code comments for PR #6509 * fix: UI issue in the stop button dialog box (fine-tuning) * Studio pt-BR: translate 18 new keys from main merge (password dialog, GGUF export, dataset streaming) for PR #6509 * Rounding to GB * Fix/adjust System resources tab for PR #6509 * Fix/adjust GPU monitor review items for PR #6509 * Fix/adjust remaining GPU monitor review items for PR #6509 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix/adjust MLX resource fallback for PR #6509 * floating window implementation * resize for floating window * Fix resource monitor review items * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore frontend optional dependency lock entries * Make GPU selection tests hermetic * Fix GPU monitor CI test failures * Bound MLX GGUF reload smoke * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix MLX GGUF reload smoke exit --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: wasimysaid --- studio/backend/main.py | 131 ++- .../backend/tests/test_anthropic_messages.py | 8 +- studio/backend/tests/test_gpu_selection.py | 231 ++++- studio/backend/utils/hardware/hardware.py | 272 +++-- .../src/components/floating-monitor.tsx | 160 +++ .../features/hub/catalog/models-header.tsx | 6 +- studio/frontend/src/features/hub/hub-page.tsx | 11 +- .../settings/components/usage-examples.tsx | 149 ++- .../src/features/settings/settings-dialog.tsx | 208 ++-- .../settings/stores/monitor-overlay-store.ts | 24 + .../settings/stores/settings-dialog-store.ts | 2 + .../features/settings/tabs/general-tab.tsx | 1 + .../features/settings/tabs/resources-tab.tsx | 477 +++++++++ .../studio/sections/progress-section.tsx | 128 ++- studio/frontend/src/hooks/index.ts | 2 + studio/frontend/src/hooks/use-gpu-info.ts | 47 +- .../frontend/src/hooks/use-gpu-utilization.ts | 6 +- studio/frontend/src/hooks/use-system.ts | 130 +++ studio/frontend/src/i18n/AGENTS.md | 3 +- studio/frontend/src/i18n/check-parity.ts | 6 +- studio/frontend/src/i18n/locales/en.ts | 55 ++ studio/frontend/src/i18n/locales/pt-br.ts | 934 ++++++++++++++++++ studio/frontend/src/i18n/messages.ts | 13 +- tests/studio/install/test_rocm_support.py | 6 +- tests/studio/run_real_mlx_smoke.py | 66 +- 25 files changed, 2691 insertions(+), 385 deletions(-) create mode 100644 studio/frontend/src/components/floating-monitor.tsx create mode 100644 studio/frontend/src/features/settings/stores/monitor-overlay-store.ts create mode 100644 studio/frontend/src/features/settings/tabs/resources-tab.tsx create mode 100644 studio/frontend/src/hooks/use-system.ts create mode 100644 studio/frontend/src/i18n/locales/pt-br.ts diff --git a/studio/backend/main.py b/studio/backend/main.py index 5402e5eb7b..0613a5ae53 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -12,6 +12,8 @@ from pathlib import Path as _Path import asyncio from dataclasses import asdict +from typing import Any, Optional + # Suppress C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" @@ -36,6 +38,10 @@ if sys.platform == "win32": pass del _win_stream +_SYSTEM_GPU_CACHE_TTL_SECONDS = 10.0 +_system_gpu_cache_lock = threading.Lock() +_system_gpu_cache: Optional[tuple[float, dict[str, Any]]] = None + # ── Windows AMD ROCm DLL injection ────────────────────────────────────────── # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. @@ -226,7 +232,6 @@ import shutil import warnings from contextlib import asynccontextmanager from importlib.metadata import PackageNotFoundError, version as package_version -from typing import Optional from urllib.parse import urlparse @@ -1078,8 +1083,57 @@ async def shutdown_server(request: Request, current_subject: str = Depends(get_c return {"status": "shutting_down"} +def _get_cached_system_gpu_info(logger) -> dict[str, Any]: + """Return merged GPU visibility/utilization with bounded live-probe churn.""" + import time + from utils.hardware import get_backend_visible_gpu_info, get_visible_gpu_utilization + + global _system_gpu_cache + now = time.monotonic() + with _system_gpu_cache_lock: + if _system_gpu_cache is not None: + cached_at, cached_gpu_info = _system_gpu_cache + if now - cached_at < _SYSTEM_GPU_CACHE_TTL_SECONDS: + return cached_gpu_info + + try: + visibility_info = get_backend_visible_gpu_info() or {"available": False, "devices": []} + except Exception as e: + logger.debug(f"Failed to get GPU visibility info: {e}") + visibility_info = {"available": False, "devices": []} + + try: + utilization_info = get_visible_gpu_utilization() or {"devices": []} + except Exception as e: + logger.debug(f"Failed to get GPU utilization info: {e}") + utilization_info = {"devices": []} + + util_devices = {d.get("index"): d for d in utilization_info.get("devices", [])} + enriched_devices = [] + + for dev in visibility_info.get("devices", []): + idx = dev.get("index") + util = util_devices.get(idx, {}) + + total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0 + used_vram = util.get("vram_used_gb") or 0 + + enriched_dev = dict(dev) + enriched_dev["vram_used_gb"] = used_vram + enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0 + enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") + enriched_devices.append(enriched_dev) + + gpu_info = { + "available": visibility_info.get("available", False), + "devices": enriched_devices, + } + _system_gpu_cache = (time.monotonic(), gpu_info) + return gpu_info + + @app.get("/api/system") -async def get_system_info(current_subject: str = Depends(get_current_subject)): +def get_system_info(current_subject: str = Depends(get_current_subject)): """Get system information. Auth-gated: the response (platform, Python/GPU, memory, ML packages) can @@ -1088,31 +1142,82 @@ async def get_system_info(current_subject: str = Depends(get_current_subject)): """ import platform import psutil + import os + import time + import logging from utils.hardware import get_device from utils.hardware.hardware import _backend_label - visibility_info = get_backend_visible_gpu_info() - gpu_info = { - "available": visibility_info["available"], - "devices": visibility_info["devices"], - } + logger = logging.getLogger(__name__) + + gpu_info = _get_cached_system_gpu_info(logger) - # CPU & Memory memory = psutil.virtual_memory() + try: + cpu_freq = psutil.cpu_freq() + except Exception as e: + logger.debug(f"Failed to get CPU frequency: {e}") + cpu_freq = None + + try: + disk = psutil.disk_usage(os.path.abspath(os.sep)) + except Exception as e: + logger.debug(f"Failed to get disk usage: {e}") + disk = None + + try: + current_process = psutil.Process(os.getpid()) + process_used_mb = round(current_process.memory_info().rss / 1024**2) + except Exception as e: + logger.debug(f"Failed to get current process memory: {e}") + process_used_mb = 0 + + try: + boot_time = psutil.boot_time() + except Exception as e: + logger.debug(f"Failed to get boot time: {e}") + boot_time = None + + # Read versions from metadata so a 3s poll never imports heavy ML libs (or 500s on their import errors). + from importlib.metadata import PackageNotFoundError, version as pkg_version + + ml_packages = {} + for pkg in ("torch", "transformers"): + try: + ml_packages[pkg] = pkg_version(pkg) + except PackageNotFoundError: + pass + except Exception as e: + logger.debug(f"Failed to read {pkg} version: {e}") + return { "platform": platform.platform(), "python_version": platform.python_version(), - # _backend_label so /api/system reports "rocm" (not "cuda") on AMD, - # matching /api/hardware and /api/gpu-visibility. "device_backend": _backend_label(get_device()), - "cpu_count": psutil.cpu_count(), + "cpu_count": psutil.cpu_count(logical = True), + "uptime_seconds": max(0, round(time.time() - boot_time)) if boot_time else None, + "cpu": { + "logical_count": psutil.cpu_count(logical = True), + "physical_count": psutil.cpu_count(logical = False), + "usage_percent": psutil.cpu_percent(interval = None), + "frequency_mhz": round(cpu_freq.current, 2) + if cpu_freq and cpu_freq.current is not None + else None, + }, "memory": { - "total_gb": round(memory.total / 1e9, 2), - "available_gb": round(memory.available / 1e9, 2), + "total_gb": round(memory.total / 1024**3, 2), + "available_gb": round(memory.available / 1024**3, 2), "percent_used": memory.percent, + "process_used_mb": process_used_mb, + }, + "disk": { + "total_gb": round(disk.total / 1e9, 2) if disk else 0, + "free_gb": round(disk.free / 1e9, 2) if disk else 0, + "percent_used": disk.percent if disk else 0, }, "gpu": gpu_info, + "ml_packages": ml_packages, } diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 92a87ce045..a6c1fcda9c 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -128,13 +128,7 @@ class TestToolActionNudge: assert "call render_html once" in nudge def test_balanced_nudge_empty_without_known_tool_categories(self): - assert ( - _build_tool_action_nudge( - tools = [], - model_name = "Llama-3.1-8B-Instruct", - ) - == "" - ) + assert _build_tool_action_nudge(tools = [], model_name = "Llama-3.1-8B-Instruct") == "" # ===================================================================== diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d96f88e4a6..69ad560788 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -5,9 +5,11 @@ import asyncio import importlib.util import os import re +import sys import unittest +from contextlib import nullcontext from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import patch from fastapi import HTTPException @@ -22,6 +24,7 @@ from utils.hardware import ( estimate_required_model_memory_gb, get_backend_visible_gpu_info, get_device_map, + get_gpu_utilization, get_offloaded_device_map_entries, get_parent_visible_gpu_ids, get_visible_gpu_utilization, @@ -33,6 +36,24 @@ import utils.hardware.hardware as _hw_module _BACKEND_ROOT = Path(__file__).resolve().parent.parent +async def _inline_to_thread(func, /, *args, **kwargs): + return func(*args, **kwargs) + + +def _fake_unsloth_attention_modules(resolver): + unsloth_module = ModuleType("unsloth") + models_module = ModuleType("unsloth.models") + utils_module = ModuleType("unsloth.models._utils") + utils_module.resolve_attention_implementation = resolver + models_module._utils = utils_module + unsloth_module.models = models_module + return { + "unsloth": unsloth_module, + "unsloth.models": models_module, + "unsloth.models._utils": utils_module, + } + + def _load_route_module(name: str, relative_path: str): spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path) module = importlib.util.module_from_spec(spec) @@ -122,6 +143,139 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase): + def test_gpu_utilization_preserves_primary_shape_with_devices(self): + devices = [ + { + "index": 5, + "visible_ordinal": 0, + "gpu_utilization_pct": 11.0, + "temperature_c": 40.0, + "vram_used_gb": 4.0, + "vram_total_gb": 24.0, + "vram_utilization_pct": 16.7, + "power_draw_w": 80.0, + "power_limit_w": 300.0, + "power_utilization_pct": 26.7, + }, + { + "index": 3, + "visible_ordinal": 1, + "gpu_utilization_pct": 22.0, + "temperature_c": 50.0, + "vram_used_gb": 8.0, + "vram_total_gb": 24.0, + "vram_utilization_pct": 33.3, + "power_draw_w": 120.0, + "power_limit_w": 300.0, + "power_utilization_pct": 40.0, + }, + ] + + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch.object(_hw_module, "IS_ROCM", False), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = {"raw": "5,3", "numeric_ids": [5, 3]}, + ), + patch( + "utils.hardware.hardware._smi_query", + return_value = { + "available": True, + "devices": devices, + "backend_cuda_visible_devices": "5,3", + "parent_visible_gpu_ids": [5, 3], + "index_kind": "physical", + }, + ), + ): + result = get_gpu_utilization() + + self.assertIsInstance(result, dict) + self.assertTrue(result["available"]) + self.assertEqual(result["backend"], "cuda") + self.assertEqual(result["index"], 5) + self.assertEqual(result["visible_ordinal"], 0) + self.assertEqual(result["vram_total_gb"], 24.0) + self.assertEqual(result["parent_visible_gpu_ids"], [5, 3]) + self.assertEqual([device["index"] for device in result["devices"]], [5, 3]) + + def test_gpu_utilization_cpu_returns_legacy_unavailable_object(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU): + result = get_gpu_utilization() + + self.assertEqual(result, {"available": False, "backend": "cpu", "devices": []}) + + def test_gpu_utilization_mlx_stays_available_without_agx_stats(self): + fake_psutil = ModuleType("psutil") + fake_psutil.virtual_memory = lambda: SimpleNamespace(total = 64 * 1024**3) + + with ( + patch.dict(sys.modules, {"psutil": fake_psutil}), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX), + patch("utils.hardware.hardware._read_apple_gpu_stats", return_value = {}), + patch( + "core.training.get_training_backend", + return_value = SimpleNamespace(_progress = None), + ), + patch("utils.hardware.apple.read_gpu_temperature_c", return_value = None), + patch("utils.hardware.apple.read_gpu_power_w", return_value = None), + ): + result = get_gpu_utilization() + + self.assertTrue(result["available"]) + self.assertEqual(result["backend"], "mlx") + self.assertIsNone(result["gpu_utilization_pct"]) + self.assertEqual(result["vram_used_gb"], 0) + self.assertEqual(result["vram_total_gb"], 64.0) + self.assertEqual(len(result["devices"]), 1) + + def test_gpu_utilization_xpu_uses_visible_devices(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = { + "available": True, + "backend": "xpu", + "parent_visible_gpu_ids": [2, 0], + "index_kind": "physical", + "devices": [ + { + "index": 2, + "visible_ordinal": 1, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": 3.0, + "vram_total_gb": 16.0, + "vram_utilization_pct": 18.8, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + }, + { + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": 1.0, + "vram_total_gb": 16.0, + "vram_utilization_pct": 6.3, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + }, + ], + }, + ), + ): + result = get_gpu_utilization() + + self.assertEqual(result["backend"], "xpu") + self.assertEqual(result["index"], 0) + self.assertEqual(result["visible_ordinal"], 0) + self.assertEqual([device["index"] for device in result["devices"]], [0, 2]) + def test_visible_gpu_utilization_filters_to_parent_visible_ids(self): smi_output = "\n".join( [ @@ -272,6 +426,14 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_get_offloaded_device_map_entries_handles_models_without_device_map(self): self.assertEqual(get_offloaded_device_map_entries(SimpleNamespace()), {}) + @patch( + "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate", + new = lambda model_name, **_: model_name, + ) + @patch( + "utils.hardware.hardware._load_config_for_gpu_estimate", + new = lambda *_args, **_kwargs: None, + ) def test_estimate_required_memory_formulas(self): eight_gb = 8 * (1024**3) @@ -432,6 +594,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_explicit_ids_without_auto_selection(self): with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.resolve_requested_gpu_ids", return_value = [2, 3], @@ -464,6 +627,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self): with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", return_value = ( @@ -582,6 +746,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue], @@ -709,14 +874,23 @@ class TestRouteErrors(unittest.TestCase): has_audio_input = False, ) - with patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -835,9 +1009,9 @@ class TestRouteErrors(unittest.TestCase): with ( patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), patch.object( inference_route, @@ -849,6 +1023,13 @@ class TestRouteErrors(unittest.TestCase): "get_llama_cpp_backend", return_value = SimpleNamespace(is_loaded = False), ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -856,7 +1037,7 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -899,9 +1080,9 @@ class TestRouteErrors(unittest.TestCase): with ( patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), patch.object( inference_route, @@ -913,6 +1094,13 @@ class TestRouteErrors(unittest.TestCase): "get_llama_cpp_backend", return_value = SimpleNamespace(is_loaded = False), ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -920,7 +1108,7 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -1102,10 +1290,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): cfg._attn_implementation = "eager" return "eager" - with patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ): + with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)): hardware_module._determine_attention_impl_for_gpu_estimate(config) self.assertFalse(hasattr(config, "_attn_implementation")) @@ -1133,10 +1318,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): with ( patch.object(AutoModelForCausalLM, "_model_mapping", new = None), patch.object(AutoModel, "_model_mapping", new = None), - patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ), + patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)), ): result = hardware_module._determine_attention_impl_for_gpu_estimate(config) @@ -1173,10 +1355,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): inner._attn_implementation = "eager" return "eager" - with patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ): + with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)): hardware_module._determine_attention_impl_for_gpu_estimate(config) self.assertFalse(hasattr(config, "_attn_implementation")) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 88a1784b88..cde7070075 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -710,82 +710,159 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa return None, None +def _gpu_utilization_payload( + device: DeviceType, devices: list[Dict[str, Any]], **metadata: Any +) -> Dict[str, Any]: + """Keep the legacy primary-GPU shape and append all visible devices.""" + backend = _backend_label(device) + normalized = [] + for ordinal, raw in enumerate(devices): + dev = dict(raw) + dev.setdefault("available", True) + dev.setdefault("backend", backend) + if dev.get("visible_ordinal") is None: + dev["visible_ordinal"] = ordinal + normalized.append(dev) + + normalized.sort(key = lambda dev: dev.get("visible_ordinal", dev.get("index", 0))) + payload: Dict[str, Any] = { + "available": bool(normalized), + "backend": backend, + "devices": normalized, + } + payload.update(metadata) + if normalized: + payload.update(normalized[0]) + payload["available"] = True + payload["backend"] = normalized[0].get("backend", backend) + payload["devices"] = normalized + return payload + + def get_gpu_utilization() -> Dict[str, Any]: - """Return a live snapshot of device utilization information.""" + """Live utilization snapshot for the primary GPU plus all visible GPUs.""" device = get_device() + if device == DeviceType.XPU: + result = get_visible_gpu_utilization() + return _gpu_utilization_payload( + device, + result.get("devices", []), + parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []), + index_kind = result.get("index_kind"), + ) + if device == DeviceType.CUDA: - result = _smi_query("get_primary_gpu_utilization") - if result is not None: - result["backend"] = _backend_label(device) - if IS_ROCM: - # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.). - _reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec()) - return result - # SMI unavailable. On Windows, use Performance Counters (Task Manager - # source) for system-wide VRAM, covering cross-process usage torch can't see. + parent_visible_spec = _get_parent_visible_gpu_spec() + result = _smi_query( + "get_visible_gpu_utilization", + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result is not None and "devices" in result: + devices = result["devices"] + numeric_ids = parent_visible_spec.get("numeric_ids") + if IS_ROCM and numeric_ids is not None: + _reconcile_rocm_unified_memory(result, numeric_ids) + + return _gpu_utilization_payload( + device, + devices, + backend_cuda_visible_devices = result.get("backend_cuda_visible_devices"), + parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []), + index_kind = result.get("index_kind"), + ) + + # Fallback Windows ROCm if IS_ROCM and platform.system() == "Windows": _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() if _win_used is not None and _win_total is not None: _win_util = _rocm_windows_perf_counter_gpu_util_pct() - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": _win_util, - "temperature_c": None, - "vram_used_gb": _win_used, - "vram_total_gb": _win_total, - "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) - if _win_total > 0 - else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } - # Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed. + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": _win_util, + "temperature_c": None, + "vram_used_gb": _win_used, + "vram_total_gb": _win_total, + "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) + if _win_total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) + + # Fallback Linux ROCm if IS_ROCM and platform.system() == "Linux": _linux_used, _linux_total = _rocm_linux_sysfs_vram_gb() if _linux_used is not None and _linux_total is not None: _linux_util = _rocm_linux_sysfs_gpu_busy_pct() _linux_temp = _rocm_linux_sysfs_temp_c() _linux_power = _rocm_linux_sysfs_power_w() - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": _linux_util, - "temperature_c": _linux_temp, - "vram_used_gb": _linux_used, - "vram_total_gb": _linux_total, - "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) - if _linux_total > 0 - else None, - "power_draw_w": _linux_power, - "power_limit_w": None, - "power_utilization_pct": None, - } - # Last resort: torch mem_get_info (process-local). - _visible_spec = _get_parent_visible_gpu_spec() - _numeric_ids = _visible_spec.get("numeric_ids") or [0] - _primary_idx = [_numeric_ids[0]] if _numeric_ids else [0] - _torch_devices = _torch_get_per_device_info(_primary_idx) - if _torch_devices: - _td = _torch_devices[0] - _total = _td["total_gb"] - _used = _td["used_gb"] - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": None, - "temperature_c": None, - "vram_used_gb": _used, - "vram_total_gb": _total, - "vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": _linux_util, + "temperature_c": _linux_temp, + "vram_used_gb": _linux_used, + "vram_total_gb": _linux_total, + "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) + if _linux_total > 0 + else None, + "power_draw_w": _linux_power, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) - # MLX: _read_apple_gpu_stats() carries both VRAM-used and GPU util%. + # Last resort: torch mem_get_info (process-local) for all visible GPUs + _visible_spec = _get_parent_visible_gpu_spec() + _numeric_ids = _visible_spec.get("numeric_ids") or [] + if not _numeric_ids: + visible_count = _torch_get_physical_gpu_count() or 0 + _numeric_ids = list(range(visible_count)) + + _torch_devices = _torch_get_per_device_info(_numeric_ids) + if _torch_devices: + gpu_array = [] + for _td in _torch_devices: + _total = _td["total_gb"] + _used = _td["used_gb"] + gpu_array.append( + { + "available": True, + "backend": _backend_label(device), + "index": _td["index"], + "name": _td.get("name", "Unknown"), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": _used, + "vram_total_gb": _total, + "vram_utilization_pct": round((_used / _total) * 100, 1) + if _total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ) + return _gpu_utilization_payload(device, gpu_array) + + # MLX if device == DeviceType.MLX: try: import psutil @@ -793,9 +870,8 @@ def get_gpu_utilization() -> Dict[str, Any]: total_bytes = psutil.virtual_memory().total except Exception as e: logger.error(f"Error getting MLX GPU utilization: {e}") - return {"available": False, "backend": device.value, "error": str(e)} - if not agx: - return {"available": False, "backend": device.value} + return {"available": False, "backend": device.value, "devices": [], "error": str(e)} + allocated_bytes = agx.get("vram_used_bytes", 0) or 0 vram_used_gb = allocated_bytes / (1024**3) total_gb = total_bytes / (1024**3) @@ -814,37 +890,51 @@ def get_gpu_utilization() -> Dict[str, Any]: from . import apple - return { - "available": True, - "backend": device.value, - "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, - "temperature_c": apple.read_gpu_temperature_c(), - "vram_used_gb": round(vram_used_gb, 2), - "vram_total_gb": round(total_gb, 2), - "vram_utilization_pct": ( - round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None - ), - "power_draw_w": apple.read_gpu_power_w(), - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": device.value, + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, + "temperature_c": apple.read_gpu_temperature_c(), + "vram_used_gb": round(vram_used_gb, 2), + "vram_total_gb": round(total_gb, 2), + "vram_utilization_pct": round((vram_used_gb / total_gb) * 100, 1) + if total_gb > 0 + else None, + "power_draw_w": apple.read_gpu_power_w(), + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) mem = get_gpu_memory_info() if device != DeviceType.CPU and mem.get("available"): - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": None, - "temperature_c": None, - "vram_used_gb": round(mem.get("allocated_gb", 0), 2), - "vram_total_gb": round(mem.get("total_gb", 0), 2), - "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": mem.get("device", 0), + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": round(mem.get("allocated_gb", 0), 2), + "vram_total_gb": round(mem.get("total_gb", 0), 2), + "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) - return {"available": False, "backend": _backend_label(device)} + return {"available": False, "backend": _backend_label(device), "devices": []} def _apply_unified_memory_correction( diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx new file mode 100644 index 0000000000..f02da6612e --- /dev/null +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store"; +import { useSystemInfo } from "@/hooks/use-system"; +import { useT } from "@/i18n"; +import { cn } from "@/lib/utils"; +import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react"; +import { motion } from "motion/react"; +import { useRef } from "react"; + +function clampPercent(value: number): number { + return Math.max(0, Math.min(100, value)); +} + +function usageIndicatorClass(percent: number): string { + if (percent >= 90) return "bg-destructive"; + if (percent >= 70) return "bg-amber-500"; + return "bg-primary"; +} + +function usageTextClass(percent: number): string { + if (percent >= 90) return "text-destructive"; + if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + return "text-primary"; +} + +function formatGb(value: number): string { + const digits = value >= 10 ? 1 : 2; + return `${value.toFixed(digits)} GB`; +} + +export function FloatingMonitor() { + const t = useT(); + const { isOpen, setIsOpen } = useMonitorOverlayStore(); + const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 }); + + const constraintsRef = useRef(null); + + if (!isOpen) return null; + + const ramTotal = systemInfo.memory?.total_gb ?? 0; + const ramAvailable = systemInfo.memory?.available_gb ?? 0; + const ramUsed = Math.max(0, ramTotal - ramAvailable); + const ramPercent = clampPercent(systemInfo.memory?.percent_used ?? 0); + + const devices = systemInfo.gpu?.devices ?? []; + const vramTotal = devices.reduce( + (sum, device) => sum + (device.memory_total_gb ?? 0), + 0, + ); + const vramUsed = devices.reduce( + (sum, device) => sum + (device.vram_used_gb ?? 0), + 0, + ); + const vramPercent = clampPercent( + vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0, + ); + + const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0; + + return ( +
+ +
+
+ + + {t("settings.resources.liveMonitor.title")} + +
+
+
+ +
+ + +
+
+ + +
+
+ {t("settings.resources.liveMonitor.ram")} + + {Math.round(ramPercent)}% + +
+
+ {formatGb(ramUsed)} / {formatGb(ramTotal)} +
+ +
+ + {hasGpu && ( +
+
+ + {t("settings.resources.liveMonitor.vram")}{" "} + {devices.length > 1 + ? `(${devices.length} GPUs)` + : `(${devices[0].name ?? "GPU"})`} + + + {Math.round(vramPercent)}% + +
+
+ {formatGb(vramUsed)} / {formatGb(vramTotal)} +
+ +
+ )} + + +
+ ); +} diff --git a/studio/frontend/src/features/hub/catalog/models-header.tsx b/studio/frontend/src/features/hub/catalog/models-header.tsx index f5afce5c59..10d5a80e5e 100644 --- a/studio/frontend/src/features/hub/catalog/models-header.tsx +++ b/studio/frontend/src/features/hub/catalog/models-header.tsx @@ -15,6 +15,7 @@ import { PackageIcon, RamMemoryIcon, RemoveCircleIcon, + CpuIcon } from "@hugeicons/core-free-icons"; import type { IconSvgElement } from "@hugeicons/react"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -43,6 +44,7 @@ export function ModelsHeader({ isDataset, gpuLabel, ramLabel, + coreLabel, activeCheckpoint, activeGgufVariant, onTitleClick, @@ -53,6 +55,7 @@ export function ModelsHeader({ isDataset: boolean; gpuLabel: string; ramLabel: string; + coreLabel: string; activeCheckpoint: string | null; activeGgufVariant: string | null; onTitleClick: () => void; @@ -84,7 +87,8 @@ export function ModelsHeader({ value={String(localCount)} /> - + + {activeCheckpoint && (
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index 56aa07335d..c3920e18f6 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -1085,11 +1085,15 @@ export function ModelsPage() { const { vramInfo, minMemory } = useHubModelVram(selectedModel, gpu); const gpuLabel = gpu.available - ? `${Math.floor(gpu.memoryTotalGb)} GB` + ? `${Math.round(gpu.memoryTotalGb)} GB` : "Unavailable"; const ramLabel = - gpu.systemRamAvailableGb > 0 - ? `${Math.floor(gpu.systemRamAvailableGb)} GB` + gpu.systemRamTotalGb > 0 + ? `${Math.round(gpu.systemRamTotalGb)} GB` + : "Unavailable"; + const coreLabel = + gpu.cpuCore > 0 && gpu.cpuThread > 0 + ? `${gpu.cpuCore}/${gpu.cpuThread}` : "Unavailable"; const openNewChat = useCallback(() => { @@ -1453,6 +1457,7 @@ export function ModelsPage() { isDataset={isDatasetMode} gpuLabel={gpuLabel} ramLabel={ramLabel} + coreLabel={coreLabel} activeCheckpoint={activeCheckpoint} activeGgufVariant={activeGgufVariant} onTitleClick={handleResetToDiscover} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index aef2e7ffe6..d66ca6105d 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -33,44 +33,58 @@ import { updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; -// API call type; OS axis applies to curl only (Python is OS-identical). type ExampleType = | "curl" | "python" + | "javascript" | "curlTools" | "pythonTools" + | "javascriptTools" | "curlAdvanced" - | "pythonAdvanced"; + | "pythonAdvanced" + | "javascriptAdvanced"; type Os = "unix" | "windows"; -// plain = bare call; tools = server-side tools; advanced = sampling + thinking + tools. type Variant = "plain" | "tools" | "advanced"; const TYPE_TABS: { id: ExampleType; label: string }[] = [ { id: "curl", label: "curl" }, { id: "python", label: "Python" }, + { id: "javascript", label: "JavaScript" }, { id: "curlTools", label: "curl + tools" }, { id: "pythonTools", label: "Python + tools" }, + { id: "javascriptTools", label: "JavaScript + tools" }, { id: "curlAdvanced", label: "curl + advanced" }, { id: "pythonAdvanced", label: "Python + advanced" }, + { id: "javascriptAdvanced", label: "JavaScript + advanced" }, ]; const TYPE_LABEL_KEY: Partial> = { curlTools: "settings.apiKeys.exampleCurlTools", pythonTools: "settings.apiKeys.examplePythonTools", + javascriptTools: "settings.apiKeys.exampleJavaScriptTools", curlAdvanced: "settings.apiKeys.exampleCurlAdvanced", pythonAdvanced: "settings.apiKeys.examplePythonAdvanced", + javascriptAdvanced: "settings.apiKeys.exampleJavaScriptAdvanced", }; const OS_AWARE: Record = { curl: true, python: false, + javascript: false, curlTools: true, pythonTools: false, + javascriptTools: false, curlAdvanced: true, pythonAdvanced: false, + javascriptAdvanced: false, }; const CURL_TYPES = new Set(["curl", "curlTools", "curlAdvanced"]); +const JAVASCRIPT_TYPES = new Set([ + "javascript", + "javascriptTools", + "javascriptAdvanced", +]); const PROMPT = "Can Unsloth Studio do API calling?"; // Auto-switch demo: a second call naming a different downloaded GGUF so the @@ -82,7 +96,6 @@ const SWITCH_MODEL = "your-other-downloaded-GGUF"; const SWITCH_PROMPT = "Now answer as a different model."; // web_search + python + terminal are the reliable built-in tools. const TOOLS = ["web_search", "python", "terminal"]; -// Sampling/thinking knobs for the "+ advanced" examples. const ADV = { temperature: 0.7, top_p: 0.8, @@ -93,37 +106,18 @@ const ADV = { } as const; const DOC_LINKS = [ - { - label: "Claude Code", - href: "https://unsloth.ai/docs/basics/claude-code", - }, - { - label: "Codex", - href: "https://unsloth.ai/docs/basics/codex", - }, - { - label: "OpenClaw", - href: "https://unsloth.ai/docs/integrations/openclaw", - }, - { - label: "OpenCode", - href: "https://unsloth.ai/docs/integrations/opencode", - }, - { - label: "Hermes Agent", - href: "https://unsloth.ai/docs/integrations/hermes-agent", - }, + { label: "Claude Code", href: "https://unsloth.ai/docs/basics/claude-code" }, + { label: "Codex", href: "https://unsloth.ai/docs/basics/codex" }, + { label: "OpenClaw", href: "https://unsloth.ai/docs/integrations/openclaw" }, + { label: "OpenCode", href: "https://unsloth.ai/docs/integrations/opencode" }, + { label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" }, ]; -// JSON-encode; also a valid Python literal, so odd model names never break output. const j = (s: string): string => JSON.stringify(s); -// Embed in a POSIX single-quoted string: close, escaped quote, reopen. const shSingle = (s: string): string => s.replace(/'/g, "'\\''"); -// Embed in a PowerShell single-quoted string: '' is a literal quote. const psSingle = (s: string): string => s.replace(/'/g, "''"); const toolsJson = TOOLS.map(j).join(", "); -// Shared body fields (after model/messages, before stream) per variant. function bodyExtraLines(variant: Variant, indent: string): string[] { const lines: string[] = []; if (variant === "advanced") { @@ -152,7 +146,6 @@ function curlBodyPretty(model: string, variant: Variant): string { return `{\n${lines.join("\n")}\n }`; } -// One-line JSON for the Windows body file (PowerShell mangles inline quotes to curl.exe). function winBody(model: string, variant: Variant): string { const body: Record = { model, @@ -172,7 +165,7 @@ function winBody(model: string, variant: Variant): string { body.enabled_tools = TOOLS; } body.stream = true; - return JSON.stringify(body); + return JSON.stringify(body, null, 2); } // A leading comment (valid in both bash and PowerShell) noting the model field @@ -193,7 +186,6 @@ function curlUnix( -d '${shSingle(curlBodyPretty(model, variant))}'`; } -// Windows PowerShell: curl aliases to Invoke-WebRequest, so use curl.exe + body file. function curlWindows( base: string, key: string, @@ -233,7 +225,6 @@ function pythonSnippet( variant: Variant, autoSwitch: boolean, ): string { - // Standard OpenAI args are named; Unsloth extensions go through extra_body. const named = variant === "advanced" ? ` @@ -258,7 +249,6 @@ function pythonSnippet( ${extra.join("\n")} },` : ""; - // With tools, some chunks are tool-lifecycle events with no choices; guard it. const loop = variant !== "plain" ? `for chunk in response: @@ -281,6 +271,70 @@ response = client.chat.completions.create( ${loop}${autoSwitch ? pythonSwitchDemo() : ""}`; } +function javascriptSnippet( + base: string, + key: string, + model: string, + variant: Variant, + autoSwitch: boolean, +): string { + const options: string[] = []; + if (variant === "advanced") { + options.push(` temperature: ${ADV.temperature},`); + options.push(` top_p: ${ADV.top_p},`); + options.push(` max_tokens: ${ADV.max_tokens},`); + } + + // The JS SDK forwards unknown options into the request body, so these go at the + // top level (the Python SDK needs them under extra_body instead). + if (variant === "advanced") { + options.push(` top_k: ${ADV.top_k},`); + options.push(` min_p: ${ADV.min_p},`); + options.push(` repetition_penalty: ${ADV.repetition_penalty},`); + options.push(` enable_thinking: true,`); + } + if (variant !== "plain") { + options.push(` enable_tools: true,`); + options.push(` enabled_tools: [${toolsJson}],`); + } + + const trailingOptions = options.length ? `\n${options.join("\n")}` : ""; + + return `import OpenAI from "openai"; + +const client = new OpenAI({ + baseURL: ${j(`${base}/v1`)}, + apiKey: ${j(key)}, +}); + +const response = await client.chat.completions.create({ + model: ${j(model)}, + messages: [{ role: "user", content: ${j(PROMPT)} }],${trailingOptions} + stream: true, +}); + +for await (const chunk of response) { + process.stdout.write(chunk.choices?.[0]?.delta?.content || ""); +}${autoSwitch ? javascriptSwitchDemo() : ""}`; +} + +function javascriptSwitchDemo(): string { + return ` + +// "Switch model by request" is on: replace the model below with another GGUF you +// have downloaded and Studio loads it before serving. Unknown names keep serving +// the current model. +const switchResponse = await client.chat.completions.create({ + model: ${j(SWITCH_MODEL)}, + messages: [{ role: "user", content: ${j(SWITCH_PROMPT)} }], + stream: true, +}); + +for await (const chunk of switchResponse) { + process.stdout.write(chunk.choices?.[0]?.delta?.content || ""); +}`; +} + function buildSnippets( base: string, key: string, @@ -292,17 +346,24 @@ function buildSnippets( return { curl: curl(base, key, model, "plain", autoSwitch), python: pythonSnippet(base, key, model, "plain", autoSwitch), + javascript: javascriptSnippet(base, key, model, "plain", autoSwitch), curlTools: curl(base, key, model, "tools", autoSwitch), pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch), + javascriptTools: javascriptSnippet(base, key, model, "tools", autoSwitch), curlAdvanced: curl(base, key, model, "advanced", autoSwitch), pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch), + javascriptAdvanced: javascriptSnippet( + base, + key, + model, + "advanced", + autoSwitch, + ), }; } const KEY_PLACEHOLDER = "sk-unsloth-YOUR_KEY"; const MODEL_FALLBACK = "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL"; - -// Default ON: when a tunnel exists, examples should show the public base_url. const USE_TUNNEL_KEY = "unsloth_api_use_tunnel"; function readUseTunnelPref(): boolean { @@ -319,11 +380,10 @@ function writeUseTunnelPref(value: boolean): void { try { window.localStorage.setItem(USE_TUNNEL_KEY, value ? "true" : "false"); } catch { - // Non-fatal: the toggle still applies for this session. + // Non-fatal } } -// Active local checkpoint as repo[:variant]; external/none falls back to a default. function useLoadedModelName(): string { const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint); const ggufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); @@ -338,7 +398,6 @@ function useLoadedModelName(): string { }, [checkpoint, ggufVariant]); } -// shiki highlighting via the app's shared code plugin + themes (same as chat). const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [ typeof unslothLightTheme, typeof unslothDarkTheme, @@ -352,7 +411,6 @@ function HighlightedCode({ code: string; language: string; }) { - // Fence so Streamdown's shiki plugin highlights it (no markdown inside a fence). const markdown = useMemo( () => `\`\`\`${language}\n${code}\n\`\`\``, [code, language], @@ -390,7 +448,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); const [savingAutoSwitch, setSavingAutoSwitch] = useState(false); - // Tunnel may start after the first /api/health read; refresh so it surfaces here. useEffect(() => { void fetchDeviceType({ force: true }); }, []); @@ -410,10 +467,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { }, []); const model = useLoadedModelName(); - // Real key while revealed (before "Done"); otherwise a placeholder. const key = apiKey || KEY_PLACEHOLDER; - // Toggle on + tunnel up: public tunnel URL. Off: backend direct host:port - // (origin is only a last-resort fallback). const origin = typeof window !== "undefined" ? window.location.origin : ""; const base = useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); @@ -429,7 +483,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ? os === "windows" ? "powershell" : "bash" - : "python"; + : JAVASCRIPT_TYPES.has(lang) + ? "javascript" + : "python"; const handleCopy = async () => { if (await copyToClipboard(snippets[lang])) { @@ -539,8 +595,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { )}
- {/* Always rendered (dimmed when off) so toggling never changes the - row height and shifts the code block below. */} - {/* key on the snippet so Streamdown remounts and re-highlights when - only a substring (e.g. the base URL) changes; its block memo - otherwise keeps the stale render. */} ; case "appearance": return ; + case "resources": + return ; case "chat": return ; case "connections": @@ -100,6 +110,7 @@ export function SettingsDialog() { general: null, profile: null, appearance: null, + resources: null, chat: null, connections: null, "api-keys": null, @@ -115,110 +126,113 @@ export function SettingsDialog() { }, [open, activeTab]); return ( - !o && closeDialog()}> - { - // Restore focus to the element that triggered openDialog(). Radix's - // FocusScope races our rAF-scheduled tab focus and loses the - // previous-focus reference, so restore it by hand. - if (opener && opener.isConnected) { - e.preventDefault(); - opener.focus({ preventScroll: true }); - } - }} - className={cn( - // Cap at 820px but shrink to the viewport so it doesn't clip on - // iPad-portrait widths (640-820px) where fixed `w-[820px]` overflows. - "settings-surface !max-w-[min(820px,calc(100vw-2rem))] h-[560px] w-[min(820px,calc(100vw-2rem))] p-0 overflow-hidden", - // Soft shadow, no outline ring. Pin --radius to the light value so - // corner rounding matches in dark mode. - "shadow-border rounded-xl ring-0 [--radius:1.1rem]", - "max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none", - )} - > - - {t("settings.dialog.title")} - - - {t("settings.dialog.description")} - -
- + {tab.badgeKey ? ( + + {t(tab.badgeKey)} + + ) : null} + + ); + })} + + -
- -
- {renderTab(activeTab)} -
-
-
- - +
+ +
+ {renderTab(activeTab)} +
+
+ + + + + ); } diff --git a/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts b/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts new file mode 100644 index 0000000000..1af804a19e --- /dev/null +++ b/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { create } from "zustand"; +import { persist } from "zustand/middleware"; + +interface MonitorOverlayState { + isOpen: boolean; + isMinimized: boolean; + setIsOpen: (open: boolean) => void; + toggleMinimized: () => void; +} + +export const useMonitorOverlayStore = create()( + persist( + (set) => ({ + isOpen: false, + isMinimized: false, + setIsOpen: (isOpen) => set({ isOpen }), + toggleMinimized: () => set((state) => ({ isMinimized: !state.isMinimized })), + }), + { name: "unsloth_monitor_overlay" } + ) +); \ No newline at end of file diff --git a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts index 7fab580f3c..234e92b3d0 100644 --- a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts +++ b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts @@ -7,6 +7,7 @@ export type SettingsTab = | "general" | "profile" | "appearance" + | "resources" | "chat" | "connections" | "api-keys" @@ -60,6 +61,7 @@ function loadInitialTab(): SettingsTab { "general", "profile", "appearance", + "resources", "chat", "connections", "api-keys", diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index ce69f3d910..df684b7752 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -104,6 +104,7 @@ const PREFS_KEYS: string[] = [ "tour:studio:v1", // Update notifications "unsloth_show_llama_update_banner", + "unsloth_monitor_overlay", ]; // Set by resetAllPrefs so the unmount-commit effect skips writing back the diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx new file mode 100644 index 0000000000..d5e19cc51c --- /dev/null +++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx @@ -0,0 +1,477 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { Switch } from "@/components/ui/switch"; +import { openModelsDir } from "@/features/native-intents"; +import { useSystemInfo, type GpuDevice } from "@/hooks/use-system"; +import { isTauri } from "@/lib/api-base"; +import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { toast } from "@/lib/toast"; +import { cn } from "@/lib/utils"; +import { useT } from "@/i18n"; +import { useEffect, useMemo, useState } from "react"; +import { loadModelsFolder, type ModelsFolder } from "../api/models-folder"; +import { SettingsRow } from "../components/settings-row"; +import { SettingsSection } from "../components/settings-section"; +import { useMonitorOverlayStore } from "../stores/monitor-overlay-store"; +import { LayersIcon } from "lucide-react"; + +const POLL_MS = 3000; + +function isFiniteNumber(value: number | null | undefined): value is number { + return typeof value === "number" && Number.isFinite(value); +} + +function clampPercent(value: number | null | undefined): number { + if (!isFiniteNumber(value)) return 0; + return Math.max(0, Math.min(100, value)); +} + +function usageIndicatorClass(percent: number): string { + if (percent >= 90) return "bg-destructive"; + if (percent >= 70) return "bg-amber-500"; + return "bg-primary"; +} + +function usageTextClass(percent: number): string { + if (percent >= 90) return "text-destructive"; + if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + return "text-primary"; +} + +function formatGb(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + const digits = safe >= 10 ? 1 : 2; + return `${safe.toFixed(digits)} GB`; +} + +function formatMb(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + return `${Math.round(safe).toLocaleString()} MB`; +} + +function formatPercent(value: number | null | undefined): string { + return `${Math.round(clampPercent(value))}%`; +} + +function formatFrequency(mhz: number | null | undefined): string | null { + if (!isFiniteNumber(mhz) || mhz <= 0) return null; + if (mhz >= 1000) return `${(mhz / 1000).toFixed(2)} GHz`; + return `${Math.round(mhz)} MHz`; +} + +function formatUptime(seconds: number | null | undefined): string { + if (!isFiniteNumber(seconds) || seconds <= 0) return "0m"; + const minutes = Math.floor(seconds / 60); + const hours = Math.floor(minutes / 60); + const days = Math.floor(hours / 24); + if (days > 0) return `${days}d ${hours % 24}h`; + if (hours > 0) return `${hours}h ${minutes % 60}m`; + return `${Math.max(1, minutes)}m`; +} + +function MetricTile({ + label, + value, + detail, + percent, +}: { + label: string; + value: string; + detail: string; + percent: number; +}) { + const safePercent = clampPercent(percent); + return ( +
+
+ + {label} + + + {formatPercent(safePercent)} + +
+
+
+ {value} +
+
+ {detail} +
+
+ +
+ ); +} + +function InfoRow({ + label, + value, + detail, +}: { + label: string; + value: string; + detail?: string; +}) { + return ( +
+ + {label} + + + {detail ? `${value} (${detail})` : value} + +
+ ); +} + +function deviceOrdinal(device: GpuDevice): number | undefined { + return device.visible_ordinal ?? device.index; +} + +export function ResourcesTab() { + const t = useT(); + const [liveUpdates, setLiveUpdates] = useState(true); + const { isOpen, setIsOpen } = useMonitorOverlayStore(); + const systemInfo = useSystemInfo({ + enabled: liveUpdates, + pollMs: liveUpdates ? POLL_MS : undefined, + }); + const [modelsFolder, setModelsFolder] = useState(null); + const [modelsFolderLoaded, setModelsFolderLoaded] = useState(false); + + useEffect(() => { + let cancelled = false; + void loadModelsFolder() + .then((folder) => { + if (cancelled) return; + setModelsFolder(folder); + setModelsFolderLoaded(true); + }) + .catch(() => { + if (cancelled) return; + setModelsFolderLoaded(true); + }); + return () => { + cancelled = true; + }; + }, []); + + const metrics = useMemo(() => { + const devices = systemInfo.gpu?.devices ?? []; + const ramTotal = systemInfo.memory?.total_gb ?? 0; + const ramAvailable = systemInfo.memory?.available_gb ?? 0; + const ramUsed = Math.max(0, ramTotal - ramAvailable); + const diskTotal = systemInfo.disk?.total_gb ?? 0; + const diskFree = systemInfo.disk?.free_gb ?? 0; + const diskUsed = Math.max(0, diskTotal - diskFree); + const vramTotal = devices.reduce( + (sum, device) => sum + (device.memory_total_gb ?? 0), + 0, + ); + const vramUsed = devices.reduce( + (sum, device) => sum + (device.vram_used_gb ?? 0), + 0, + ); + const vramFree = devices.reduce( + (sum, device) => + sum + + (device.vram_free_gb ?? + Math.max(0, (device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0))), + 0, + ); + const vramPercent = vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0; + + return { + devices, + ramTotal, + ramUsed, + diskTotal, + diskFree, + diskUsed, + vramTotal, + vramUsed, + vramFree, + vramPercent, + }; + }, [systemInfo]); + + const handleModelsFolder = async () => { + const folder = modelsFolder; + if (!folder) return; + if (isTauri) { + try { + await openModelsDir(folder.path); + } catch (error) { + toast.error(t("settings.resources.storage.openError"), { + description: error instanceof Error ? error.message : undefined, + }); + } + return; + } + if (await copyToClipboard(folder.path)) { + toast.success(t("settings.resources.storage.copied")); + } else { + toast.error(t("settings.resources.storage.copyError")); + } + }; + + const cpuCoresLabel = + systemInfo.cpu?.logical_count && systemInfo.cpu?.physical_count + ? t("settings.resources.liveMonitor.cpuCores", { + logical: systemInfo.cpu.logical_count, + physical: systemInfo.cpu.physical_count, + }) + : t("settings.resources.environment.unknown"); + const cpuFrequencyLabel = formatFrequency(systemInfo.cpu?.frequency_mhz); + const hasGpu = + (systemInfo.gpu?.available ?? false) && metrics.devices.length > 0; + const backendLabel = ( + systemInfo.gpu?.backend ?? systemInfo.device_backend ?? "cpu" + ).toUpperCase(); + const modelsFolderPath = modelsFolder + ? modelsFolder.path + : modelsFolderLoaded + ? t("settings.resources.environment.unknown") + : t("common.loading"); + + return ( +
+
+
+

+ {t("settings.resources.title")} +

+

+ {t("settings.resources.description")} +

+
+
+ + +
+ {t("settings.resources.liveUpdates")} + +
+
+
+ + +
+ + + + +
+ + + + {hasGpu ? ( + metrics.devices.map((device, index) => { + const ordinal = deviceOrdinal(device); + const total = device.memory_total_gb ?? 0; + const used = device.vram_used_gb ?? 0; + const free = device.vram_free_gb ?? Math.max(0, total - used); + const percent = + device.vram_utilization_pct ?? + (total > 0 ? (used / total) * 100 : null); + const safePercent = clampPercent(percent); + return ( +
+
+
+
+ {device.name ?? + t("settings.resources.gpu.unknownDevice")} +
+
+ {ordinal === undefined + ? backendLabel + : `${t("settings.resources.gpu.deviceWithIndex", { + index: ordinal, + })}, ${backendLabel}`} +
+
+
+ + {formatPercent(safePercent)}{" "} + {t("settings.resources.gpu.vramUtilization")} + +
+
+
+ + {t("settings.resources.gpu.used", { + value: formatGb(used), + })} + + + {t("settings.resources.gpu.free", { + value: formatGb(free), + })} + + + {t("settings.resources.gpu.total", { + value: formatGb(total), + })} + +
+ +
+ ); + }) + ) : ( +
+ {t("settings.resources.gpu.noGpu")} +
+ )} + + + + + +
+ + {modelsFolderPath} + + +
+ + + + + + + + + + + +
+ ); +} diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index 6aaca86e0b..abab35db93 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -29,6 +29,7 @@ import { import { getTrainingMethodLabel } from "@/features/training/lib/training-methods"; import type { TrainingViewData } from "@/features/training"; import { useGpuUtilization } from "@/hooks"; +import type { GpuUtilization } from "@/hooks/use-gpu-utilization"; import { cn } from "@/lib/utils"; import { ChartAverageIcon, @@ -42,7 +43,7 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Link, useNavigate } from "@tanstack/react-router"; -import { type ReactElement, type ReactNode, useState } from "react"; +import { type ReactElement, type ReactNode, useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { ChartSettingsSheet } from "./charts/chart-settings-sheet"; import { @@ -123,18 +124,17 @@ export function ProgressSection({ const [stopDialogOpen, setStopDialogOpen] = useState(false); const [stopRequestedLocal, setStopRequestedLocal] = useState(false); - // Auto-resets when training stops; no useEffect needed const stopRequested = data.isTrainingRunning && stopRequestedLocal; const pct = data.totalSteps > 0 ? Math.min( - 100, - Math.max( - 0, - Math.round((data.currentStep / data.totalSteps) * 100), - ), - ) + 100, + Math.max( + 0, + Math.round((data.currentStep / data.totalSteps) * 100), + ), + ) : Math.round(data.progressPercent); const elapsed = data.elapsedSeconds; @@ -214,16 +214,16 @@ export function ProgressSection({ }, ...(data.trainingMethod !== "full" ? [ - { - section: "LoRA", - rows: [ - configRow(t("studio.progress.rank"), cfgLoraRank), - configRow(t("studio.progress.alpha"), cfgLoraAlpha), - configRow(t("studio.progress.dropout"), cfgLoraDropout), - configRow(t("studio.progress.variant"), cfgLoraVariant), - ], - }, - ] + { + section: "LoRA", + rows: [ + configRow(t("studio.progress.rank"), cfgLoraRank), + configRow(t("studio.progress.alpha"), cfgLoraAlpha), + configRow(t("studio.progress.dropout"), cfgLoraDropout), + configRow(t("studio.progress.variant"), cfgLoraVariant), + ], + }, + ] : []), ]; @@ -350,8 +350,8 @@ export function ProgressSection({ {stepsPerSecond == null ? t("studio.progress.noStepsPerSecond") : t("studio.progress.stepsPerSecond", { - value: stepsPerSecond.toFixed(2), - })} + value: stepsPerSecond.toFixed(2), + })} {data.currentNumTokens != null && ( {t("studio.progress.tokens", { value: data.currentNumTokens })} @@ -373,14 +373,50 @@ function LiveGpuPanel({ isTrainingRunning: boolean; }): ReactElement { const t = useT(); - const gpu = useGpuUtilization(isTrainingRunning); + const [selectedGpu, setSelectedGpu] = useState(0); + const gpuData = useGpuUtilization(isTrainingRunning); + const gpus: GpuUtilization[] = + Array.isArray(gpuData?.devices) && gpuData.devices.length > 0 + ? gpuData.devices + : gpuData && Object.keys(gpuData).length > 0 + ? [gpuData] + : []; + + useEffect(() => { + if (selectedGpu > 0 && selectedGpu >= gpus.length) { + setSelectedGpu(0); + } + }, [gpus.length, selectedGpu]); + + const gpuCount = gpus.length; + const currentGpu: Partial = gpus[selectedGpu] || gpus[0] || {}; return (
-
-

- {t("studio.progress.gpuMonitor")} -

+
+
+

+ {t("studio.progress.gpuMonitor")} +

+ {gpuCount > 1 && ( + + )} +
{t("studio.progress.live")} @@ -388,51 +424,44 @@ function LiveGpuPanel({
- } + icon={} value={ - gpu.gpu_utilization_pct != null - ? `${gpu.gpu_utilization_pct}%` + currentGpu.gpu_utilization_pct != null + ? `${currentGpu.gpu_utilization_pct}%` : "--" } - pct={gpu.gpu_utilization_pct ?? 0} + pct={currentGpu.gpu_utilization_pct ?? 0} /> - } + icon={} value={ - gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--" + currentGpu.temperature_c != null ? `${currentGpu.temperature_c}°C` : "--" } - pct={gpu.temperature_c ?? 0} + pct={currentGpu.temperature_c ?? 0} max={100} /> } value={ - gpu.vram_used_gb != null && gpu.vram_total_gb != null - ? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB` + currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null + ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB` : "--" } - pct={gpu.vram_utilization_pct ?? 0} + pct={currentGpu.vram_utilization_pct ?? 0} /> } value={ - gpu.power_draw_w != null - ? gpu.power_limit_w != null - ? `${gpu.power_draw_w} / ${gpu.power_limit_w} W` - : `${gpu.power_draw_w} W` + currentGpu.power_draw_w != null + ? currentGpu.power_limit_w != null + ? `${currentGpu.power_draw_w} / ${currentGpu.power_limit_w} W` + : `${currentGpu.power_draw_w} W` : "--" } - pct={gpu.power_utilization_pct ?? 0} + pct={currentGpu.power_utilization_pct ?? 0} />
@@ -560,7 +589,10 @@ function TrainingHeaderActions({ {stopRequested ? t("studio.training.stopping") : t("studio.training.stopAction")} - + {t("studio.training.stopTitle")} diff --git a/studio/frontend/src/hooks/index.ts b/studio/frontend/src/hooks/index.ts index d5c923d57d..5b2fa43ae7 100644 --- a/studio/frontend/src/hooks/index.ts +++ b/studio/frontend/src/hooks/index.ts @@ -1,6 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + export { useDebouncedValue } from "./use-debounced-value"; export { useGpuInfo } from "./use-gpu-info"; export { useGpuUtilization } from "./use-gpu-utilization"; @@ -9,3 +10,4 @@ export { useHfDatasetSplits } from "./use-hf-dataset-splits"; export { useHfTokenValidation } from "./use-hf-token-validation"; export { useTauriBackend } from "./use-tauri-backend"; export { useCollapseScrollLock } from "./use-collapse-scroll-lock"; +export { useSystemInfo } from "./use-system"; diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts index eb4d89abd8..1e313acdf3 100644 --- a/studio/frontend/src/hooks/use-gpu-info.ts +++ b/studio/frontend/src/hooks/use-gpu-info.ts @@ -3,19 +3,26 @@ import { authFetch } from "@/features/auth"; import { useEffect, useState } from "react"; +import type { SystemInfoResponse } from "./use-system"; export interface GpuInfo { available: boolean; name: string; memoryTotalGb: number; + cpuCore: number; + cpuThread: number; systemRamAvailableGb: number; + systemRamTotalGb: number } const DEFAULT_GPU: GpuInfo = { available: false, name: "Unknown", memoryTotalGb: 0, + cpuCore: 0, + cpuThread: 0, systemRamAvailableGb: 0, + systemRamTotalGb: 0 }; // Module-level cache so multiple components share one fetch. @@ -30,24 +37,30 @@ async function fetchGpuOnce(): Promise { try { const res = await authFetch("/api/system"); if (!res.ok) throw new Error(`HTTP ${res.status}`); - const data = await res.json(); - const ramAvailableGb = data?.memory?.available_gb ?? 0; + + const data = await res.json() as SystemInfoResponse; const gpuData = data?.gpu; - if (!gpuData?.available || !gpuData.devices?.length) { - // No discrete GPU (e.g. Mac): still surface system RAM so memory math - // (unified memory) has a budget to work with. - const info: GpuInfo = { ...DEFAULT_GPU, systemRamAvailableGb: ramAvailableGb }; - cachedGpu = info; - return info; - } - const devices = gpuData.devices as Array<{ name?: string; memory_total_gb?: number }>; - const totalGb = devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0); - const info: GpuInfo = { - available: true, - name: devices[0]?.name ?? "Unknown", - memoryTotalGb: totalGb, - systemRamAvailableGb: ramAvailableGb, + + // CPU/RAM exist even on hosts without a GPU, so populate them on every path. + // No discrete GPU (e.g. Mac): still surface system RAM so memory math + // (unified memory) has a budget to work with. + const base = { + cpuCore: data?.cpu?.physical_count ?? 0, + cpuThread: data?.cpu?.logical_count ?? 0, + systemRamAvailableGb: data?.memory?.available_gb ?? 0, + systemRamTotalGb: data?.memory?.total_gb ?? 0, }; + + const devices = gpuData?.devices ?? []; + const info: GpuInfo = + gpuData?.available && devices.length + ? { + ...base, + available: true, + name: devices[0]?.name ?? "Unknown", + memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), + } + : { ...DEFAULT_GPU, ...base }; cachedGpu = info; return info; } catch { @@ -78,4 +91,4 @@ export function useGpuInfo(): GpuInfo { }, []); return gpu; -} +} \ No newline at end of file diff --git a/studio/frontend/src/hooks/use-gpu-utilization.ts b/studio/frontend/src/hooks/use-gpu-utilization.ts index 1a9f6102dd..be16647a4d 100644 --- a/studio/frontend/src/hooks/use-gpu-utilization.ts +++ b/studio/frontend/src/hooks/use-gpu-utilization.ts @@ -7,6 +7,9 @@ import { useEffect, useRef, useState } from "react"; export interface GpuUtilization { available: boolean; backend: string | null; + devices?: GpuUtilization[]; + index?: number; + visible_ordinal?: number; gpu_utilization_pct: number | null; temperature_c: number | null; vram_used_gb: number | null; @@ -57,11 +60,10 @@ export function useGpuUtilization( const json = (await res.json()) as GpuUtilization; if (!cancelled) setData(json); } catch { - // Silently ignore — next poll will retry + // Retry on the next poll. } } - // Fetch immediately, then set up interval void poll(); timerRef.current = setInterval(() => void poll(), intervalMs); diff --git a/studio/frontend/src/hooks/use-system.ts b/studio/frontend/src/hooks/use-system.ts new file mode 100644 index 0000000000..a135cce86e --- /dev/null +++ b/studio/frontend/src/hooks/use-system.ts @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { useEffect, useState } from "react"; + +export interface GpuDevice { + index?: number; + index_kind?: string; + visible_ordinal?: number; + name?: string; + memory_total_gb?: number; + vram_used_gb?: number; + vram_free_gb?: number; + vram_utilization_pct?: number | null; +} + +export interface SystemInfoResponse { + platform: string; + python_version: string; + device_backend: "cuda" | "rocm" | "cpu" | "mlx" | "xpu"; + uptime_seconds: number | null; + cpu: { + logical_count: number; + physical_count: number; + usage_percent: number; + frequency_mhz: number | null; + }; + memory: { + total_gb: number; + available_gb: number; + percent_used: number; + process_used_mb: number; + }; + disk: { + total_gb: number; + free_gb: number; + percent_used: number; + }; + gpu: { + available: boolean; + backend?: string; + backend_cuda_visible_devices?: string | null; + parent_visible_gpu_ids?: number[]; + index_kind?: string; + devices: GpuDevice[]; + }; + ml_packages: { + torch?: string; + transformers?: string; + }; +} + +let cachedSystem: SystemInfoResponse | null = null; +let systemFetchPromise: Promise | null = null; + +const DEFAULT_SYSTEM: SystemInfoResponse = { + platform: "Unknown", + python_version: "Unknown", + device_backend: "cpu", + uptime_seconds: 0, + cpu: { logical_count: 0, physical_count: 0, usage_percent: 0, frequency_mhz: null }, + memory: { total_gb: 0, available_gb: 0, percent_used: 0, process_used_mb: 0 }, + disk: { total_gb: 0, free_gb: 0, percent_used: 0 }, + gpu: { available: false, devices: [] }, + ml_packages: {} +}; + +async function fetchSystemOnce({ + force = false, +}: { force?: boolean } = {}): Promise { + if (systemFetchPromise) return systemFetchPromise; + if (!force && cachedSystem) return cachedSystem; + + systemFetchPromise = (async () => { + try { + const res = await authFetch("/api/system"); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const data = await res.json(); + + cachedSystem = data as SystemInfoResponse; + return cachedSystem; + } catch { + cachedSystem = null; + return DEFAULT_SYSTEM; + } finally { + systemFetchPromise = null; + } + })(); + + return systemFetchPromise; +} + +interface UseSystemInfoOptions { + pollMs?: number; + enabled?: boolean; +} + +export function useSystemInfo({ + pollMs, + enabled = true, +}: UseSystemInfoOptions = {}): SystemInfoResponse { + const [systemInfo, setSystemInfo] = useState(cachedSystem ?? DEFAULT_SYSTEM); + + useEffect(() => { + if (!enabled) return; + + let cancelled = false; + let timeoutId: number | null = null; + + const update = (force: boolean) => { + void fetchSystemOnce({ force }) + .then((info) => { + if (!cancelled) setSystemInfo(info); + }) + .finally(() => { + if (cancelled || !pollMs) return; + timeoutId = window.setTimeout(() => update(true), pollMs); + }); + }; + + update(Boolean(pollMs)); + return () => { + cancelled = true; + if (timeoutId !== null) window.clearTimeout(timeoutId); + }; + }, [enabled, pollMs]); + + return systemInfo; +} diff --git a/studio/frontend/src/i18n/AGENTS.md b/studio/frontend/src/i18n/AGENTS.md index 42d964acca..35c025cd0b 100644 --- a/studio/frontend/src/i18n/AGENTS.md +++ b/studio/frontend/src/i18n/AGENTS.md @@ -2,10 +2,11 @@ - `locales/en.ts` is the complete baseline message file. - Non-English locale files may be partial. Missing keys must fall back to English at runtime. -- Use BCP 47 locale tags for new languages, for example `zh-CN`, `ja-JP`, and `ko-KR`. +- Use BCP 47 locale tags for new languages, for example `zh-CN`, `pt-BR`, `ja-JP`, and `ko-KR`. - Do not change fallback logic to hide missing translations. - Do not add automatic DOM translation, MutationObserver text replacement, or runtime guess-based translation. - Preserve interpolation variables exactly, for example `{count}`, `{model}`, and `{provider}`. - Keep product and technical names unchanged unless there is an established localized name, for example `Unsloth Studio`, `LoRA`, `GGUF`, and `Hugging Face`. - Keep translation changes small and reviewable. Prefer separate commits for runtime changes, UI migration, and locale text. - When adding user-facing Studio UI text, add the English message key first and add non-English overrides only when the translation is clear. +- Run `npx tsx src/i18n/check-parity.ts` before committing to ensure there are no shape mismatches or placeholder discrepancies in the non-English overlays. \ No newline at end of file diff --git a/studio/frontend/src/i18n/check-parity.ts b/studio/frontend/src/i18n/check-parity.ts index 8c027f9ce3..e2b66f5cdc 100644 --- a/studio/frontend/src/i18n/check-parity.ts +++ b/studio/frontend/src/i18n/check-parity.ts @@ -3,13 +3,14 @@ // Parity check between en.ts and every non-English locale. // - Locale files may be partial; missing keys must fall back to English. -// - All zh-CN keys must exist in en (no extras). +// - All non-English keys must exist in en (no extras). // - Placeholder set must match per leaf between en and the overlay. // // Run: npx tsx src/i18n/check-parity.ts import { en } from "./locales/en.ts"; import { zhCN } from "./locales/zh-CN.ts"; +import { ptBR } from "./locales/pt-br.ts"; import { ja } from "./locales/ja.ts"; type Tree = { readonly [k: string]: string | Tree }; @@ -90,6 +91,7 @@ function checkExtras( const overlays: Record = { "zh-CN": zhCN as unknown as Tree, + "pt-BR": ptBR as unknown as Tree, "ja": ja as unknown as Tree, }; let anyError = false; @@ -112,4 +114,4 @@ for (const [locale, overlay] of Object.entries(overlays)) { } if (anyError) process.exit(1); -console.log("\nAll locale overlays pass parity."); +console.log("\nAll locale overlays pass parity."); \ No newline at end of file diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 136e8523ba..92abc222a0 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -92,6 +92,7 @@ export const en = { general: "General", profile: "Profile", appearance: "Appearance", + resources: "System", chat: "Chat", connections: "Connections", apiKeys: "API", @@ -275,6 +276,58 @@ export const en = { "Keep the sidebar expanded instead of collapsing to icons.", }, }, + resources: { + title: "System", + description: "Monitor this Studio server's hardware and storage.", + liveUpdates: "Live updates", + floatingWindow: "Floating window", + disableOverlay: "Disable overlay", + liveMonitor: { + title: "Live monitor", + cpu: "CPU", + ram: "RAM", + disk: "Disk", + vram: "VRAM", + cpuCores: "{logical} logical / {physical} physical cores", + currentLoad: "Current load", + free: "{value} free", + noGpu: "No visible GPU", + }, + gpu: { + title: "GPU devices", + noGpu: "No visible GPU detected. CPU-only resources are shown above.", + unknownDevice: "Unknown GPU", + deviceWithIndex: "GPU {index}", + vramUtilization: "VRAM", + used: "{value} used", + free: "{value} free", + total: "{value} total", + }, + storage: { + title: "Storage", + systemDisk: "System disk", + diskUsage: "{used} used / {total}", + diskFree: "{free} free", + modelsFolder: "Models folder", + modelsFolderDescription: "Where downloaded models are stored.", + openAction: "Open", + copyAction: "Copy path", + copied: "Path copied", + openError: "Couldn't open the folder", + copyError: "Couldn't copy the path", + }, + environment: { + title: "Environment", + backend: "Backend", + python: "Python", + torch: "Torch", + transformers: "Transformers", + uptime: "Uptime", + processMemory: "Process memory", + notInstalled: "Not installed", + unknown: "Unknown", + }, + }, chat: { title: "Chat", description: "Manage chat history stored on this device.", @@ -373,8 +426,10 @@ export const en = { usageTools: "Tools", exampleCurlTools: "curl + tools", examplePythonTools: "Python + tools", + exampleJavaScriptTools: "JavaScript + tools", exampleCurlAdvanced: "curl + advanced", examplePythonAdvanced: "Python + advanced", + exampleJavaScriptAdvanced: "JavaScript + advanced", osUnix: "Linux / macOS / WSL", osWindows: "Windows", secureHttps: "Secure HTTPS", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts new file mode 100644 index 0000000000..c261e4ed0c --- /dev/null +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -0,0 +1,934 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +export const ptBR = { + common: { + cancel: "Cancelar", + close: "Fechar", + delete: "Excluir", + done: "Concluído", + error: "Erro", + export: "Exportar", + help: "Ajuda", + loading: "Carregando...", + new: "Novo", + rename: "Renomear", + save: "Salvar", + saving: "Salvando...", + search: "Buscar", + shutdown: "Desligar", + }, + shell: { + beta: "BETA", + brand: "unsloth", + product: "Unsloth Studio", + accountMenu: "Menu de conta {name}", + updateAvailable: "Atualização disponível", + aria: { + home: "Início do Unsloth", + closeSidebar: "Fechar barra lateral", + openSidebar: "Abrir barra lateral", + chatOptions: "Opções de chat", + runOptions: "Opções de execução", + }, + navigation: { + newChat: "Novo Chat", + returnToChat: "Retornar ao Chat", + compare: "Comparar", + search: "Buscar", + hub: "Hub", + train: "Treinar", + recipes: "Receitas", + export: "Exportar", + recents: "Recentes", + settings: "Configurações", + api: "API", + lightMode: "Modo Claro", + darkMode: "Modo Escuro", + guidedTour: "Tour Guiado", + help: "Ajuda", + logOut: "Sair", + shutdown: "Desligar", + }, + notFound: { + title: "Página não encontrada", + description: "{path} não existe.", + backToChat: "Voltar para o chat", + }, + dialog: { + deleteChat: { + title: "Excluir chat", + description: 'Tem certeza de que deseja excluir este chat "{name}"?', + }, + deleteRun: { + title: "Excluir execução de treino", + description: 'Tem certeza de que deseja excluir esta execução "{name}"?', + }, + renameChat: { + title: "Renomear chat", + placeholder: "Título do chat", + }, + renameRun: { + title: "Renomear execução", + placeholder: "Nome da execução", + }, + }, + toast: { + cannotDeleteRunningRun: "Não é possível excluir uma execução de treino em andamento", + failedToDeleteChat: "Falha ao excluir o chat", + failedToDeleteRun: "Falha ao excluir a execução", + failedToRenameChat: "Falha ao renomear o chat", + failedToRenameRun: "Falha ao renomear a execução", + }, + }, + settings: { + title: "Configurações", + dialog: { + title: "Configurações", + description: "Gerencie suas preferências do Unsloth.", + closeAriaLabel: "Fechar configurações", + }, + tabs: { + general: "Geral", + profile: "Perfil", + appearance: "Aparência", + resources: "Sistema", + chat: "Chat", + connections: "Conexões", + apiKeys: "API", + about: "Sobre", + }, + general: { + title: "Geral", + description: "Preferências globais do Unsloth.", + account: "Conta", + huggingFaceToken: "Token do Hugging Face", + huggingFaceTokenDescription: + "Usado para carregar modelos restritos e enviar artefatos.", + tokenSaved: "Token salvo", + hideToken: "Ocultar token", + showToken: "Mostrar token", + password: "Senha", + passwordDescription: "Altere a senha desta conta do Studio.", + passwordDialog: { + trigger: "Alterar senha", + title: "Alterar senha", + description: + "Insira sua senha atual e escolha uma nova (no mínimo {minLength} caracteres).", + currentPassword: "Senha atual", + newPassword: "Nova senha", + confirmPassword: "Confirmar nova senha", + currentTooShort: + "A senha atual deve ter no mínimo {minLength} caracteres.", + newTooShort: "A nova senha deve ter no mínimo {minLength} caracteres.", + mismatch: "As senhas não coincidem.", + samePassword: + "A nova senha deve ser diferente da senha atual.", + update: "Atualizar senha", + updating: "Atualizando...", + updated: "Senha atualizada.", + updateFailed: "Falha ao atualizar a senha.", + }, + chatDefaults: "Padrões do chat", + autoTitleNewChats: "Gerar título automático para novos chats", + autoTitleNewChatsDescription: + "Gera um título curto a partir da primeira mensagem.", + helperLlm: { + sectionTitle: "LLM Auxiliar", + preloadOnStartup: "Pré-carregar LLM Auxiliar na inicialização", + preloadOnStartupDescription: + "Baixa o modelo auxiliar do Assistente de IA em segundo plano ao iniciar. Desativado por padrão; o Assistente de IA ainda pode buscá-lo sob demanda.", + disabledByEnv: + "Desativado por UNSLOTH_HELPER_MODEL_DISABLE no ambiente de backend.", + loadError: "Falha ao carregar as configurações do LLM Auxiliar.", + saveError: "Falha ao salvar as configurações do LLM Auxiliar.", + }, + notifications: { + sectionTitle: "Notificações", + showLlamaUpdates: "Notificações de atualização do llama.cpp", + showLlamaUpdatesDescription: + "Notifica quando uma nova versão do llama.cpp estiver disponível. Desative se você apenas realiza treinos.", + }, + gettingStarted: "Primeiros passos", + startOnboarding: "Iniciar integração", + startOnboardingDescription: + "Reabre o assistente de configuração sem alterar sua conta.", + startOnboardingAction: "Iniciar integração", + uploads: { + sectionTitle: "Uploads", + maxUploadSize: "Limite de upload do dataset de treino", + maxUploadSizeDescription: + "O padrão é {defaultSize} MB.", + }, + storage: { + sectionTitle: "Armazenamento", + modelsFolder: "Pasta de modelos", + modelsFolderDescription: + "Onde os modelos baixados são armazenados.", + openAction: "Abrir", + copyAction: "Copiar caminho", + copied: "Caminho copiado", + openError: "Não foi possível abrir a pasta", + copyError: "Não foi possível copiar o caminho", + }, + resetPreferences: { + sectionTitle: "Zona de perigo", + label: "Redefinir todas as preferências locais", + description: + "Limpa apenas as preferências locais. Chats, acesso à API e configurações salvas no banco de dados são mantidos.", + action: "Redefinir preferências", + confirmTitle: "Redefinir todas as preferências locais?", + confirmDescription: + "Limpa as preferências locais e recarrega o Unsloth. Chats, acesso à API e configurações salvas no banco de dados são mantidos.", + confirmAction: "Redefinir e recarregar", + }, + }, + profile: { + title: "Perfil", + description: "Como seu perfil aparece no Unsloth.", + changePicture: "Alterar foto de perfil", + displayName: "Nome de exibição", + nickname: "Como o Unsloth deve chamar você?", + nicknamePlaceholder: "Apelido", + nicknameSaved: "Nome preferido salvo", + avatarShape: "Formato da foto de perfil", + avatarShapeCircle: "Círculo", + avatarShapeRounded: "Arredondado", + chooseSloth: "Ou escolha uma preguiça", + nameSaved: "Nome de perfil salvo", + namePersistErrorTitle: "Não foi possível persistir o nome de perfil", + namePersistErrorDescription: + "Nome atualizado para esta sessão, mas pode não persistir após recarregar.", + photoUpdated: "Foto de perfil atualizada", + photoPersistErrorTitle: "Não foi possível persistir a foto de perfil", + photoPersistErrorDescription: + "Foto atualizada para esta sessão, mas pode não persistir após recarregar.", + photoUpdateErrorTitle: "Não foi possível atualizar a foto de perfil", + imageUseError: "Não foi possível usar esta imagem.", + }, + appearance: { + title: "Aparência", + description: "Como o Unsloth Studio se parece neste dispositivo.", + theme: { + title: "Tema", + label: "Esquema de cores", + description: "Claro, escuro ou seguir o sistema.", + system: "Sistema", + light: "Claro", + dark: "Escuro", + }, + language: { + title: "Idioma", + label: "Idioma de exibição", + description: "O idioma utilizado pelo Unsloth.", + }, + layout: { + title: "Layout", + compactSidebar: "Fixar barra lateral por padrão", + compactSidebarDescription: + "Mantém a barra lateral expandida em vez de recolhê-la em ícones.", + }, + }, + resources: { + title: "Sistema", + description: "Monitore o hardware e o armazenamento deste servidor Studio.", + liveUpdates: "Atualizações ao vivo", + floatingWindow: "Janela flutuante", + disableOverlay: "Desativar sobreposição", + liveMonitor: { + title: "Monitor ao vivo", + cpu: "CPU", + ram: "RAM", + disk: "Disco", + vram: "VRAM", + cpuCores: "{logical} lógicos / {physical} físicos", + currentLoad: "Carga atual", + free: "{value} livres", + noGpu: "Nenhuma GPU visível", + }, + gpu: { + title: "Dispositivos GPU", + noGpu: "Nenhuma GPU visível detectada. Os recursos somente CPU aparecem acima.", + unknownDevice: "GPU desconhecida", + deviceWithIndex: "GPU {index}", + vramUtilization: "VRAM", + used: "{value} usados", + free: "{value} livres", + total: "{value} total", + }, + storage: { + title: "Armazenamento", + systemDisk: "Disco do sistema", + diskUsage: "{used} usados / {total}", + diskFree: "{free} livres", + modelsFolder: "Pasta de modelos", + modelsFolderDescription: "Onde os modelos baixados são armazenados.", + openAction: "Abrir", + copyAction: "Copiar caminho", + copied: "Caminho copiado", + openError: "Não foi possível abrir a pasta", + copyError: "Não foi possível copiar o caminho", + }, + environment: { + title: "Ambiente", + backend: "Backend", + python: "Python", + torch: "Torch", + transformers: "Transformers", + uptime: "Tempo ativo", + processMemory: "Memória do processo", + notInstalled: "Não instalado", + unknown: "Desconhecido", + }, + }, + chat: { + title: "Chat", + description: "Gerencie o histórico de chat armazenado neste dispositivo.", + modelDisclaimer: "Mostrar aviso do modelo", + modelDisclaimerDescription: + 'Mostra "LLMs podem cometer erros" abaixo da caixa de chat.', + artifacts: { + title: "Canvas", + collapseHtmlBlocks: "Recolher blocos HTML", + collapseHtmlBlocksDescription: + "O modo Canvas recolhe o HTML completo automaticamente. Ative isso para também recolher documentos HTML delimitados quando o Canvas estiver desativado.", + allowNetworkAccess: "Permitir acesso à rede no canvas", + allowNetworkAccessDescription: + "Permite que as pré-visualizações do canvas carreguem scripts, estilos, fontes, mídia e recursos de rede de CDNs. Mantenha desativado para pré-visualizações totalmente offline.", + }, + data: "Dados", + exportHistory: "Exportar histórico de chat", + exportHistoryDescription: + "Baixe todos os chats e mensagens em formato JSON.", + exportAction: "Exportar", + exportingAction: "Exportando...", + exportConversations: "Exportar Recentes e Projetos", + exportConversationsDescription: + "Baixe os Recentes ou Recentes mais chats de projetos como JSONL bruto, CSV ou ShareGPT JSONL, combinados ou por chat.", + exportConversationsAction: "Exportar", + exportScopeRecents: "Recentes", + exportScopeAll: "Recentes + Projetos", + exportCombinedSuffix: "(combinado)", + exportPerChatSuffix: "(por chat)", + importChats: "Importar chats", + importChatsDescription: + "Importe um arquivo exportado em JSONL, NDJSON ou CSV para os Recentes.", + importChatsAction: "Importar", + importNoConversations: "Nenhuma conversa encontrada no arquivo.", + importedOneChat: "Importada 1 conversa para os Recentes.", + importedChatCount: "Importadas {count} conversas para os Recentes.", + importFailed: "Falha na importação.", + clearHistory: "Limpar histórico de chat", + clearHistoryDescription: "Exclui o histórico de chat deste dispositivo.", + clearAction: "Limpar", + clearAllChats: "Limpar todos os chats", + clearAllChatsDescription: "Exclui permanentemente todos os chats deste dispositivo.", + noChatsToClear: "Nenhum chat para limpar.", + clearOneChatDescription: + "Exclui permanentemente o único chat deste dispositivo.", + clearChatCountDescription: + "Exclui permanentemente todos os {count} chats deste dispositivo.", + clearChatsAction: "Limpar chats", + clearOneChatTitle: "Limpar 1 chat?", + clearChatsTitle: "Limpar {count} chats?", + clearChatsConfirmDescription: + "Exclui permanentemente todos os chats deste dispositivo. Esta ação não pode ser desfeita.", + clearingAction: "Limpando...", + clearOneChatAction: "Limpar 1 chat", + clearChatCountAction: "Limpar {count} chats", + clearedAllChats: "Todos os chats foram limpos", + clearedOneChat: "1 chat foi limpo", + clearedChatCount: "{count} chats foram limpos", + someChatsCouldNotBeCleared: "Não foi possível limpar alguns chats", + chatsClearedRemainOne: + "{clearedCount} chats limpos; 1 chat restante. Por favor, tente novamente.", + chatsClearedRemain: + "{clearedCount} chats limpos; {remainingCount} chats restantes. Por favor, tente novamente.", + oneChatClearedRemain: + "1 chat limpo; {remainingCount} chats restantes. Por favor, tente novamente.", + oneChatClearedRemainOne: "1 chat limpo; 1 chat restante. Por favor, tente novamente.", + storageClearFailedOne: + "Falha ao limpar o armazenamento; 1 chat pode ter restado. Por favor, tente novamente.", + storageClearFailed: + "Falha ao limpar o armazenamento; {count} chats podem ter restado. Por favor, tente novamente.", + failedToClearChats: "Falha ao limpar os chats", + }, + connections: { + title: "Conexões", + description: "Gerencie provedores e conexões externas.", + }, + apiKeys: { + title: "API", + description: + "Acesse o Unsloth por meio da API compatível com OpenAI.", + readDocs: "Leia a documentação da API", + noAccess: "Nenhum acesso à API ainda.", + newBadge: "Novo", + accessTokens: "Tokens de acesso", + loadError: "Não foi possível carregar o acesso à API.", + createError: "Não foi possível criar o token de acesso.", + revokeError: "Não foi possível revogar o token de acesso.", + never: "Nunca", + tokenNamePlaceholder: "Nome do token (ex: producao)", + newAccessTokenName: "Nome do novo token de acesso", + createToken: "Criar token", + creating: "Criando...", + newTokenCreated: "Novo token de acesso criado", + accessTokenCopied: "Token de acesso copiado", + copyAccessToken: "Copiar token de acesso", + copyNow: "Copie agora - isto não será exibido novamente.", + usageExamples: "Exemplos de uso", + usageTools: "Ferramentas", + exampleCurlTools: "curl + ferramentas", + examplePythonTools: "Python + ferramentas", + exampleJavaScriptTools: "JavaScript + ferramentas", + exampleCurlAdvanced: "curl + avançado", + examplePythonAdvanced: "Python + avançado", + exampleJavaScriptAdvanced: "JavaScript + avançado", + osUnix: "Linux / macOS / WSL", + osWindows: "Windows", + secureHttps: "HTTPS Seguro", + secureHttpsHint: + "A porta 0.0.0.0 ainda está acessível globalmente. Para segurança total, inicie o Unsloth Studio com --secure para expor apenas este link HTTPS.", + copyTunnelUrl: "Copiar URL do túnel", + copySnippet: "Copiar trecho de código", + copy: "Copiar", + copied: "Copiado", + setupDocs: "Docs de configuração:", + relativeNever: "nunca", + relativeJustNow: "agora mesmo", + relativeHoursAgo: "há {count}h", + relativeDaysAgo: "há {count}d", + relativeMonthsAgo: "há {count} meses", + relativeYearsAgo: "há {count} anos", + expired: "expirado", + today: "hoje", + inDays: "em {count}d", + created: "Criado {value}", + used: "Usado {value}", + expires: "Expira {value}", + actionsFor: "Ações para {name}", + copyPrefix: "Copiar prefixo", + revokeToken: "Revogar token", + revokeTitle: 'Revogar token de acesso "{name}"?', + revokeDescription: + "Aplicativos que usam este token perderão o acesso imediatamente. Esta ação não pode ser desfeita.", + revokeAction: 'Revogar "{name}"', + revoking: "Revogando...", + }, + about: { + title: "Sobre", + description: + "Documentação, notas de lançamento, feedback e informações da build.", + studioVersion: "Versão do Unsloth", + packageVersion: "Versão do Pacote", + llamaCppVersion: "Versão do llama.cpp", + hardware: "Hardware", + gpu: "GPU", + cuda: "CUDA", + rocm: "ROCm", + updates: "Atualização", + help: "Ajuda", + documentation: "Documentação", + releaseNotes: "Notas de lançamento", + whatsNew: "O que há de novo", + feedback: "Feedback", + reportIssue: "Reportar um problema", + license: { + sectionTitle: "Licença", + studioLabel: "Unsloth Studio", + studioLicense: "AGPL-3.0", + studioDescription: + "Código aberto sob a licença GNU AGPL v3.0.", + libraryLabel: "Unsloth Core", + libraryLicense: "Apache-2.0", + libraryDescription: "Licenciado sob Apache 2.0.", + }, + dangerZone: "Zona de perigo", + shutDownStudio: "Desligar Unsloth Studio", + shutDownStudioDescription: + "Interrompe o servidor Unsloth e encerra sua sessão.", + shutDown: "Desligar", + update: { + title: "Atualizar Unsloth Studio", + commandText: "Texto de {label}", + copied: "Copiado", + copyCommand: "Copiar comando", + commandCopied: "{label} copiado", + copyNamedCommand: "Copiar {label}", + checkingInstall: "Verificando como o Unsloth foi instalado...", + installIntro: "Para instalar ou atualizar o Unsloth:", + localUpdateHeading: "Atualização local", + installCommandUnix: "Comando de instalação para macOS/Linux", + installCommandWindows: "Comando de instalação para Windows", + localInstallDetected: + "Instalação local detectada. Atualize a partir do seu repositório original para evitar substituí-lo pelo PyPI.", + pullThenUpdate: "Puxe as últimas alterações (git pull) e depois execute o instalador local:", + gitPullCommand: "comando git pull", + localInstallerCommand: "comando do instalador local", + sourceInstallDetected: + "Instalação do pacote por código-fonte ou VCS detectada. Reinstale a partir do caminho local original ou URL do Git.", + repoCheckoutFallback: + "Se você ainda tiver o repositório baixado, execute o instalador local a partir dele:", + restartAfterUpdate: "Reinicie o Unsloth após a atualização.", + desktopManaged: + "O aplicativo de desktop mantém seu backend integrado atualizado e avisará quando uma nova versão estiver disponível.", + unknownInstall: + "Não foi possível detectar como o Unsloth foi instalado. Para instalações via instalador ou PyPI, use os comandos acima.", + localCheckout: + "Para instalações de repositório local, execute o instalador local a partir desse diretório:", + docs: "Docs de instalação:", + docsInstall: "Instalação", + docsUpdating: "Atualização", + docsMac: "Mac", + docsWindows: "Windows", + }, + }, + }, + studio: { + routeTitle: "Treinar", + title: "Estúdio de Fine-tuning", + subtitles: { + configure: "Configure e inicie o treinamento", + trainingInProgress: "Treinamento em andamento", + viewPastRuns: "Visualizar execuções de treino anteriores", + viewingPastRun: "Visualizando execução anterior", + }, + tabs: { + configure: "Configurar", + currentRun: "Execução Atual", + history: "Histórico", + }, + loadingRuntime: "Carregando ambiente de execução de treino...", + backToHistory: "Voltar ao histórico", + sections: { + model: "Modelo", + dataset: "Dataset", + params: "Parâmetros", + training: "Treinamento", + charts: "Gráficos", + progress: "Progresso do Treinamento", + }, + configure: { + title: "Configurar", + description: "Escolha um modelo, dataset e configurações de treinamento.", + startTraining: "Iniciar Treinamento", + starting: "Iniciando...", + loadingModel: "Carregando modelo...", + checkingDataset: "Verificando dataset...", + trainingConfig: "Configuração de Treino", + }, + model: { + title: "Modelo", + description: "Selecione o modelo base e o método de treinamento", + fasterTrainingBadge: "Treinamento 2x Mais Rápido", + baseModel: "Modelo base", + localModel: "Modelo Local", + localModelTooltip: + "Caminho para um modelo baixado localmente ou um repositório HF customizado.", + scanningLocalAndCachedModels: "Escaneando modelos locais e em cache...", + scanning: "Escaneando...", + scanningLocalModels: "Escaneando modelos locais...", + noLocalModelsFound: "Nenhum modelo local encontrado", + noLocalModelsFoundManual: "Nenhum modelo local encontrado. Insira o caminho manualmente.", + failedToLoadLocalModels: "Falha ao carregar modelos locais", + hfCache: "Cache do HF", + customFolders: "Pastas Customizadas", + localDir: "Diretório local", + huggingFaceModel: "Modelo do Hugging Face", + huggingFaceModelTooltip: + "Busque modelos no Hugging Face ou escolha da nossa lista recomendada.", + searchModels: "Buscar modelos...", + searching: "Buscando...", + noModelsFound: "Nenhum modelo encontrado", + needsVram: "Precisa de ~{vram}GB de VRAM (GPU: {gpu}GB)", + tightVram: "~{vram}GB de VRAM (limite na {gpu}GB)", + vramEstimate: "~{vram}GB de VRAM", + method: "Método", + methodTooltip: + "O QLoRA usa quantização de 4 bits para menor uso de VRAM. O LoRA usa 16 bits. O Full atualiza todos os pesos. O CPT (Continued Pretraining) treina em texto bruto para adaptar o modelo a um novo domínio sem formatação de chat.", + readMore: "Leia mais", + fullFineTune: "Fine-tune Completo (Full)", + checkingToken: "Verificando token...", + getOrUpdateToken: "Obter ou atualizar token", + huggingFaceTokenOptional: "Token do Hugging Face (Opcional)", + continuedPretraining: "Pré-treinamento Contínuo (CPT)", + localModels: "Modelos locais", + localModelsFound: "{count} modelos locais/em cache encontrados", + loadingLocalModels: "Carregando modelos locais...", + }, + dataset: { + title: "Dataset", + description: "Selecione ou envie os dados de treinamento", + source: "Origem do dataset", + chooseDataset: "Escolher dataset", + chooseDatasetTooltip: + "Use as abas do pop-up para alternar entre o Hugging Face e as saídas de receitas locais.", + localTab: "Local", + searchHuggingFaceDatasets: "Buscar datasets no Hugging Face...", + searchLocalDatasets: "Buscar datasets locais...", + searching: "Buscando...", + noDatasetsFound: "Nenhum dataset encontrado", + loadingLocalDatasets: "Carregando datasets locais...", + failedToLoadLocalDatasets: "Falha ao carregar datasets locais.", + noLocalDatasetsYet: "Nenhum dataset local ainda.", + noLocalDatasetsMatchSearch: "Nenhum dataset local corresponde à busca.", + openDataRecipes: "Abrir Receitas de Dados", + browsingSource: "Navegando em {browsing}. A seleção atual permanece {current}.", + localDatasets: "Datasets locais", + localDataset: "Dataset local", + localDatasetRows: " / {count} linhas", + huggingFaceDataset: "Dataset do Hugging Face", + localDatasetMetadata: "Metadados do dataset local", + dataRecipeOutput: "Saída da Receita de Dados.", + rows: "Linhas", + columns: "Colunas", + batches: "Lotes", + updated: "Atualizado", + evalDataset: "Dataset de validação (Eval)", + uploading: "Enviando...", + upload: "Upload", + uploadEvalFile: "Enviar arquivo de validação", + evalDatasetDescription: + "Opcional. Se não for fornecido, uma pequena parte será dividida a partir dos dados de treinamento.", + advanced: "Avançado", + targetFormat: "Formato de Destino", + targetFormatTooltip: + "Formato dos seus dados de treinamento. A detecção automática funciona para a maioria dos datasets.", + auto: "Auto", + rawText: "Texto Bruto", + trainSplitStart: "Início da Divisão de Treino", + trainSplitStartTooltip: + "Treine apenas em um subconjunto da sua divisão de treino especificando um índice de linha inicial (inclusivo, baseado em 0). Deixe em branco para começar da primeira linha.", + trainSplitEnd: "Fim da Divisão de Treino", + trainSplitEndTooltip: + "Último índice de linha a ser incluído da divisão de treino (inclusivo, baseado em 0). Por exemplo, defina o Início como 0 e o Fim como 99 para treinar nas primeiras 100 linhas. Deixe em branco para usar todas as linhas restantes.", + endPlaceholder: "Fim", + clear: "Limpar", + dropFileOrClick: "Solte 1 arquivo aqui ou clique para fazer upload", + viewDataset: "Visualizar dataset", + uploadFailed: "Falha no envio", + unknownError: "Erro desconhecido", + unsupportedFileType: "Tipo de arquivo não suportado", + uploadOneFileType: "Envie um arquivo do tipo {types}.", + datasetUploaded: "Dataset enviado", + evalDatasetUploaded: "Dataset de validação enviado", + uploadOneFileAtATime: "Envie um arquivo por vez", + uploadSingleFileDescription: + "O upload do dataset de treinamento aceita apenas um único arquivo.", + checkingToken: "Verificando token...", + getOrUpdateToken: "Obter ou atualizar token", + preview: "Pré-visualizar dataset", + split: "Divisão (Split)", + subset: "Subconjunto (Subset)", + s3: { + title: "Configuração do S3", + description: "Carregue datasets em .parquet, .json, .jsonl ou .csv do Amazon S3", + bucket: "Nome do Bucket", + bucketPlaceholder: "meu-bucket-de-dados-de-treino", + region: "Região da AWS", + regionPlaceholder: "us-east-1", + prefix: "Prefixo do Caminho", + prefixPlaceholder: "datasets/whisper/", + prefixTooltip: "Caminho opcional dentro do bucket para os arquivos do seu dataset", + accessKeyId: "ID da Chave de Acesso", + accessKeyIdPlaceholder: "AKIAIOSFODNN7EXAMPLE", + secretAccessKey: "Chave de Acesso Secreta", + secretAccessKeyPlaceholder: "Sua chave de acesso secreta da AWS", + useIamRole: "Usar Função IAM", + useIamRoleTooltip: "Usa credenciais de função IAM em vez de chaves de acesso (recomendado para EC2/SageMaker)", + testConnection: "Testar Conexão", + connectionSuccess: "Conectado com sucesso ao bucket S3", + connectionFailed: "Falha ao conectar ao bucket S3", + comingSoon: "Integração com S3 em breve", + comingSoonDescription: "O carregamento de datasets do S3 requer o boto3. Este recurso está em desenvolvimento.", + }, + }, + params: { + title: "Parâmetros", + description: "Configure os hiperparâmetros de treinamento", + loraSettings: "Configurações do LoRA", + trainingHyperparameters: "Hiperparâmetros de Treinamento", + maxSteps: "Passos Máximos (Max Steps)", + epochs: "Épocas (Epochs)", + useMaxSteps: "Usar Passos Máximos", + useEpochs: "Usar Épocas", + maxStepsTooltip: "Sobrescreve o total de passos do otimizador.", + epochsTooltip: "Número de passagens completas pelo dataset.", + epochsDescription: "Cada época é uma passagem completa pelo seu dataset.", + maxStepsDescription: + "Limita o treinamento a um número fixo de passos do otimizador.", + contextLength: "Comprimento do Contexto", + contextLengthTooltip: "Número máximo de tokens por amostra de treinamento.", + customContextLength: "Insira um valor personalizado", + contextLengthDescription: "Comprimento máximo de sequência para amostras de treino", + learningRate: "Taxa de Aprendizado (Learning Rate)", + learningRateTooltip: + "Tamanho do passo para atualizações de peso. Valores menores treinam mais lentamente, mas com mais estabilidade.", + learningRateDescription: + "Recomendado: 2e-4 para LoRA, 5e-5 para CPT, 2e-5 para fine-tune completo", + embeddingLearningRate: "Taxa de Aprendizado do Embedding", + embeddingLearningRateTooltip: + "Usado apenas quando o CPT está treinando embed_tokens. Os embeddings são mais fáceis de desestabilizar do que os pesos LoRA, por isso geralmente precisam de um LR menor. Deixe em branco para usar lr/10; a faixa típica de funcionamento é de 2x a 10x menor que o LR principal. Aumente apenas se a adaptação de vocabulário ou de tokens de domínio estiver muito lenta.", + embeddingLearningRateDescription: + "Deixe em branco para usar lr/10 (recomendado). A faixa típica é de 2x a 10x menor que a taxa de aprendizado principal.", + rank: "Rank", + rankTooltip: + "Dimensão das matrizes de baixo rank. Maior = mais capacidade.", + alpha: "Alpha", + alphaTooltip: "Fator de escala para atualizações LoRA. Geralmente o dobro do rank.", + dropout: "Dropout", + dropoutTooltip: + "Probabilidade de dropout para as camadas LoRA para reduzir o overfitting.", + visionLayers: "Camadas de visão", + languageLayers: "Camadas de linguagem", + attentionModules: "Módulos de atenção", + mlpModules: "Módulos MLP", + targetModules: "Módulos de Destino", + enableLora: "Ativar LoRA", + trainWithLora: "Treinar com LoRA", + stableRank: "Stable Rank", + memoryEfficient: "Eficiente em Memória", + optimization: "Otimização", + schedule: "Cronograma", + memory: "Memória", + optimizer: "Otimizador", + optimizerTooltip: + "Algoritmo de otimização. Variantes de 8 bits reduzem o uso de memória. Fused é recomendado para modelos de visão.", + lrScheduler: "Agendador de LR", + lrSchedulerTooltip: + "Como a taxa de aprendizado muda ao longo do treino. Linear decai de forma constante; cosine decai em curva.", + optimizerOptions: { + adamw8bit: "AdamW 8-bit", + pagedAdamw8bit: "Paged AdamW 8-bit", + adamwBnb8bit: "AdamW BNB 8-bit", + pagedAdamw32bit: "Paged AdamW 32-bit", + adamwTorch: "AdamW (PyTorch)", + adamwTorchFused: "AdamW (PyTorch Fused)", + }, + lrSchedulerOptions: { + linear: "Linear", + cosine: "Cosine", + }, + batchSize: "Tamanho do Lote (Batch Size)", + batchSizeTooltip: "Amostras processadas por passo. Maior consome mais VRAM.", + gradAccum: "Acúmulo de Gradiente", + gradAccumTooltip: "Simula tamanhos de lote maiores sem gastar VRAM extra.", + weightDecay: "Decaimento de Peso", + weightDecayTooltip: "Regularização L2 para evitar overfitting.", + warmupSteps: "Passos de Aquecimento (Warmup)", + warmupStepsTooltip: + "Aumenta gradualmente a LR no início do treino para garantir estabilidade.", + scheduleEpochsTooltip: + "Número de passagens completas pelo dataset. Defina 0 para rodar por passos máximos.", + saveSteps: "Passos para Salvar", + saveStepsTooltip: "Salva um checkpoint a cada N passos. 0 para desativar.", + evalSteps: "Passos de Validação", + evalStepsTooltip: + "Fração dos passos totais de treino entre as validações (0-1). Defina como 0 para desativar. Ex: 0.01 = valida a cada 1% dos passos.", + seed: "Seed", + seedTooltip: "Semente aleatória para reprodutibilidade.", + gradCheckpoint: "Grad Checkpoint", + gradCheckpointTooltip: + "Troca processamento por memória recalculando as ativações.", + none: "Nenhum", + standard: "Padrão", + enablePacking: "Ativar empacotamento (packing)", + assistantCompletionsOnly: "Apenas respostas do assistente", + readMore: "Leia mais", + }, + training: { + title: "Treinamento", + description: "Monitore e controle o treinamento", + chartNoDataTitle: "Nenhum dado de treinamento ainda", + chartNoDataDescription: "Inicie o treinamento para ver o progresso da loss", + startTraining: "Iniciar Treinamento", + starting: "Iniciando...", + loadingModel: "Carregando modelo...", + checkingDataset: "Verificando dataset...", + configLabel: "Configuração de Treino", + upload: "Upload", + uploadConfigTooltip: "Carregar uma configuração YAML salva", + save: "Salvar", + saveConfigTooltip: "Baixar configuração atual como YAML", + reset: "Redefinir", + resetConfigTooltip: "Redefinir para os padrões do modelo", + configLoaded: "Configuração carregada", + failedToLoadConfig: "Falha ao carregar a configuração", + invalidYamlFile: "Arquivo YAML inválido", + failedToReadFile: "Falha ao ler o arquivo", + parametersReset: "Parâmetros redefinidos para os padrões do modelo", + audioIncompatible: + "Este modelo não suporta áudio. Mude para um modelo compatível com áudio ou escolha um dataset sem áudio.", + visionIncompatible: + "O modelo de texto não é compatível com um dataset multimodal. Mude para um modelo de visão ou escolha um dataset apenas de texto.", + cancelTitle: "Cancelar Treinamento", + cancelDescription: "Deseja cancelar a execução de treinamento atual?", + continueAction: "Continuar Treinamento", + cancelAction: "Cancelar Treinamento", + stopTitle: "Interromper Treinamento", + stopDescription: "Escolha como você deseja interromper a execução de treinamento atual.", + stopAction: "Interromper", + stopping: "Interrompendo...", + stopAndSave: "Interromper e Salvar", + compareInChat: "Comparar no Chat", + exportModel: "Exportar Modelo", + milestone: "Marco", + halfwayDone: "Metade concluída. O treinamento passou de 50%.", + doneNextStep: + "Treinamento concluído. Próximo passo: comparar as saídas do modelo base vs fine-tuned.", + }, + history: { + title: "Histórico", + emptyTitle: "Nenhuma execução de treino ainda", + emptyDescription: + "Nenhuma execução de treino ainda. Inicie sua primeira execução na aba Configurar.", + loadError: "Falha ao carregar as execuções de treino", + deleteError: "Falha ao excluir a execução de treino. Por favor, tente novamente.", + retry: "Tentar novamente", + loadMore: "Carregar mais", + loading: "Carregando...", + loadingRun: "Carregando execução de treino...", + runNotFound: "Execução não encontrada", + deleteTitle: "Excluir execução de treino?", + deleteDescription: + "Isso excluirá permanentemente esta execução de treino e todas as suas métricas. Esta ação não pode ser desfeita.", + runCount: "{count} execuções", + oneRun: "1 execução", + resume: "Retomar", + resumeTraining: "Retomar treinamento", + resuming: "Retomando...", + deleteRun: "Excluir execução", + loss: "Loss", + steps: "Passos", + lossTrendSparkline: "Minigráfico de tendência da loss", + relativeJustNow: "agora mesmo", + relativeMinutesAgo: "há {count}m", + relativeHoursAgo: "há {count}h", + relativeDaysAgo: "há {count}d", + status: { + completed: "Concluído", + stopped: "Interrompido", + error: "Erro", + running: "Em andamento", + continued: "Continuado", + }, + message: { + completed: "Treinamento concluído", + stopped: "Treinamento interrompido", + running: "Treinamento em andamento", + errored: "Treinamento com erro", + }, + }, + charts: { + settings: "Configurações do Gráfico", + settingsDescription: + "Ajuste a apresentação do gráfico enquanto o treinamento continua rodando.", + openSettings: "Abrir configurações do gráfico", + viewWindow: "Janela de visualização", + viewWindowDescription: "Mostra apenas os passos mais recentes ou o histórico completo.", + window: "Janela", + all: "Tudo", + trainingLoss: "Loss de Treinamento", + trainingLossDescription: "Controle as sobreposições e a suavização EMA.", + smoothing: "Suavização", + smoothingDescription: "Mova para a direita para mais suavização. `0` = bruto.", + showRawLoss: "Mostrar loss bruta", + showSmoothedLoss: "Mostrar loss suavizada", + showAverageLine: "Mostrar linha média", + scaleAndCleanup: "Escala e limpeza", + linear: "Linear", + log: "Log", + noClip: "Sem corte", + clipP99: "Cortar p99", + clipP95: "Cortar p95", + lossAxis: "Eixo da loss", + gradientNormAxis: "Eixo da norma do gradiente", + learningRateAxis: "Eixo da taxa de aprendizado", + resetDefaults: "Redefinir padrões", + loss: "Loss", + smoothed: "Suavizado", + evalLoss: "Loss de Validação", + learningRate: "Taxa de Aprendizado", + lr: "LR", + gradNorm: "Norma do Grad.", + gradientNorm: "Norma do Gradiente", + step: "Passo {step}", + averageValue: "média {value}", + waitingForFirstEvaluationStep: "Aguardando o primeiro passo de validação...", + evaluationNotConfigured: "Validação não configurada", + evalChartWillAppear: "O gráfico aparecerá assim que o eval_steps for alcançado", + setEvalDatasetAndSteps: + "Defina o dataset de validação e eval_steps para acompanhar a loss de validação", + }, + progress: { + title: "Progresso do Treinamento", + liveMetrics: "Métricas de treino em tempo real", + exportGguf: "Exportar para GGUF", + openConfig: "Abrir configuração de treino", + configLabel: "Configuração de Treino", + hyperparams: "Hiperparâmetros", + epochs: "Épocas", + batchSize: "Tamanho do lote", + learningRate: "Taxa de aprendizado", + optimizer: "Otimizador", + maxSteps: "Passos máximos", + contextLength: "Comprimento do contexto", + warmupSteps: "Passos de warmup", + rank: "Rank", + alpha: "Alpha", + dropout: "Dropout", + variant: "Variante", + epoch: "Época {value}", + percentComplete: "{percent}% completo", + stepProgress: "Passo {current} / {total}", + loss: "Loss", + lr: "LR", + gradNorm: "Norma do Grad.", + model: "Modelo", + method: "Método", + elapsed: "Decorrido: {value}", + eta: "ETA: {value}", + stepsPerSecond: "{value} passos/s", + noStepsPerSecond: "-- passos/s", + tokens: "Tokens: {value}", + gpuMonitor: "Monitor da GPU", + live: "Ao vivo", + utilization: "Utilização", + temperature: "Temperatura", + vram: "VRAM", + power: "Energia", + phase: { + idle: "Ocioso", + downloadingModel: "Baixando modelo", + downloadingDataset: "Baixando dataset", + loadingModel: "Carregando modelo", + loadingDataset: "Carregando dataset", + configuring: "Configurando", + training: "Treinando", + completed: "Concluído", + error: "Erro", + stopped: "Interrompido", + }, + }, + trainingStart: { + ready: "Pronto", + downloading: "Baixando", + preparing: "Preparando", + left: "restam {eta}", + downloaded: "{size} baixados", + terminalStart: "> treinamento do unsloth iniciado...", + preparingResources: "> Preparando modelo e dataset...", + gettingReady: "> Estamos deixando tudo pronto para a sua execução...", + waitingForFirstStep: "> {message} | aguardando o primeiro passo... ({step})", + resumingTraining: "Retomando treinamento...", + startingTraining: "iniciando treinamento...", + dataset: "Dataset", + datasetStreaming: "Dataset: streaming (sem download completo)", + modelWeights: "Pesos do modelo", + }, + tour: { + guidedTour: "Tour Guiado", + }, + }, +} as const; diff --git a/studio/frontend/src/i18n/messages.ts b/studio/frontend/src/i18n/messages.ts index ef58ca23fe..23db2ce326 100644 --- a/studio/frontend/src/i18n/messages.ts +++ b/studio/frontend/src/i18n/messages.ts @@ -4,19 +4,26 @@ import { getLocale } from "./locale-store"; import { en } from "./locales/en"; import { zhCN } from "./locales/zh-CN"; +import { ptBR } from "./locales/pt-br"; import { ja } from "./locales/ja"; import type { InterpolationValues, MessageKey } from "./types"; export const LOCALES = { en: { label: "English", nativeLabel: "English" }, "zh-CN": { label: "Chinese (Simplified)", nativeLabel: "简体中文" }, - ja: { label: "Japanese", nativeLabel: "日本語" }, + "pt-BR": { label: "Portuguese (Brazil)", nativeLabel: "Português (Brasil)" }, + "ja": { label: "Japanese", nativeLabel: "日本語" }, } as const; export type Locale = keyof typeof LOCALES; export type TranslationKey = MessageKey; -export const messages = { en, "zh-CN": zhCN, ja } as const; +export const messages = { + en, + "zh-CN": zhCN, + "pt-BR": ptBR, + ja +} as const; const PLACEHOLDER_PATTERN = /\{([a-zA-Z0-9_]+)\}/g; @@ -75,4 +82,4 @@ export function isSupportedLocale(value: unknown): value is Locale { typeof value === "string" && Object.prototype.hasOwnProperty.call(LOCALES, value) ); -} +} \ No newline at end of file diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 4c50fa1c8e..092e1803f8 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1473,12 +1473,14 @@ class TestHardwareAmdBranching: assert "from . import amd" in source def test_hardware_branches_on_is_rocm_for_utilization(self): - """get_gpu_utilization dispatches to amd.py via _smi_query when IS_ROCM.""" + """get_gpu_utilization dispatches visible metrics through amd.py on ROCm.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] - assert '_smi_query("get_primary_gpu_utilization"' in func_body + assert "_smi_query(" in func_body + assert '"get_visible_gpu_utilization"' in func_body + assert "_reconcile_rocm_unified_memory" in func_body smi = source[ source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1) ] diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 0843a2b447..7a63dcfb85 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -565,30 +565,50 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: raise SystemExit(f"no .gguf files in {save_dir}") gguf_path = gguf_files[0] + # This is a save/reload-integrity smoke; a few generated tokens are enough. + # Keep llama.cpp bounded on macOS runners where BF16 GGUF decode is CPU-bound. + n_predict = os.environ.get("UNSLOTH_GGUF_RELOAD_N", "8") + n_threads = os.environ.get("UNSLOTH_GGUF_RELOAD_THREADS", str(os.cpu_count() or 4)) + reload_timeout = int(os.environ.get("UNSLOTH_GGUF_RELOAD_TIMEOUT", "420")) + with Phase("reload_gguf", metrics): - proc = subprocess.run( - [ - str(llama_cli), - "-m", - str(gguf_path), - "-p", - PROMPT, - "-n", - "24", - "--temp", - "0", - "--seed", - str(SEED), - "-no-cnv", - "--no-warmup", - ], - capture_output = True, - text = True, - timeout = 300, - # Hand llama-cli an immediate EOF; without it -no-cnv can still leave the - # process blocked reading stdin, which times out instead of generating. - stdin = subprocess.DEVNULL, - ) + argv = [ + str(llama_cli), + "-m", + str(gguf_path), + "-p", + PROMPT, + "-n", + n_predict, + "-t", + n_threads, + "--temp", + "0", + "--seed", + str(SEED), + "-c", + "256", + "--no-warmup", + ] + try: + proc = subprocess.run( + argv, + capture_output = True, + text = True, + timeout = reload_timeout, + # Newer llama.cpp keeps llama-cli in chat mode; exit after one reply. + input = "/exit\n", + ) + except subprocess.TimeoutExpired as exc: + + def _decode(stream) -> str: + if isinstance(stream, bytes): + return stream.decode("utf-8", errors = "replace") + return stream or "" + + print(f" [reload:gguf] TIMEOUT stdout:\n{_decode(exc.stdout)[:1000]}", flush = True) + print(f" [reload:gguf] TIMEOUT stderr:\n{_decode(exc.stderr)[:1000]}", flush = True) + raise metrics["llama_cli_returncode"] = proc.returncode metrics["generation"] = (proc.stdout or "")[:1500] From d33a7a7a1a536fea1b14af26652ac8e4e4de096d Mon Sep 17 00:00:00 2001 From: Ayushman <139611211+InfoSage05@users.noreply.github.com> Date: Fri, 3 Jul 2026 00:11:10 +0530 Subject: [PATCH 08/23] Fix: skip fp16/bf16 validation for full finetuning in RL trainers (#6813) --------- Co-authored-by: Ayushman Paul --- unsloth/models/rl.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 53668d14d8..602de69d3f 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1015,6 +1015,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "dtype = _get_dtype(dtype)\n" "float16 = dtype == torch.float16\n" "bfloat16 = dtype == torch.bfloat16\n" + "if full_finetuning:\n" + " if bfloat16 and use_fp16: use_fp16 = False\n" + " if float16 and use_bf16: use_bf16 = False\n" "if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n" "if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n" "if force_float32:\n" From 2bfeb47c92201ebb1c9ab304130f03e3f5e6f092 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 11:49:56 -0700 Subject: [PATCH 09/23] studio/frontend: drop developer-only /grid-test route (#5662) --------- Co-authored-by: danielhanchen --- studio/frontend/src/app/router.tsx | 2 - studio/frontend/src/app/routes/grid-test.tsx | 69 -------------------- 2 files changed, 71 deletions(-) delete mode 100644 studio/frontend/src/app/routes/grid-test.tsx diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index 5c18e637e2..586c03d5df 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -10,7 +10,6 @@ import { Route as dataRecipesRoute } from "./routes/data-recipes"; import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; import { Route as chatRoute } from "./routes/chat"; import { Route as exportRoute } from "./routes/export"; -import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; import { Route as loginRoute } from "./routes/login"; import { Route as hubRoute } from "./routes/hub"; @@ -25,7 +24,6 @@ const routeTree = rootRoute.addChildren([ onboardingRoute, loginRoute, changePasswordRoute, - gridTestRoute, hubRoute, settingsRoute, studioRoute, diff --git a/studio/frontend/src/app/routes/grid-test.tsx b/studio/frontend/src/app/routes/grid-test.tsx deleted file mode 100644 index c4b6b505a1..0000000000 --- a/studio/frontend/src/app/routes/grid-test.tsx +++ /dev/null @@ -1,69 +0,0 @@ -// SPDX-License-Identifier: AGPL-3.0-only -// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -import { DashboardGrid, DashboardLayout } from "@/components/layout"; -import { - Card, - CardContent, - CardDescription, - CardHeader, - CardTitle, -} from "@/components/ui/card"; -import { createRoute } from "@tanstack/react-router"; -import { requireAuth } from "../auth-guards"; -import { Route as rootRoute } from "./__root"; - -export const Route = createRoute({ - getParentRoute: () => rootRoute, - path: "/grid-test", - beforeLoad: () => requireAuth(), - component: GridTestPage, -}); - -function GridTestPage() { - return ( - -
-
-

Grid Test - 3 Columns

-

- max-w-7xl, gap-6, responsive 1→2→3 -

-
- - - {[1, 2, 3].map((i) => ( - - - Card {i} - ~400px at 1280px viewport - - -
- - - ))} - - -
-

4 Columns

-

~296px per card at 1280px

-
- - - {[1, 2, 3, 4].map((i) => ( - - - Card {i} - Smaller cards - - -
- - - ))} - -
- - ); -} From 73e8245ee857b0afa7750870896662bd1ee5dcee Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Thu, 2 Jul 2026 16:11:20 -0500 Subject: [PATCH 10/23] [Studio] Add --with-llama-cpp-dir installer flag to reuse a local llama.cpp (#6472) * Add --with-llama-cpp-dir flag to install.ps1 and install.sh Users can now pass --with-llama-cpp-dir /path/to/llama.cpp to the installer to skip downloading or building llama.cpp and use a local directory instead. A junction (Windows) or symlink (Linux/macOS) is created at the canonical install location, bypassing both the prebuilt download (Phase 3) and source build (Phase 4) steps in setup.ps1/setup.sh. The path is passed via UNSLOTH_LOCAL_LLAMA_CPP_DIR env var which setup.ps1 and setup.sh read directly. Ported from the idea in unslothai/unsloth#4384, reimplemented against current Studio architecture. * test: add static wiring test for --with-llama-cpp-dir flag Cross-checks install.sh, install.ps1, studio/setup.sh and studio/setup.ps1 so the flag's contract (parse -> UNSLOTH_LOCAL_LLAMA_CPP_DIR env var -> link local dir, skip prebuilt download and source build) can't silently regress. Wired into studio-backend-ci.yml alongside the other tests/sh installer tests. * Address review feedback on --with-llama-cpp-dir flag - setup.ps1: delete an existing junction/symlink via DirectoryInfo.Delete() instead of a recursive remove, which can traverse the link and wipe the user's real llama.cpp directory on PowerShell 5.1. - setup.ps1: short-circuit the build chain when a local dir is linked so CMake never runs inside the user's checkout when it lacks a Windows-layout binary. - install.sh / setup.sh: resolve paths with CDPATH= cd -P so a set CDPATH cannot corrupt the resolved path. - install.sh: seed _WITH_LLAMA_CPP_DIR from UNSLOTH_LOCAL_LLAMA_CPP_DIR so an exported env var (piped-install style) is honored instead of being clobbered. - setup.sh: create the root llama-quantize shim when linking a local source build so GGUF export's check_llama_cpp() still finds it. - setup.sh / setup.ps1: drop a stale link before the custom-home ownership assert so re-runs with the flag stay idempotent. - test: pin the new linked-dir build short-circuit. * Harden --with-llama-cpp-dir against Codex/Gemini review findings - install.sh: error when --with-llama-cpp-dir is the final arg with no path, matching the existing --package/--python post-loop guards (was a silent fallback to the normal prebuilt/source install). - studio/setup.sh: canonicalize LLAMA_CPP_DIR before the self-link no-op compare. _RESOLVED_LOCAL is fully resolved while LLAMA_CPP_DIR was textual, so a symlinked $HOME made the guard miss and the rm -rf could wipe the user's real llama.cpp tree. - studio/setup.sh: make the llama-quantize shim non-fatal; it writes through the link into the user's tree, which may be read-only (shared/CI cache), and under set -e a failed ln aborted an otherwise-good reuse. - studio/setup.ps1: detect a broken junction via Get-Item -Force instead of Test-Path so a dangling link from a prior run is removed and mklink can relink to a new valid directory. - studio/setup.ps1: use Copy-Item -LiteralPath so a source path containing [ ] isn't treated as a wildcard in the junction copy fallback. - tests: update the wiring assertions for the LiteralPath copy and the canonicalized compare. * Validate/reuse local llama.cpp tree and guard the in-use case Addresses the second Codex pass on the --with-llama-cpp-dir flag: - Validate the linked tree before disabling installs (setup.sh + setup.ps1): reusing a local dir skips BOTH the prebuilt download and the source build, so the dir must already contain a runnable llama-server (build/bin on Linux/macOS, build\bin\Release\llama-server.exe on Windows). Bail out with a clear message instead of linking an unbuilt/wrong-platform checkout and leaving Studio with no usable binary. - Treat a canonical-path target as already linked when it holds a build (setup.sh + setup.ps1): point the flag at ~/.unsloth/llama.cpp itself and an existing build is reused (skip prebuilt + source) rather than clobbered by the staged prebuilt installer (which uses os.replace()/replace). An empty canonical dir still falls through to the normal in-place install. - Abort when an in-use llama.cpp can't be removed on Windows (setup.ps1): Remove-Item -ErrorAction SilentlyContinue can silently leave a locked tree in place; detect that and stop with the same active-process message + exit 3 the prebuilt path uses, instead of junctioning over a half-present dir. Left as follow-up (already tracked by the PR author as a non-blocker): the in-app "Update llama.cpp" updater does not yet recognize a local-link install as externally managed; that fix belongs in studio/backend/utils/llama_cpp_update.py. * Accept all backend llama-server layouts in --with-llama-cpp-dir validation The linked-tree validation only accepted build/bin[/Release]/llama-server, but LlamaCppBackend._layout_candidates() resolves a root-level llama-server first, then build/bin, then build/bin/Release on Windows. A `make` build or a flat release extract (binary at the dir root) was therefore rejected with a hard installer failure even though Studio would have run it. Validate the same candidate set the backend uses in both setup scripts, and add wiring-test assertions so the check can't silently narrow again. * Treat --with-llama-cpp-dir local links as externally managed A --with-llama-cpp-dir install junctions/symlinks the canonical llama.cpp dir to the user's own checkout, but two backend paths still treated it as a Studio-owned tree: - The in-app updater (llama_cpp_update) offered and could apply an official prebuilt over the link, writing through it into the user's checkout (or failing) and silently dropping the link the flag created. - Orphan cleanup (LlamaCppBackend._kill_orphaned_servers) resolved the linked root into its kill allowlist, so a llama-server the user launched from the same checkout was classified as ours and killed on startup. Detect the canonical dir being a symlink/junction (reparse point) and treat the install as unmanaged: get_update_status reports unsupported, start_update refuses with reason "local_link", and the linked root is left out of the orphan allowlist. Adds behavioral tests (link vs plain dir, updater refusal, and the spared-vs-killed orphan control). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add behavioral shell test for --with-llama-cpp-dir linking The existing tests/sh/test_with_llama_cpp_dir_flag.sh is a static grep of the scripts. This adds a behavioral test that extracts the real link block from studio/setup.sh (by content anchors, with a self-validating extraction) and runs it against hermetic fake dirs, asserting the outcomes that matter: - an external CMake build links and arms neither the prebuilt download nor the source build - a flat / make tree (root-level llama-server, no build/bin) is accepted too - an unbuilt tree is rejected with a non-zero exit and no link left behind - relinking over a stale link preserves the target's contents (no data loss) - pointing at the canonical path is a no-op reuse, not a self-referential link Symlink-identity checks run only where real symlinks exist (skipped on Windows git-bash copy-mode); the link/skip/no-data-loss checks run everywhere. Wired into studio-backend-ci.yml next to the static test. * Install psutil in backend CI so orphan-cleanup tests run The new orphan-cleanup tests import psutil for the process scan, but the Backend CI deps step installed studio.txt plus a fixed extras list that omits it, so the two tests failed with ModuleNotFoundError. Add psutil to both backend pytest dep steps (kept in shared shape), and guard the import with pytest.importorskip so a minimal env without psutil skips these tests instead of erroring. --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/studio-backend-ci.yml | 11 +- install.ps1 | 17 ++ install.sh | 24 +++ studio/backend/core/inference/llama_cpp.py | 26 +++ .../tests/test_local_llama_cpp_link.py | 137 ++++++++++++++ studio/backend/utils/llama_cpp_update.py | 75 ++++++++ studio/setup.ps1 | 92 +++++++++- studio/setup.sh | 85 ++++++++- tests/sh/test_with_llama_cpp_dir_flag.sh | 172 ++++++++++++++++++ .../test_with_llama_cpp_dir_link_behavior.sh | 132 ++++++++++++++ 10 files changed, 763 insertions(+), 8 deletions(-) create mode 100644 studio/backend/tests/test_local_llama_cpp_link.py create mode 100644 tests/sh/test_with_llama_cpp_dir_flag.sh create mode 100644 tests/sh/test_with_llama_cpp_dir_link_behavior.sh diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index bce355458a..3022127a2b 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -68,9 +68,10 @@ jobs: pip install -r studio/backend/requirements/studio.txt # Extras that studio.txt does not list but the import chain needs # (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography - # for the auth DB, yaml/jinja2 for utils.models.model_config, etc.): + # for the auth DB, yaml/jinja2 for utils.models.model_config, psutil for + # the orphan-cleanup process scan, etc.): pip install \ - python-multipart aiofiles sqlalchemy cryptography \ + python-multipart aiofiles sqlalchemy cryptography psutil \ pyyaml jinja2 mammoth unpdf requests \ 'numpy<3' pytest pytest-asyncio httpx # Torch CPU + transformers are required by a chunk of the backend test @@ -133,7 +134,7 @@ jobs: python -m pip install --upgrade pip pip install -r studio/backend/requirements/studio.txt pip install \ - python-multipart aiofiles sqlalchemy cryptography \ + python-multipart aiofiles sqlalchemy cryptography psutil \ pyyaml jinja2 mammoth unpdf requests typer \ 'numpy<3' pytest pytest-asyncio httpx # torchvision: unsloth_zoo.vision_utils imports it at module scope. @@ -229,7 +230,9 @@ jobs: tests/sh/test_resolve_cuda_archs.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ - tests/sh/test_torch_flavor.sh; do + tests/sh/test_torch_flavor.sh \ + tests/sh/test_with_llama_cpp_dir_flag.sh \ + tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do echo "::group::$s" bash "$s" echo "::endgroup::" diff --git a/install.ps1 b/install.ps1 index f7f9540970..8c667df079 100644 --- a/install.ps1 +++ b/install.ps1 @@ -99,6 +99,7 @@ function Install-UnslothStudio { $TauriMode = $false $SkipTorch = $false $ShortcutsOnly = $false + $WithLlamaCppDir = "" $argList = $args for ($i = 0; $i -lt $argList.Count; $i++) { switch ($argList[$i]) { @@ -116,6 +117,14 @@ function Install-UnslothStudio { } $PackageName = $argList[$i] } + "--with-llama-cpp-dir" { + $i++ + if ($i -ge $argList.Count) { + Write-Host "[ERROR] --with-llama-cpp-dir requires a path argument." -ForegroundColor Red + return (Exit-InstallFailure "--with-llama-cpp-dir requires a path argument.") + } + $WithLlamaCppDir = $argList[$i] + } } } @@ -2430,6 +2439,13 @@ exit 0 } $studioArgs = @('studio', 'setup') if ($script:UnslothVerbose) { $studioArgs += '--verbose' } + if ($WithLlamaCppDir) { + if (-not (Test-Path -LiteralPath $WithLlamaCppDir -PathType Container)) { + Write-Host "[ERROR] --with-llama-cpp-dir path does not exist: $WithLlamaCppDir" -ForegroundColor Red + return (Exit-InstallFailure "--with-llama-cpp-dir path does not exist.") + } + $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR = (Resolve-Path -LiteralPath $WithLlamaCppDir).Path + } $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" # Hand the venv interpreter to setup.ps1 so it reuses the Python we already # resolved and built the venv with, instead of re-probing the system (which @@ -2445,6 +2461,7 @@ exit 0 } else { Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue } + Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue } diff --git a/install.sh b/install.sh index 7a9f0be87f..0370559540 100755 --- a/install.sh +++ b/install.sh @@ -53,6 +53,11 @@ _VERBOSE=false _SHORTCUTS_ONLY=false _next_is_package=false _next_is_python=false +_next_is_llama_cpp_dir=false +# Seed from the environment so a caller who exports UNSLOTH_LOCAL_LLAMA_CPP_DIR +# (the documented piped-install style) is honored; the --with-llama-cpp-dir +# flag below overrides it when given. +_WITH_LLAMA_CPP_DIR="${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" for arg in "$@"; do if [ "$_next_is_package" = true ]; then PACKAGE_NAME="$arg" @@ -64,6 +69,11 @@ for arg in "$@"; do _next_is_python=false continue fi + if [ "$_next_is_llama_cpp_dir" = true ]; then + _WITH_LLAMA_CPP_DIR="$arg" + _next_is_llama_cpp_dir=false + continue + fi case "$arg" in --local) STUDIO_LOCAL_INSTALL=true ;; --package) _next_is_package=true ;; @@ -72,6 +82,7 @@ for arg in "$@"; do --no-torch) _NO_TORCH_FLAG=true ;; --verbose|-v) _VERBOSE=true ;; --shortcuts-only) _SHORTCUTS_ONLY=true ;; + --with-llama-cpp-dir) _next_is_llama_cpp_dir=true ;; esac done @@ -255,6 +266,10 @@ if [ "$_next_is_python" = true ]; then echo "❌ ERROR: --python requires a version argument (e.g. --python 3.12)." >&2 exit 1 fi +if [ "$_next_is_llama_cpp_dir" = true ]; then + echo "❌ ERROR: --with-llama-cpp-dir requires a path argument." >&2 + exit 1 +fi # Validate --package to prevent injection into shell/Python commands. # Must start with a letter/digit (rejects leading dashes that uv would parse as flags). @@ -3023,6 +3038,13 @@ _run_setup_with_studio_home() { "$@" fi } +if [ -n "$_WITH_LLAMA_CPP_DIR" ]; then + if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then + echo "[ERROR] --with-llama-cpp-dir path does not exist: $_WITH_LLAMA_CPP_DIR" >&2 + exit 1 + fi + _WITH_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_WITH_LLAMA_CPP_DIR" && pwd -P)" +fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then _run_setup_with_studio_home env \ SKIP_STUDIO_BASE="$_SKIP_BASE" \ @@ -3031,6 +3053,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then STUDIO_LOCAL_INSTALL=1 \ STUDIO_LOCAL_REPO="$_REPO_ROOT" \ UNSLOTH_NO_TORCH="$SKIP_TORCH" \ + UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \ bash "$SETUP_SH" bool: + """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink + or a Windows directory junction / reparse point. Such a link resolves into + the user's own llama.cpp checkout, which Studio does not own.""" + try: + if os.path.islink(path): + return True + except OSError: + return False + if os.name == "nt": + try: + import stat + attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] + return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) + except (OSError, AttributeError): + return False + return False + + class LlamaCppBackend: """Manages a llama-server subprocess for GGUF model inference. @@ -7166,6 +7185,13 @@ class LlamaCppBackend: resolved_roots: list[Path] = [] for root in install_roots: try: + # A --with-llama-cpp-dir local link (symlink/junction) + # resolves into the user's own checkout. Adding it would let + # us treat the user's externally-launched llama-server as our + # orphan and kill it, so leave such roots out of the + # allowlist (we forgo orphan-reaping for local-link installs). + if _is_external_link(root): + continue resolved_roots.append(root.resolve()) except OSError: pass diff --git a/studio/backend/tests/test_local_llama_cpp_link.py b/studio/backend/tests/test_local_llama_cpp_link.py new file mode 100644 index 0000000000..c78c029d91 --- /dev/null +++ b/studio/backend/tests/test_local_llama_cpp_link.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Behavioral tests for the --with-llama-cpp-dir 'unmanaged local link' contract. + +When the canonical llama.cpp dir is a symlink (POSIX) / junction (Windows) to a +user's own checkout, Studio must treat it as externally managed: + - the in-app updater must not offer or apply a prebuilt over the link + - orphan cleanup must not kill a llama-server the user launched from that tree + +These exercise real link behavior rather than grepping the scripts. +""" + +import os +import subprocess +from pathlib import Path + +import pytest + +from utils import llama_cpp_update as u +from core.inference.llama_cpp import LlamaCppBackend + + +def _make_link(link: Path, target: Path) -> None: + """Create a directory junction (Windows) / symlink (POSIX); neither needs + elevation.""" + target.mkdir(parents = True, exist_ok = True) + if os.name == "nt": + subprocess.run( + ["cmd", "/c", "mklink", "/J", str(link), str(target)], + check = True, + capture_output = True, + text = True, + ) + else: + link.symlink_to(target, target_is_directory = True) + + +def _server_subpath() -> Path: + return Path( + "build/bin/Release/llama-server.exe" if os.name == "nt" else "build/bin/llama-server" + ) + + +class _FakeProc: + def __init__(self, pid: int, exe: str) -> None: + self.info = {"pid": pid, "name": "llama-server", "exe": exe} + self.killed = False + + def kill(self) -> None: + self.killed = True + + +def test_is_external_link_detects_link_vs_plain_dir(tmp_path: Path) -> None: + plain = tmp_path / "plain" + plain.mkdir() + assert u._is_external_link(plain) is False + + link = tmp_path / "link" + _make_link(link, tmp_path / "tgt") + assert u._is_external_link(link) is True + + +def test_active_install_is_local_link(tmp_path: Path) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + binary = str(link / _server_subpath()) + assert u._active_install_is_local_link(binary) is True + + # A plain (non-link) llama.cpp dir is Studio-managed, not a local link. + plain = tmp_path / "plain" / "llama.cpp" + plain.mkdir(parents = True) + assert u._active_install_is_local_link(str(plain / _server_subpath())) is False + + +def test_get_update_status_reports_local_link(tmp_path: Path, monkeypatch) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath())) + st = u.get_update_status() + assert st["supported"] is False + assert st["update_available"] is False + assert st["local_link"] is True + + +def test_start_update_refuses_local_link(tmp_path: Path, monkeypatch) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath())) + res = u.start_update() + assert res["started"] is False + assert res["reason"] == "local_link" + + +def _run_orphan_scan(monkeypatch, studio_root: Path, fake: _FakeProc) -> int: + # psutil drives the cross-platform process scan; skip (rather than error) if a + # minimal test env lacks it. CI installs it so these tests actually run. + psutil = pytest.importorskip("psutil") + + monkeypatch.setattr( + LlamaCppBackend, + "_resolved_studio_root_and_is_legacy", + staticmethod(lambda: (studio_root.resolve(), False)), + ) + monkeypatch.setattr(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)) + monkeypatch.setattr(psutil, "process_iter", lambda attrs = None: iter([fake])) + return LlamaCppBackend._kill_orphaned_servers() + + +def test_orphan_cleanup_spares_local_link_tree(tmp_path: Path, monkeypatch) -> None: + studio_root = tmp_path / "studio-home" + studio_root.mkdir() + external = tmp_path / "external" + (external / _server_subpath().parent).mkdir(parents = True) + (external / _server_subpath()).write_text("x") + _make_link(studio_root / "llama.cpp", external) + + exe_under_link = str((external / _server_subpath()).resolve()) + fake = _FakeProc(os.getpid() + 777, exe_under_link) + killed = _run_orphan_scan(monkeypatch, studio_root, fake) + assert killed == 0 + assert fake.killed is False + + +def test_orphan_cleanup_kills_under_real_root(tmp_path: Path, monkeypatch) -> None: + # Control: a real (non-link) managed root still gets its orphan reaped, so + # the spare-the-link test above is meaningful (not a no-op). + studio_root = tmp_path / "studio-home" + bin_dir = studio_root / "llama.cpp" / _server_subpath().parent + bin_dir.mkdir(parents = True) + exe = studio_root / "llama.cpp" / _server_subpath() + exe.write_text("x") + + fake = _FakeProc(os.getpid() + 888, str(exe.resolve())) + killed = _run_orphan_scan(monkeypatch, studio_root, fake) + assert killed == 1 + assert fake.killed is True diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 8648b053d5..c16ae91467 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -324,12 +324,74 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: } +def _is_external_link(path: Optional[Path]) -> bool: + """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink + or a Windows directory junction / reparse point. Such a link resolves into + the user's own llama.cpp checkout, so Studio must never auto-update it.""" + if path is None: + return False + try: + if os.path.islink(path): + return True + except OSError: + return False + if os.name == "nt": + try: + import stat + attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] + return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) + except (OSError, AttributeError): + return False + return False + + +def _active_install_is_local_link(binary: Optional[str]) -> bool: + """True when the active llama-server resolves through a --with-llama-cpp-dir + local link at the canonical llama.cpp directory. An update would write + through that link into the user's own checkout (or fail), so the install is + treated as externally managed: no update is offered or applied. Checks only + up to and including the ``llama.cpp`` dir so a symlinked HOME / studio root + above it can't trip a false positive.""" + if not binary: + return False + for parent in Path(binary).parents: + if _is_external_link(parent): + return True + if parent.name == "llama.cpp": + break + return False + + +def _local_link_status() -> dict: + """Status payload for a local-link install: unmanaged, no update offered.""" + with _job_lock: + job = dict(_job) + return { + "supported": False, + "update_available": False, + "stale": False, + "installed_tag": None, + "latest_tag": None, + "published_repo": None, + "installed_at_utc": None, + "age_days": None, + "source_build": False, + "local_link": True, + "update_size_bytes": None, + "job": job, + } + + def get_update_status(*, force_refresh: bool = False) -> dict: """Report whether a newer prebuilt exists plus the current job state. force_refresh bypasses the 24h release cache for an explicit "check now". """ binary = _find_binary() + # A --with-llama-cpp-dir local link is the user's own tree; never offer to + # replace it. Bail before any network/freshness work. + if _active_install_is_local_link(binary): + return _local_link_status() marker = read_install_marker(binary) with _job_lock: @@ -537,6 +599,19 @@ def start_update() -> dict: """Kick off a background update. Idempotent: a second call while one is running returns the in-flight job rather than starting another.""" binary = _find_binary() + # Refuse to update a --with-llama-cpp-dir local link: installing a prebuilt + # here would write through the link into the user's own checkout (or fail) + # and silently drop the link the flag created. + if _active_install_is_local_link(binary): + return { + "started": False, + "reason": "local_link", + "message": ( + "llama.cpp is a local directory linked with --with-llama-cpp-dir; " + "Studio won't replace it. Update your own llama.cpp checkout instead." + ), + "job": get_update_status()["job"], + } marker = read_install_marker(binary) script = _installer_script() if script is None: diff --git a/studio/setup.ps1 b/studio/setup.ps1 index ae4e8464ec..bb1e88cc4e 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -3180,7 +3180,86 @@ if ($LlamaPr) { $SkipPrebuiltInstall = $true } -if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { +$LocalLlamaCppLinked = $false +$LocalLlamaCppSrc = $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR +if ($LocalLlamaCppSrc) { + if (-not (Test-Path -LiteralPath $LocalLlamaCppSrc -PathType Container)) { + step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $LocalLlamaCppSrc" "Red" + exit 1 + } + $ResolvedLocal = (Resolve-Path -LiteralPath $LocalLlamaCppSrc).Path + # Reusing a local dir disables both the prebuilt download and the source + # build, so a runnable llama-server.exe must already be present. Accept any + # layout LlamaCppBackend._layout_candidates() resolves (root-level, build\bin, + # or build\bin\Release) so the flag never rejects a tree Studio could run. + $LocalLlamaServerFound = $false + foreach ($_cand in @( + (Join-Path $ResolvedLocal "llama-server.exe"), + (Join-Path $ResolvedLocal "build\bin\llama-server.exe"), + (Join-Path $ResolvedLocal "build\bin\Release\llama-server.exe"))) { + if (Test-Path -LiteralPath $_cand) { $LocalLlamaServerFound = $true; break } + } + if ($ResolvedLocal -eq $LlamaCppDir) { + # Points at the canonical install location itself: never delete-then-link + # onto itself. Reuse an existing build here (skip prebuilt + source) so the + # staged prebuilt installer can't replace a build the user asked to reuse; + # if nothing is built yet, fall through to the normal install. + if ($LocalLlamaServerFound) { + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" "Yellow" + $LocalLlamaCppLinked = $true + $NeedLlamaSourceBuild = $false + } else { + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" "Yellow" + } + } else { + # Fail clearly rather than junction an unbuilt or wrong-platform checkout + # and leave Studio with no usable binary. + if (-not $LocalLlamaServerFound) { + step "llama.cpp" "no llama-server.exe under $ResolvedLocal (looked for .\llama-server.exe, .\build\bin and .\build\bin\Release) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "Red" + exit 1 + } + # If the target is already a junction/symlink (e.g. a previous + # --with-llama-cpp-dir run), delete only the link via DirectoryInfo.Delete(). + # Remove-Item -Recurse -Force on a reparse point can traverse the link and + # wipe the user's real llama.cpp directory on PowerShell 5.1. Dropping the + # stale link here also keeps the custom-home ownership check below idempotent. + # Use Get-Item -Force (not Test-Path): a *broken* junction whose target was + # moved/deleted makes Test-Path return false, which would leave the dangling + # link in place and make mklink below fail; Get-Item still resolves it so we + # can remove it and relink to a new valid directory. + $existing = Get-Item -LiteralPath $LlamaCppDir -Force -ErrorAction SilentlyContinue + if ($existing -and ($existing.Attributes -band [System.IO.FileAttributes]::ReparsePoint)) { + $existing.Delete() + } + if ($StudioHomeIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + } + if (Test-Path -LiteralPath $LlamaCppDir) { + Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue + # A locked/in-use tree can silently survive removal (SilentlyContinue + # masks it). Don't then junction/copy over a half-present dir; mirror the + # prebuilt path's active-process handling and stop with a clear message. + if (Test-Path -LiteralPath $LlamaCppDir) { + step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" + substep "Close Studio or other llama.cpp users and retry" "Yellow" + exit 3 + } + } + cmd /c "mklink /J `"$LlamaCppDir`" `"$ResolvedLocal`"" 2>&1 | Out-Null + if ($LASTEXITCODE -ne 0) { + substep "Could not create directory junction; copying instead..." "Yellow" + Copy-Item -Recurse -LiteralPath $ResolvedLocal -Destination $LlamaCppDir + } + Write-Host "" + step "llama.cpp" "linked local directory: $ResolvedLocal" + $LocalLlamaCppLinked = $true + $NeedLlamaSourceBuild = $false + } +} + +if ($LocalLlamaCppLinked) { + # local directory linked above; skip prebuilt install +} elseif ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { Write-Host "" substep "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" "Yellow" $NeedLlamaSourceBuild = $true @@ -3390,7 +3469,8 @@ if (Test-Path -LiteralPath $LlamaServerBin) { # Install build tools now (last resort) rather than eagerly in Phase 1, so the # prebuilt path stays fast. Same condition as the if/elseif chain below: a source -# build runs only when needed and no usable binary is already present. +# build runs only when needed and no usable binary is already present. A linked +# local dir sets $NeedLlamaSourceBuild = $false, so this no-ops for that path. $WillBuildLlamaFromSource = $NeedLlamaSourceBuild -and ` -not ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") if ($WillBuildLlamaFromSource) { @@ -3399,7 +3479,13 @@ if ($WillBuildLlamaFromSource) { $HasCmakeForBuild = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) } -if (-not $NeedLlamaSourceBuild) { +if ($LocalLlamaCppLinked) { + # Local dir linked above -- honor the flag's contract: skip BOTH the prebuilt + # download and the source build. Falling through here would run CMake inside + # the user's checkout (via the junction) when it lacks build\bin\Release\llama-server.exe. + Write-Host "" + step "llama.cpp" "linked (skipping build)" +} elseif (-not $NeedLlamaSourceBuild) { Write-Host "" step "llama.cpp" "prebuilt (validated)" } elseif ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { diff --git a/studio/setup.sh b/studio/setup.sh index 22a922355d..6a74cd2296 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1271,7 +1271,90 @@ fi verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)" -if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then +# GGUF export's check_llama_cpp() looks for a llama-quantize shim at the root of +# the install dir, but a source build keeps the binary under build/bin/. Mirror +# the source-build-reuse step and create the shim when the reused tree has one +# but no root shim yet. Best-effort: the tree may be read-only (shared/CI cache), +# and under `set -e` a failed ln would otherwise abort an good reuse. +_link_local_llama_quantize_shim() { + if [ -x "$1/build/bin/llama-quantize" ] && [ ! -e "$1/llama-quantize" ]; then + ln -sf build/bin/llama-quantize "$1/llama-quantize" 2>/dev/null || \ + substep "could not create llama-quantize shim in linked dir (read-only?); GGUF export may be unavailable" + fi +} + +# Accept any layout LlamaCppBackend._layout_candidates() resolves so the flag +# never rejects a tree Studio could actually run: a root-level llama-server (a +# `make` build or a flat-extracted release) or the CMake build/bin/llama-server. +_has_local_llama_server() { + [ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ] +} + +_LOCAL_LLAMA_CPP_LINKED=false +if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" ]; then + if [ ! -d "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" ]; then + step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $UNSLOTH_LOCAL_LLAMA_CPP_DIR" "$C_ERR" + exit 1 + fi + _RESOLVED_LOCAL="$(CDPATH= cd -P -- "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" && pwd -P)" + # Canonicalize the install path the same way before comparing: _RESOLVED_LOCAL + # is fully resolved, but LLAMA_CPP_DIR is textual ($UNSLOTH_HOME/llama.cpp). If + # $HOME (or UNSLOTH_HOME) contains a symlink, the two never match even when the + # user pointed the flag at the canonical install itself -- and the rm -rf below + # would then wipe the very tree they asked to reuse. Resolve via the parent so + # this works whether or not the leaf currently exists. + _CANON_LLAMA_CPP_DIR="$LLAMA_CPP_DIR" + _LLAMA_CPP_PARENT="$(dirname "$LLAMA_CPP_DIR")" + if [ -d "$_LLAMA_CPP_PARENT" ]; then + _CANON_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_LLAMA_CPP_PARENT" && pwd -P)/$(basename "$LLAMA_CPP_DIR")" + fi + if [ "$_RESOLVED_LOCAL" = "$_CANON_LLAMA_CPP_DIR" ]; then + # Points at the canonical install location itself: never delete-then-link + # it onto itself. If a usable build is already there, reuse it and skip + # both the prebuilt download and the source build -- the prebuilt installer + # uses os.replace() and would otherwise clobber an existing source build at + # this path. If nothing is built there yet, fall through to the normal + # install so it gets built in place exactly as it would without the flag. + if _has_local_llama_server "$LLAMA_CPP_DIR"; then + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" + _link_local_llama_quantize_shim "$LLAMA_CPP_DIR" + _LOCAL_LLAMA_CPP_LINKED=true + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=true + else + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" + fi + else + # Reusing disables BOTH the prebuilt download and the source build, so the + # linked tree must already contain a runnable llama-server in one of the + # layouts the backend resolves (root-level or build/bin/). Fail clearly + # rather than link an unbuilt or wrong-platform checkout and leave Studio + # with no usable binary. + if ! _has_local_llama_server "$_RESOLVED_LOCAL"; then + step "llama.cpp" "no llama-server under $_RESOLVED_LOCAL (looked for ./llama-server and ./build/bin/llama-server) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "$C_ERR" + exit 1 + fi + # A stale link from a previous --with-llama-cpp-dir run isn't Studio-owned + # content; drop it before the ownership check so re-runs stay idempotent + # for a custom UNSLOTH_STUDIO_HOME (the assert would otherwise follow the + # link into the user's dir and reject it as unowned). + [ -L "$LLAMA_CPP_DIR" ] && rm -f "$LLAMA_CPP_DIR" + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + fi + rm -rf "$LLAMA_CPP_DIR" + ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR" + _link_local_llama_quantize_shim "$LLAMA_CPP_DIR" + step "llama.cpp" "linked local directory: $_RESOLVED_LOCAL" + _LOCAL_LLAMA_CPP_LINKED=true + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=true + fi +fi + +if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then + : # local directory linked above; skip prebuilt install +elif [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then step "llama.cpp" "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt" "$C_WARN" _NEED_LLAMA_SOURCE_BUILD=true elif [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then diff --git a/tests/sh/test_with_llama_cpp_dir_flag.sh b/tests/sh/test_with_llama_cpp_dir_flag.sh new file mode 100644 index 0000000000..cee158bf40 --- /dev/null +++ b/tests/sh/test_with_llama_cpp_dir_flag.sh @@ -0,0 +1,172 @@ +#!/bin/bash +# Static analysis: the --with-llama-cpp-dir flag must be wired consistently +# across both installers (install.sh / install.ps1) and both setup scripts +# (studio/setup.sh / studio/setup.ps1). +# +# The flag lets a user point the installer at a local llama.cpp directory so it +# skips BOTH the prebuilt download (Phase 3) and the source build (Phase 4), +# linking the local dir into the canonical install location instead. The path +# crosses the installer->setup boundary via the UNSLOTH_LOCAL_LLAMA_CPP_DIR env +# var. These checks pin that contract so a future refactor of either side can't +# silently break it (e.g. installer parses the flag but setup never reads the +# env var, or setup links the dir but still runs the build). +# +# This is a shape/wiring test, not a behavioral one: it greps the committed +# scripts. It needs no Python, no GPU, no network. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +INSTALL_PS1="$SCRIPT_DIR/../../install.ps1" +SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh" +SETUP_PS1="$SCRIPT_DIR/../../studio/setup.ps1" +ENV_VAR="UNSLOTH_LOCAL_LLAMA_CPP_DIR" +PASS=0 +FAIL=0 + +assert_contains() { + _label="$1"; _file="$2"; _needle="$3" + if grep -qF -- "$_needle" "$_file"; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected to find '$_needle' in $(basename "$_file"))" + FAIL=$((FAIL + 1)) + fi +} + +# Count of distinct lines matching a regex, used to assert a guard appears +# in more than one place (e.g. env var forwarded on both setup invocations). +assert_min_count() { + _label="$1"; _file="$2"; _pattern="$3"; _min="$4" + _n=$(grep -cE -- "$_pattern" "$_file" || true) + if [ "$_n" -ge "$_min" ]; then + echo " PASS: $_label (found $_n, need >= $_min)" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (found $_n in $(basename "$_file"), need >= $_min)" + FAIL=$((FAIL + 1)) + fi +} + +echo "" +echo "=== install.sh: parses --with-llama-cpp-dir and forwards the env var ===" + +assert_contains \ + "install.sh: accepts --with-llama-cpp-dir flag" \ + "$INSTALL_SH" "--with-llama-cpp-dir" +assert_contains \ + "install.sh: validates the path exists before forwarding" \ + "$INSTALL_SH" 'if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then' +# The path must be forwarded to setup.sh on BOTH the local and the +# non-local setup invocations, else --local users (the documented path) +# would silently lose the flag. +assert_min_count \ + "install.sh: forwards $ENV_VAR on both setup invocations" \ + "$INSTALL_SH" "$ENV_VAR=\"\\\$_WITH_LLAMA_CPP_DIR\"" 2 + +echo "" +echo "=== install.ps1: parses --with-llama-cpp-dir and forwards the env var ===" + +assert_contains \ + "install.ps1: accepts --with-llama-cpp-dir flag" \ + "$INSTALL_PS1" '"--with-llama-cpp-dir"' +assert_contains \ + "install.ps1: errors when flag is given with no path argument" \ + "$INSTALL_PS1" "--with-llama-cpp-dir requires a path argument" +assert_contains \ + "install.ps1: validates the path exists before forwarding" \ + "$INSTALL_PS1" "--with-llama-cpp-dir path does not exist" +assert_contains \ + "install.ps1: exports $ENV_VAR for setup.ps1" \ + "$INSTALL_PS1" "\$env:$ENV_VAR =" +# The exported env var must be cleaned up so a later setup invocation in the +# same shell session doesn't inherit a stale local-dir link. +assert_contains \ + "install.ps1: clears $ENV_VAR after the setup run" \ + "$INSTALL_PS1" "Remove-Item Env:$ENV_VAR" + +echo "" +echo "=== studio/setup.sh: reads the env var, links, and skips download+build ===" + +assert_contains \ + "setup.sh: reads $ENV_VAR" \ + "$SETUP_SH" "$ENV_VAR" +assert_contains \ + "setup.sh: symlinks the local dir into the canonical install location" \ + "$SETUP_SH" 'ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR"' +assert_contains \ + "setup.sh: disables the source build when the local dir is linked" \ + "$SETUP_SH" "_NEED_LLAMA_SOURCE_BUILD=false" +assert_contains \ + "setup.sh: skips the prebuilt download when the local dir is linked" \ + "$SETUP_SH" "_SKIP_PREBUILT_INSTALL=true" +# The link branch must short-circuit the FORCE_COMPILE / prebuilt chain rather +# than fall through into it. +assert_contains \ + "setup.sh: link branch gates the prebuilt/compile chain" \ + "$SETUP_SH" 'if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then' + +echo "" +echo "=== studio/setup.ps1: reads the env var, junctions, and skips download+build ===" + +assert_contains \ + "setup.ps1: reads $ENV_VAR" \ + "$SETUP_PS1" "\$env:$ENV_VAR" +assert_contains \ + "setup.ps1: creates a directory junction into the canonical location" \ + "$SETUP_PS1" "mklink /J" +assert_contains \ + "setup.ps1: falls back to a copy when the junction can't be created" \ + "$SETUP_PS1" "Copy-Item -Recurse -LiteralPath \$ResolvedLocal -Destination \$LlamaCppDir" +assert_contains \ + "setup.ps1: disables the source build when the local dir is linked" \ + "$SETUP_PS1" '$NeedLlamaSourceBuild = $false' +# The link branch must gate the prebuilt-install chain (the elseif on +# FORCE_COMPILE), and the linked-dir case must short-circuit the build chain +# so neither a prebuilt download nor a source build runs against it. +assert_contains \ + "setup.ps1: link branch gates the prebuilt/compile chain" \ + "$SETUP_PS1" 'if ($LocalLlamaCppLinked) {' +assert_contains \ + "setup.ps1: linked-dir case short-circuits the build chain" \ + "$SETUP_PS1" 'step "llama.cpp" "linked (skipping build)"' + +echo "" +echo "=== both setup scripts: validate against every layout the backend resolves ===" + +# The linked tree is accepted only if it already holds a runnable llama-server, +# but the check must match LlamaCppBackend._layout_candidates() (root-level +# first, then build/bin, then build/bin/Release on Windows). A narrower check +# would reject a make/flat-release tree the backend could run. +assert_contains \ + "setup.sh: accepts root-level or build/bin llama-server layouts" \ + "$SETUP_SH" '[ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ]' +assert_contains \ + "setup.ps1: accepts the build\\bin (non-Release) llama-server.exe layout" \ + "$SETUP_PS1" 'Join-Path $ResolvedLocal "build\bin\llama-server.exe"' +assert_contains \ + "setup.ps1: accepts the root-level llama-server.exe layout" \ + "$SETUP_PS1" 'Join-Path $ResolvedLocal "llama-server.exe"' + +echo "" +echo "=== both setup scripts: a local dir pointing at the canonical path is a no-op ===" + +# Guard against the self-link footgun: if the user passes the canonical install +# dir itself, neither script should delete-then-link it onto itself. +assert_contains \ + "setup.sh: ignores a local dir equal to the canonical install location" \ + "$SETUP_SH" 'if [ "$_RESOLVED_LOCAL" = "$_CANON_LLAMA_CPP_DIR" ]; then' +assert_contains \ + "setup.ps1: ignores a local dir equal to the canonical install location" \ + "$SETUP_PS1" 'if ($ResolvedLocal -eq $LlamaCppDir) {' + +echo "" +echo "=== Results ===" +echo " PASS: $PASS" +echo " FAIL: $FAIL" +if [ "$FAIL" -gt 0 ]; then + echo "FAILED" + exit 1 +fi +echo "ALL PASSED" diff --git a/tests/sh/test_with_llama_cpp_dir_link_behavior.sh b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh new file mode 100644 index 0000000000..fb09e56c51 --- /dev/null +++ b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# Behavioral test for the --with-llama-cpp-dir linking block in studio/setup.sh. +# The companion test_with_llama_cpp_dir_flag.sh is a static wiring check; this one +# actually RUNS the real link logic (extracted from setup.sh by content anchors, +# not line numbers) against hermetic fake dirs and asserts the outcomes Lee asked +# for: an external built dir gets linked, neither the prebuilt download nor the +# source build is armed, an unbuilt dir is rejected, a relink doesn't destroy the +# target, and pointing at the canonical path is a no-op. POSIX symlinks here; +# the Windows junction path is covered by the backend test suite. +set -u +HERE="$(CDPATH= cd -P -- "$(dirname "$0")" && pwd -P)" +SETUP="$HERE/../../studio/setup.sh" +fails=0 +check() { # name expected actual + if [ "$2" = "$3" ]; then printf ' PASS %s\n' "$1" + else printf ' FAIL %s : expected [%s] got [%s]\n' "$1" "$2" "$3"; fails=$((fails+1)); fi +} + +# Extract the two helpers + the whole `UNSLOTH_LOCAL_LLAMA_CPP_DIR` if-block. +# Starts at the quantize-shim helper, ends at the first column-0 `fi` after the +# `if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR..` guard (inner ifs are indented). +block="$(awk ' + /^_link_local_llama_quantize_shim\(\) \{/ {grab=1} + grab {print} + /^if \[ -n "\$\{UNSLOTH_LOCAL_LLAMA_CPP_DIR/ {inif=1} + inif && /^fi$/ {exit} +' "$SETUP")" + +# Self-validate the extraction so a future setup.sh refactor fails loudly here. +case "$block" in *'ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR"'*) : ;; + *) echo "FAIL: link block extraction broke (no ln -sfn)"; exit 1 ;; esac +case "$block" in *'_has_local_llama_server'*) : ;; + *) echo "FAIL: link block extraction broke (no _has_local_llama_server)"; exit 1 ;; esac + +# Stub setup.sh's logging + ownership helpers, seed the vars the block reads, +# then run the extracted block and print the resulting state. +PREAMBLE=' +set -u +step() { :; }; substep() { :; }; verbose_substep() { :; } +_assert_studio_owned_or_absent() { :; } +C_ERR="" +_STUDIO_HOME_IS_CUSTOM=false +_NEED_LLAMA_SOURCE_BUILD=UNSET +_SKIP_PREBUILT_INSTALL=UNSET +' +EPILOGUE=' +echo "LINKED=$_LOCAL_LLAMA_CPP_LINKED" +echo "NEED_BUILD=$_NEED_LLAMA_SOURCE_BUILD" +echo "SKIP_PREBUILT=$_SKIP_PREBUILT_INSTALL" +if [ -L "$LLAMA_CPP_DIR" ]; then echo "ISLINK=1"; echo "TARGET=$(readlink "$LLAMA_CPP_DIR")"; else echo "ISLINK=0"; fi +' +SNIP="$PREAMBLE"$'\n'"$block"$'\n'"$EPILOGUE" + +# run_link -> prints state lines; RC in $RC +run_link() { + OUT="$(env -i PATH="$PATH" HOME="$T" \ + UNSLOTH_LOCAL_LLAMA_CPP_DIR="$1" LLAMA_CPP_DIR="$2" \ + bash -c "$SNIP" 2>/dev/null)" + RC=$? +} +val() { printf '%s\n' "$OUT" | grep "^$1=" | head -1 | cut -d= -f2-; } + +T="$(mktemp -d)" +trap 'rm -rf "$T"' EXIT + +# Some environments (Windows git-bash without native symlinks) make `ln -s` copy +# instead of link. The symlink-identity assertions (ISLINK / readlink target) +# only run where real symlinks exist; the link/skip/no-data-loss assertions run +# everywhere, including CI (Linux), where the link path is the real one. +ln -s "$T" "$T/.symprobe" 2>/dev/null +if [ -L "$T/.symprobe" ]; then SYMLINKS=1; else SYMLINKS=0; fi +rm -rf "$T/.symprobe" + +# A built external tree (CMake layout) + a flat/`make` tree (root-level binary). +# The fake binary is a shebang script so the `-x` test in _has_local_llama_server +# holds on both Linux (chmod +x) and Windows git-bash (MSYS treats #!-files as +# executable), without needing a real platform binary. +mk_exe() { printf '#!/bin/sh\necho fake\n' > "$1"; chmod +x "$1"; } +mk_built() { mkdir -p "$1/build/bin"; mk_exe "$1/build/bin/llama-server"; } +mk_flat() { mkdir -p "$1"; mk_exe "$1/llama-server"; } + +# 1. External CMake build -> linked, and BOTH install paths disarmed. +EXT1="$T/ext_cmake"; mk_built "$EXT1"; : > "$EXT1/keep.txt" +CANON1="$T/home1/llama.cpp"; mkdir -p "$(dirname "$CANON1")" +run_link "$EXT1" "$CANON1" +check "cmake build: linked" "true" "$(val LINKED)" +check "cmake build: source build off" "false" "$(val NEED_BUILD)" +check "cmake build: prebuilt skipped" "true" "$(val SKIP_PREBUILT)" +if [ "$SYMLINKS" = 1 ]; then + check "cmake build: canonical is a symlink" "1" "$(val ISLINK)" + check "cmake build: link points at external" "$(CDPATH= cd -P -- "$EXT1" && pwd -P)" "$(val TARGET)" +else + printf ' SKIP cmake build: symlink-identity (no real symlinks here)\n' +fi + +# 2. Flat / make tree (root-level llama-server, no build/bin) -> still linked +# (the new layout-candidate acceptance; the old check rejected this). +EXT2="$T/ext_flat"; mk_flat "$EXT2" +CANON2="$T/home2/llama.cpp"; mkdir -p "$(dirname "$CANON2")" +run_link "$EXT2" "$CANON2" +check "flat build: linked (root-level llama-server accepted)" "true" "$(val LINKED)" + +# 3. Unbuilt tree -> rejected (non-zero exit, no link created). +EXT3="$T/ext_empty"; mkdir -p "$EXT3" +CANON3="$T/home3/llama.cpp"; mkdir -p "$(dirname "$CANON3")" +run_link "$EXT3" "$CANON3" +check "unbuilt tree: rejected (exit != 0)" "yes" "$([ "$RC" -ne 0 ] && echo yes || echo no)" +check "unbuilt tree: no link left behind" "no" "$([ -L "$CANON3" ] && echo yes || echo no)" + +# 4. Relink over a stale link must NOT destroy the (new) target's contents. +OLD="$T/ext_old"; mk_built "$OLD" +NEW="$T/ext_new"; mk_built "$NEW"; : > "$NEW/precious.txt" +CANON4="$T/home4/llama.cpp"; mkdir -p "$(dirname "$CANON4")" +ln -sfn "$OLD" "$CANON4" # simulate a prior --with-llama-cpp-dir run +run_link "$NEW" "$CANON4" +if [ "$SYMLINKS" = 1 ]; then + check "relink: now points at the new external" "$(CDPATH= cd -P -- "$NEW" && pwd -P)" "$(val TARGET)" +fi +check "relink: new target's contents preserved" "yes" "$([ -f "$NEW/precious.txt" ] && echo yes || echo no)" +check "relink: old target's contents preserved" "yes" "$([ -f "$OLD/build/bin/llama-server" ] && echo yes || echo no)" + +# 5. Pointing at the canonical path itself is a no-op reuse: linked, not turned +# into a self-referential symlink, contents untouched. +CANON5="$T/home5/llama.cpp"; mk_built "$CANON5"; : > "$CANON5/keep.txt" +run_link "$CANON5" "$CANON5" +check "canonical no-op: linked" "true" "$(val LINKED)" +check "canonical no-op: not made a symlink" "0" "$(val ISLINK)" +check "canonical no-op: contents preserved" "yes" "$([ -f "$CANON5/keep.txt" ] && echo yes || echo no)" + +echo "" +if [ "$fails" -ne 0 ]; then echo "$fails check(s) failed"; exit 1; fi +echo "All checks passed" From d91824583452f8d1faf3973a15d3dc4ef5a334ac Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Fri, 3 Jul 2026 06:02:26 +0800 Subject: [PATCH 11/23] Add MLX-aware public Unsloth trainer API (#6462) * feat: add mlx public trainer api * test: cover mlx public trainer api * fix: preserve mlx epoch trainer configs * fix: pass mlx warmup ratio through config * fix: align mlx trainer dataset order * fix: keep mlx chat templates import-light * fix: infer mlx trainer context length * fix: mirror cuda mlx context defaults * fix: align mlx notebook trainer defaults * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: keep mlx public helpers import-light * refactor: reuse mlx optimizer normalization * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: address mlx review feedback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: tighten mlx training argument parity * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: align mlx trainer eos default * Fix MLX trainer to accept DataCollatorForSeq2Seq and handle TokenizerWrapper in get_chat_template * Trim redundant docstrings on internal MLX helpers * MLX review fixes: Studio optimizer import-safe on non-MLX hosts, preserve explicit max_length, skip MLX tests before import * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX review round 2: defer max_length to model context, optimizer alias fallback for older zoo, skip non-MLX test on missing GPU deps * MLX review round 3: keep chat_templates importable without torch on MLX * fix: preserve MLX trainer notebook shims * fix: ignore CUDA tokenizer moves on MLX * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: harden MLX trainer shims * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: unwrap MLX scheduler enum args * fix: coerce integral MLX epoch counts * fix: spoof CUDA compatibility APIs on MLX * fix: harden MLX notebook compatibility shims * MLX: add torch.cuda.mem_get_info to the compatibility shim Notebook memory cells call torch.cuda.mem_get_info()[0] directly (not gated by is_available), so on MLX it raises without a shim. Return (free, total) bytes from the MLX device stats, consistent with the other torch.cuda compat helpers, and add a matching assertion to the compat-API test. * MLX: use active memory for mem_get_info; fix BatchEncoding.to keyword device Address review on the MLX compatibility shim: - torch.cuda.mem_get_info() now derives free bytes from current active MLX memory instead of the peak high-water mark, so a capacity check stays accurate after a transient spike or a prior run. - BatchEncoding.to(device=...) passed by keyword no longer forwards a positional None alongside the keyword (which raised "multiple values for 'device'"), so non-CUDA keyword moves like .to(device="cpu") delegate correctly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX: accept preserve_dataset_order; stub RL trainers with a clear error Two fixes so unmigrated notebooks behave predictably on MLX (torch present): - preserve_dataset_order is a real MLXTrainingConfig field but was missing from the extra-argument allowlist, so passing it (as a config or trainer kwarg) could be rejected as unknown on a zoo without the field. Add it to _MLX_IMPLEMENTED_EXTRA_ARGUMENTS so the documented no-shuffle path is reachable. - GRPO/DPO/ORPO (and KTO/PPO/Reward) have no MLX trainer yet. Retarget the ones the installed trl exposes to a stub that raises a clear 'not supported on MLX' error instead of importing the real torch/CUDA trainer and crashing deep inside it. Only existing trainers are retargeted (no invented attributes), idempotent across re-imports. * MLX: make RL-trainer stubbing import-safe; back current-memory APIs with active memory Address review on the MLX shims: - The RL-trainer stub loop probed trl with getattr(_trl, name), which triggers trl's lazy trainer import and pulls torch -- that can crash import unsloth on a torch-free MLX install just to check existence. Decide what to stub from trl.__all__ + already-materialized attrs (vars) instead; never resolve the real trainer. All trl trainer names are in __all__, so they are still stubbed (even torch-free), and the probe no longer imports torch. - torch.cuda.memory_reserved / memory_allocated (the current, non-max APIs) were aliased to peak max_memory_reserved. Back them with current active MLX memory so cleanup / capacity checks see live usage; max_* keep the peak high-water mark. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MLX: keep TRL's SFTConfig epoch default under the trl.SFTConfig alias Unmigrated notebooks import SFTConfig from trl, which the MLX build aliases to the public training-args class. TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the native MLX config defaults to max_steps=60. So an SFTConfig built without an explicit length silently ran 60 MLX steps instead of TRL's 3 epochs under the alias. Alias trl.SFTConfig to a thin subclass that seeds the TRL epoch default only when neither max_steps nor num_train_epochs is given; explicit lengths pass through untouched, and the native public args class keeps its MLX default. Epoch mode is supported by the MLX trainer. * MLX CI: keep the GGUF reload smoke under the job timeout The RELOAD-GGUF-via-llama-cli step timed out at 300s. BF16 GGUF decode is CPU-bound on the macOS runner (~10s+/token), so generating 24 tokens landed right on the 300s cliff and killed the process. This step is a save/reload integrity smoke (it only needs a few chars of output), so the token count is incidental: generate 8 tokens with explicit threads and a small headroom on the subprocess timeout, all env-tunable (UNSLOTH_GGUF_RELOAD_N / _THREADS / _TIMEOUT). Cuts the reload well under the 25 minute job budget. * MLX: broaden trainer stubs, real peak-memory reset, fix shim tests Address review on the MLX public API: - The SFTConfig identity tests asserted trl.SFTConfig is UnslothTrainingArguments, but the alias now points at the _MLXSFTConfig subclass that preserves TRL's epoch default, so the MLX suite failed before testing the shim. Assert issubclass instead. - torch.cuda.reset_peak_memory_stats was a no-op, so max_memory_reserved kept earlier model-load peaks across a scoped run. Wire it to mx.reset_peak_memory with the same core/metal fallback used for the reads. - The unsupported-trainer stubs were a fixed list, so trainers outside it (a newer RLOOTrainer) still routed to the real torch trainer. Derive the set from trl.__all__ (every non-SFT *Trainer) so all non-SFT surfaces fail with a clear MLX message; names come from __all__ so trl is never resolved. - The non-MLX export smoke skipped only on missing bitsandbytes/triton; other absent GPU deps (numpy/torch/unsloth-zoo, or _gpu_init re-raising ImportError) made it fail on CPU hosts. Skip on any ImportError. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: keep MLX notebook compatibility minimal * MLX CI: force CPU + small context for the GGUF reload smoke The RELOAD-GGUF-via-llama-cli step timed out even at 8 tokens (>420s), so it is a fixed hang, not per-token cost: on the paravirtual macOS runner GPU llama.cpp's Metal backend stalls, and the gemma3 GGUF advertises a 32768 context that llama-cli would otherwise fully allocate. Run llama-cli CPU-only (-ngl 0) with a small context (-c 256); keep generation short. All env-tunable (UNSLOTH_GGUF_RELOAD_NGL / _CTX / _N / _THREADS / _TIMEOUT). Also print llama.cpp's partial stdout/stderr on timeout so a future hang is diagnosable instead of an opaque TimeoutExpired. * MLX CI: export the reload-smoke GGUF as q8_0, not bf16 The GGUF reload via llama-cli timed out on the runner even CPU-only with a tiny context and 8 tokens. Root cause is the format, not the flags: the smoke exported quantization_method='not_quantized', which maps to a bf16 GGUF, and llama.cpp's bf16 CPU decode is unusably slow on the paravirtual macOS runner. Export q8_0 (fast_quantized, the exporter default and what users deploy) instead -- llama.cpp has optimized q8_0 CPU kernels, so the fresh-process reload loads and generates in seconds. The reload stays CPU-only (-ngl 0) with a small context. * test: clear TRL shim before availability check --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 --- studio/backend/core/training/worker.py | 58 +- .../tests/test_mlx_training_worker_config.py | 2 +- tests/python/test_mlx_public_trainer_api.py | 1206 ++++++++++++++++ tests/studio/run_real_mlx_smoke.py | 54 +- unsloth/__init__.py | 1270 ++++++++++++++++- unsloth/chat_templates.py | 67 +- 6 files changed, 2595 insertions(+), 62 deletions(-) create mode 100644 tests/python/test_mlx_public_trainer_api.py diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 610af2472e..17dc1299ca 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1253,32 +1253,48 @@ def _adapt_for_mlx_vlm( return adapted -_MLX_STUDIO_OPTIM_MAP = { - "adamw_8bit": "adamw", - "paged_adamw_8bit": "adamw", - "adamw_bnb_8bit": "adamw", - "paged_adamw_32bit": "adamw", - "adamw_torch": "adamw", - "adamw_torch_fused": "adamw", - "adamw": "adamw", - "adafactor": "adafactor", - "sgd": "sgd", - "adam": "adam", - "muon": "muon", - "lion": "lion", -} _MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"} +# Fallback alias map mirroring unsloth_zoo._normalize_mlx_optimizer_name, used +# only when mlx (Apple Silicon) is not importable so Studio config validation +# still works on non-MLX hosts. The zoo function stays the source of truth. +_MLX_STUDIO_ADAMW_ALIASES = frozenset( + ( + "adamw_8bit", + "paged_adamw_8bit", + "adamw_bnb_8bit", + "paged_adamw_32bit", + "adamw_torch", + "adamw_torch_fused", + "paged_adamw", + "adamw_32bit", + "adamw_hf", + "adamw_anyprecision", + "adamw_apex_fused", + ) +) +_MLX_STUDIO_NATIVE_OPTIMIZERS = ("adafactor", "adamw", "adam", "sgd", "muon", "lion") + + def _normalize_mlx_studio_optimizer(value): - raw = str(value or "adamw_8bit").strip().lower() try: - return _MLX_STUDIO_OPTIM_MAP[raw] - except KeyError: - supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP)) - raise ValueError( - f"Unsupported optimizer for MLX training: {value!r}. " f"Supported values: {supported}." - ) + from unsloth_zoo.mlx.trainer import _normalize_mlx_optimizer_name + return _normalize_mlx_optimizer_name(value or "adamw_8bit") + except (ImportError, ValueError): + # Missing mlx, or an older unsloth-zoo whose normalizer lacks CUDA/TRL + # aliases: map common adamw_* names locally so notebook defaults work. + opt = str(getattr(value, "value", value) or "adamw_8bit").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_STUDIO_ADAMW_ALIASES: + opt = "adamw" + if opt not in _MLX_STUDIO_NATIVE_OPTIMIZERS: + supported = ", ".join(_MLX_STUDIO_NATIVE_OPTIMIZERS) + raise ValueError( + f"Unsupported optimizer for MLX training: {value!r}. " + f"Supported optimizers: {supported}." + ) + return opt def _normalize_mlx_studio_scheduler(value): diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index dce5e27c08..14fc0933d0 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -76,7 +76,7 @@ def test_mlx_studio_optimizer_aliases_are_explicit(): def test_mlx_studio_rejects_unknown_optimizer(): - with pytest.raises(ValueError, match = "Unsupported optimizer for MLX training"): + with pytest.raises(ValueError, match = "Supported"): _normalize_mlx_studio_optimizer("adamw_typo") diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py new file mode 100644 index 0000000000..2c33f86af1 --- /dev/null +++ b/tests/python/test_mlx_public_trainer_api.py @@ -0,0 +1,1206 @@ +"""Tests for the MLX public trainer compatibility surface.""" + +from __future__ import annotations + +import builtins +import importlib +import importlib.util +import platform +import sys +import types +import warnings + +import pytest + +_MLX_SKIP_REASON = "MLX public trainer API is only active on the MLX backend" + + +def _import_mlx_unsloth(): + """Import unsloth and skip when the current platform is not using MLX.""" + # Skip before importing unsloth so non-MLX hosts missing optional GPU deps + # (e.g. bitsandbytes) skip cleanly instead of erroring at collection. + if not ( + platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None + ): + pytest.skip(_MLX_SKIP_REASON) + unsloth = importlib.import_module("unsloth") + if getattr(unsloth, "DEVICE_TYPE", None) != "mlx": + pytest.skip(_MLX_SKIP_REASON) + return unsloth + + +class _DummyModel: + """Small model stub that satisfies MLXTrainer constructor probes.""" + + def trainable_parameters(self): + """Return no trainable parameters for constructor-only tests.""" + return {} + + +class _DummyVLMModel(_DummyModel): + """Small VLM model stub for MLX vision trainer constructor probes.""" + + _is_vlm_model = True + + +def test_mlx_exports_unsloth_trainer_api(): + """MLX imports should expose the public Unsloth trainer API.""" + unsloth = _import_mlx_unsloth() + from unsloth import ( + RawTextDataLoader, + TextPreprocessor, + UnslothTrainer, + UnslothTrainingArguments, + clear_gpu_memory, + get_gpu_memory_stats, + ) + + assert RawTextDataLoader is unsloth.RawTextDataLoader + assert TextPreprocessor is unsloth.TextPreprocessor + assert UnslothTrainer is unsloth.UnslothTrainer + assert UnslothTrainingArguments is unsloth.UnslothTrainingArguments + assert get_gpu_memory_stats is unsloth.get_gpu_memory_stats + assert clear_gpu_memory is unsloth.clear_gpu_memory + assert issubclass(UnslothTrainer, unsloth.MLXTrainer) + assert issubclass(UnslothTrainingArguments, unsloth.MLXTrainingConfig) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_non_mlx_exports_public_trainer_api_when_available(): + """GPU/ROCm imports should keep exporting the public Unsloth trainer API.""" + try: + unsloth = importlib.import_module("unsloth") + except ImportError as exc: + # Non-MLX import pulls the optional GPU stack (numpy/torch/unsloth-zoo, + # bitsandbytes/triton, and _gpu_init can re-raise missing deps as + # ImportError). Skip when any of it is unavailable rather than failing + # collection on CPU/ROCm/XPU review hosts. + pytest.skip(f"non-MLX import dependency unavailable: {exc}") + if getattr(unsloth, "DEVICE_TYPE", None) == "mlx": + pytest.skip("non-MLX export smoke test only runs on GPU/ROCm backends") + + assert unsloth.UnslothTrainer is not None + assert unsloth.UnslothTrainingArguments is not None + assert callable(unsloth.get_gpu_memory_stats) + assert callable(unsloth.clear_gpu_memory) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_mlx_training_arguments_accept_trl_style_kwargs(): + """TRL/SFTConfig-style kwargs should normalize without breaking MLX config.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "bf16.*dataset_kwargs"): + args = unsloth.UnslothTrainingArguments( + max_length = 123, + max_steps = 10, + warmup_ratio = 0.2, + remove_unused_columns = False, + dataset_kwargs = {"skip_prepare_dataset": True}, + bf16 = True, + ) + + assert args.max_seq_length == 123 + assert args.warmup_steps == 2 + assert args.remove_unused_columns is False + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.bf16 is True + assert args.warmup_ratio == 0.2 + assert args._unsloth_mlx_max_seq_length_explicit is False + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): + """Implemented and falsey inert compatibility kwargs should stay quiet.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + args = unsloth.UnslothTrainingArguments( + warmup_ratio = 0.2, + max_steps = 10, + padding_free = False, + remove_unused_columns = False, + assistant_only_loss = False, + completion_only_loss = False, + ) + + assert args.warmup_steps == 2 + assert args.padding_free is False + assert args.remove_unused_columns is False + assert args.completion_only_loss is False + assert caught == [] + + +def test_mlx_training_arguments_prefer_canonical_max_seq_length(): + """Canonical MLX config fields should win over compatibility aliases.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_seq_length = 456, max_length = 123) + dict_args = unsloth.UnslothTrainingArguments( + {"max_length": 123, "max_seq_length": 456}, + ) + + assert args.max_seq_length == 456 + assert args.max_length == 456 + assert args._unsloth_mlx_max_length_value == 456 + assert dict_args.max_seq_length == 456 + assert dict_args.max_length == 456 + assert dict_args._unsloth_mlx_max_length_value == 456 + assert args._unsloth_mlx_max_seq_length_explicit is True + assert dict_args._unsloth_mlx_max_seq_length_explicit is True + + +def test_mlx_training_arguments_preserve_explicit_positive_warmup_steps(): + """Explicit warmup_steps should take precedence over warmup_ratio.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments( + max_steps = 10, + warmup_steps = 5, + warmup_ratio = 0.1, + ) + + assert args.warmup_steps == 5 + assert args._unsloth_mlx_warmup_steps_explicit is True + + +def test_mlx_clear_gpu_memory_uses_metal_fallback(monkeypatch): + """Older MLX releases expose cache clearing under mx.metal.clear_cache.""" + unsloth = _import_mlx_unsloth() + import mlx.core as mx + + called = [] + metal = getattr(mx, "metal", None) or type("Metal", (), {})() + monkeypatch.delattr(mx, "clear_cache", raising = False) + monkeypatch.setattr(mx, "metal", metal, raising = False) + monkeypatch.setattr(metal, "clear_cache", lambda: called.append("metal"), raising = False) + + unsloth.clear_gpu_memory() + + assert called == ["metal"] + + +def test_mlx_training_arguments_preserve_explicit_epoch_training(): + """Epoch-based configs should not inherit the MLX max_steps default.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(num_train_epochs = 1, warmup_ratio = 0.1) + default_args = unsloth.UnslothTrainingArguments() + + assert args.num_train_epochs == 1 + assert args.max_steps == -1 + assert args.warmup_ratio == 0.1 + assert args._unsloth_mlx_warmup_steps_explicit is False + assert default_args.max_steps == unsloth.MLXTrainingConfig.max_steps + + +def test_mlx_training_arguments_keep_mlx_dataset_order_default(): + """Training arguments alone should not override MLX's native data order.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_steps = 1) + explicit_default = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ) + + assert args.dataset_order == "default" + assert args._unsloth_mlx_dataset_order_explicit is False + assert args._unsloth_mlx_max_seq_length_explicit is False + assert explicit_default.dataset_order == "default" + assert explicit_default._unsloth_mlx_dataset_order_explicit is True + + +def test_mlx_training_arguments_warn_on_meaningful_inert_kwargs(): + """Unsupported TrainingArguments knobs should not be silently ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "push_to_hub.*save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "steps", + push_to_hub = True, + padding_free = False, + ) + + assert args.save_strategy == "steps" + assert args.push_to_hub is True + assert args.padding_free is False + + +def test_mlx_training_arguments_reject_unknown_kwargs(): + """Unknown SFTConfig flags should fail instead of becoming inert attributes.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth.UnslothTrainingArguments(assistant_only_loss = True) + + completion_args = unsloth.UnslothTrainingArguments(completion_only_loss = True) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_reject_unsupported_object_flags(): + """Object-style SFTConfig flags should not be silently dropped.""" + unsloth = _import_mlx_unsloth() + + class ArgsObject: + max_steps = 1 + assistant_only_loss = True + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth._coerce_mlx_training_args(ArgsObject()) + + class CompletionArgsObject: + max_steps = 1 + completion_only_loss = True + + completion_args = unsloth._coerce_mlx_training_args(CompletionArgsObject()) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_accept_output_dir_positional(): + """A single positional output_dir should match TrainingArguments behavior.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("custom-outputs", max_steps = 3) + + assert args.output_dir == "custom-outputs" + assert args.max_steps == 3 + + +def test_mlx_training_arguments_normalize_optim_and_object_aliases(): + """Common notebook optimizer names and object aliases should normalize.""" + unsloth = _import_mlx_unsloth() + + class Scheduler: + value = "cosine" + + class ArgsObject: + optim = "adamw_8bit" + eval_steps = None + lr_scheduler_type = Scheduler() + max_length = 321 + max_steps = 10 + num_train_epochs = 3.0 + save_steps = 500 + save_strategy = "no" + warmup_ratio = 0.1 + warmup_steps = 0 + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth._coerce_mlx_training_args(ArgsObject()) + + assert args.optim == "adamw" + assert args.eval_steps == 0 + assert args.lr_scheduler_type == "cosine" + assert args.max_seq_length == 321 + assert args.num_train_epochs == 3 + assert type(args.num_train_epochs) is int + assert args.save_steps == 0 + assert args.warmup_steps == 1 + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_accept_supported_notebook_kwargs(): + """Supported SFT notebooks should be able to pass their current args.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns( + RuntimeWarning, + match = "bf16.*dataset_kwargs.*gradient_checkpointing_kwargs.*save_strategy", + ): + args = unsloth.UnslothTrainingArguments( + bf16 = True, + dataset_kwargs = {"skip_prepare_dataset": True}, + dataset_num_proc = 4, + dataset_text_field = "text", + embedding_learning_rate = 5e-5, + fp16 = False, + gradient_accumulation_steps = 8, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = {"use_reentrant": False}, + learning_rate = 1e-4, + logging_steps = 2, + lr_scheduler_type = "cosine", + max_grad_norm = 0.3, + max_length = 1024, + max_steps = 10, + num_train_epochs = 1, + optim = "paged_adamw_8bit", + output_dir = "outputs", + padding_free = False, + per_device_train_batch_size = 1, + remove_unused_columns = False, + report_to = "none", + save_strategy = "steps", + seed = 123, + warmup_ratio = 0.1, + weight_decay = 0.01, + ) + + assert args.dataset_num_proc == 4 + assert args.dataset_text_field == "text" + assert args.embedding_learning_rate == 5e-5 + assert args.gradient_accumulation_steps == 8 + assert args.gradient_checkpointing is True + assert args.learning_rate == 1e-4 + assert args.logging_steps == 2 + assert args.lr_scheduler_type == "cosine" + assert args.max_grad_norm == 0.3 + assert args.max_seq_length == 1024 + assert args.max_steps == 10 + assert args.num_train_epochs == 1 + assert args.optim == "adamw" + assert args.output_dir == "outputs" + assert args.per_device_train_batch_size == 1 + assert args.report_to == "none" + assert args.seed == 123 + assert args.warmup_ratio == 0.1 + assert args.warmup_steps == 1 + assert args.weight_decay == 0.01 + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.gradient_checkpointing_kwargs == {"use_reentrant": False} + assert args.save_strategy == "steps" + + +def test_mlx_training_arguments_honor_direct_no_save_strategy(): + """Direct kwargs should map save_strategy=no to save_steps=0.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "no", + save_steps = 500, + ) + + assert args.save_steps == 0 + + +def test_mlx_trainer_accepts_common_sft_kwargs(): + """UnslothTrainer should accept common SFTTrainer kwargs on MLX.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + dataset_num_proc = 8, + max_length = 456, + optim = "adamw_bnb_8bit", + processing_class = object(), + ) + + assert trainer.args.max_steps == 1 + assert trainer.args.dataset_num_proc == 8 + assert trainer.args.max_seq_length == 456 + assert trainer.args.max_grad_norm == 1.0 + assert trainer.args.optim == "adamw" + assert trainer.args.dataset_order == "torch_randperm" + assert trainer._unsloth_mlx_ignored_trainer_kwargs == {} + assert caught == [] + + +def test_mlx_trainer_preserves_explicit_dataset_order(): + """UnslothTrainer should only set torch_randperm when order is implicit.""" + unsloth = _import_mlx_unsloth() + + explicit_default = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ), + ) + explicit_sequential = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "sequential", + ), + ) + implicit_with_override = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + dataset_num_proc = 4, + ) + implicit_streaming = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, streaming = True), + ) + explicit_no_clip = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + max_grad_norm = 0.0, + ), + ) + + assert explicit_default.args.dataset_order == "default" + assert explicit_sequential.args.dataset_order == "sequential" + assert implicit_with_override.args.dataset_order == "torch_randperm" + assert implicit_streaming.args.dataset_order == "default" + assert implicit_with_override.args.max_grad_norm == 1.0 + assert explicit_no_clip.args.max_grad_norm == 0.0 + + +def test_mlx_trainer_uses_model_context_length_when_implicit(): + """UnslothTrainer should mirror CUDA's max_length bridge precedence.""" + unsloth = _import_mlx_unsloth() + model = _DummyModel() + model.max_seq_length = 321 + max_length_model = _DummyModel() + max_length_model.max_seq_length = 321 + none_model = _DummyModel() + none_model.max_seq_length = 321 + explicit_seq_model = _DummyModel() + explicit_seq_model.max_seq_length = 321 + clamped_seq_model = _DummyModel() + clamped_seq_model.max_seq_length = 321 + model_max_length = _DummyModel() + model_max_length.max_length = 777 + metadata_model = _DummyModel() + metadata_model.config = type("Config", (), {"max_position_embeddings": 888})() + metadata_tokenizer = type("Tokenizer", (), {"model_max_length": 999})() + explicit_max_length_no_model = _DummyModel() + trainer_override_model = _DummyModel() + trainer_override_model.max_seq_length = 321 + config_override_model = _DummyModel() + config_override_model.max_seq_length = 432 + + implicit = unsloth.UnslothTrainer( + model = model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + max_length_args = unsloth.UnslothTrainer( + model = max_length_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + none_args = unsloth.UnslothTrainer( + model = none_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = None), + ) + explicit_seq = unsloth.UnslothTrainer( + model = explicit_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 123), + ) + clamped_seq = unsloth.UnslothTrainer( + model = clamped_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 654), + ) + model_max_length_only = unsloth.UnslothTrainer( + model = model_max_length, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + metadata_ignored = unsloth.UnslothTrainer( + model = metadata_model, + tokenizer = metadata_tokenizer, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + explicit_max_length = unsloth.UnslothTrainer( + model = explicit_max_length_no_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + trainer_override = unsloth.UnslothTrainer( + model = trainer_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + max_seq_length = 654, + ) + config_with_override = unsloth.UnslothTrainer( + model = config_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.MLXTrainingConfig(max_steps = 1), + dataset_num_proc = 4, + ) + + assert implicit.args.max_seq_length == 321 + assert implicit.args.max_length == 321 + assert max_length_args.args.max_seq_length == 321 + assert max_length_args.args.max_length == 321 + assert none_args.args.max_seq_length == 321 + assert none_args.args.max_length == 321 + assert explicit_seq.args.max_seq_length == 123 + assert explicit_seq.args.max_length == 123 + assert clamped_seq.args.max_seq_length == 321 + assert clamped_seq.args.max_length == 321 + assert model_max_length_only.args.max_seq_length == 777 + assert model_max_length_only.args.max_length == 777 + assert metadata_ignored.args.max_seq_length == 1024 + assert metadata_ignored.args.max_length == 1024 + assert explicit_max_length.args.max_seq_length == 123 + assert explicit_max_length.args.max_length == 123 + assert trainer_override.args.max_seq_length == 654 + assert trainer_override.args.max_length == 654 + assert config_with_override.args.max_seq_length == 432 + assert config_with_override.args.max_length == 432 + + +def test_mlx_trainer_processing_class_overrides_explicit_none_tokenizer(): + """TRL passes tokenizer=None while processing_class carries the tokenizer.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = processor, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_vision_collator_processor_overrides_processing_class(): + """Vision notebooks pass the tokenizer as processing_class and processor in collator.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_preserves_explicit_processor_over_vision_collator(): + """Explicit processor kwargs should stay authoritative over collator metadata.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + explicit_processor = object() + + class Processor: + pass + + collator_processor = Processor() + collator_processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), collator_processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processor = explicit_processor, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is explicit_processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_forwards_vision_collator_positional_defaults(): + """Vision collator CUDA-style positionals should route into MLX args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator( + _DummyVLMModel(), + object(), + 2048, + None, + "max", + -100, + False, + None, + None, + True, + None, + False, + ) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + data_collator = collator, + ) + + assert trainer.args.max_seq_length == 2048 + assert trainer.args.image_size == "max" + assert trainer.args.completion_only_loss is False + + +def test_mlx_vision_collator_default_does_not_override_explicit_args(): + """Implicit collator defaults should not override explicit trainer args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), object()) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = False, + ), + data_collator = collator, + ) + + assert trainer.args.completion_only_loss is False + + +def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs(): + """Unsupported kwargs that change training semantics should fail on MLX.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "peft_config"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + peft_config = object(), + ) + + +def test_mlx_trainer_rejects_metrics_and_callbacks(): + """Trainer hooks should fail because MLXTrainer cannot honor them yet.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "callbacks"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [object()], + ) + with pytest.raises(NotImplementedError, match = "compute_metrics"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + compute_metrics = lambda *_: None, + ) + + +def test_mlx_trainer_rejects_custom_data_collator(): + """MLXTrainer owns batching; custom SFT data collators must not be ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "data_collator"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + data_collator = object(), + ) + + +def test_mlx_trainer_rejects_text_completion_only_loss(): + """Text MLX training should not silently ignore completion_only_loss=True.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "completion_only_loss=True"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + +def test_mlx_trainer_allows_vlm_completion_only_loss(): + """VLM MLX training supports completion_only_loss during collation.""" + unsloth = _import_mlx_unsloth() + + class VLMModel(_DummyModel): + _is_vlm_model = True + + trainer = unsloth.UnslothTrainer( + model = VLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + assert trainer.args.completion_only_loss is True + + +def test_mlx_trainer_accepts_trl_style_positional_args(): + """TRL-style positional `(model, args, ...)` should not be read as tokenizer.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("trl-outputs", max_steps = 2) + trainer = unsloth.UnslothTrainer( + _DummyModel(), + args, + train_dataset = [], + tokenizer = None, + ) + + assert trainer.args is args + assert trainer.args.output_dir == "trl-outputs" + assert trainer.train_dataset == [] + + +def test_mlx_trainer_accepts_trl_none_placeholder_positionals(): + """Explicit TRL default placeholders should preserve later positional args.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + processing_class = object() + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + None, + processing_class, + ) + + assert getattr(trainer.train_dataset, "_dataset", trainer.train_dataset) is dataset + assert getattr(trainer, "_mlx_train_dataset_for_batches", dataset) is dataset + assert trainer.tokenizer is processing_class + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_none_placeholder_positionals(): + """Short TRL placeholder calls should keep the fourth arg as train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_placeholders_with_keyword_dataset(): + """Short TRL placeholders should not conflict with keyword train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + train_dataset = dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_preserves_mlx_positional_schema_with_none_tokenizer(): + """MLX-style `(model, tokenizer, train_dataset, ...)` should still work.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + dataset, + None, + ) + + assert trainer.tokenizer is None + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + + +def test_mlx_compatibility_shims_are_installed(): + """Old notebook imports should resolve to the MLX public API after unsloth import.""" + unsloth = _import_mlx_unsloth() + + trl = importlib.import_module("trl") + trainer_module = importlib.import_module("unsloth.trainer") + chat_templates = importlib.import_module("unsloth.chat_templates") + dataset_utils = importlib.import_module("unsloth_zoo.dataset_utils") + + assert importlib.util.find_spec("trl") is not None + assert importlib.util.find_spec("unsloth.trainer") is not None + assert unsloth.trainer is trainer_module + assert unsloth.chat_templates is chat_templates + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trainer_module.UnslothTrainer is unsloth.UnslothTrainer + assert trainer_module.UnslothVisionDataCollator is unsloth.UnslothVisionDataCollator + assert chat_templates.train_on_responses_only is dataset_utils.train_on_responses_only + assert callable(unsloth.train_on_responses_only) + + +def test_mlx_trl_shim_preserves_existing_trl_module(monkeypatch): + """The MLX TRL shim should patch, not replace, an already-loaded TRL module.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.existing_marker = object() + trl.ExistingExport = object() + trl.__all__ = ["ExistingExport", "BrokenExport"] + + def _raise_for_broken_export(name): + if name == "BrokenExport": + raise RuntimeError("optional dependency missing") + raise AttributeError(name) + + trl.__getattr__ = _raise_for_broken_export + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + assert sys.modules["trl"] is trl + assert trl.__path__ == ["real-trainer-package"] + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + assert "ExistingExport" in trl.__all__ + assert "BrokenExport" not in trl.__all__ + assert "SFTTrainer" in trl.__all__ + assert "SFTConfig" in trl.__all__ + + +def test_mlx_trl_shim_installs_real_trl_or_stub(monkeypatch): + """The MLX TRL shim should prefer real TRL and stub only if unavailable.""" + unsloth = _import_mlx_unsloth() + monkeypatch.delitem(sys.modules, "trl", raising = False) + real_trl_available = importlib.util.find_spec("trl") is not None + + unsloth._install_mlx_trl_sft_shim() + trl = importlib.import_module("trl") + + if real_trl_available: + assert trl.__version__ != "0.0.0+unsloth-mlx" + else: + assert trl.__version__ == "0.0.0+unsloth-mlx" + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + + +def test_mlx_trl_star_import_exports_public_shims(): + """Existing `from trl import *` callers should receive MLX SFT shims.""" + unsloth = _import_mlx_unsloth() + namespace = {} + + exec("from trl import *", namespace) + + assert namespace["SFTTrainer"] is unsloth.UnslothTrainer + assert issubclass(namespace["SFTConfig"], unsloth.UnslothTrainingArguments) + + +def test_mlx_rl_trainers_stub_with_clear_error(monkeypatch): + """GRPO/DPO/ORPO trainers have no MLX path, so the shim retargets the ones trl + exposes to a clear NotImplementedError instead of a confusing CUDA crash, and + never invents trainers trl does not have.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + + class _RealTrainer: + def __init__(self, *args, **kwargs): + raise AssertionError("the real torch/CUDA trainer must not run on MLX") + + trl.GRPOTrainer = _RealTrainer + trl.DPOTrainer = _RealTrainer + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + for name in ("GRPOTrainer", "DPOTrainer"): + assert getattr(trl, name) is not _RealTrainer + with pytest.raises(NotImplementedError) as exc: + getattr(trl, name)(model = None, args = None) + assert "MLX" in str(exc.value) and name in str(exc.value) + # trainers trl never exposed must not be invented + assert not hasattr(trl, "PPOTrainer") + # idempotent: a second install keeps the same stub + stub = trl.GRPOTrainer + unsloth._install_mlx_trl_sft_shim() + assert trl.GRPOTrainer is stub + + +def test_mlx_rl_trainer_stub_is_lazy_import_safe(monkeypatch): + """Stubbing unsupported trl trainers must not resolve them: trl lazy-imports + pull torch, so on a torch-free MLX install a getattr probe would crash + `import unsloth`. The shim reads __all__/vars metadata and never triggers + trl's __getattr__ for a trainer it is about to replace.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "GRPOTrainer", "DPOTrainer"] + resolved = [] + + def _lazy_getattr(name): + resolved.append(name) + raise ImportError(f"lazy import of {name} would pull torch") + + trl.__getattr__ = _lazy_getattr + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() # must not raise despite the lazy trl + + # trainers declared in __all__ are stubbed WITHOUT ever resolving the real one + assert resolved == [] + for name in ("GRPOTrainer", "DPOTrainer"): + with pytest.raises(NotImplementedError): + getattr(trl, name)(model = None) + + +def test_mlx_stubs_trl_trainers_outside_fixed_set(monkeypatch): + """Any non-SFT trainer trl exports (e.g. a newer RLOOTrainer not in the fixed + list) must be stubbed too, so no torch trainer slips through on MLX.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "RLOOTrainer"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + with pytest.raises(NotImplementedError) as exc: + trl.RLOOTrainer(model = None) + assert "MLX" in str(exc.value) and "RLOOTrainer" in str(exc.value) + # SFT stays usable; only non-SFT trainers are stubbed + assert trl.SFTTrainer is unsloth.UnslothTrainer + + +def test_mlx_preserve_dataset_order_is_accepted(): + """preserve_dataset_order=True must be accepted (it is a real MLX config field), + not rejected as an unknown/unsupported argument.""" + unsloth = _import_mlx_unsloth() + args = unsloth.UnslothTrainingArguments( + output_dir = "mlx-out", + max_steps = 10, + preserve_dataset_order = True, + ) + assert getattr(args, "preserve_dataset_order", False) is True + + +def test_mlx_sftconfig_alias_keeps_trl_epoch_default(monkeypatch): + """`trl.SFTConfig` (aliased on MLX) keeps TRL's default training length: with + no explicit max_steps/num_train_epochs it runs TRL's 3 epochs, not the native + MLX 60-step default. An explicit length is authoritative and untouched.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + # no explicit length -> TRL epoch default (3 epochs, step cap disabled) + cfg = trl.SFTConfig(output_dir = "mlx-out") + assert cfg.num_train_epochs == 3 + assert cfg.max_steps == -1 + # explicit step / epoch counts stay exactly as written + assert trl.SFTConfig(output_dir = "mlx-out", max_steps = 17).max_steps == 17 + assert trl.SFTConfig(output_dir = "mlx-out", num_train_epochs = 2).num_train_epochs == 2 + + +def test_mlx_vision_collator_is_constructor_compatible(): + """Vision notebooks should be able to instantiate the collator placeholder.""" + unsloth = _import_mlx_unsloth() + + collator = unsloth.UnslothVisionDataCollator("model", "processor", flag = True) + + assert collator.model == "model" + assert collator.processor == "processor" + assert collator.kwargs == {"completion_only_loss": True, "flag": True} + + +def test_mlx_train_on_responses_only_returns_shared_mask_function(): + """The MLX public shim should expose the shared response-mask helper.""" + unsloth = _import_mlx_unsloth() + + class Tokenizer: + def __call__( + self, + text, + add_special_tokens = False, + ): + return types.SimpleNamespace( + input_ids = { + "": [1], + "": [2], + }[text] + ) + + def convert_tokens_to_ids(self, token): + return token + + mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + ) + masked = mask_fn( + { + "input_ids": [[1, 10, 2, 20, 21, 1, 11]], + } + ) + + assert masked == {"labels": [[-100, -100, -100, 20, 21, -100, -100]]} + + last_mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + last_response_only = True, + ) + last_masked = last_mask_fn( + { + "input_ids": [[1, 10, 2, 20, 1, 11, 2, 30]], + } + ) + + assert last_masked == {"labels": [[-100, -100, -100, -100, -100, -100, -100, 30]]} + + +def test_mlx_get_chat_template_uses_light_tokenizer_patch(monkeypatch): + """MLX notebooks should not import CUDA-heavy tokenizer/save helpers.""" + _import_mlx_unsloth() + from unsloth.chat_templates import get_chat_template + import unsloth_zoo.tokenizer_utils as tokenizer_utils + + class Tokenizer: + is_fast = True + padding_side = "right" + eos_token = "" + bos_token = "" + unk_token = "" + pad_token = "" + added_tokens_decoder = {} + + def fake_patch_tokenizer(model, tokenizer): + return model, tokenizer + + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name.startswith("unsloth.models") or name.startswith("unsloth.save"): + raise AssertionError(f"unexpected CUDA-heavy import: {name}") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(tokenizer_utils, "patch_tokenizer", fake_patch_tokenizer) + monkeypatch.setattr(builtins, "__import__", guarded_import) + + tokenizer = get_chat_template( + Tokenizer(), + chat_template = ("{{ messages }}", ""), + ) + + assert tokenizer.chat_template == "{{ messages }}" + assert tokenizer.padding_side == "right" + + +def test_mlx_gpu_memory_stats_helper_shape(): + """The portable memory helper should return CUDA-shaped values.""" + unsloth = _import_mlx_unsloth() + + stats, used, total = unsloth.get_gpu_memory_stats() + + assert isinstance(stats.name, str) + assert hasattr(stats, "total_memory") + assert isinstance(used, float) + assert total > 0 + + +def test_mlx_torch_cuda_compatibility_shim(): + """Existing CUDA memory and move calls should run on MLX.""" + unsloth = _import_mlx_unsloth() + torch = pytest.importorskip("torch") + from transformers.tokenization_utils_base import BatchEncoding + + stats, used, total = unsloth.get_gpu_memory_stats() + cuda_stats = torch.cuda.get_device_properties(0) + + assert cuda_stats.name == stats.name + assert cuda_stats.total_memory == stats.total_memory + assert torch.cuda.get_device_name(0) == stats.name + assert torch.cuda.max_memory_reserved() == int(used * 1024 * 1024 * 1024) + assert torch.cuda.max_memory_allocated() == torch.cuda.max_memory_reserved() + # current (non-max) APIs report live active memory, not the peak high-water + # mark, and never exceed it. + assert 0 <= torch.cuda.memory_reserved() <= torch.cuda.max_memory_reserved() + assert torch.cuda.memory_allocated() == torch.cuda.memory_reserved() + assert torch.cuda.device_count() == 1 + assert torch.cuda.current_device() == 0 + assert torch.cuda.get_device_capability() == (0, 0) + assert total > 0 + + free_bytes, total_bytes = torch.cuda.mem_get_info() + assert total_bytes == int(total * 1024 * 1024 * 1024) + assert 0 <= free_bytes <= total_bytes + + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + torch.cuda.set_device(0) + + tensor = torch.tensor([1, 2, 3]) + assert tensor.to("cuda") is tensor + assert tensor.cuda() is tensor + assert tensor.to(device = "cuda") is tensor + assert tensor.to("cuda", dtype = torch.float32).dtype == torch.float32 + + batch = BatchEncoding({"input_ids": tensor}) + assert batch.to("cuda") is batch + assert batch.to(device = "cuda") is batch diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 7a63dcfb85..275fe7ac57 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -403,10 +403,14 @@ def cmd_train(args) -> int: metrics["gguf_dir"] = str(gguf_dir) with Phase("save_gguf", metrics): try: + # q8_0 (the exporter default), not bf16: llama.cpp has optimized q8_0 + # CPU kernels, whereas bf16 CPU decode is unusably slow on the runner + # and made the fresh-process llama-cli reload below time out. q8_0 is + # also what users deploy by default. model.save_pretrained_gguf( str(gguf_dir), tokenizer = tokenizer, - quantization_method = "not_quantized", + quantization_method = "fast_quantized", ) gguf_files = sorted(gguf_dir.glob("*.gguf")) if not gguf_files: @@ -565,31 +569,36 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: raise SystemExit(f"no .gguf files in {save_dir}") gguf_path = gguf_files[0] - # This is a save/reload-integrity smoke; a few generated tokens are enough. - # Keep llama.cpp bounded on macOS runners where BF16 GGUF decode is CPU-bound. + # Save/reload-integrity smoke (assert below only needs a few chars). The GGUF is + # exported q8_0 (see save_gguf) because llama.cpp bf16 CPU decode is unusably slow + # on the runner. Run CPU-only (-ngl 0), cap the context (-c 256, the model + # advertises 32768), and keep generation short; all env-tunable. n_predict = os.environ.get("UNSLOTH_GGUF_RELOAD_N", "8") n_threads = os.environ.get("UNSLOTH_GGUF_RELOAD_THREADS", str(os.cpu_count() or 4)) + n_ctx = os.environ.get("UNSLOTH_GGUF_RELOAD_CTX", "256") + n_gpu_layers = os.environ.get("UNSLOTH_GGUF_RELOAD_NGL", "0") reload_timeout = int(os.environ.get("UNSLOTH_GGUF_RELOAD_TIMEOUT", "420")) - + argv = [ + str(llama_cli), + "-m", + str(gguf_path), + "-p", + PROMPT, + "-n", + n_predict, + "-t", + n_threads, + "-c", + n_ctx, + "-ngl", + n_gpu_layers, + "--temp", + "0", + "--seed", + str(SEED), + "--no-warmup", + ] with Phase("reload_gguf", metrics): - argv = [ - str(llama_cli), - "-m", - str(gguf_path), - "-p", - PROMPT, - "-n", - n_predict, - "-t", - n_threads, - "--temp", - "0", - "--seed", - str(SEED), - "-c", - "256", - "--no-warmup", - ] try: proc = subprocess.run( argv, @@ -606,6 +615,7 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: return stream.decode("utf-8", errors = "replace") return stream or "" + print(f" [reload:gguf] TIMEOUT running: {' '.join(argv)}", flush = True) print(f" [reload:gguf] TIMEOUT stdout:\n{_decode(exc.stdout)[:1000]}", flush = True) print(f" [reload:gguf] TIMEOUT stderr:\n{_decode(exc.stderr)[:1000]}", flush = True) raise diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 8202195ca8..04cc600725 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -33,6 +33,27 @@ if platform.system() == "Windows": pass +class _UnslothDeviceStats: + """Portable device metadata used by backend memory-reporting helpers.""" + + def __init__( + self, + name, + total_memory = 0, + ): + """Store a display name and total memory in bytes.""" + self.name = name + self.total_memory = int(total_memory or 0) + self.major = 0 + self.minor = 0 + self.multi_processor_count = 0 + + +def _bytes_to_gb(value): + """Convert byte counts to GiB rounded""" + return round(float(value or 0) / 1024 / 1024 / 1024, 3) + + def _is_mlx_available(): # Transitional import barrier: keep non-Apple-Silicon imports from touching # unsloth_zoo until unsloth_zoo.mlx is import-safe on GPU hosts. Then this @@ -66,7 +87,12 @@ if _IS_MLX: # mlx.trainer / mlx.loader submodules. Surface a friendly install hint # instead of a raw ImportError on the submodule path. try: - from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx.trainer import ( + MLXTrainer, + MLXTrainingConfig, + _is_vlm_model, + _normalize_mlx_optimizer_name, + ) from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as _e: raise ImportError( @@ -75,6 +101,53 @@ if _IS_MLX: "`pip install -U unsloth-zoo` or rerun install.sh." ) from _e + import dataclasses as _dataclasses + import importlib.machinery as _machinery + import sys as _sys + import types as _types + import warnings as _warnings + + __version__ = unsloth_zoo.__version__ + DEVICE_TYPE = "mlx" + + def _is_mlx_cuda_device_target(device): + """Return True when a torch .to/.cuda target asks for CUDA on MLX.""" + if device is None: + return False + return str(device).lower().startswith("cuda") + + def _patch_mlx_batch_encoding_to_cuda(): + """Treat tokenizer_output.to("cuda") as a no-op on the MLX backend.""" + try: + from transformers.tokenization_utils_base import BatchEncoding + except Exception: + return + + original_to = getattr(BatchEncoding, "to", None) + if original_to is None or getattr(original_to, "_unsloth_mlx_cuda_noop", False): + return + + def batch_encoding_to( + self, + device = None, + *args, + **kwargs, + ): + target = kwargs.get("device", device) + if _is_mlx_cuda_device_target(target): + return self + # device given by keyword: don't also pass the positional None, or the + # original raises "multiple values for 'device'" (e.g. .to(device="cpu")). + if "device" in kwargs: + return original_to(self, *args, **kwargs) + return original_to(self, device, *args, **kwargs) + + batch_encoding_to._unsloth_mlx_cuda_noop = True + batch_encoding_to._unsloth_original_to = original_to + BatchEncoding.to = batch_encoding_to + + _patch_mlx_batch_encoding_to_cuda() + # Load raw_text helpers without executing dataprep/__init__.py, which # imports synthetic.py -> torch and would defeat the torch-free MLX path. from pathlib import Path as _Path @@ -89,9 +162,6 @@ if _IS_MLX: TextPreprocessor = _raw_text.TextPreprocessor del _raw_text, _raw_text_spec, _raw_text_path, _Path - __version__ = unsloth_zoo.__version__ - DEVICE_TYPE = "mlx" - class FastLanguageModel: @staticmethod def from_pretrained(*args, **kwargs): @@ -141,14 +211,1202 @@ if _IS_MLX: is_bf16_supported = is_bfloat16_supported + def get_gpu_memory_stats(): + """Return MLX device stats, peak memory, and total memory in GiB.""" + import mlx.core as mx + + info = mx.device_info() + total = info.get("memory_size") or info.get("max_recommended_working_set_size") or 0 + get_peak_memory = getattr(mx, "get_peak_memory", None) + if get_peak_memory is None and hasattr(mx, "metal"): + get_peak_memory = getattr(mx.metal, "get_peak_memory", None) + peak = get_peak_memory() if callable(get_peak_memory) else 0 + stats = _UnslothDeviceStats(info.get("device_name", "Apple GPU"), total) + max_memory = _bytes_to_gb(total) or 1.0 + return stats, _bytes_to_gb(peak), max_memory + + def clear_gpu_memory(): + """Clear MLX's cached GPU memory for compatibility cleanup helpers.""" + import mlx.core as mx + + clear_cache = getattr(mx, "clear_cache", None) + if clear_cache is None and hasattr(mx, "metal"): + clear_cache = getattr(mx.metal, "clear_cache", None) + if callable(clear_cache): + clear_cache() + + def _patch_mlx_torch_cuda_compat_api(): + """Expose CUDA-shaped torch helpers for compatibility callers on MLX.""" + try: + import torch + except Exception: + return + + cuda = getattr(torch, "cuda", None) + if cuda is not None and not getattr(cuda, "_unsloth_mlx_cuda_compat_api", False): + + def get_device_properties(device = None): + """Return MLX device stats through torch.cuda's compatibility API.""" + return get_gpu_memory_stats()[0] + + def get_device_name(device = None): + """Return the MLX device name through torch.cuda's compatibility API.""" + return get_device_properties(device).name + + def max_memory_reserved(device = None): + """Return MLX peak memory in bytes for torch.cuda compatibility API.""" + return int(get_gpu_memory_stats()[1] * 1024 * 1024 * 1024) + + def empty_cache(): + """Clear MLX cache through torch.cuda.empty_cache().""" + clear_gpu_memory() + + def _mlx_active_memory_bytes(): + """Current active MLX memory in bytes (not the peak high-water mark).""" + import mlx.core as mx + + get_active = getattr(mx, "get_active_memory", None) + if get_active is None and hasattr(mx, "metal"): + get_active = getattr(mx.metal, "get_active_memory", None) + return int(get_active()) if callable(get_active) else 0 + + def memory_current(device = None): + """Return CURRENT MLX memory in bytes. torch.cuda.memory_reserved / + memory_allocated report live usage, not the peak (that is max_*).""" + return _mlx_active_memory_bytes() + + def mem_get_info(device = None): + """Return (free, total) bytes for torch.cuda compatibility API. + Free uses CURRENT active memory, not the peak high-water mark, so + a capacity check stays accurate after a transient spike.""" + total = int(get_gpu_memory_stats()[2] * 1024 * 1024 * 1024) + return (max(total - _mlx_active_memory_bytes(), 0), total) + + def reset_peak_memory_stats(device = None): + """Reset MLX's peak-memory counter so a later max_memory_reserved / + max_memory_allocated scopes to the run, not earlier model-load peaks.""" + import mlx.core as mx + + reset = getattr(mx, "reset_peak_memory", None) + if reset is None and hasattr(mx, "metal"): + reset = getattr(mx.metal, "reset_peak_memory", None) + if callable(reset): + reset() + + def synchronize(device = None): + """Wait for queued MLX work when torch.cuda.synchronize() is called.""" + import mlx.core as mx + + sync = getattr(mx, "synchronize", None) + if callable(sync): + sync() + + cuda.get_device_properties = get_device_properties + cuda.get_device_name = get_device_name + cuda.max_memory_reserved = max_memory_reserved + cuda.max_memory_allocated = max_memory_reserved + cuda.memory_reserved = memory_current + cuda.memory_allocated = memory_current + cuda.empty_cache = empty_cache + cuda.mem_get_info = mem_get_info + cuda.reset_peak_memory_stats = reset_peak_memory_stats + cuda.synchronize = synchronize + cuda.current_device = lambda: 0 + cuda.device_count = lambda: 1 + cuda.set_device = lambda device = None: None + cuda.get_device_capability = lambda device = None: (0, 0) + cuda.is_bf16_supported = lambda *args, **kwargs: is_bfloat16_supported() + cuda._unsloth_mlx_cuda_compat_api = True + + tensor_to = getattr(torch.Tensor, "to", None) + if tensor_to is not None and not getattr(tensor_to, "_unsloth_mlx_cuda_noop", False): + + def _coerce_mlx_dtype_to_torch(value): + """Map MLX dtype objects to their torch dtype equivalents.""" + try: + import mlx.core as mx + except Exception: + return value + dtype_map = { + mx.bool_: torch.bool, + mx.int8: torch.int8, + mx.int16: torch.int16, + mx.int32: torch.int32, + mx.int64: torch.int64, + mx.uint8: torch.uint8, + mx.float16: torch.float16, + mx.float32: torch.float32, + mx.bfloat16: torch.bfloat16, + } + mapped = dtype_map.get(value, None) + if mapped is not None: + return mapped + dtype_name = str(value).rsplit(".", 1)[-1] + name_map = { + "bool_": torch.bool, + "int8": torch.int8, + "int16": torch.int16, + "int32": torch.int32, + "int64": torch.int64, + "uint8": torch.uint8, + "float16": torch.float16, + "float32": torch.float32, + "bfloat16": torch.bfloat16, + } + return name_map.get(dtype_name, value) + + def mlx_tensor_to(self, *args, **kwargs): + """Ignore CUDA device targets while preserving dtype conversions.""" + args = list(args) + kwargs = dict(kwargs) + removed_cuda_device = False + if args and _is_mlx_cuda_device_target(args[0]): + args.pop(0) + removed_cuda_device = True + if _is_mlx_cuda_device_target(kwargs.get("device", None)): + kwargs.pop("device", None) + removed_cuda_device = True + if removed_cuda_device and not args: + cuda_only_kwargs = ("non_blocking", "copy", "memory_format") + if all(key in cuda_only_kwargs for key in kwargs): + return self + if removed_cuda_device and not args and not kwargs: + return self + if args: + args[0] = _coerce_mlx_dtype_to_torch(args[0]) + if "dtype" in kwargs: + kwargs["dtype"] = _coerce_mlx_dtype_to_torch(kwargs["dtype"]) + return tensor_to(self, *args, **kwargs) + + mlx_tensor_to._unsloth_mlx_cuda_noop = True + mlx_tensor_to._unsloth_original_to = tensor_to + torch.Tensor.to = mlx_tensor_to + + tensor_cuda = getattr(torch.Tensor, "cuda", None) + if tensor_cuda is not None and not getattr(tensor_cuda, "_unsloth_mlx_cuda_noop", False): + + def mlx_tensor_cuda(self, *args, **kwargs): + """Treat tensor.cuda() as a no-op on MLX.""" + return self + + mlx_tensor_cuda._unsloth_mlx_cuda_noop = True + mlx_tensor_cuda._unsloth_original_cuda = tensor_cuda + torch.Tensor.cuda = mlx_tensor_cuda + + _patch_mlx_torch_cuda_compat_api() + + _MLX_TRAINING_CONFIG_FIELDS = {_field.name for _field in _dataclasses.fields(MLXTrainingConfig)} + _MLX_TRAINING_ARGUMENT_ALIASES = { + "max_length": "max_seq_length", + } + _MLX_COMPAT_EXTRA_ARGUMENTS = frozenset( + ( + "bf16", + "dataloader_num_workers", + "dataloader_pin_memory", + "dataset_kwargs", + "ddp_find_unused_parameters", + "disable_tqdm", + "eval_strategy", + "evaluation_strategy", + "fp16", + "full_determinism", + "gradient_checkpointing_kwargs", + "hub_model_id", + "hub_token", + "log_level", + "logging_strategy", + "neftune_noise_alpha", + "optim_args", + "padding_free", + "push_to_hub", + "remove_unused_columns", + "save_on_each_node", + "save_safetensors", + "save_strategy", + "torch_compile", + ) + ) + _MLX_IMPLEMENTED_EXTRA_ARGUMENTS = frozenset( + ( + "image_size", + "preserve_dataset_order", + "warmup_ratio", + ) + ) + _MLX_ALLOWED_EXTRA_ARGUMENTS = _MLX_COMPAT_EXTRA_ARGUMENTS | _MLX_IMPLEMENTED_EXTRA_ARGUMENTS + _MLX_UNSUPPORTED_TASK_ARGUMENTS = frozenset( + ( + "assistant_only_loss", + "completion_only_loss", + ) + ) + + def _is_mlx_no_save_strategy(value): + if hasattr(value, "value"): + value = value.value + strategy = str(value or "").strip().lower() + strategy = strategy.rsplit(".", 1)[-1] + return strategy in ("no", "none", "false") + + _MLX_ADAMW_OPTIMIZER_ALIASES = frozenset( + ( + "adamw_8bit", + "paged_adamw_8bit", + "adamw_bnb_8bit", + "paged_adamw_32bit", + "adamw_torch", + "adamw_torch_fused", + "paged_adamw", + "adamw_32bit", + "adamw_hf", + "adamw_anyprecision", + "adamw_apex_fused", + ) + ) + + def _normalize_mlx_training_value(key, value): + if key == "eval_steps" and value is None: + return 0 + if key == "num_train_epochs" and value is not None and not isinstance(value, bool): + try: + epochs = float(value) + except (TypeError, ValueError): + pass + else: + if epochs.is_integer(): + return int(epochs) + if key == "lr_scheduler_type" and hasattr(value, "value"): + return value.value + if key != "optim": + return value + try: + return _normalize_mlx_optimizer_name(value) + except ValueError: + # Older unsloth-zoo lacks CUDA/TRL optimizer aliases; map common + # adamw_* names so notebook defaults (optim="adamw_8bit") still work. + opt = str(getattr(value, "value", value) or "adamw").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_ADAMW_OPTIMIZER_ALIASES: + return "adamw" + raise + + def _mlx_training_argument_values(args): + values = {} + for field in _dataclasses.fields(MLXTrainingConfig): + if hasattr(args, field.name): + values[field.name] = _normalize_mlx_training_value( + field.name, + getattr(args, field.name), + ) + for alias, target in _MLX_TRAINING_ARGUMENT_ALIASES.items(): + if target not in values and hasattr(args, alias): + values[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else alias] = getattr( + args, alias + ) + for name in _MLX_ALLOWED_EXTRA_ARGUMENTS: + if hasattr(args, name): + values[name] = getattr(args, name) + for name in _MLX_UNSUPPORTED_TASK_ARGUMENTS: + if hasattr(args, name): + value = getattr(args, name) + if ( + name == "completion_only_loss" + and value is not None + and name in _MLX_TRAINING_CONFIG_FIELDS + ): + values[name] = value + elif value is not None and value is not False: + values[name] = value + if _is_mlx_no_save_strategy(values.get("save_strategy", None)): + values["save_steps"] = 0 + return values + + def _split_mlx_trainer_kwargs(kwargs): + trainer_kwargs = {} + config_kwargs = {} + ignored_kwargs = {} + for key, value in kwargs.items(): + if key in _MLX_TRAINER_KWARGS: + trainer_kwargs[key] = value + continue + target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) + if target in _MLX_TRAINING_CONFIG_FIELDS or key in _MLX_ALLOWED_EXTRA_ARGUMENTS: + config_kwargs[key] = value + else: + ignored_kwargs[key] = value + return trainer_kwargs, config_kwargs, ignored_kwargs + + def _is_mlx_training_args_like(value): + if isinstance(value, (MLXTrainingConfig, dict, str, os.PathLike)): + return True + return any( + hasattr(value, name) + for name in ( + "output_dir", + "per_device_train_batch_size", + "gradient_accumulation_steps", + "max_steps", + "learning_rate", + ) + ) + + def _should_use_trl_positional_schema(args): + if len(args) < 2: + return False + if _is_mlx_training_args_like(args[1]): + return True + # TRL callers often pass explicit defaults: + # SFTTrainer(model, None, None, train_dataset, ...) + return len(args) >= 3 and args[1] is None and (args[2] is None or callable(args[2])) + + def _assign_mlx_positional_kwarg(kwargs, name, value): + if name in kwargs: + raise TypeError( + f"UnslothTrainer.__init__() got multiple values for argument " f"{name!r}" + ) + kwargs[name] = value + + def _normalize_mlx_trainer_init_args(args, kwargs): + kwargs = dict(kwargs) + if len(args) == 0: + return kwargs + + use_trl_schema = _should_use_trl_positional_schema(args) + positional_names = ( + _TRL_SFT_TRAINER_POSITIONAL_KWARGS if use_trl_schema else _MLX_TRAINER_POSITIONAL_KWARGS + ) + if len(args) > len(positional_names): + raise TypeError( + f"UnslothTrainer.__init__() takes at most " + f"{len(positional_names)} positional arguments on MLX " + f"({len(args)} given)" + ) + for name, value in zip(positional_names, args): + _assign_mlx_positional_kwarg(kwargs, name, value) + return kwargs + + def _is_meaningful_mlx_extra_value(value): + if value is None or value is False: + return False + if isinstance(value, (str, bytes)) and len(value) == 0: + return False + if isinstance(value, (dict, list, tuple, set, frozenset)) and len(value) == 0: + return False + return True + + def _warn_ignored_mlx_training_args(extra_kwargs): + names = sorted( + key + for key, value in extra_kwargs.items() + if (key in _MLX_COMPAT_EXTRA_ARGUMENTS and _is_meaningful_mlx_extra_value(value)) + ) + if not names: + return + _warnings.warn( + "Unsloth MLX: accepting but not applying unsupported " + "TrainingArguments kwargs: " + f"{', '.join(names)}. These options are not implemented by " + "MLXTrainer yet.", + RuntimeWarning, + stacklevel = 3, + ) + + def _is_meaningful_mlx_trainer_kwarg(key, value): + if key == "optimizers" and value == (None, None): + return False + return _is_meaningful_mlx_extra_value(value) + + def _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs): + names = sorted( + key + for key, value in ignored_kwargs.items() + if _is_meaningful_mlx_trainer_kwarg(key, value) + ) + if not names: + return + raise NotImplementedError( + "Unsloth MLX: unsupported SFTTrainer kwargs cannot be ignored safely: " + f"{', '.join(names)}. Remove these kwargs or use a supported MLX " + "trainer configuration." + ) + + def _raise_unknown_mlx_training_args(extra_kwargs): + names = sorted(key for key in extra_kwargs if key not in _MLX_ALLOWED_EXTRA_ARGUMENTS) + if not names: + return + raise NotImplementedError( + "Unsloth MLX: unsupported TrainingArguments/SFTConfig kwargs: " + f"{', '.join(names)}. Remove these kwargs or use fields implemented " + "by MLXTrainingConfig." + ) + + def _positive_mlx_context_length(value): + if value is None or isinstance(value, bool): + return None + try: + length = int(value) + except (TypeError, ValueError, OverflowError): + return None + if length <= 0: + return None + return length + + def _positive_mlx_training_number(value): + if value is None or isinstance(value, bool): + return None + try: + number = float(value) + except (TypeError, ValueError, OverflowError): + return None + if number <= 0: + return None + return number + + def _set_mlx_cuda_style_context_length(args, length): + args.max_seq_length = length + args.max_length = length + args._unsloth_mlx_max_length_value = length + return args + + class UnslothTrainingArguments(MLXTrainingConfig): + """MLX-compatible public training arguments for Unsloth notebooks.""" + + def __init__(self, *args, **kwargs): + if len(args) == 1 and isinstance(args[0], dict): + kwargs = {**args[0], **kwargs} + elif len(args) == 1 and isinstance(args[0], (str, os.PathLike)): + kwargs = {"output_dir": os.fspath(args[0]), **kwargs} + elif args: + raise TypeError( + "UnslothTrainingArguments on MLX accepts keyword arguments, " + "a dict, or a single positional output_dir." + ) + + max_length_value = kwargs.get("max_length", None) + # Only the canonical max_seq_length marks context length explicit; TRL + # max_length stays a compatibility alias and defers to the model's + # context length when one is available. + max_seq_length_explicit = ( + _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None + ) + if "max_length" in kwargs and "max_seq_length" not in kwargs: + kwargs["max_seq_length"] = kwargs["max_length"] + elif ( + "max_length" in kwargs + and _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None + ): + max_length_value = kwargs["max_seq_length"] + if "num_train_epochs" in kwargs and "max_steps" not in kwargs: + kwargs["max_steps"] = -1 + + dataset_order_explicit = "dataset_order" in kwargs or bool( + kwargs.get("preserve_dataset_order", False) + ) + append_eos_explicit = "append_eos" in kwargs + grad_clip_explicit = any( + name in kwargs for name in ("max_grad_norm", "max_grad_value", "max_grad_leaf_norm") + ) + warmup_ratio = kwargs.get("warmup_ratio", None) + warmup_steps_supplied = "warmup_steps" in kwargs + warmup_steps_value = kwargs.get("warmup_steps", None) + warmup_steps_explicit = False + if warmup_steps_supplied: + try: + warmup_steps_explicit = int(warmup_steps_value) > 0 + except (TypeError, ValueError): + warmup_steps_explicit = True + filtered_kwargs = {} + extra_kwargs = {} + for key, value in kwargs.items(): + target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) + if key != target and target in kwargs: + continue + value = _normalize_mlx_training_value(target, value) + if target in _MLX_UNSUPPORTED_TASK_ARGUMENTS: + if ( + target == "completion_only_loss" + and value is not None + and target in _MLX_TRAINING_CONFIG_FIELDS + ): + filtered_kwargs[target] = value + elif _is_meaningful_mlx_extra_value(value): + extra_kwargs[key] = value + continue + if target in _MLX_TRAINING_CONFIG_FIELDS: + filtered_kwargs[target] = value + else: + extra_kwargs[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else key] = value + + _raise_unknown_mlx_training_args(extra_kwargs) + + if _is_mlx_no_save_strategy(extra_kwargs.get("save_strategy", None)): + filtered_kwargs["save_steps"] = 0 + + if warmup_ratio is not None and not warmup_steps_explicit: + import math as _math + max_steps = filtered_kwargs.get( + "max_steps", + getattr(MLXTrainingConfig, "max_steps", 60), + ) + try: + if int(max_steps) > 0: + filtered_kwargs["warmup_steps"] = max( + 0, + _math.ceil(int(max_steps) * float(warmup_ratio)), + ) + except (TypeError, ValueError): + pass + + super().__init__(**filtered_kwargs) + self._unsloth_mlx_dataset_order_explicit = dataset_order_explicit + self._unsloth_mlx_append_eos_explicit = append_eos_explicit + self._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit + self._unsloth_mlx_max_length_value = max_length_value + if "max_length" in kwargs: + self.max_length = max_length_value + self._unsloth_mlx_grad_clip_explicit = grad_clip_explicit + self._unsloth_mlx_warmup_steps_explicit = warmup_steps_explicit + self._unsloth_mlx_extra_args = extra_kwargs + for key, value in extra_kwargs.items(): + setattr(self, key, value) + _warn_ignored_mlx_training_args(extra_kwargs) + + def _resolve_mlx_cuda_style_max_seq_length(args, model = None): + model_max_seq_length = _positive_mlx_context_length( + getattr(model, "max_seq_length", None), + ) + args_max_seq_length = _positive_mlx_context_length( + getattr(args, "max_seq_length", None), + ) + args_max_seq_length_explicit = getattr( + args, + "_unsloth_mlx_max_seq_length_explicit", + None, + ) + if args_max_seq_length_explicit is None: + default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) + args_max_seq_length_explicit = ( + args_max_seq_length is not None and args_max_seq_length != default_max_seq_length + ) + if not args_max_seq_length_explicit: + args_max_seq_length = None + + if args_max_seq_length is None and model_max_seq_length is not None: + args_max_seq_length = model_max_seq_length + elif ( + args_max_seq_length is not None + and model_max_seq_length is not None + and args_max_seq_length > model_max_seq_length + ): + print( + "Unsloth: You set `max_seq_length` as " + f"{args_max_seq_length} but the maximum the model supports is " + f"{model_max_seq_length}. We shall reduce it." + ) + args_max_seq_length = model_max_seq_length + + if args_max_seq_length is not None: + _set_mlx_cuda_style_context_length(args, args_max_seq_length) + return args + + model_max_length = model_max_seq_length + if model_max_length is None: + model_max_length = _positive_mlx_context_length( + getattr(model, "max_length", None), + ) + if model_max_length is not None: + _set_mlx_cuda_style_context_length(args, model_max_length) + return args + + args_max_length = _positive_mlx_context_length( + getattr(args, "max_length", None), + ) + if args_max_length is None: + args_max_length = _positive_mlx_context_length( + getattr(args, "_unsloth_mlx_max_length_value", None), + ) + if args_max_length is not None: + _set_mlx_cuda_style_context_length(args, args_max_length) + if model is not None: + setattr(model, "max_seq_length", args_max_length) + return args + + _set_mlx_cuda_style_context_length(args, 1024) + return args + + def _apply_unsloth_trainer_mlx_defaults( + args, + model = None, + max_seq_length_explicit = False, + ): + if ( + not getattr(args, "streaming", False) + and not getattr(args, "preserve_dataset_order", False) + and not getattr(args, "_unsloth_mlx_dataset_order_explicit", False) + ): + default_order = getattr(MLXTrainingConfig, "dataset_order", "default") + if getattr(args, "dataset_order", default_order) in (None, default_order): + args.dataset_order = "torch_randperm" + + if isinstance(args, UnslothTrainingArguments) and not getattr( + args, "_unsloth_mlx_append_eos_explicit", False + ): + args.append_eos = False + + if isinstance(args, UnslothTrainingArguments) and not getattr( + args, "_unsloth_mlx_grad_clip_explicit", False + ): + max_grad_norm = _positive_mlx_training_number( + getattr(args, "max_grad_norm", None), + ) + max_grad_value = _positive_mlx_training_number( + getattr(args, "max_grad_value", None), + ) + max_grad_leaf_norm = _positive_mlx_training_number( + getattr(args, "max_grad_leaf_norm", None), + ) + if max_grad_norm is None and max_grad_value is None and max_grad_leaf_norm is None: + args.max_grad_norm = 1.0 + + if not max_seq_length_explicit: + _resolve_mlx_cuda_style_max_seq_length(args, model = model) + return args + + def _coerce_mlx_training_args(args, overrides = None): + overrides = overrides or {} + if isinstance(args, MLXTrainingConfig) and not overrides: + return args + dataset_order_explicit = None + append_eos_explicit = None + max_seq_length_explicit = None + max_length_value = None + grad_clip_explicit = None + if args is None: + values = {} + elif isinstance(args, dict): + values = dict(args) + elif isinstance(args, (str, os.PathLike)): + values = {"output_dir": os.fspath(args)} + else: + dataset_order_explicit = getattr( + args, + "_unsloth_mlx_dataset_order_explicit", + False, + ) + append_eos_explicit = getattr( + args, + "_unsloth_mlx_append_eos_explicit", + None, + ) + max_seq_length_explicit = getattr( + args, + "_unsloth_mlx_max_seq_length_explicit", + None, + ) + if max_seq_length_explicit is None: + args_max_seq_length = _positive_mlx_context_length( + getattr(args, "max_seq_length", None), + ) + default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) + max_seq_length_explicit = ( + args_max_seq_length is not None + and args_max_seq_length != default_max_seq_length + ) + max_length_value = getattr( + args, + "_unsloth_mlx_max_length_value", + getattr(args, "max_length", None), + ) + grad_clip_explicit = getattr( + args, + "_unsloth_mlx_grad_clip_explicit", + None, + ) + values = _mlx_training_argument_values(args) + if hasattr(args, "max_length"): + values["max_length"] = getattr(args, "max_length") + values.update(overrides) + coerced = UnslothTrainingArguments(**values) + if ( + dataset_order_explicit is not None + and "dataset_order" not in overrides + and "preserve_dataset_order" not in overrides + ): + coerced._unsloth_mlx_dataset_order_explicit = dataset_order_explicit + if append_eos_explicit is not None and "append_eos" not in overrides: + coerced._unsloth_mlx_append_eos_explicit = append_eos_explicit + if ( + max_seq_length_explicit is not None + and "max_seq_length" not in overrides + and "max_length" not in overrides + ): + coerced._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit + if max_length_value is not None and "max_length" not in overrides: + coerced._unsloth_mlx_max_length_value = max_length_value + coerced.max_length = max_length_value + if ( + grad_clip_explicit is not None + and "max_grad_norm" not in overrides + and "max_grad_value" not in overrides + and "max_grad_leaf_norm" not in overrides + ): + coerced._unsloth_mlx_grad_clip_explicit = grad_clip_explicit + return coerced + + _MLX_TRAINER_POSITIONAL_KWARGS = ( + "model", + "tokenizer", + "train_dataset", + "eval_dataset", + "dataset_text_field", + "max_seq_length", + "packing", + "data_collator", + "args", + "formatting_func", + "processor", + ) + _TRL_SFT_TRAINER_POSITIONAL_KWARGS = ( + "model", + "args", + "data_collator", + "train_dataset", + "eval_dataset", + "processing_class", + "compute_loss_func", + "compute_metrics", + "callbacks", + "optimizers", + "optimizer_cls_and_kwargs", + "preprocess_logits_for_metrics", + "peft_config", + "formatting_func", + ) + _MLX_TRAINER_KWARGS = frozenset(_MLX_TRAINER_POSITIONAL_KWARGS) + + def _is_mlx_native_text_collator(collator): + """HF pad/copy collators are redundant on MLX; match by class name.""" + for klass in type(collator).__mro__: + name = klass.__name__ + if name in ( + "DataCollatorForSeq2Seq", + "DataCollatorWithPadding", + "DefaultDataCollator", + ): + return True + if name == "DataCollatorForLanguageModeling": + # Plain causal padding is fine; MLM masking changes semantics. + return not bool(getattr(collator, "mlm", False)) + return False + + _MLX_VISION_COLLATOR_FORWARDED_KWARGS = frozenset( + ("completion_only_loss", "formatting_func", "max_seq_length") + ) + _MLX_VISION_COLLATOR_IMAGE_KWARGS = frozenset(("image_size", "resize")) + _MLX_VISION_COLLATOR_POSITIONAL_KWARGS = ( + "max_seq_length", + "formatting_func", + "resize", + "ignore_index", + "train_on_responses_only", + "instruction_part", + "response_part", + "force_match", + "num_proc", + "completion_only_loss", + "pad_to_multiple_of", + "resize_dimension", + "snap_to_patch_size", + "last_response_only", + ) + _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS = { + "ignore_index": -100, + "train_on_responses_only": False, + "instruction_part": None, + "response_part": None, + "force_match": True, + "num_proc": None, + "pad_to_multiple_of": None, + "resize_dimension": 0, + "snap_to_patch_size": False, + "last_response_only": False, + } + + def _is_default_mlx_vision_collator_value(key, value): + """Return whether an unsupported collator value is the CUDA default.""" + if key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS: + return False + default = _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS[key] + if default is None: + return value is None + if isinstance(default, bool): + return value is default + return value == default and type(value) is type(default) + + def _has_mlx_training_arg_value(args, key): + """Return whether training args already carry an explicit config value.""" + if args is None or isinstance(args, (str, os.PathLike)): + return False + if isinstance(args, dict): + return key in args + return getattr(args, key, None) is not None + + def _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs): + """Reject VLM collator kwargs that cannot be ignored safely on MLX.""" + unsupported = sorted( + key + for key, value in collator_kwargs.items() + if ( + key not in _MLX_VISION_COLLATOR_FORWARDED_KWARGS + and key not in _MLX_VISION_COLLATOR_IMAGE_KWARGS + and ( + ( + key in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS + and not _is_default_mlx_vision_collator_value(key, value) + ) + or ( + key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS + and _is_meaningful_mlx_extra_value(value) + ) + ) + ) + ) + if unsupported: + raise NotImplementedError( + "Unsloth MLX: unsupported UnslothVisionDataCollator kwargs " + f"cannot be ignored safely: {', '.join(unsupported)}." + ) + + class UnslothTrainer(MLXTrainer): + """Backend-aware public trainer that routes supported SFT notebooks to MLX.""" + + def __init__(self, *args, **kwargs): + kwargs = _normalize_mlx_trainer_init_args(args, kwargs) + processing_class = kwargs.pop("processing_class", None) + processor_from_processing_class = False + if processing_class is not None: + if kwargs.get("processor", None) is None: + kwargs["processor"] = processing_class + processor_from_processing_class = True + if kwargs.get("tokenizer", None) is None: + kwargs["tokenizer"] = getattr( + processing_class, + "tokenizer", + processing_class, + ) + kwargs.setdefault("tokenizer", None) + + data_collator = kwargs.pop("data_collator", None) + if data_collator is not None: + if isinstance(data_collator, UnslothVisionDataCollator): + collator_processor = getattr(data_collator, "processor", None) + if collator_processor is not None and ( + kwargs.get("processor", None) is None or processor_from_processing_class + ): + kwargs["processor"] = collator_processor + if kwargs.get("tokenizer", None) is None: + kwargs["tokenizer"] = getattr( + collator_processor, + "tokenizer", + collator_processor, + ) + collator_kwargs = getattr(data_collator, "kwargs", None) or {} + collator_explicit_kwargs = getattr( + data_collator, + "_unsloth_mlx_explicit_kwargs", + set(collator_kwargs), + ) + collator_image_size = collator_kwargs.get( + "image_size", + collator_kwargs.get("resize", None), + ) + if isinstance(collator_image_size, list): + collator_image_size = tuple(collator_image_size) + if ( + isinstance(collator_image_size, str) + and collator_image_size.lower() == "max" + ): + collator_image_size = "max" + if "image_size" not in kwargs and ( + isinstance(collator_image_size, int) + or collator_image_size == "max" + or ( + isinstance(collator_image_size, tuple) + and len(collator_image_size) == 2 + and all(isinstance(x, int) for x in collator_image_size) + ) + ): + kwargs["image_size"] = collator_image_size + for collator_key in _MLX_VISION_COLLATOR_FORWARDED_KWARGS: + collator_defaulted_value = collator_key not in collator_explicit_kwargs + if collator_defaulted_value and _has_mlx_training_arg_value( + kwargs.get("args"), collator_key + ): + continue + if ( + collator_key in collator_kwargs + and collator_key not in kwargs + and collator_kwargs[collator_key] is not None + ): + kwargs[collator_key] = collator_kwargs[collator_key] + _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs) + elif _is_mlx_native_text_collator(data_collator): + pass # redundant on MLX; MLXTrainer batches/masks/pads natively + else: + raise NotImplementedError( + "Unsloth MLX: custom data_collator is not supported by " + "MLXTrainer. Pass the dataset directly or use the MLX " + "trainer's native batching path." + ) + + trainer_kwargs, config_kwargs, ignored_kwargs = _split_mlx_trainer_kwargs(kwargs) + _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs) + trainer_kwargs["args"] = _coerce_mlx_training_args( + trainer_kwargs.get("args"), + config_kwargs, + ) + if getattr( + trainer_kwargs["args"], "completion_only_loss", None + ) is True and not _is_vlm_model(trainer_kwargs.get("model")): + raise NotImplementedError( + "Unsloth MLX: completion_only_loss=True is only supported " + "for VLM training. For text SFT, call train_on_responses_only " + "after constructing the trainer." + ) + if getattr( + trainer_kwargs["args"], "train_on_completions", None + ) is True and not _is_vlm_model(trainer_kwargs.get("model")): + raise NotImplementedError( + "Unsloth MLX: train_on_completions=True is only supported " + "for VLM training. For text SFT, call train_on_responses_only " + "after constructing the trainer." + ) + trainer_kwargs["args"] = _apply_unsloth_trainer_mlx_defaults( + trainer_kwargs["args"], + model = trainer_kwargs.get("model"), + max_seq_length_explicit = (trainer_kwargs.get("max_seq_length") is not None), + ) + + super().__init__(**trainer_kwargs) + self.processing_class = ( + processing_class + if processing_class is not None + else self.processor or self.tokenizer + ) + if trainer_kwargs.get("max_seq_length") is not None: + _set_mlx_cuda_style_context_length( + self.args, + self.args.max_seq_length, + ) + self._unsloth_mlx_ignored_trainer_kwargs = ignored_kwargs + class UnslothVisionDataCollator: + def __init__( + self, + model = None, + processor = None, + *args, + **kwargs, + ): + explicit_kwargs = set(kwargs) + if len(args) > len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS): + raise TypeError( + "UnslothVisionDataCollator on MLX accepts at most " + f"{len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS)} positional " + "options after model and processor." + ) + for key, value in zip(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS, args): + if key in kwargs: + raise TypeError( + f"UnslothVisionDataCollator got multiple values for argument {key!r}" + ) + kwargs[key] = value + explicit_kwargs.add(key) + if "completion_only_loss" not in kwargs: + kwargs["completion_only_loss"] = True + self.model = model + self.processor = processor + self.args = () + self.kwargs = kwargs + self._unsloth_mlx_explicit_kwargs = explicit_kwargs + + def __call__(self, features): + raise NotImplementedError( + "Unsloth: UnslothVisionDataCollator is a compatibility placeholder " + "on MLX. Pass the dataset to UnslothTrainer; MLXTrainer performs " + "vision batching internally." + ) + + def get_chat_template(*args, **kwargs): + """Apply an Unsloth chat template through a lazy MLX-safe import.""" + from .chat_templates import get_chat_template as _get_chat_template + return _get_chat_template(*args, **kwargs) + + def apply_chat_template(*args, **kwargs): + """Format a dataset with an Unsloth chat template through a lazy import.""" + from .chat_templates import apply_chat_template as _apply_chat_template + return _apply_chat_template(*args, **kwargs) + + def standardize_data_formats(*args, **kwargs): + """Normalize ShareGPT-style datasets through the shared zoo helper.""" + from unsloth_zoo.dataset_utils import standardize_data_formats as _standardize_data_formats + return _standardize_data_formats(*args, **kwargs) + + def standardize_sharegpt(*args, **kwargs): + """Alias ShareGPT standardization to the shared dataset-format helper.""" + return standardize_data_formats(*args, **kwargs) + + def train_on_responses_only(*args, **kwargs): + """Mask non-response tokens through the shared zoo dataset helper.""" + from unsloth_zoo.dataset_utils import train_on_responses_only as _train_on_responses_only + return _train_on_responses_only(*args, **kwargs) + + def _safe_mlx_trl_star_exports(_trl): + """Return importable TRL star exports plus the MLX SFT shims.""" + exports = list(getattr(_trl, "__all__", ())) + safe_exports = [] + for name in exports: + try: + getattr(_trl, name) + except Exception: + continue + safe_exports.append(name) + for name in ("SFTConfig", "SFTTrainer"): + if name not in safe_exports: + safe_exports.append(name) + return safe_exports + + # trl trainers with no MLX implementation yet. Swap them for stubs that fail + # with a clear message instead of importing the real torch/CUDA trainer and + # crashing deep inside it, so an unmigrated GRPO/DPO/ORPO notebook is legible. + _MLX_UNSUPPORTED_TRL_TRAINERS = ( + "GRPOTrainer", + "DPOTrainer", + "ORPOTrainer", + "KTOTrainer", + "PPOTrainer", + "RewardTrainer", + ) + + def _make_mlx_unsupported_trl_trainer(name): def __init__(self, *args, **kwargs): raise NotImplementedError( - "Unsloth: UnslothVisionDataCollator is not used on MLX. " - "Use the MLX trainer/data path instead." + f"Unsloth: {name} is not yet supported on the MLX (Apple Silicon) " + f"backend. Only SFT training runs on MLX today; use a CUDA/ROCm GPU " + f"for {name}." ) + return type(name, (), {"__init__": __init__, "_unsloth_mlx_unsupported": True}) + + class _MLXSFTConfig(UnslothTrainingArguments): + """`trl.SFTConfig` alias that keeps TRL's default training length. + + TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the + native MLX config defaults to max_steps=60. An unmigrated notebook that + builds SFTConfig without an explicit length would otherwise silently run + 60 MLX steps under this alias, so seed the TRL epoch default when neither + max_steps nor num_train_epochs is given (epoch mode is MLX-supported). + """ + + def __init__(self, *args, **kwargs): + keys = set(kwargs) + if len(args) == 1 and isinstance(args[0], dict): + keys |= set(args[0]) + if not ({"max_steps", "num_train_epochs"} & keys): + kwargs.setdefault("num_train_epochs", 3) + super().__init__(*args, **kwargs) + + def _install_mlx_trl_sft_shim(): + """Install MLX-backed TRL SFT shims without replacing the TRL module.""" + _trl = _sys.modules.get("trl") + if _trl is None: + try: + import trl as _trl + except ImportError: + _trl = _types.ModuleType("trl") + _trl.__version__ = "0.0.0+unsloth-mlx" + _trl.__package__ = "trl" + _trl.__path__ = [] + _trl.__spec__ = _machinery.ModuleSpec("trl", loader = None, is_package = True) + _sys.modules["trl"] = _trl + + _trl.SFTTrainer = UnslothTrainer + _trl.SFTConfig = _MLXSFTConfig + # Only retarget trainers the installed trl actually exposes (don't invent + # attributes); idempotent so re-importing unsloth is a no-op. + # Decide what to stub from trl's declared exports (__all__) and already + # materialized attrs only. A getattr probe here would trigger trl's lazy + # trainer import, pulling torch and breaking `import unsloth` on torch-free + # MLX just to check existence. + _trl_exports = set(getattr(_trl, "__all__", ()) or ()) + # Stub every non-SFT trainer trl exposes, not just a fixed list, so newer + # trainers (RLOOTrainer, ...) also fail with a clear MLX message instead + # of importing the real torch trainer. Names come from __all__ so we never + # resolve them (that would trigger trl's lazy import and pull torch). + _unsupported = set(_MLX_UNSUPPORTED_TRL_TRAINERS) | { + _n for _n in _trl_exports if _n.endswith("Trainer") and _n != "SFTTrainer" + } + for _name in _unsupported: + _current = vars(_trl).get(_name) + if getattr(_current, "_unsloth_mlx_unsupported", False): + continue + if _name in _trl_exports or _current is not None: + setattr(_trl, _name, _make_mlx_unsupported_trl_trainer(_name)) + _trl.__all__ = _safe_mlx_trl_star_exports(_trl) + _trl.__UNSLOTH_MLX_COMPAT__ = True + + def _install_mlx_unsloth_trainer_shim(): + module_name = f"{__name__}.trainer" + _trainer = _types.ModuleType(module_name) + _trainer.__package__ = __name__ + _trainer.__spec__ = _machinery.ModuleSpec(module_name, loader = None) + _trainer.MLXTrainer = MLXTrainer + _trainer.MLXTrainingConfig = MLXTrainingConfig + _trainer.UnslothTrainer = UnslothTrainer + _trainer.UnslothTrainingArguments = UnslothTrainingArguments + _trainer.UnslothVisionDataCollator = UnslothVisionDataCollator + _sys.modules[module_name] = _trainer + globals()["trainer"] = _trainer + + _install_mlx_trl_sft_shim() + _install_mlx_unsloth_trainer_shim() + else: # GPU path: load everything from _gpu_init from ._gpu_init import * from ._gpu_init import __version__ + + def get_gpu_memory_stats(): + """Return CUDA/ROCm/XPU device stats, peak memory, and total memory in GiB.""" + try: + import torch + if hasattr(torch, "xpu") and torch.xpu.is_available(): + props = torch.xpu.get_device_properties(0) + peak = ( + torch.xpu.max_memory_reserved() + if hasattr(torch.xpu, "max_memory_reserved") + else torch.xpu.max_memory_allocated() + ) + total = getattr(props, "total_memory", 0) + return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 + if hasattr(torch, "cuda") and torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + peak = torch.cuda.max_memory_reserved() + total = getattr(props, "total_memory", 0) + return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 + except Exception: + pass + stats = _UnslothDeviceStats("Unknown GPU", 0) + return stats, 0.0, 1.0 + + def clear_gpu_memory(): + """Clear cached GPU memory on CUDA, ROCm, or XPU when available.""" + try: + import torch + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() + elif hasattr(torch, "cuda") and torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 60eb8de5d7..169b2dbd0e 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -27,18 +27,25 @@ __all__ = [ "test_construct_chat_template", ] -from transformers import StoppingCriteria, StoppingCriteriaList -from torch import LongTensor, FloatTensor -from transformers.models.llama.modeling_llama import logger +from transformers.utils import logging +try: + from torch import LongTensor, FloatTensor +except ImportError: + LongTensor = FloatTensor = None +logger = logging.get_logger(__name__) import os import shutil -from .tokenizer_utils import * import re from .ollama_template_mappers import OLLAMA_TEMPLATES -from unsloth_zoo.dataset_utils import ( - train_on_responses_only, - standardize_data_formats, -) +try: + from unsloth_zoo.dataset_utils import ( + train_on_responses_only, + standardize_data_formats, + ) +except ImportError: + # dataset_utils pulls torch; keep chat_templates importable on torch-free + # (MLX) hosts, which expose these via the backend-specific wrappers instead. + train_on_responses_only = standardize_data_formats = None standardize_sharegpt = standardize_data_formats CHAT_TEMPLATES = {} DEFAULT_SYSTEM_MESSAGE = {} @@ -1838,11 +1845,24 @@ def get_chat_template( map_eos_token = True, system_message = None, patch_saving = True, - use_zoo_tokenizer_patch = False, + use_zoo_tokenizer_patch = None, ): assert(type(map_eos_token) is bool) + import sys + is_mlx_backend = getattr(sys.modules.get("unsloth"), "DEVICE_TYPE", None) == "mlx" + if use_zoo_tokenizer_patch is None: + use_zoo_tokenizer_patch = is_mlx_backend old_tokenizer = tokenizer + # mlx-lm's TokenizerWrapper._tokenizer is the HF tokenizer, not the Rust + # backend the vocab-edit paths below need; unwrap here, re-wrap before return. + _mlx_tokenizer_wrapper = None + if is_mlx_backend and tokenizer.__class__.__name__ == "TokenizerWrapper": + _inner_tokenizer = getattr(tokenizer, "_tokenizer", None) + if _inner_tokenizer is not None and hasattr(_inner_tokenizer, "is_fast"): + _mlx_tokenizer_wrapper = tokenizer + tokenizer = _inner_tokenizer + IS_GEMMA = False if tokenizer.__class__.__name__.startswith("Gemma"): if chat_template == "chatml": chat_template = "gemma_chatml" @@ -1952,6 +1972,7 @@ def get_chat_template( pass # Must fix the sentence piece tokenizer since there's no tokenizer.model file! + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) else: pass @@ -1997,6 +2018,7 @@ def get_chat_template( # Must fix the sentence piece tokenizer since there's no tokenizer.model file! token_mapping = { old_eos_token : stop_word, } + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) pass @@ -2057,13 +2079,25 @@ def get_chat_template( # stopping_criteria = create_stopping_criteria(tokenizer, stop_word) # Patch saving functions - if patch_saving: + if patch_saving and not is_mlx_backend: from .save import patch_saving_functions tokenizer = patch_saving_functions(tokenizer) # Add Ollama tokenizer._ollama_modelfile = ollama_modelfile tokenizer._system_message = system_message + + # Re-wrap so the trainer gets the same TokenizerWrapper type back. + if _mlx_tokenizer_wrapper is not None: + _mlx_tokenizer_wrapper._tokenizer = tokenizer + eos_token_id = getattr(tokenizer, "eos_token_id", None) + if eos_token_id is not None: + _mlx_tokenizer_wrapper._eos_token_ids = {eos_token_id} + _mlx_tokenizer_wrapper._chat_template = None + _mlx_tokenizer_wrapper.has_chat_template = ( + getattr(tokenizer, "chat_template", None) is not None + ) + tokenizer = _mlx_tokenizer_wrapper return tokenizer#, stopping_criteria @@ -2749,6 +2783,15 @@ extra_eos_tokens = None, def create_stopping_criteria(tokenizer, stop_word = "eos_token"): + try: + import torch + from transformers import StoppingCriteria, StoppingCriteriaList + except ImportError as exc: + raise ImportError( + "Unsloth: create_stopping_criteria requires PyTorch and is only " + "supported on Torch backends." + ) from exc + class StoppingCriteriaSub(StoppingCriteria): __slots__ = "stop_token", "single_match", "length", @@ -2828,10 +2871,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) @@ -2845,10 +2888,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_old_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) # We add ourselves From abdc968e8d7ec82cab3349a5ae7edf5c41936ab4 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Fri, 3 Jul 2026 16:13:54 +0300 Subject: [PATCH 12/23] report a complete load once llama-server is healthy (#6790) * report a complete load once llama-server is healthy load_progress() derived its fraction purely from the llama-server's VmRSS over the GGUF shard total. With layers offloaded to VRAM (-ngl) the process releases the mmap'd weight pages after upload, so VmRSS sinks back well below the shard total: the fraction climbs toward ~1.0 during mmap, then collapses to a small value (~8%) once the weights are on the GPU. A fraction-driven progress bar therefore restarts and sticks there indefinitely even though the model is loaded and serving, which reads as a hang at "Starting model...". Once the server is healthy the load is complete by definition, so report fraction 1.0 (and bytes_loaded == bytes_total) in the ready phase regardless of resident set size. The VmRSS read is factored into _read_rss_bytes() with its original semantics preserved (0 on a missing VmRSS line, None when /proc is unavailable) so it can be unit-tested off Linux. Fixes #5740 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * stub heavy deps in the load-progress test and guard a valueless VmRSS Two review fixes: 1. The new test imported core.inference.llama_cpp at module top, which pulls in loggers/structlog/httpx and fails collection with ModuleNotFoundError in the lightweight backend test env when the file is run on its own. Stub loggers, structlog and httpx via sys.modules.setdefault before the import, mirroring test_llama_cpp_load_progress_matrix.py; setdefault keeps the real modules when installed. Verified the file now collects and passes with only pytest present. 2. Catch IndexError in _read_rss_bytes: a "VmRSS:" line with no value column would make line.split()[1] raise and crash a load-progress poll. Return None instead, with a test for the valueless line. * Hold load-progress high-water mark and explain a never-healthy load (#5740) load_progress() now holds a per-process VmRSS high-water mark, so the bar no longer regresses to ~8% when -ngl offloads the weights and frees the mmap pages mid-load. A live server that never returns 200 on /health now gets a specific error (context/VRAM too large, or a local proxy/VPN intercepting the loopback probe) instead of the generic invalid-GGUF/out-of-memory message. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Hakan Baysal Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 61 +++++-- ..._llama_cpp_start_failure_classification.py | 10 ++ .../tests/test_llama_cpp_wait_for_health.py | 9 + .../test_load_progress_ready_fraction.py | 166 ++++++++++++++++++ 4 files changed, 236 insertions(+), 10 deletions(-) create mode 100644 studio/backend/tests/test_load_progress_ready_fraction.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8b984c0bfe..035e5d12c7 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1273,6 +1273,7 @@ class LlamaCppBackend: self._is_diffusion: bool = False self._diffusion_visual_bin: Optional[str] = None self._healthy = False + self._load_rss_hwm = (None, 0) # (pid, peak VmRSS) for load_progress self._stats_logger = None # vLLM-style engine-stats poller, set on load # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None @@ -1480,6 +1481,21 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + @staticmethod + def _read_rss_bytes(pid: int) -> Optional[int]: + """Resident set size of ``pid`` in bytes, from /proc//status (Linux). + 0 when the status has no VmRSS line (zombie / kernel thread); None where + /proc is unavailable (macOS/Windows) or the value is unreadable.""" + try: + with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: + for line in f: + if line.startswith("VmRSS:"): + # IndexError guards a "VmRSS:" line with no value column. + return int(line.split()[1]) * 1024 # kB -> bytes + except (FileNotFoundError, PermissionError, ValueError, IndexError, OSError): + return None + return 0 # readable but no VmRSS line + def load_progress(self) -> Optional[dict]: """Return live model-load progress, or None if not loading. @@ -1539,22 +1555,32 @@ class LlamaCppBackend: except OSError: pass - # Read VmRSS from /proc//status (kilobytes on Linux). - bytes_loaded = 0 - try: - with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: - for line in f: - if line.startswith("VmRSS:"): - kb = int(line.split()[1]) - bytes_loaded = kb * 1024 - break - except (FileNotFoundError, PermissionError, ValueError, OSError): + # VmRSS of the llama-server; None where /proc is unavailable. + bytes_loaded = LlamaCppBackend._read_rss_bytes(pid) + if bytes_loaded is None: return None + # RSS climbs as weights page in, then drops once -ngl offloads them to + # VRAM and the mmap pages are freed. Hold a per-process high-water mark + # so the bar never regresses to ~8% mid-load (#5740). + hwm_pid, hwm = getattr(self, "_load_rss_hwm", (None, 0)) + hwm = bytes_loaded if hwm_pid != pid else max(hwm, bytes_loaded) + self._load_rss_hwm = (pid, hwm) + bytes_loaded = hwm + phase = "ready" if self._healthy else "mmap" fraction = 0.0 if bytes_total > 0: fraction = min(1.0, bytes_loaded / bytes_total) + # Once llama-server is healthy the load is complete by definition. With + # layers offloaded to VRAM (-ngl) the process releases the mmap'd weight + # pages, so VmRSS sinks back well below the shard total; the raw RSS + # fraction would then report a partial (~8%) load indefinitely and freeze + # a fraction-driven progress bar even though the model is ready (#5740). + if self._healthy: + if bytes_total > 0: + bytes_loaded = bytes_total + fraction = 1.0 return { "phase": phase, "bytes_loaded": bytes_loaded, @@ -4232,6 +4258,17 @@ class LlamaCppBackend: "expected; otherwise check the llama-server log for the cause." ) + # A live server that never answered 200 on /health is not a bad GGUF: + # the load is too large for VRAM/context, or a local proxy/VPN grabbed + # the loopback probe (#5740). + if "health check timed out" in lowered: + return ( + "llama-server started but never became healthy on its local " + "/health endpoint. Try a smaller context length or a more " + "quantized GGUF, and if you use a VPN or HTTP proxy make sure " + "localhost bypasses it (NO_PROXY=127.0.0.1,localhost)." + ) + # Fallback: genuinely unknown failure (OOM, missing binary ...). return ( "llama-server failed to start. " @@ -7501,6 +7538,10 @@ class LlamaCppBackend: time.sleep(interval) + # Leave a marker so _classify_llama_start_failure tells a live but + # never-healthy load (too large, or a proxy hijacking the loopback + # probe) apart from a bad GGUF (#5740). + self._stdout_lines.append(f"llama-server health check timed out after {timeout}s") logger.error(f"llama-server health check timed out after {timeout}s") return False diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py index 6b26121cf8..246d810602 100644 --- a/studio/backend/tests/test_llama_cpp_start_failure_classification.py +++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py @@ -140,6 +140,16 @@ class TestOllamaAndFallback: msg = _classify("", None, None) assert "llama-server failed to start" in msg + def test_health_timeout_names_probe_not_generic(self): + # A live server that never returns 200 on /health must name the probe and + # proxy/context causes, not blame a bad GGUF (#5740). + msg = _classify( + "llama-server health check timed out after 600.0s", "/models/x.gguf", "local/x" + ) + assert "/health" in msg + assert "NO_PROXY" in msg + assert "GGUF file is valid" not in msg + class TestOsKillReturncode: """SIGKILL (-9) with no diagnostic output is the OOM killer and gets a named, diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py index 1ba6c9f7b5..82c5b4931a 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -67,6 +67,15 @@ class TestWaitForHealthResilience: monkeypatch.setattr(httpx, "get", lambda *a, **kw: ok_resp) assert b._wait_for_health(timeout = 1.0, interval = 0.01) is True + def test_timeout_records_marker_for_classification(self, monkeypatch): + """A live-but-never-healthy server leaves a marker so the failure is + classified as a /health timeout, not a bad GGUF (#5740).""" + b = _make_backend() + b._process.poll.return_value = None + monkeypatch.setattr(httpx, "get", lambda *a, **kw: mock.Mock(status_code = 503)) + assert b._wait_for_health(timeout = 0.02, interval = 0.01) is False + assert any("health check timed out" in ln for ln in b._stdout_lines) + def test_read_error_loops_to_subprocess_poll(self, monkeypatch): """WinError 10054 (httpx.ReadError) must be swallowed; the next iteration sees the dead subprocess and returns False with a structured exit-code log.""" b = _make_backend() diff --git a/studio/backend/tests/test_load_progress_ready_fraction.py b/studio/backend/tests/test_load_progress_ready_fraction.py new file mode 100644 index 0000000000..2e499cd8c6 --- /dev/null +++ b/studio/backend/tests/test_load_progress_ready_fraction.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""load_progress() must report a complete load once llama-server is healthy. + +With layers offloaded to VRAM (-ngl) the server releases the mmap'd weight pages +after upload, so its VmRSS sinks back well below the shard total. The raw RSS +fraction would then sit at a partial (~8%) value forever and freeze a +fraction-driven progress bar even though the model is ready -- the "stuck around +8% on the second pass" symptom in #5740. In the ready phase the fraction must be +1.0 regardless of resident set size. +""" + +from __future__ import annotations + +import io +import sys +import types +from pathlib import Path +from unittest.mock import patch + +import pytest + +# Stub heavy/unavailable deps before importing the module under test, so a +# targeted run in the lightweight backend env (no structlog/httpx) still +# collects. setdefault keeps the real modules when they are installed. Mirrors +# test_llama_cpp_load_progress_matrix.py. +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +sys.modules.setdefault("structlog", types.ModuleType("structlog")) + +_httpx_stub = types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + + +def _backend( + gguf_path, + *, + healthy, + pid = 4321, +): + # Bare instance: exercise load_progress() without the heavy real __init__. + be = object.__new__(LlamaCppBackend) + be._process = types.SimpleNamespace(pid = pid) + be._gguf_path = str(gguf_path) + be._healthy = healthy + return be + + +def _gguf(tmp_path, size_bytes): + f = tmp_path / "model-Q4_K_M.gguf" + f.write_bytes(b"\0" * size_bytes) + return f + + +def test_ready_reports_complete_despite_low_rss(tmp_path, monkeypatch): + # Healthy, but VmRSS has dropped to ~8% of the shard total after VRAM upload. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(_gguf(tmp_path, 10000), healthy = True) + p = be.load_progress() + assert p["phase"] == "ready" + assert p["fraction"] == 1.0 # not 0.08 + assert p["bytes_loaded"] == p["bytes_total"] == 10000 + + +def test_mmap_phase_reports_raw_rss_fraction(tmp_path, monkeypatch): + # Still loading: the bar should track real residency, not jump to 1.0. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(_gguf(tmp_path, 10000), healthy = False) + p = be.load_progress() + assert p["phase"] == "mmap" + assert p["fraction"] == 0.08 + assert p["bytes_loaded"] == 800 + assert p["bytes_total"] == 10000 + + +def test_progress_fraction_is_monotonic(tmp_path, monkeypatch): + # RSS peaks during page-in, then drops after -ngl offload; the bar must hold + # its high-water mark instead of collapsing back to ~8% (#5740). + be = _backend(_gguf(tmp_path, 10000), healthy = False) + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 9000)) + assert be.load_progress()["fraction"] == 0.9 + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + p = be.load_progress() + assert p["fraction"] == 0.9 + assert p["bytes_loaded"] == 9000 + + +def test_ready_without_shard_size_still_completes(tmp_path, monkeypatch): + # bytes_total unknown (file unstattable): fraction must still read complete. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(tmp_path / "missing.gguf", healthy = True) + p = be.load_progress() + assert p["phase"] == "ready" + assert p["fraction"] == 1.0 + assert p["bytes_total"] == 0 + + +def test_none_when_no_process(tmp_path): + be = _backend(_gguf(tmp_path, 10000), healthy = True) + be._process = None + assert be.load_progress() is None + + +def test_none_when_rss_unreadable(tmp_path, monkeypatch): + # /proc unavailable (macOS/Windows) or unreadable -> no progress payload. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: None)) + be = _backend(_gguf(tmp_path, 10000), healthy = False) + assert be.load_progress() is None + + +def test_read_rss_bytes_absent_pid_is_none(): + # A pid with no readable /proc entry (or no /proc at all) yields None, never + # raises. + assert LlamaCppBackend._read_rss_bytes(2**31 - 1) is None + + +def test_read_rss_bytes_valueless_line_is_none(): + # A "VmRSS:" line with no value column must not raise (IndexError) -> None. + def fake_open(path, *a, **kw): + if str(path).startswith("/proc/"): + return io.StringIO("Name:\ttest\nVmRSS:\n") + return open(path, *a, **kw) + + with patch("builtins.open", side_effect = fake_open): + assert LlamaCppBackend._read_rss_bytes(4321) is None + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason = "/proc is Linux-only") +def test_read_rss_bytes_reads_self_on_linux(): + rss = LlamaCppBackend._read_rss_bytes(__import__("os").getpid()) + assert isinstance(rss, int) and rss > 0 From 9fd4a503e841b9eb7aa7b19e840d9feb705add6c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 08:16:32 -0700 Subject: [PATCH 13/23] fast_generate: clear error for vLLM-style inputs when fast_inference=False (#6786) * fast_generate: clear error for vLLM-style inputs when fast_inference=False When fast_inference=False, fast_generate falls back to HuggingFace generate, and the wrapper already rejects vLLM-only usage (a sampling_params or lora_request kwarg, or a string prompt). A vLLM prompt dict ({'prompt':..., 'multi_modal_data':...}) or a SamplingParams passed positionally slipped through and hit transformers.generate, raising a cryptic 'SamplingParams object has no attribute update'. Detect both and raise the same clear 'only supported with fast_inference=True' error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fast_generate: also reject positional list of SamplingParams and list of vLLM prompt dicts Address review feedback: the slow-mode guard missed SamplingParams passed inside a positional list and a list of {"prompt": ...} dicts, both valid vLLM batched shapes that leaked into transformers.generate. Fold the checks into small predicates and extend the GPU-free test (now 7 reject + 3 pass). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test_fast_generate_slow_guard: expose assertions via a test_ function so pytest collects them The assertions lived in run(), only called from __main__, so pytest reported no tests collected and CI skipped the coverage. Rename to test_fast_generate_slow_guard; the standalone script entrypoint still works. * fast_generate: reject vLLM tokenized/embeds prompt dicts in the slow-mode guard vLLM also accepts prompt dicts keyed by prompt_token_ids or prompt_embeds, not just prompt/multi_modal_data. Those slipped past the slow-mode guard and fell through to HuggingFace generate with a cryptic error. Recognize all vLLM prompt-dict keys and add a TokensPrompt test case (now 8 reject + 3 pass). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fast_generate slow-mode guard: catch vLLM prompts= keyword form vLLM's generate names its first argument `prompts`, so a slow-mode call like fast_generate(prompts="hi") or prompts=[{"prompt": ...}] bypassed the guard and leaked into HuggingFace generate as an unexpected kwarg. Check kwargs["prompts"] with the same _is_vllm_prompt predicate and add two test cases. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fast_generate slow-mode guard: reject vLLM tokenized prompt kwargs vLLM's legacy call shape passes tokens as prompt_token_ids= (and prompt_embeds=), which are not HuggingFace generate arguments. In slow mode these bypassed the guard and leaked into HF generate as unexpected kwargs. Reject their presence with the same tokenize-first message and add a test case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fast_generate slow-mode guard: treat prompts= as vLLM-only prompts is a vLLM keyword, not a HuggingFace generate argument, so any value passed as prompts= (including a bare token-id list, which _is_vllm_prompt deliberately ignores for positional HF token ids) is a vLLM-style call. Reject prompts= / prompt_token_ids= / prompt_embeds= on presence, and keep the conservative _is_vllm_prompt check only for the positional arg. * fast_generate slow-mode guard: reject vLLM prompt kwargs on presence prompts / prompt_token_ids / prompt_embeds are vLLM-only keyword names that HuggingFace generate does not accept, so a defaulted call like prompts=None should raise the actionable slow-mode error instead of leaking a None kwarg into HF generate. Check membership in kwargs rather than a non-None value. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- tests/test_fast_generate_slow_guard.py | 95 ++++++++++++++++++++++++++ unsloth/models/_utils.py | 70 +++++++++++-------- 2 files changed, 136 insertions(+), 29 deletions(-) create mode 100644 tests/test_fast_generate_slow_guard.py diff --git a/tests/test_fast_generate_slow_guard.py b/tests/test_fast_generate_slow_guard.py new file mode 100644 index 0000000000..6bfc561e54 --- /dev/null +++ b/tests/test_fast_generate_slow_guard.py @@ -0,0 +1,95 @@ +"""GPU-free test for the fast_generate slow-mode guard in _utils.py. + +When fast_inference=False, model.fast_generate falls back to HuggingFace generate, so vLLM-only +inputs must be rejected with a clear message instead of leaking into transformers.generate. Covers +a string prompt, a vLLM {"prompt":..., "multi_modal_data":...} dict, SamplingParams passed both +positionally and as a kwarg, and a normal tokenized call passing through. +""" + +import ast, functools, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +UTILS = os.path.join(HERE, "unsloth", "models", "_utils.py") + + +def _load_factory(): + src = open(UTILS).read() + for node in ast.parse(src).body: + if isinstance(node, ast.FunctionDef) and node.name == "make_fast_generate_wrapper": + ns = {"functools": functools} + exec(ast.get_source_segment(src, node), ns) + return ns["make_fast_generate_wrapper"] + raise AssertionError("make_fast_generate_wrapper not found in _utils.py") + + +make_fast_generate_wrapper = _load_factory() + + +class _SamplingParams: + pass + + +_SamplingParams.__name__ = "SamplingParams" # match by class name, no vllm import needed + + +def _wrapper(): + state = {} + + def original_generate(*a, **k): + state["hit"] = True + return "ok" + + return make_fast_generate_wrapper(original_generate), state + + +def _rejects(fn, needle): + try: + fn() + except ValueError as e: + assert needle in str(e), str(e) + return True + raise AssertionError("expected ValueError") + + +def test_fast_generate_slow_guard(): + w, _ = _wrapper() + # reject every vLLM-only shape + assert _rejects(lambda: w("hello"), "fast_inference=True") + assert _rejects( + lambda: w({"prompt": "hi", "multi_modal_data": {"image": None}}), "fast_inference=True" + ) + assert _rejects(lambda: w(["a", "b"]), "fast_inference=True") + assert _rejects(lambda: w([{"prompt": "hi"}]), "fast_inference=True") # list of prompt dicts + assert _rejects( + lambda: w({"prompt_token_ids": [1, 2, 3]}), "fast_inference=True" + ) # vLLM TokensPrompt + assert _rejects(lambda: w(prompts = "hello"), "fast_inference=True") # vLLM `prompts` kwarg + assert _rejects( + lambda: w(prompts = [{"prompt": "hi"}]), "fast_inference=True" + ) # vLLM `prompts` kwarg list + assert _rejects( + lambda: w(prompt_token_ids = [1, 2, 3]), "fast_inference=True" + ) # vLLM legacy tokenized kwarg + assert _rejects( + lambda: w(prompts = [1, 2, 3]), "fast_inference=True" + ) # token-id list via vLLM-only `prompts` kwarg + assert _rejects( + lambda: w(prompts = None), "fast_inference=True" + ) # vLLM-only kwarg present even if None + assert _rejects(lambda: w({"prompt": "hi"}, _SamplingParams()), "sampling_params") + assert _rejects( + lambda: w({"prompt": "hi"}, [_SamplingParams()]), "sampling_params" + ) # list of SamplingParams + assert _rejects(lambda: w(sampling_params = object()), "sampling_params") + + # pass normal tokenized calls with no false positives + w, state = _wrapper() + assert w(input_ids = "TOKENS", max_new_tokens = 8) == "ok" and state.get("hit") + assert w([1, 2, 3], max_new_tokens = 8) == "ok" # positional token ids + assert w([], max_new_tokens = 8) == "ok" # empty positional + print("13 reject + 3 pass fast_generate slow-mode guard cases passed") + + +if __name__ == "__main__": + test_fast_generate_slow_guard() + print("OK: fast_generate rejects vLLM-style inputs when fast_inference=False") diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 599a5c0262..047783c35e 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3602,8 +3602,27 @@ def make_fast_generate_wrapper(original_generate): @functools.wraps(original_generate) def _fast_generate_wrapper(*args, **kwargs): - # Check for vLLM-specific arguments - if "sampling_params" in kwargs: + def _has_sampling_params(a): + # SamplingParams passed directly or inside a positional list/tuple + return type(a).__name__ == "SamplingParams" or ( + isinstance(a, (list, tuple)) + and any(type(i).__name__ == "SamplingParams" for i in a) + ) + + def _is_vllm_prompt(a): + # str prompt, a vLLM prompt dict (prompt / prompt_token_ids / prompt_embeds / + # multi_modal_data), or a list/tuple of those + head = a[0] if isinstance(a, (list, tuple)) and len(a) > 0 else a + return isinstance(head, str) or ( + isinstance(head, dict) + and any( + k in head + for k in ("prompt", "prompt_token_ids", "prompt_embeds", "multi_modal_data") + ) + ) + + # vLLM-only; also catch SamplingParams passed positionally (fast_generate(prompt, params)) + if "sampling_params" in kwargs or any(_has_sampling_params(a) for a in args): raise ValueError( "Unsloth: `sampling_params` is only supported when `fast_inference=True` (vLLM). " "Since `fast_inference=False`, use HuggingFace generate arguments instead:\n" @@ -3616,33 +3635,26 @@ def make_fast_generate_wrapper(original_generate): "Since `fast_inference=False`, LoRA weights are already merged into the model." ) - # Check if first positional argument is a string or list of strings - if len(args) > 0: - first_arg = args[0] - is_string_input = False - - if isinstance(first_arg, str): - is_string_input = True - elif isinstance(first_arg, (list, tuple)) and len(first_arg) > 0: - if isinstance(first_arg[0], str): - is_string_input = True - - if is_string_input: - raise ValueError( - "Unsloth: Passing text strings to `fast_generate` is only supported " - "when `fast_inference=True` (vLLM). Since `fast_inference=False`, you must " - "tokenize the input first:\n\n" - " messages = tokenizer.apply_chat_template(\n" - ' [{"role": "user", "content": "Your prompt here"}],\n' - " tokenize=True, add_generation_prompt=True,\n" - ' return_tensors="pt", return_dict=True\n' - " )\n" - " output = model.fast_generate(\n" - " **messages.to('cuda'),\n" - " max_new_tokens=64,\n" - " temperature=1.0,\n" - " )" - ) + # A vLLM-style prompt (string, {"prompt":..., "multi_modal_data":...} dict, or a list/tuple + # of either) only works under vLLM; tokenize first when fast_inference=False. A positional + # arg may be HF token ids, so check it conservatively with _is_vllm_prompt. The `prompts` / + # `prompt_token_ids` / `prompt_embeds` keywords are vLLM-only names that HuggingFace generate + # does not accept, so any of them being present is a vLLM-style call (even a bare token list, + # or an explicit None from a defaulted kwargs dict), hence membership rather than a value check. + vllm_prompt_kwarg = any( + k in kwargs for k in ("prompts", "prompt_token_ids", "prompt_embeds") + ) + if (len(args) > 0 and _is_vllm_prompt(args[0])) or vllm_prompt_kwarg: + raise ValueError( + "Unsloth: Passing vLLM-style prompts to `fast_generate` is only supported when " + "`fast_inference=True` (vLLM). Since `fast_inference=False`, tokenize first:\n\n" + " inputs = tokenizer.apply_chat_template(\n" + ' [{"role": "user", "content": "Your prompt here"}],\n' + " tokenize=True, add_generation_prompt=True,\n" + ' return_tensors="pt", return_dict=True,\n' + " )\n" + " output = model.fast_generate(**inputs.to('cuda'), max_new_tokens=64, temperature=1.0)" + ) # Call original generate return original_generate(*args, **kwargs) From b8400f40df39a3c006284cedfe31b8e4fb0a5131 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Fri, 3 Jul 2026 20:47:27 +0530 Subject: [PATCH 14/23] CLI: Rename unsloth connect to unsloth start (#6613) * replaced connect with start * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Studio: build the coding-agent command from the selected server The API keys panel showed a hardcoded `unsloth start claude`. `unsloth start` defaults to 127.0.0.1:8888 and only mints a key for a loopback server, so a non-default port or a tunnel/remote base would target the wrong server or fail to mint. Build the command from the panel base/key (and emit a key for non-loopback), matching the other snippets in the panel. * CLI: keep `unsloth connect` as a hidden alias for `unsloth start` Avoids breaking existing scripts and docs that still call `unsloth connect`. * Tests: stub _unstarted_cleanup in same-task disconnect test The test builds _SameTaskStreamingResponse via __new__, so set the attribute that __call__ now reads. * Match coding-agent command loopback check to the CLI 127.0.0.0/8 rule (#6613) * Keep unsloth_cli.commands.connect importable as a deprecated shim (#6613) * Format the new coding-agents panel strings and import per biome (#6613) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop the unsloth connect alias and shim; unsloth start is the only command (#6613) * Route unsloth connect to unsloth start as a hidden backward-compatible alias (#6613) * Forward unsloth run model-load flags to unsloth start (gguf-variant, context-length, load-in-4bit, tensor-parallel) (#6613) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Session-scope coding agent config in unsloth start Configure each agent for the current session instead of writing the Studio endpoint, key, and default model into the user's own config. Codex, OpenCode, OpenClaw, and Hermes get a private config relocated through their config-path env vars (CODEX_HOME, OPENCODE_CONFIG overlay, OPENCLAW_CONFIG_PATH plus OPENCLAW_STATE_DIR, HERMES_HOME). Claude Code suppresses the attribution header for the session via the CLAUDE_CODE_ATTRIBUTION_HEADER env var plus a --settings overlay, with no ~/.claude write. --launch uses an ephemeral temp dir removed after the agent exits; --no-launch uses a stable Unsloth-owned dir and prints the matching export lines. * Read relocated agent session config in Local Agent Guides CI The contract crosscheck and the openclaw/hermes patch helpers now read each agent's config from the relocated path printed by unsloth start --no-launch (CODEX_HOME, OPENCODE_CONFIG, OPENCLAW_CONFIG_PATH, HERMES_HOME) instead of fixed home paths. The Claude attribution A/B toggles the header for the session only (shipped-config HIT vs vanilla MISS) instead of editing ~/.claude/settings.json. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Skip the POSIX-only --no-launch parser test on Windows test_no_launch_output_is_parseable mirrors the #6547 bash CI parser, which greps export/unset lines and only runs on Linux/macOS runners. On Windows --no-launch prints PowerShell ($env: / Remove-Item), so the export-line assertion does not apply there. Cross-OS staging CI surfaced this. * Size Claude Code's auto-compact window to the loaded model's context Claude Code auto-compacts against its native (~600k token) window, so against a smaller local model it overflows the server's context (silent truncation) long before it compacts. Set CLAUDE_CODE_AUTO_COMPACT_WINDOW to the loaded model's real context length (the value codex/openclaw already get via model_context_window / contextWindow). Omitted when the model reports no context length. * Pin OpenCode/Hermes context window and set 90% compaction across agents Feed every agent the server-determined sequence length (the value /v1/models reports from runtime_context_length) and a ~90% compaction threshold. OpenCode: a custom-provider model with no limit defaults to context 0, which silently disables auto-compaction, so set limit.context/output and scale the compaction buffer to 10% of the window. Hermes: pin model.context_length (it otherwise falls back to a 256k default when the server's /v1/models omits the field) and set compression.threshold 0.9. Claude: add CLAUDE_AUTOCOMPACT_PCT_OVERRIDE=90 alongside the window. Codex (model_context_window) and OpenClaw (contextWindow) already carried the window and auto-manage off it. * Add `unsloth start pi` recipe Pi was the only agent without a built-in recipe, so the agent-guides CI hand-wrote ~/.pi/agent/models.json. Add a first-class `pi` command mirroring the others: - write_pi_config writes the session-scoped OpenAI-compatible provider config (key in the config, like openclaw/opencode). - pi() launches `pi --provider unsloth --model ` (Pi defaults to the google provider, so the provider/model are pinned on the command line) with HOME relocated for the session. Pi has no config-dir env var and resolves ~/.pi off $HOME, so HOME-scoping keeps the user's ~/.pi untouched. Migrate the agent-guides CI off the hand-written config onto the `unsloth start pi --no-launch` path (connection + file-edit), with a crosscheck for the provider api, so the documented recipe is exercised. * Harden unsloth start for Windows and WSL agent launches Address the Codex review on PR 6613: - write_pi_config now pins the loaded contextWindow and a sane maxTokens so Pi compacts instead of overflowing a small Studio context (it otherwise assumes its 128000 default), matching the other agents. - pi() sets USERPROFILE (and HOMEDRIVE/HOMEPATH when present) alongside HOME on native Windows, where Node resolves ~/.pi via USERPROFILE rather than HOME, so the session no longer reads or writes the user's real ~/.pi. - The WSLENV bridge flags path-valued vars with /p so a Windows npm shim under /mnt receives translated paths, while scalar vars (the numeric context window) pass through untranslated. WSLENV is deduped on the bare name. - _print_env prints the launch command with PowerShell-safe quoting so the inline --settings JSON survives copy-paste on native Windows --no-launch. Add tests for the WSLENV path flagging, PowerShell quoting, the Pi context window, and the Pi USERPROFILE relocation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Set CLAUDE_CODE_NO_FLICKER for the Claude session A local server streams in bursts, so Claude Code's full-screen TUI redraw flickers between tokens. Disable it for the session via CLAUDE_CODE_NO_FLICKER, alongside the other CLAUDE_CODE_* session env knobs. * Add a normalized --yolo flag routed to each agent's auto-approve mode It is easy to forget which agent spells "run tools without prompting" which way, so `unsloth start` now accepts all three spellings as one option (--yolo, --dangerously-skip-permissions, --dangerously-bypass-approvals-and-sandbox) and routes to the agent's own mechanism: - claude: --dangerously-skip-permissions - codex: --dangerously-bypass-approvals-and-sandbox - hermes: --yolo - pi: --approve (Pi's only approval gate is project trust) - opencode: a permission allow block in opencode.json (no CLI flag exists) - openclaw: tools.exec security=full / ask=off / host=gateway (no CLI flag exists) Because the option is parsed by `unsloth start`, the "wrong" spelling for an agent still routes correctly instead of leaking through to the agent and erroring. IS_SANDBOX is deliberately left unset for Claude so its root/sandbox safety gate still applies. Adds routing, cross-routing, and per-config tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix review findings: IPv6 loopback command, pi USERPROFILE under WSL, yolo guard From a 10-reviewer pass over the PR: - studio/frontend agent-command.ts: normalize bracketed IPv6 hosts. URL.hostname returns "[::1]" for http://[::1]:8888, which never matched the "::1" loopback checks, so the copied command embedded the placeholder API key for a local IPv6 server instead of the bare auto-minting command. Now [::1] is treated as loopback like the CLI's is_loopback_url, so the command matches the CLI contract. - pi(): also relocate USERPROFILE (and HOMEDRIVE/HOMEPATH) when running under WSL against a /mnt Windows shim, not just on native Windows. Windows Node resolves ~/.pi via USERPROFILE, and the WSLENV bridge translates the path, so pi no longer falls back to the user's real ~/.pi in that case. - _yolo_command_flags: use .get so a config-based agent (or a typo) yields no flag instead of a latent KeyError. Adds tests for the WSL pi USERPROFILE relocation, the yolo unmapped-agent guard, and that opencode/openclaw --yolo stays config-only (no argv flag). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix round-2 review findings: WSLENV /p upgrade, agent help text - _merge_wslenv now upgrades a user's pre-existing unflagged WSLENV entry (e.g. a bare HOME or USERPROFILE) to the path-translated form (HOME/p) instead of leaving it as-is, so a Windows agent shim under WSL receives the translated session path rather than the raw Linux path. - Generalize the `unsloth start` registration help to list all six agents (was only "Claude Code, Codex"). Adds a test for the WSLENV unflagged-entry upgrade. * Fix round-3 review findings: complete openclaw --yolo, refresh stale copy - openclaw --yolo now also writes the host approvals file (exec-approvals.json with defaults security=full / ask=off / askFallback=full) alongside the tools.exec config. OpenClaw gates tool execution on both layers (the stricter wins), so the config alone could still leave it prompting or denying. Mirrors `openclaw exec-policy preset yolo`. ask=off means nothing is ever prompted, so the runtime socket block is unnecessary. - Studio API panel copy: clarify that a local server auto-mints the key while a remote one embeds it in the command, and add pi to the swap hint. - Local Agent Guides CI: drop the stale "pi has no start.py recipe" note now that all six agents are driven via `unsloth start --no-launch`. Adds the openclaw approvals-file assertions and a no-yolo openclaw test. * start: parse claude --version with a regex so a format change does not drop optimization flags * start: offer to install a missing agent (prompt then run its install command) * start: auto-start a Studio server for --model when none is running, and stop it on exit * inference: surface an actionable message when llama-server cannot compile a tool grammar * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix review findings: kill the auto-started server tree on Windows; apply the tool-grammar message to the OpenAI passthrough too * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * start: split --model org/repo:variant so a running session is not evicted `unsloth start --model org/repo:QUANT` failed against an already-running Studio server and, worse, killed whatever model another session had loaded. /v1/models lists a loaded GGUF under its bare repo id (e.g. unsloth/Qwen3-1.7B-GGUF), so _resolve_model never matched the `:QUANT`-suffixed request. It then POSTed /api/inference/load with model_path=org/repo:QUANT, which (a) Hugging Face rejects ("Repo id must use alphanumeric chars, '-', '_' or '.'") and (b) evicts the model the other session was using, so a second 'unsloth start' in a new tmux/terminal tore down the first. Re-running the command then attached to the now-empty server, which is why it 'worked the second time'. Mirror the org/repo:QUANT -> org/repo + --gguf-variant QUANT shorthand that 'unsloth run' and llama.cpp already accept, splitting it in _connect before we match or serve. Matching now resolves against the loaded bare repo id (no spurious reload, no eviction), and any real load uses a valid repo id plus gguf_variant. An explicit --gguf-variant still wins; local paths and Windows drive letters pass through untouched. The auto-serve path likewise spawns 'unsloth run --model org/repo --gguf-variant QUANT'. * start: harden auth-key handling, codex teardown, and CI transcript redaction Three review findings: 1. CI could leak a live key. agent-guides-drive.sh printed the raw 'unsloth start --no-launch' transcript (which carries export UNSLOTH_API_KEY / ANTHROPIC_AUTH_TOKEN lines) to the Actions log on both the failure path and the success path before redact() ran. Add cat_redacted() and use it for those two prints, so the key is scrubbed on the way to the log while the on-disk file stays intact for the env parsing that follows. 2. Outages masqueraded as bad keys. _key_accepted caught a broad Exception and returned False, so a 5xx or timeout while checking a cached key looked like a rejection: it discarded a good key and minted extra ones (local) or reported 'no saved key' (remote). Only treat HTTP 401/403 as a rejection; let other errors propagate so a real outage surfaces. 3. Codex preflight could leave the auto-started server up. _require_gguf_for_codex runs after _connect may have auto-started Studio but before _run installs its teardown finally, so a preflight rejection (e.g. a transformers-backend model) left the server holding the port/GPU until the atexit backstop. Tear it down explicitly at the point of failure. Tests: a 5xx on a saved key surfaces without minting; a non-GGUF codex preflight tears down the auto-served server. * start: fix IPv6/portless studio URLs, Pi config-dir isolation, and Pi install recipe Four review findings: 1. Pi ignored the session config when PI_CODING_AGENT_DIR was already set. Pi's getAgentDir() reads process.env.PI_CODING_AGENT_DIR before falling back to $HOME/.pi/agent, so a value inherited from the user's shell sent Pi to their real config and skipped our provider/key (the HOME relocation alone was not enough). Pin PI_CODING_AGENT_DIR at the session's .pi/agent dir; it is path-valued so the WSL bridge translates it automatically. 2. Pi install hint dropped Pi's documented --ignore-scripts. Pi's README installs with 'npm install -g --ignore-scripts @earendil-works/pi-coding-agent' and notes it needs no install scripts, so accepting the prompt now follows that safe recipe. 3. Auto-start ignored a portless UNSLOTH_STUDIO_URL. unsloth run binds to 'parsed.port or 8888', so http://127.0.0.1 launched the child on 8888 but the health poll (and the returned base) still used port 80, stalling until the startup timeout. Normalize the base to host:8888 (IPv6-safe) before starting and polling. 4. API-panel command mistook IPv6 loopback for the bare default. The bare 'unsloth start' only probes 127.0.0.1:8888 on the IPv4 stack, so http://[::1]:8888 must carry an explicit UNSLOTH_STUDIO_URL. Drop ::1 from the bare-default host set while keeping it a loopback host (URL emitted, no key needed). Tests: PI_CODING_AGENT_DIR is set to the session dir; _effective_base normalizes portless/IPv6 bases; a portless UNSLOTH_STUDIO_URL auto-serves on :8888. * start: apply fresh-review findings across CLI, CI, and the API-panel command From a fresh multi-reviewer pass over the merged head plus the latest Codex bot review: 1. Load knobs now always consult the server. _resolve_model matched on model id alone, so --gguf-variant / --context-length / --no-load-in-4bit / --tensor-parallel were silently ignored whenever the id was already loaded (asking for UD-Q4_K_XL kept a Q8_0 serving). With any explicit knob the CLI defers to /api/inference/load, whose already-loaded dedup answers without reloading when variant and settings match, so a second session running the same command still attaches without evicting the first. 2. OpenCode --yolo and the session model pin now ride in OPENCODE_CONFIG_CONTENT. A project's own opencode.json outranks OPENCODE_CONFIG, so a repo config could silently override the session model and the --yolo permission block; OPENCODE_CONFIG_CONTENT outranks project config. The API key stays in the private file, never in printed env. 3. The --no-launch recipe's last line is a self-contained one-liner (inline VAR=value assignments before the command, conflicting vars blanked). People copy just the last line, and a bare codex/claude there ran against the user's real ~/.codex or Anthropic credentials with zero isolation, e.g. inheriting a pre-existing damaged ~/.codex state DB and blaming the recipe. The CI drive script scrubs the key from the one 'invoking:' echo this adds. 4. The auto-serve log is 0600 and the parent handle is closed. It sat world-readable in the shared tempdir under a predictable name while carrying the minted sk-unsloth- key from the unsloth run banner. 5. _key_accepted fails with a clean message on outages. Non-auth errors (5xx, network, timeout) surfaced as a raw traceback; 401/403 still mean a rejected key. 6. _effective_base strips URL paths, and https loopback targets never auto-serve. http://127.0.0.1:8888/studio polled /studio/api/health (404) and https://127.0.0.1 polled the wrong scheme, both spinning until the 15-minute startup timeout. 7. API-panel command: only literal 127.0.0.1:8888 earns the bare command. localhost can resolve to ::1, which the bare CLI never probes, so it keeps UNSLOTH_STUDIO_URL. 8. CI artifact sweep covers redacted-configs/ and agent-workdir/, not just logs/. Tests: 125 CLI tests pass (new coverage for each fix), 156 backend tests pass, ruff clean. Adds an unsloth connect alias regression test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * start: hand Pi a clean screen at launch Pi paints inline from wherever the cursor sits: its first render assumes a clean screen instead of clearing or entering the alternate screen itself (current Pi never emits a clear at startup). Launched under unsloth start, that left the session starting mid-scroll beneath the connection output. Clear the screen (click.clear, cross-platform, no-op without a TTY) right before the Studio banner so Pi opens exactly one line down on a clean viewport. Launch path only: --no-launch recipes and piped output are never wiped, and alternate-screen agents are left alone. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * start: auto-override hermes' 64K context floor for small model windows Hermes refuses to initialize when the served model's context window is under 64,000 tokens, and a second copy of the same check rejects the compression model mid-session. write_hermes_config previously pinned the real window, so any small local model (e.g. 40,960) failed at startup with manual config.yaml instructions. For windows below the floor the recipe now claims 65,536 in model.context_length, scales compression.threshold so compaction still fires at 90% of the real window, and sets auxiliary.compression.context_length to cover the mid-session check. Windows at or above the floor keep the exact previous behavior. * ci: install pi with --ignore-scripts, matching the start.py hint The pi cell predates the pi recipe in start.py and still installed the package with lifecycle scripts enabled, so CI stopped exercising the exact command users are prompted to run. npm_retry now passes extra flags through, the pi branch mirrors the install hint verbatim, and the stale no-recipe comment is refreshed. * ci: fail loudly when a relocation var is missing from connect output The empty-string guards ran after appending /config.toml or /config.yaml, so they could never fire: crosscheck_contract silently skipped its contract checks and patch_hermes_tools died on the root path with a bare traceback. Check the raw variable first and guide_fail with the real cause. * staging: 6613 round 6 (https elision, no-launch home reuse, auto-start key fallback) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: shimmyshimmer <107991372+shimmyshimmer@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- .github/scripts/agent-guides-drive.sh | 274 +-- .github/scripts/agent-guides-install.sh | 33 +- .github/scripts/serve-unsloth-run.sh | 2 +- .github/workflows/local-agent-guides-ci.yml | 50 +- studio/backend/routes/inference.py | 32 +- .../tests/test_openai_tool_passthrough.py | 27 + .../settings/components/agent-command.ts | 73 + .../settings/components/usage-examples.tsx | 41 + studio/frontend/src/i18n/locales/en.ts | 4 + unsloth_cli/__init__.py | 13 +- unsloth_cli/commands/connect.py | 777 ------- unsloth_cli/commands/start.py | 1497 +++++++++++++ unsloth_cli/tests/test_connect.py | 954 --------- unsloth_cli/tests/test_start.py | 1848 +++++++++++++++++ 14 files changed, 3722 insertions(+), 1903 deletions(-) create mode 100644 studio/frontend/src/features/settings/components/agent-command.ts delete mode 100644 unsloth_cli/commands/connect.py create mode 100644 unsloth_cli/commands/start.py delete mode 100644 unsloth_cli/tests/test_connect.py create mode 100644 unsloth_cli/tests/test_start.py diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 3c7cea919c..9b85b20177 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -6,12 +6,12 @@ # Local Agent Guides CI. All failures from here are failure class (c) # "guide drift": the server preflight already passed and the agent CLI # already installed, so a failure here means the documented recipe in -# unsloth_cli/commands/connect.py no longer produces a working flow. +# unsloth_cli/commands/start.py no longer produces a working flow. # -# Self-updating: for the 5 agents with a connect.py recipe we obtain the -# exact env + command from `unsloth connect --no-launch` and run -# THAT, so a recipe change is exercised automatically. Pi (no connect.py -# command at HEAD) is driven by a hand-written recipe. +# Self-updating: for all six agents (claude, codex, hermes, openclaw, +# opencode, pi) we obtain the exact env + command from +# `unsloth start --no-launch` and run THAT, so a recipe change is +# exercised automatically. # # Every agent invocation is wrapped in `timeout` so a headless-TTY prompt # can never hang the runner -- a timeout is reported as guide drift with a @@ -53,14 +53,14 @@ REDACTED_DIR="$REPO_ROOT/redacted-configs" WORKDIR_BASE="$REPO_ROOT/agent-workdir" CACHE_HELPER="$SCRIPT_DIR/assert-prompt-cache.sh" mkdir -p "$LOGS_DIR" "$REDACTED_DIR" -CONNECT_REF="unsloth_cli/commands/connect.py" +CONNECT_REF="unsloth_cli/commands/start.py" # Prefill-shrinking flags for Claude Code. The heavyweight agents send # multi-thousand-token system prompts + full tool schemas, which on a CPU-only # runner is minutes of prefill per model round-trip (~16 tok/s for a 4B model). # Replacing the ~5.7k default system prompt with a tiny one (--system-prompt-file) # and restricting tools cuts the prefill to a few hundred tokens so it completes -# quickly on CPU. These only shape the request size; the connect.py recipe +# quickly on CPU. These only shape the request size; the start.py recipe # (endpoint, auth, model) is still exercised end to end. # # The bulk of Claude Code's prompt is the built-in tool JSON schemas: measured @@ -105,6 +105,13 @@ redact() { done } +# Print a file to the log with the key scrubbed, without mutating it (the raw file is +# still needed to parse the real env). Use this instead of `cat` for any transcript that +# carries an `export UNSLOTH_API_KEY=...` line, so a live key never reaches Actions logs. +cat_redacted() { + sed "s#${UNSLOTH_API_KEY}##g" "$1" +} + # A reply must be non-empty and free of connection/auth errors. assert_reply() { local out="$1" @@ -131,45 +138,29 @@ run_timed() { # $1=outfile, rest=command return "$rc" } -# ── Pi: no connect.py command at HEAD -> hand-written recipe ────────────── -write_pi_config() { - if unsloth connect pi --help >/dev/null 2>&1; then - # Tripwire: once a real recipe exists, the hand-written config would mask any - # drift in it, defeating the point of this CI. Fail hard so the cell is - # migrated to the self-updating `unsloth connect pi --no-launch` path. - guide_fail "connect.py now ships a 'pi' command -- migrate this CI cell to the 'unsloth connect pi --no-launch' path so the documented recipe is exercised (the hand-written Pi config no longer reflects it)" - fi - mkdir -p "$HOME/.pi/agent" - python3 - "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" "$UNSLOTH_MODEL_ID" <<'PY' -import json, os, sys -base, key, model = sys.argv[1], sys.argv[2], sys.argv[3] -cfg = {"providers": {"unsloth": { - "api": "openai-completions", - "baseUrl": f"{base}/v1", - "apiKey": key, - "models": [{"id": model}], -}}} -path = os.path.expanduser("~/.pi/agent/models.json") -with open(path, "w") as fh: - json.dump(cfg, fh, indent=2) -PY - cp "$HOME/.pi/agent/models.json" "$REDACTED_DIR/pi-models.json" 2>/dev/null || true - redact "$REDACTED_DIR/pi-models.json" +# Read a value from an `export VAR=...` line in the connect --no-launch output. +# `unsloth start` writes each agent's session config off the user's ~ and points +# at it through a relocation env var (CODEX_HOME / OPENCODE_CONFIG / +# OPENCLAW_CONFIG_PATH), so the contract checks read the path from here. +raw_env() { # $1 = var name -> value (one shlex-quote layer stripped) + local raw="$LOGS_DIR/connect-${AGENT}.txt" + local v; v="$(sed -n "s/^export $1=//p" "$raw" | tail -1)" + v="${v#\'}"; v="${v%\'}"; printf '%s' "$v" } -# ── 5-agent connect.py path: parse env + command from --no-launch ───────── +# ── 5-agent start.py path: parse env + command from --no-launch ───────── # Populates globals CONNECT_ENV (export/unset lines) and CONNECT_CMD (the -# launch command on the last printed line), and runs connect.py's config -# writers as a side effect (it writes ~/.codex, ~/.claude, etc.). +# launch command on the last printed line), and runs start.py's config +# writers as a side effect (it writes each agent's relocated session config). parse_connect() { local raw="$LOGS_DIR/connect-${AGENT}.txt" - if ! unsloth connect "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then - cat "$raw" - guide_fail "'unsloth connect ${AGENT} --no-launch' exited non-zero" + if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + cat_redacted "$raw" + guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero" fi - echo "[$AGENT] connect --no-launch printed:"; cat "$raw" + echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw" CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)" - # The launch command is the last non-export, non-status line. connect.py + # The launch command is the last non-export, non-status line. start.py # prints "Studio · model " and "Updated ..." status lines first. CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \ | grep -E '[^[:space:]]' | tail -1)" @@ -177,45 +168,63 @@ parse_connect() { redact "$raw" } -# Cross-check the documented contract knobs so silent connect.py changes +# Cross-check the documented contract knobs so silent start.py changes # (env-var rename, wire_api flip, attribution setting drop) also fail/flag. crosscheck_contract() { local raw="$LOGS_DIR/connect-${AGENT}.txt" + local cfg home case "$AGENT" in codex) grep -q 'UNSLOTH_STUDIO_AUTH_TOKEN' "$raw" \ - || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (connect.py _CODEX_ENV_KEY)" - if [ -f "$HOME/.codex/config.toml" ]; then - grep -q 'wire_api = "responses"' "$HOME/.codex/config.toml" \ - || guide_fail "Codex wire_api is no longer \"responses\" in ~/.codex/config.toml" - cp "$HOME/.codex/config.toml" "$REDACTED_DIR/codex-config.toml" + || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (start.py _CODEX_ENV_KEY)" + home="$(raw_env CODEX_HOME)" + # An empty relocation var would make cfg "/config.toml" and silently + # skip the [ -f ] contract check below; fail loudly instead. + [ -n "$home" ] || guide_fail "CODEX_HOME missing from connect output (start.py codex())" + cfg="$home/config.toml" + if [ -f "$cfg" ]; then + grep -q 'wire_api = "responses"' "$cfg" \ + || guide_fail "Codex wire_api is no longer \"responses\" in \$CODEX_HOME/config.toml" + cp "$cfg" "$REDACTED_DIR/codex-config.toml" fi grep -q 'codex --oss --profile unsloth_api' "$raw" \ || echo "::warning::Codex launch command changed from 'codex --oss --profile unsloth_api'" ;; claude) grep -q 'ANTHROPIC_AUTH_TOKEN' "$raw" \ - || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (connect.py claude())" - if [ -f "$HOME/.claude/settings.json" ]; then - grep -q '"CLAUDE_CODE_ATTRIBUTION_HEADER"' "$HOME/.claude/settings.json" \ - || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER not written to ~/.claude/settings.json (ensure_claude_attribution_header)" - cp "$HOME/.claude/settings.json" "$REDACTED_DIR/claude-settings.json" - fi + || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (start.py claude())" + grep -q 'CLAUDE_CODE_ATTRIBUTION_HEADER' "$raw" \ + || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER no longer set for the session (start.py claude())" ;; hermes) grep -q 'UNSLOTH_API_KEY' "$raw" \ - || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (connect.py _HERMES_ENV_KEY)" - [ -f "$HOME/.hermes/config.yaml" ] && cp "$HOME/.hermes/config.yaml" "$REDACTED_DIR/hermes-config.yaml" + || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (start.py _HERMES_ENV_KEY)" + home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "HERMES_HOME missing from connect output (start.py hermes())" + cfg="$home/config.yaml" + [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/hermes-config.yaml" ;; openclaw) - if [ -f "$HOME/.openclaw/openclaw.json" ]; then - grep -q '"openai-completions"' "$HOME/.openclaw/openclaw.json" \ + cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + if [ -n "$cfg" ] && [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ || echo "::warning::OpenClaw provider api is no longer 'openai-completions' (write_openclaw_config)" - cp "$HOME/.openclaw/openclaw.json" "$REDACTED_DIR/openclaw.json" + cp "$cfg" "$REDACTED_DIR/openclaw.json" fi ;; opencode) - [ -f "$HOME/.config/opencode/opencode.json" ] && cp "$HOME/.config/opencode/opencode.json" "$REDACTED_DIR/opencode.json" + cfg="$(raw_env OPENCODE_CONFIG)" + [ -n "$cfg" ] && [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/opencode.json" + ;; + pi) + # Pi has no config-dir env var; the session is HOME-relocated, and the + # provider config lives at $HOME/.pi/agent/models.json. + cfg="$(raw_env HOME)/.pi/agent/models.json" + if [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ + || echo "::warning::Pi provider api is no longer 'openai-completions' (write_pi_config)" + cp "$cfg" "$REDACTED_DIR/pi-models.json" + fi ;; esac redact "$REDACTED_DIR"/* 2>/dev/null || true @@ -229,16 +238,23 @@ crosscheck_contract() { # Hermes: an explicit empty cli toolset disables all tools (and drops the # tool-gated guidance blocks), so -z sends ~300 tokens instead of thousands. -# hermes ships a DEFAULT config.yaml that already has a populated -# platform_toolsets, and `unsloth connect` merges into it, so we must override -# cli (not just append). That needs a YAML parser, and the runner's bare -# python3 has no PyYAML -- but the venv that ships `unsloth` does (connect.py -# imports yaml), so run the patch with that interpreter. +# Hermes enables its default cli toolset when the session config does not pin one, +# so we must set platform_toolsets.cli explicitly to [] (not just append) to get +# zero tools. That needs a YAML parser, and the runner's bare python3 has no +# PyYAML -- but the venv that ships `unsloth` does (start.py imports yaml), so run +# the patch with that interpreter. We patch the relocated $HERMES_HOME/config.yaml +# that `unsloth start` printed, not the user's ~/.hermes. # (-z reads platform_toolsets.cli; --ignore-rules is a no-op under -z.) patch_hermes_tools() { # $1 = none|default + # Check the raw var BEFORE appending /config.yaml: the joined path is never + # empty, so the old guard could not fire and the patcher would die on + # "/config.yaml" with a bare traceback instead of this clear failure. + local home; home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "Hermes HERMES_HOME missing from connect output (start.py hermes())" + local cfg; cfg="$home/config.yaml" # Find a python that can import yaml. The runner's bare python3 cannot, but the # interpreter in the `unsloth` console-script shebang provably can (it runs - # connect.py's write_hermes_config, which imports yaml). Try that first, then + # start.py's write_hermes_config, which imports yaml). Try that first, then # any python on PATH, then the venv sibling, picking the first with PyYAML. local cand py="" shebang shebang="$(head -1 "$(command -v unsloth)" 2>/dev/null | sed -n 's/^#![[:space:]]*//p' | awk '{print $1}')" @@ -247,13 +263,13 @@ patch_hermes_tools() { # $1 = none|default { [ -x "$cand" ] || command -v "$cand" >/dev/null 2>&1; } || continue if "$cand" -c 'import yaml' 2>/dev/null; then py="$cand"; break; fi done - [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch ~/.hermes/config.yaml" - echo "[hermes] patching config with $py" - "$py" - "$1" <<'PY' + [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch the hermes session config" + echo "[hermes] patching $cfg with $py" + "$py" - "$1" "$cfg" <<'PY' import os, sys import yaml mode = sys.argv[1] -p = os.path.expanduser("~/.hermes/config.yaml") +p = sys.argv[2] cfg = (yaml.safe_load(open(p)) or {}) if os.path.exists(p) else {} ts = cfg.get("platform_toolsets") if not isinstance(ts, dict): @@ -274,10 +290,14 @@ PY # drop the auto-injected AGENTS.md/SOUL.md bootstrap (the bulk of the prompt) for # both modes. --agent must reference a defined agent, so write it before invoking. patch_openclaw_agent() { # $1 = notools|tools - python3 - "$1" <<'PY' + # OpenClaw reads its config from the relocated OPENCLAW_CONFIG_PATH that + # `unsloth start` printed, so patch THAT file (not the user's ~/.openclaw). + local cfg; cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + [ -n "$cfg" ] || guide_fail "OpenClaw OPENCLAW_CONFIG_PATH missing from connect output (start.py openclaw())" + python3 - "$1" "$cfg" <<'PY' import os, sys, json mode = sys.argv[1] -p = os.path.expanduser("~/.openclaw/openclaw.json") +p = sys.argv[2] cfg = json.load(open(p)) if os.path.exists(p) else {} agents = cfg.setdefault("agents", {}) agents.setdefault("defaults", {})["skipBootstrap"] = True @@ -293,20 +313,24 @@ print(f"[openclaw] agent ci tools = {agent.get('tools', 'default')}") PY } -# Build an invoke script that applies connect.py's env then runs the launch +# Build an invoke script that applies start.py's env then runs the launch # command (with extra args appended) under bash. We do NOT eval connect's env # into this shell; we write it into a one-shot script so the export/unset -# semantics are exactly what connect.py printed. The script path is absolute +# semantics are exactly what start.py printed. The script path is absolute # so it is valid even when the caller has cd'd into a scratch work dir. invoke_via_connect() { # $1=outfile, rest=extra args appended to the command local out="$1"; shift local script="$LOGS_DIR/invoke-${AGENT}.sh" local real; real="$(mktemp)" + # CONNECT_ENV_EXTRA / CONNECT_CMD_OVERRIDE let a caller (attribution-ab) flip a + # session knob without editing the user's config; empty -> use what start.py emitted. + local cmd="${CONNECT_CMD_OVERRIDE:-$CONNECT_CMD}" { echo "set -uo pipefail" echo "$CONNECT_ENV" + [ -n "${CONNECT_ENV_EXTRA:-}" ] && echo "$CONNECT_ENV_EXTRA" # Append extra args (the prompt / flags) to the launch command verbatim. - printf '%s' "$CONNECT_CMD" + printf '%s' "$cmd" local a for a in "$@"; do printf ' %q' "$a"; done printf '\n' @@ -318,7 +342,9 @@ invoke_via_connect() { # $1=outfile, rest=extra args appended to the command # Writing the redacted copy up front keeps the key out of the artifact even if # the run times out (run_timed exits before returning here). cp "$real" "$script"; redact "$script" - echo "[$AGENT] invoking (timeout ${TIMEOUT}s): $CONNECT_CMD $*" + # The connect one-liner now carries the key as an inline env assignment; scrub it on + # the way to the log (the executed $real keeps the live value). + echo "[$AGENT] invoking (timeout ${TIMEOUT}s): ${cmd//${UNSLOTH_API_KEY}/} $*" run_timed "$out" bash "$real" local rc=$? rm -f "$real" @@ -332,27 +358,23 @@ case "$MODE" in connection) PROMPT='Reply with exactly the single word: pong' OUT="$LOGS_DIR/${AGENT}-connection.txt" - if [ "$AGENT" = "pi" ]; then - write_pi_config - run_timed "$OUT" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$PROMPT" - else - parse_connect - crosscheck_contract - # claude/codex run in print mode via the flags connect.py emits - # (claude -p / codex exec). For agents whose default subcommand prints - # to stdout we pass the prompt through ctx.args. - case "$AGENT" in - claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; - codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; - opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; - hermes) patch_hermes_tools none - invoke_via_connect "$OUT" -z "$PROMPT" ;; - openclaw) patch_openclaw_agent notools - invoke_via_connect "$OUT" agent --local --agent ci \ - --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; - *) invoke_via_connect "$OUT" "$PROMPT" ;; - esac - fi + parse_connect + crosscheck_contract + # claude/codex run in print mode via the flags start.py emits + # (claude -p / codex exec). For agents whose default subcommand prints + # to stdout we pass the prompt through ctx.args. + case "$AGENT" in + claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; + codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; + opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; + pi) invoke_via_connect "$OUT" -p "$PROMPT" ;; + hermes) patch_hermes_tools none + invoke_via_connect "$OUT" -z "$PROMPT" ;; + openclaw) patch_openclaw_agent notools + invoke_via_connect "$OUT" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; + *) invoke_via_connect "$OUT" "$PROMPT" ;; + esac # A non-zero exit from the documented launch command is drift even if it # printed something: a benign-looking "command not found" / usage dump would # otherwise slip past assert_reply (which only flags empty/error-keyword text). @@ -371,22 +393,18 @@ case "$MODE" in T1='Create a file named hello.py in the current directory whose entire contents are a single line: print("Hello"). Do not run it.' T2='Run hello.py with python and show me the exact output.' - # The connect.py recipe writers + crosscheck must see the repo; run them + # The start.py recipe writers + crosscheck must see the repo; run them # from the repo root BEFORE cd-ing into the scratch work dir. - if [ "$AGENT" != "pi" ]; then - parse_connect - crosscheck_contract - # File-edit needs real tools, so we cannot zero them as in connection. - # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md - # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work - # dir is empty, so no project context files are auto-loaded either. - case "$AGENT" in - hermes) patch_hermes_tools default ;; - openclaw) patch_openclaw_agent tools ;; - esac - else - write_pi_config - fi + parse_connect + crosscheck_contract + # File-edit needs real tools, so we cannot zero them as in connection. + # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md + # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work + # dir is empty, so no project context files are auto-loaded either. + case "$AGENT" in + hermes) patch_hermes_tools default ;; + openclaw) patch_openclaw_agent tools ;; + esac # Drive from inside the work dir so the agent edits files there. All log # writes use absolute $LOGS_DIR, so cwd does not matter for them. @@ -395,7 +413,14 @@ case "$MODE" in invoke_turn() { # $1=outfile $2=continue? $3=prompt local out="$1" cont="$2" prompt="$3" case "$AGENT" in - pi) run_timed "$out" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$prompt" ;; + pi) + # Pi continues the previous session with -c; provider/model come from + # the parsed `unsloth start pi` recipe (CONNECT_CMD), not hardcoded here. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" -p --continue "$prompt" + else + invoke_via_connect "$out" -p "$prompt" + fi ;; claude) # --dangerously-skip-permissions lets headless claude actually use the # Write/Bash tools (otherwise it blocks on an approval prompt and emits @@ -466,33 +491,32 @@ case "$MODE" in # right before the measured turn, so an earlier turn's reuse can't leak in. LLAMA_LOG_DIR="${UNSLOTH_LLAMA_LOG_DIR:-$HOME/.unsloth/studio/logs/llama-server}" export LLAMA_LOG_DIR - parse_connect # writes ~/.claude/settings.json (header=0) + env + parse_connect # prints session env + suppression flags (no ~/.claude write) crosscheck_contract PROMPT='Reply with exactly the single word: pong' - # Phase A: header DISABLED (=0, the documented setting) -> expect a HIT on - # the continued turn. connect.py's ensure_claude_attribution_header() set 0. + # Phase A: the suppression start.py ships (CLAUDE_CODE_ATTRIBUTION_HEADER=0 + + # --exclude-dynamic-system-prompt-sections + --settings overlay) -> expect a + # HIT on the continued turn, since the system-prompt prefix is stable. invoke_via_connect "$LOGS_DIR/claude-ab-hit-1.txt" -p "$PROMPT" # turn 1 primes FROM_HIT="$(bash "$CACHE_HELPER" mark)" # offset before turn 2 invoke_via_connect "$LOGS_DIR/claude-ab-hit-2.txt" -p --continue "$PROMPT again" CACHE_LOG_FROM="$FROM_HIT" bash "$CACHE_HELPER" log HIT - # Phase B: header ENABLED -> expect a MISS. The header prepends a - # per-request-changing attribution line to the system prompt, so the shared - # prefix changes every turn and the KV cache is invalidated (~90% slower); - # this is exactly what the guide flag prevents. - python3 - <<'PY' -import json, os -p = os.path.expanduser("~/.claude/settings.json") -s = json.load(open(p)) if os.path.exists(p) else {} -s.setdefault("env", {})["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "1" -json.dump(s, open(p, "w"), indent=2) -PY + # Phase B: vanilla Claude with the header ENABLED -> expect a MISS. We flip + # the env var to 1 and strip the suppression flags from the launch command + # (without them the dynamic attribution line is included and changes every + # turn, so the shared prefix moves and the KV cache is invalidated, ~90% + # slower). This is session-only: nothing is written to ~/.claude. + CONNECT_ENV_EXTRA='export CLAUDE_CODE_ATTRIBUTION_HEADER=1' + CONNECT_CMD_OVERRIDE="$(printf '%s' "$CONNECT_CMD" \ + | sed -E "s/ --exclude-dynamic-system-prompt-sections//; s/ --settings '[^']*'//")" invoke_via_connect "$LOGS_DIR/claude-ab-miss-1.txt" -p "$PROMPT" FROM_MISS="$(bash "$CACHE_HELPER" mark)" invoke_via_connect "$LOGS_DIR/claude-ab-miss-2.txt" -p --continue "$PROMPT again" CACHE_LOG_FROM="$FROM_MISS" bash "$CACHE_HELPER" log MISS - echo "[claude] attribution A/B OK (header=0 HIT, header=1 MISS)" + unset CONNECT_ENV_EXTRA CONNECT_CMD_OVERRIDE + echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" ;; *) diff --git a/.github/scripts/agent-guides-install.sh b/.github/scripts/agent-guides-install.sh index dfab8aec80..daf4bacd3e 100755 --- a/.github/scripts/agent-guides-install.sh +++ b/.github/scripts/agent-guides-install.sh @@ -7,7 +7,7 @@ # is the single biggest source of false reds, so installs retry with # backoff and the only ::error:: this script can emit is class (b). The # install recipes mirror the install_hint strings in -# unsloth_cli/commands/connect.py at HEAD. +# unsloth_cli/commands/start.py at HEAD. # # Usage: agent-guides-install.sh # agent in: claude codex hermes openclaw opencode pi @@ -25,13 +25,14 @@ install_fail() { } # npm registry flakiness is common in CI; retry 3x with linear backoff. +# Extra npm flags may precede the package (e.g. npm_retry --ignore-scripts pkg). npm_retry() { - local pkg="$1" i + local i for i in 1 2 3; do - if npm install -g "$pkg" >> "$LOG" 2>&1; then + if npm install -g "$@" >> "$LOG" 2>&1; then return 0 fi - echo "[install] npm install -g $pkg attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + echo "[install] npm install -g $* attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" sleep "$((i * 10))" done return 1 @@ -60,30 +61,30 @@ curl_bash() { echo "[install] agent=$AGENT (log=$LOG)" case "$AGENT" in claude) - # connect.py install_hint: curl -fsSL https://claude.ai/install.sh | bash + # start.py install_hint: curl -fsSL https://claude.ai/install.sh | bash curl_bash "https://claude.ai/install.sh" || install_fail "claude installer failed" # The installer drops the binary under ~/.local/bin. echo "$HOME/.local/bin" >> "$GITHUB_PATH" ;; codex) - # connect.py install_hint: npm install -g @openai/codex + # start.py install_hint: npm install -g @openai/codex npm_retry "@openai/codex" || install_fail "npm install -g @openai/codex failed" ;; opencode) - # connect.py install_hint: npm install -g opencode-ai + # start.py install_hint: npm install -g opencode-ai npm_retry "opencode-ai" || install_fail "npm install -g opencode-ai failed" ;; openclaw) - # connect.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash + # start.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash # npm is the more deterministic path in CI and matches the agent's docs; - # fall back to the connect.py curl installer if the npm tag is missing. + # fall back to the start.py curl installer if the npm tag is missing. if ! npm_retry "openclaw@latest"; then curl_bash "https://openclaw.ai/install.sh" || install_fail "openclaw install failed (npm + curl)" echo "$HOME/.local/bin" >> "$GITHUB_PATH" fi ;; hermes) - # connect.py install_hint: + # start.py install_hint: # curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash curl_bash "https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh" \ --non-interactive --skip-setup --skip-browser --no-skills \ @@ -91,11 +92,13 @@ case "$AGENT" in echo "$HOME/.local/bin" >> "$GITHUB_PATH" ;; pi) - # No connect.py recipe; the agent's documented package name. The CLI moved - # from the now-deprecated @mariozechner scope to @earendil-works (the old - # scope is frozen, so installing it would test a stale Pi against the API). - npm_retry "@earendil-works/pi-coding-agent" \ - || install_fail "npm install -g @earendil-works/pi-coding-agent failed" + # start.py install_hint: npm install -g --ignore-scripts @earendil-works/pi-coding-agent + # (--ignore-scripts matches Pi's documented recipe; exercising the exact hint + # catches guide drift). The CLI moved from the now-deprecated @mariozechner + # scope to @earendil-works (the old scope is frozen, so installing it would + # test a stale Pi against the API). + npm_retry --ignore-scripts "@earendil-works/pi-coding-agent" \ + || install_fail "npm install -g --ignore-scripts @earendil-works/pi-coding-agent failed" ;; *) install_fail "unknown agent '$AGENT'" diff --git a/.github/scripts/serve-unsloth-run.sh b/.github/scripts/serve-unsloth-run.sh index 34b8b962c6..6ac98ded7c 100755 --- a/.github/scripts/serve-unsloth-run.sh +++ b/.github/scripts/serve-unsloth-run.sh @@ -27,7 +27,7 @@ # # Outputs written to $GITHUB_ENV (and echoed): # UNSLOTH_API_KEY the sk-unsloth-* key minted on the banner -# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth connect` +# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth start` # finds THIS server, not the hardcoded :8888) # UNSLOTH_BASE_URL same as UNSLOTH_STUDIO_URL (alias for clarity) # UNSLOTH_MODEL_ID the canonical id reported by /v1/models diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 299ee3f18b..47f75dc1ba 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -6,29 +6,27 @@ # Detects when our local-agent setup recipes drift out of sync with # `unsloth run`. Boots a real `unsloth run --disable-tools` server and # drives the coding agents end to end through the *exact* recipes defined -# in unsloth_cli/commands/connect.py (the in-repo source of truth -- there -# is no docs/ tree). Wherever connect.py has a recipe we drive the agent -# via `unsloth connect --no-launch` and execute what it prints, so -# the test self-updates against connect.py and catches silent recipe drift. +# in unsloth_cli/commands/start.py (the in-repo source of truth -- there +# is no docs/ tree). Wherever start.py has a recipe we drive the agent +# via `unsloth start --no-launch` and execute what it prints, so +# the test self-updates against start.py and catches silent recipe drift. # # Source-of-truth files this workflow guards: -# unsloth_cli/commands/connect.py the `unsloth connect ` recipes +# unsloth_cli/commands/start.py the `unsloth start ` recipes # unsloth_cli/commands/studio.py the `unsloth run` banner (API Key line) # # Failure taxonomy (each surfaced with a distinct ::error:: + the agent name -# + the connect.py location, so a red X is immediately triageable): +# + the start.py location, so a red X is immediately triageable): # (a) Unsloth server/API regression -- the dialect HTTP preflight fails # BEFORE the agent runs (or the server never becomes healthy). # (b) Agent package install failed -- npm/curl install of the CLI failed. # (c) Guide drift -- preflight passed + install ok, but -# the documented `unsloth connect` flow produced no/garbled output. +# the documented `unsloth start` flow produced no/garbled output. # # Agents covered (6): claude, codex, hermes, openclaw, opencode, pi. -# - claude/codex/hermes/openclaw/opencode have a connect.py recipe. -# - pi has NO `unsloth connect pi` command in connect.py at HEAD; it is -# driven by a hand-written recipe and the matrix cell asserts that the -# missing connect recipe is the (known) reason, so the day connect.py -# grows a `pi` command this cell flips to the self-updating path. +# - All six have a `unsloth start ` recipe, so each cell obtains its +# env + command from `unsloth start --no-launch` and runs THAT +# (self-updating: a recipe change is exercised automatically). name: Local Agent Guides CI @@ -83,7 +81,7 @@ jobs: # ═════════════════════════════════════════════════════════════════════ # Job 1: connection # Per-agent: serve gemma-3-270m, HTTP-preflight the agent's dialect, - # install the agent, run `unsloth connect --no-launch`, execute + # install the agent, run `unsloth start --no-launch`, execute # the emitted recipe with a trivial prompt, assert a non-empty reply. # Runs on PR + weekly + dispatch. Each matrix cell is its own runner so # it serves exactly one model on its own port. @@ -103,7 +101,9 @@ jobs: env: # gemma-4-E4B (128K context, capable enough to drive every agent for a # trivial reply; the 270m model produced empty/failed responses for - # codex/openclaw and is below hermes' 64K context floor). Served as a flat + # codex/openclaw). Hermes' 64K context floor no longer constrains the model + # choice: write_hermes_config claims the floor for smaller windows and + # scales compaction back to the real window. Served as a flat # GGUF file (the -MTP- repo ships no separate draft, so this is plain 4B). GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf @@ -209,7 +209,7 @@ jobs: ;; *) # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). - # OpenClaw's connect.py recipe writes an "openai-completions" + # OpenClaw's start.py recipe writes an "openai-completions" # provider (write_openclaw_config), so it uses this path, not # /v1/messages. code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ @@ -227,13 +227,13 @@ jobs: AGENT: ${{ matrix.agent }} run: bash .github/scripts/agent-guides-install.sh "$AGENT" - # ── (c) drive the agent via connect.py and assert a reply ────────── - # For the 5 agents with a connect.py recipe we run - # `unsloth connect --no-launch`, eval its env/unset exports, + # ── (c) drive the agent via start.py and assert a reply ────────── + # For the 5 agents with a start.py recipe we run + # `unsloth start --no-launch`, eval its env/unset exports, # then run the printed command with a hard timeout (no headless-TTY # hang). Pi has no connect recipe, so it is driven by hand and the # cell asserts that absence is the (known) reason. - - name: Drive ${{ matrix.agent }} via unsloth connect (class-c isolation) + - name: Drive ${{ matrix.agent }} via unsloth start (class-c isolation) env: AGENT: ${{ matrix.agent }} run: bash .github/scripts/agent-guides-drive.sh connection "$AGENT" @@ -248,8 +248,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi @@ -438,8 +440,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi @@ -582,8 +586,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9d9db83543..a948a6eaf5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -165,6 +165,26 @@ def _friendly_error(exc: Exception) -> str: return "An internal error occurred" +def _friendly_upstream_error(text: str) -> str: + """Rewrite a raw llama-server error body into an actionable message where we can. + + The main case is a tool-calling grammar that llama-server can't compile ("failed to + parse grammar" / "failed to initialize samplers"). This surfaces to coding agents as + a hard 400 on every tool-bearing turn. It is a llama-server limitation with some + model/quant + tool-schema combinations, and recent llama.cpp builds handle the common + coding-agent tools, so point the user at updating Studio rather than the raw body. + """ + lowered = text.lower() + if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered: + return ( + "The model couldn't compile a tool-calling grammar for this request. This is a " + "llama-server limitation with some model/quant and tool-schema combinations. " + "Update Studio (it installs the latest llama.cpp, which handles the common " + "coding-agent tools) or try a different GGUF model." + ) + return f"llama-server error: {text}" + + def _clamp_finish_reason(value) -> str: """Coerce an upstream finish_reason into OpenAI's known chat values. @@ -278,8 +298,8 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException": """HTTPException for a non-200 upstream response on the OpenAI passthrough (tools / response_format). An over-context upstream error is mapped to a 400 with code="context_length_exceeded" so these paths deliver the same signal as - the non-passthrough path; any other upstream error keeps llama-server's - message verbatim.""" + the non-passthrough path; a tool-grammar compile failure gets the same actionable + guidance as the Anthropic passthrough; any other upstream error stays verbatim.""" if _classify_llama_generation_error(Exception(text)): return HTTPException( status_code = 400, @@ -292,7 +312,7 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException": ) return HTTPException( status_code = status_code, - detail = f"llama-server error: {text[:500]}", + detail = _friendly_upstream_error(text[:500]), ) @@ -8529,7 +8549,7 @@ async def _responses_stream( "output": [], "error": { "code": resp.status_code, - "message": f"llama-server error: {err_text[:500]}", + "message": _friendly_upstream_error(err_text[:500]), }, }, }, @@ -10096,7 +10116,7 @@ async def _anthropic_passthrough_stream( yield build_anthropic_sse_event( "error", anthropic_error_body( - f"llama-server error: {_err_text}", + _friendly_upstream_error(_err_text), status = resp.status_code, ), ) @@ -10199,7 +10219,7 @@ async def _anthropic_passthrough_non_streaming( if resp.status_code != 200: raise HTTPException( status_code = resp.status_code, - detail = f"llama-server error: {resp.text[:500]}", + detail = _friendly_upstream_error(resp.text[:500]), ) data = resp.json() diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index aa36c6fed4..1d725acd45 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -40,6 +40,7 @@ from routes.inference import ( _effective_max_tokens, _extract_content_parts, _friendly_error, + _friendly_upstream_error, _merge_user_content, _monitor_openai_chunk, _monitor_openai_sse_event, @@ -57,6 +58,32 @@ from routes.inference import ( from state.tool_policy import reset_tool_policy +class TestFriendlyUpstreamError: + def test_grammar_parse_failure_gets_actionable_message(self): + raw = '{"error":{"code":400,"message":"Failed to initialize samplers: failed to parse grammar","type":"invalid_request_error"}}' + msg = _friendly_upstream_error(raw) + assert "failed to parse grammar" not in msg # raw body is not surfaced verbatim + assert "tool-calling grammar" in msg and "Update Studio" in msg + + def test_failed_to_initialize_samplers_alone_matches(self): + assert "tool-calling grammar" in _friendly_upstream_error("Failed to initialize samplers") + + def test_unrelated_error_passes_through(self): + assert _friendly_upstream_error("out of memory") == "llama-server error: out of memory" + + def test_openai_passthrough_error_rewrites_grammar_failure(self): + # OpenAI-compatible agents (opencode/openclaw/hermes/pi via /v1/chat/completions) + # get the same actionable message as the Anthropic passthrough, not the raw body. + from routes.inference import _openai_passthrough_error + + exc = _openai_passthrough_error( + 400, '{"error":{"message":"Failed to initialize samplers: failed to parse grammar"}}' + ) + assert "tool-calling grammar" in exc.detail + # An unrelated upstream error still passes through verbatim. + assert "llama-server error:" in _openai_passthrough_error(500, "disk full").detail + + # ===================================================================== # ChatMessage — tool role, tool_calls, optional content # ===================================================================== diff --git a/studio/frontend/src/features/settings/components/agent-command.ts b/studio/frontend/src/features/settings/components/agent-command.ts new file mode 100644 index 0000000000..9e87922970 --- /dev/null +++ b/studio/frontend/src/features/settings/components/agent-command.ts @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// Build the `unsloth start ` command for the API-keys panel. +// `unsloth start` reads UNSLOTH_STUDIO_URL (default 127.0.0.1:8888) and only +// auto-mints a key for a loopback server, so the bare command is correct only for +// the default local server. For a non-default port or tunnel/remote base, emit the +// URL (plus a key for non-loopback) so the copy targets what the UI shows. + +const DEFAULT_STUDIO_PORT = "8888"; +const DEFAULT_AGENT = "claude"; + +// URL.hostname brackets IPv6 literals (`new URL("http://[::1]:8888").hostname` is +// "[::1]"), so strip the brackets before matching the bare "::1" loopback rules below. +function normalizeHost(host: string): string { + const lower = host.toLowerCase(); + return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower; +} + +// The bare `unsloth start` probes exactly http://127.0.0.1:8888, so only that literal +// host earns the bare command. `localhost` can resolve to ::1 (and `::1` is never +// probed), so both keep an explicit UNSLOTH_STUDIO_URL -- harmless when they alias +// 127.0.0.1, correct when they don't. +function isDefaultLocalHost(host: string): boolean { + return host === "127.0.0.1"; +} + +// Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8. +function isLoopbackHost(host: string): boolean { + if (host === "localhost" || host === "::1") return true; + const octets = host.split("."); + return ( + octets.length === 4 && + octets[0] === "127" && + octets.every((o) => /^\d{1,3}$/.test(o) && Number(o) <= 255) + ); +} + +export function buildAgentCommand( + base: string | null | undefined, + key: string | null | undefined, + os: "unix" | "windows", + agent: string = DEFAULT_AGENT, +): string { + const bare = `unsloth start ${agent}`; + + let url: URL | null = null; + try { + if (base) url = new URL(base); + } catch { + url = null; + } + // Unknown base: fall back to the bare default-local command. + if (!url) return bare; + + const host = normalizeHost(url.hostname); + const loopback = isLoopbackHost(host); + // Default local server (http://127.0.0.1/localhost:8888): bare command + // auto-discovers it. The CLI's bare default probes plain HTTP, so an HTTPS + // loopback on the same port must keep its explicit UNSLOTH_STUDIO_URL. + if (url.protocol === "http:" && isDefaultLocalHost(host) && url.port === DEFAULT_STUDIO_PORT) { + return bare; + } + + // Non-default server: set the URL; non-loopback also needs an explicit key. + let cmd = bare; + if (!loopback && key) cmd += ` --api-key ${key}`; + + const studioUrl = url.origin; + return os === "windows" + ? `$env:UNSLOTH_STUDIO_URL="${studioUrl}"; ${cmd}` + : `UNSLOTH_STUDIO_URL=${studioUrl} ${cmd}`; +} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index d66ca6105d..c43dd0f219 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -32,6 +32,7 @@ import { loadOpenAIAutoSwitchSettings, updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; +import { buildAgentCommand } from "./agent-command"; type ExampleType = | "curl" @@ -441,6 +442,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); + const [copiedAgent, setCopiedAgent] = useState(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); // null while loading; the same setting the General tab exposes (shared cache). const [autoSwitch, setAutoSwitch] = useState( @@ -477,6 +479,11 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { () => buildSnippets(base, key, model, os, autoSwitchOn), [base, key, model, os, autoSwitchOn], ); + // Agent command must target the server the panel shows, not the :8888 default. + const agentCommand = useMemo( + () => buildAgentCommand(base, key, os), + [base, key, os], + ); const osAware = OS_AWARE[lang]; const shikiLang = CURL_TYPES.has(lang) @@ -520,6 +527,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { } }; + const handleCopyAgent = async () => { + if (await copyToClipboard(agentCommand)) { + setCopiedAgent(true); + setTimeout(() => setCopiedAgent(false), 1800); + } + }; + return (

@@ -689,6 +703,33 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { language={shikiLang} />

+
+ + {t("settings.apiKeys.codingAgents")} + + + {t("settings.apiKeys.codingAgentsHint")} + +
+ + {agentCommand} + + +
+ + {t("settings.apiKeys.codingAgentsSwap")} + +
{t("settings.apiKeys.setupDocs")} {DOC_LINKS.map((link) => ( diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 92abc222a0..3d73ad4343 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -440,6 +440,10 @@ export const en = { copy: "Copy", copied: "Copied", setupDocs: "Setup docs:", + codingAgents: "Coding agents", + codingAgentsHint: + "Launch a coding agent against this server. It uses the loaded model; a local server mints an API key automatically, a remote one includes it in the command.", + codingAgentsSwap: "Swap claude for codex, openclaw, opencode, hermes, or pi.", relativeNever: "never", relativeJustNow: "just now", relativeHoursAgo: "{count}h ago", diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index 440b6276cd..b3831f5314 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -11,7 +11,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError from unsloth_cli.commands.train import train from unsloth_cli.commands.inference import inference from unsloth_cli.commands.chat import chat -from unsloth_cli.commands.connect import connect_app +from unsloth_cli.commands.start import start_app from unsloth_cli.commands.export import export, list_checkpoints from unsloth_cli.commands.studio import ( run as studio_run, @@ -79,9 +79,16 @@ app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.") app.add_typer( - connect_app, + start_app, + name = "start", + help = "Start a coding agent (Claude, Codex, OpenClaw, OpenCode, Hermes, Pi) against Studio.", +) +# Backwards-compatible hidden alias: `unsloth connect` routes to `unsloth start`. +app.add_typer( + start_app, name = "connect", - help = "Connect a coding agent (Claude Code, Codex) to Studio.", + hidden = True, + help = "Deprecated alias for `unsloth start`.", ) # Top-level `unsloth run` aliases `unsloth studio run`; same context diff --git a/unsloth_cli/commands/connect.py b/unsloth_cli/commands/connect.py deleted file mode 100644 index 096a7925e5..0000000000 --- a/unsloth_cli/commands/connect.py +++ /dev/null @@ -1,777 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -"""`unsloth connect` — launch a coding agent against a running Studio server.""" - -import json -import os -import re -import shlex -import shutil -import signal -import subprocess -import urllib.error -import urllib.request -from pathlib import Path -from typing import NoReturn, Optional - -import typer - -from unsloth_cli._inference import ( - _USER_AGENT, - _studio_token, - ensure_studio_backend_path, - find_studio_server, - is_loopback_url, - urlopen_no_redirect, - verify_studio_identity, -) - -connect_app = typer.Typer( - help = "Connect a coding agent to a running Studio server.", - no_args_is_help = True, - context_settings = {"help_option_names": ["-h", "--help"]}, -) - -_CODEX_PROFILE = "unsloth_api" -_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" -_HERMES_ENV_KEY = "UNSLOTH_API_KEY" -_HERMES_PROVIDER = "unsloth" -_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" -_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} -_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") - -# Shared by every agent command; only the config/env/command differ. -_MODEL_OPTION = typer.Option( - None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." -) -_KEY_OPTION = typer.Option( - None, - "--api-key", - envvar = "UNSLOTH_API_KEY", - help = ( - "Studio API key. For a local Studio it is minted automatically and " - "remembered per server. For a remote server, pass one with --api-key " - "(or UNSLOTH_API_KEY); it is remembered for next time." - ), -) -_LAUNCH_OPTION = typer.Option( - True, - "--launch/--no-launch", - help = "--no-launch prints the env and command instead (remote shells, WSL).", -) - - -def _fail(message: str) -> NoReturn: - typer.echo(message, err = True) - raise typer.Exit(code = 1) - - -def _http_error_detail(exc: urllib.error.HTTPError) -> str: - try: - body = json.loads(exc.read().decode()) - return body.get("detail") or body["error"]["message"] - except Exception: - return str(exc) - - -def _http_json( - method: str, - url: str, - token: str, - payload = None, - timeout = 30, - error = None, -): - """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" - request = urllib.request.Request( - url, - data = None if payload is None else json.dumps(payload).encode(), - headers = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "User-Agent": _USER_AGENT, - }, - method = method, - ) - try: - # No redirects: a 3xx would leak this bearer token to an unvetted base. - with urlopen_no_redirect(request, timeout = timeout) as response: - return json.loads(response.read().decode() or "{}") - except urllib.error.HTTPError as exc: - if error is None: - raise - _fail(f"{error}: {_http_error_detail(exc)}") - except (urllib.error.URLError, TimeoutError) as exc: - if error is None: - raise - _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") - - -def _require_studio() -> str: - base = find_studio_server() - if base is None: - expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888") - _fail( - f"No running Studio server found at {expected}. Start one with " - "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." - ) - return base - - -def _key_cache_path() -> Path: - ensure_studio_backend_path() - from utils.paths import auth_root - return auth_root() / "agent_api_key.json" - - -def _read_cache(cache: Path) -> dict: - try: - data = json.loads(cache.read_text(encoding = "utf-8")) - except Exception: - return {} - return data if isinstance(data, dict) else {} - - -def _server_buckets(servers: dict, base: str) -> dict: - # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a - # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). - entry = servers.get(base) if isinstance(servers, dict) else None - if isinstance(entry, list): - return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} - if not isinstance(entry, dict): - return {"saved": [], "minted": []} - - def _strs(name: str) -> list: - value = entry.get(name) - return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] - - return {"saved": _strs("saved"), "minted": _strs("minted")} - - -def _cached_keys(cache: Path, base: str, source: str) -> list: - # Keys are scoped per server. `source` splits user-supplied --api-key keys - # ("saved", trusted for that base) from auto-minted ones ("minted", replayed - # only after the identity check). Legacy unscoped caches are ignored. - return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] - - -def _write_private_json(path: Path, data: dict) -> None: - # O_CREAT with 0o600 so a file holding an API key is never world-readable, - # even briefly (existing files keep whatever perms the user set). - path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) - fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) - with os.fdopen(fd, "w") as handle: - handle.write(json.dumps(data, indent = 2) + "\n") - - -def _read_json_object(path: Path) -> Optional[dict]: - # {} when missing, None when it can't be parsed as an object (so the caller - # leaves a user-managed file untouched rather than clobbering it). - if not path.exists(): - return {} - try: - data = json.loads(path.read_text(encoding = "utf-8")) - except (ValueError, OSError): - return None - return data if isinstance(data, dict) else None - - -def _subdict(parent: dict, key: str) -> dict: - child = parent.get(key) - if not isinstance(child, dict): - child = parent[key] = {} - return child - - -def _remember_key(cache: Path, base: str, key: str, source: str) -> None: - data = _read_cache(cache) - servers = data.get("servers") - if not isinstance(servers, dict): - servers = data["servers"] = {} - buckets = _server_buckets(servers, base) - other = "minted" if source == "saved" else "saved" - buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] - buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance - new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} - if servers.get(base) == new_entry: - return - servers[base] = new_entry - # Collapse legacy unscoped fields. - data.pop("keys", None) - data.pop("key", None) - try: - _write_private_json(cache, data) - except OSError: - pass # worst case the next launch mints another key - - -def _key_accepted(base: str, key: str) -> bool: - try: - _http_json("GET", f"{base}/v1/models", key) - return True - except Exception: - return False - - -def _agent_api_key(base: str, explicit: Optional[str]) -> str: - cache = _key_cache_path() - if explicit: - _remember_key(cache, base, explicit, "saved") - return explicit - - # Replay a key the user saved for *this exact* server first (scoped per base, - # so it only goes back there -- including a remote/SSH-tunnelled Studio whose - # secret the local handshake can't match). Skip ones the server rejects. - for key in _cached_keys(cache, base, "saved"): - if _key_accepted(base, key): - _remember_key(cache, base, key, "saved") - return key - - # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() - # trusts a base after only a health check, so both are limited to a loopback - # server we can cryptographically confirm is ours. - if not is_loopback_url(base): - _fail( - f"No saved API key for {base} and automatic minting only runs against " - "a local Studio. Create an API key in Studio → Settings → API and " - "pass it with --api-key (it is remembered per server), or set " - "UNSLOTH_API_KEY." - ) - if not verify_studio_identity(base): - _fail( - f"Couldn't verify that {base} is your Studio (it may be running as a " - "different OS user, or another process took the port). Create an API " - "key in Studio → Settings → API and pass it with --api-key, or set " - "UNSLOTH_API_KEY." - ) - - # Identity verified: replay a previously auto-minted key, else mint a new one. - for key in _cached_keys(cache, base, "minted"): - if _key_accepted(base, key): - _remember_key(cache, base, key, "minted") - return key - - # Self-issue a JWT (signed with the local secret) and mint a key. - token = _studio_token() - if token is None: - _fail( - "Couldn't authenticate with the Studio server automatically. Create " - "an API key in Studio → Settings → API and pass it with --api-key, " - "or set UNSLOTH_API_KEY." - ) - key = _http_json( - "POST", - f"{base}/api/auth/api-keys", - token, - {"name": "Coding agents (unsloth connect)"}, - error = "Couldn't create an API key", - )["key"] - _remember_key(cache, base, key, "minted") - return key - - -def _loaded_models(base: str, key: str) -> list: - return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) - - -def _resolve_model(base: str, key: str, requested: Optional[str]) -> dict: - models = _loaded_models(base, key) - match = next((m for m in models if m["id"] == requested), None) - if requested and match is None: - typer.echo(f"Loading {requested} on the Studio server (this can take a while)…") - loaded = _http_json( - "POST", - f"{base}/api/inference/load", - key, - {"model_path": requested}, - timeout = 3600, - error = "Model load failed", - ) - # Studio registers the model under a canonical id (resolved identifier, - # casing) that /v1/models echoes but which may differ from the path we - # passed; match on the id the load reports so we don't silently fall - # through to models[0] and connect to a different loaded model. - wanted = {requested} - if isinstance(loaded, dict): - wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} - models = _loaded_models(base, key) - match = next((m for m in models if m["id"] in wanted), None) - if match is not None: - return match - if requested: - # We asked Studio to load it and it didn't surface in /v1/models; don't - # silently hand back an unrelated loaded model. - _fail( - f"Studio didn't report '{requested}' as loaded. Double-check the model " - "id, or load it from the model dropdown in the UI." - ) - if not models: - _fail( - "No model is loaded in Studio. Load one from the model dropdown in " - "the UI, or pass --model to load it from here." - ) - return models[0] - - -def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: - # Codex always streams, and Studio only streams /v1/responses from llama-server. - try: - status = _http_json("GET", f"{base}/api/inference/status", key) - except urllib.error.HTTPError as exc: - if exc.code == 404: - return # older server without the endpoint; don't block the launch - raise - if status.get("is_gguf"): - return - hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" - _fail( - f"Codex needs a GGUF model served by llama-server, but {model_id} is on " - f"the transformers backend. Try: unsloth connect codex --model {hint}" - ) - - -def claude_settings_path() -> Path: - return Path.home() / ".claude" / "settings.json" - - -def ensure_claude_attribution_header() -> None: - # The header invalidates the llama.cpp KV cache (~90% slower) and Claude - # Code only honors this setting from settings.json, not the env var. - path = claude_settings_path() - settings = {} - if path.exists(): - try: - settings = json.loads(path.read_text(encoding = "utf-8")) - except (ValueError, OSError): - settings = None - if not isinstance(settings, dict): - typer.echo( - f"Warning: couldn't parse {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " - 'to "0" in its "env" section yourself, or local inference will be much slower.', - err = True, - ) - return - env = settings.get("env") - if not isinstance(env, dict): - env = settings["env"] = {} - if str(env.get("CLAUDE_CODE_ATTRIBUTION_HEADER")) == "0": - return - env["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "0" - try: - path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(json.dumps(settings, indent = 2) + "\n", encoding = "utf-8") - except OSError: - typer.echo( - f"Warning: couldn't write {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " - 'to "0" in its "env" section yourself, or local inference will be much slower.', - err = True, - ) - return - typer.echo(f"Disabled Claude Code's attribution header in {path} (it breaks KV-cache reuse).") - - -_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" - - -def _claude_cache_flags() -> list: - # The flag moves per-machine context (cwd, env info, git status) out of - # the system prompt, where it changes every session and defeats llama.cpp - # prefix caching. As of 2.1.175 it only takes effect in print mode (`-p` - # passed through ctx.args); interactive sessions accept and ignore it. - # Claude Code < 2.1.98 aborts on the unknown flag, so check the version - # first; no local binary means a --no-launch printout for another machine. - executable = shutil.which("claude") - if executable is None: - return [_DYNAMIC_SECTIONS_FLAG] - try: - result = subprocess.run( - [executable, "--version"], capture_output = True, text = True, timeout = 10 - ) - version = tuple(int(part) for part in result.stdout.split()[0].split(".")) - except Exception: - return [] - return [_DYNAMIC_SECTIONS_FLAG] if version >= (2, 1, 98) else [] - - -def codex_home() -> Path: - return Path(os.environ.get("CODEX_HOME") or Path.home() / ".codex") - - -def _merge_codex_config(existing: str, base: str) -> str: - chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table - if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): - if chunks[0] and not chunks[0].endswith("\n"): - chunks[0] += "\n" - chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' - # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. - stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") - text = "".join(c for c in chunks if not c.startswith(stale)) - if not text.endswith("\n"): - text += "\n" - if not text.endswith("\n\n"): - text += "\n" - return text + ( - f"{_PROVIDER_HEADER}\n" - 'name = "Unsloth Studio"\n' - f"base_url = {json.dumps(base + '/v1')}\n" - f'env_key = "{_CODEX_ENV_KEY}"\n' - 'wire_api = "responses"\n' - "requires_openai_auth = false\n" - ) - - -def write_codex_config(base: str, model: dict) -> None: - home = codex_home() - home.mkdir(parents = True, exist_ok = True) - - config = home / "config.toml" - existing = config.read_text(encoding = "utf-8") if config.exists() else "" - merged = _merge_codex_config(existing, base) - if merged != existing: - config.write_text(merged, encoding = "utf-8") - typer.echo(f"Updated {config}") - - # oss_provider here too: codex --oss picks the provider from it, and the - # profile layer must beat a user-set value (e.g. "ollama") in config.toml. - profile_text = ( - f'oss_provider = "{_CODEX_PROFILE}"\n' - f'model_provider = "{_CODEX_PROFILE}"\n' - f"model = {json.dumps(model['id'])}\n" - ) - window = model.get("context_length") or model.get("max_context_length") - if window: - profile_text += f"model_context_window = {int(window)}\n" - profile = home / f"{_CODEX_PROFILE}.config.toml" - if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: - profile.write_text(profile_text, encoding = "utf-8") - typer.echo(f"Updated {profile}") - - -def _wsl_windows_executable(command: list) -> Optional[str]: - if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): - return None - executable = shutil.which(command[0]) - if executable and executable.startswith("/mnt/"): - return executable - return None - - -def _merge_wslenv(current: str, names: tuple) -> str: - entries = [entry for entry in current.split(":") if entry] - existing = {entry.split("/", 1)[0] for entry in entries} - entries.extend(name for name in names if name not in existing) - return ":".join(entries) - - -def _print_env( - env: dict, - command: list, - unset_env: tuple = (), - wsl_env_bridge: tuple = (), -) -> None: - if os.name == "nt": - for name in unset_env: - typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") - for name, value in env.items(): - # PowerShell: ` is the escape char, and $ triggers expansion inside "". - escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") - typer.echo(f'$env:{name} = "{escaped}"') - typer.echo(subprocess.list2cmdline(command)) - return - for name in unset_env: - typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") - for name, value in env.items(): - typer.echo(f"export {name}={shlex.quote(value)}") - if wsl_env_bridge: - typer.echo( - f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" - ) - typer.echo(shlex.join(command)) - - -def _launch( - command: list, - env: dict, - install_hint: str, - unset_env: tuple = (), -) -> NoReturn: - executable = shutil.which(command[0]) - if executable is None: - _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") - wsl_env_bridge = ( - tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () - ) - child_env = dict(os.environ) - if wsl_env_bridge: - child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) - for name in unset_env: - child_env[name] = "" - else: - for name in unset_env: - child_env.pop(name, None) - child_env.update(env) - # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. - previous = signal.signal(signal.SIGINT, signal.SIG_IGN) - try: - code = subprocess.run([executable, *command[1:]], env = child_env).returncode - finally: - signal.signal(signal.SIGINT, previous) - # Negative returncode means killed by signal N; shells expect 128+N. - raise typer.Exit(code = code if code >= 0 else 128 - code) - - -def _connect(api_key: Optional[str], model: Optional[str]) -> tuple: - base = _require_studio() - key = _agent_api_key(base, api_key) - return base, key, _resolve_model(base, key, model) - - -def _run( - base: str, - entry: dict, - env: dict, - command: list, - *, - launch: bool, - install_hint: str, - unset_env: tuple = (), -) -> None: - typer.echo(f"Studio {base} · model {entry['id']}") - wsl_env_bridge = ( - tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () - ) - if not launch: - _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) - return - _launch(command, env, install_hint = install_hint, unset_env = unset_env) - - -def openclaw_config_path() -> Path: - return Path.home() / ".openclaw" / "openclaw.json" - - -def write_openclaw_config(base: str, key: str, model: dict) -> None: - path = openclaw_config_path() - config = _read_json_object(path) - if config is None: - typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", - err = True, - ) - return - before = json.dumps(config, sort_keys = True) - # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). - provider_model = {"id": model["id"], "name": model["id"]} - window = model.get("context_length") or model.get("max_context_length") - if window: - provider_model["contextWindow"] = int(window) - models = _subdict(config, "models") - models.setdefault("mode", "merge") - _subdict(models, "providers")["unsloth"] = { - "baseUrl": f"{base}/v1", - "apiKey": key, - "api": "openai-completions", - "models": [provider_model], - } - # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). - defaults = _subdict(_subdict(config, "agents"), "defaults") - _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" - # Unauthenticated loopback gateway: without auth.mode=none the client won't open - # the websocket. The daemon must still be started separately (`openclaw gateway`). - gateway = _subdict(config, "gateway") - gateway.setdefault("mode", "local") - _subdict(gateway, "auth").setdefault("mode", "none") - if json.dumps(config, sort_keys = True) != before: - _write_private_json(path, config) - typer.echo(f"Updated {path}") - - -def opencode_config_path() -> Path: - config_home = os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config" - return Path(config_home) / "opencode" / "opencode.json" - - -def write_opencode_config(base: str, key: str, model: dict) -> None: - path = opencode_config_path() - config = _read_json_object(path) - if config is None: - typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", - err = True, - ) - return - before = json.dumps(config, sort_keys = True) - config.setdefault("$schema", "https://opencode.ai/config.json") - _subdict(config, "provider")["unsloth"] = { - "npm": "@ai-sdk/openai-compatible", - "name": "Unsloth Studio", - "options": {"baseURL": f"{base}/v1", "apiKey": key}, - "models": {model["id"]: {"name": model["id"]}}, - } - # OpenCode selects a model by "/". - config["model"] = f"unsloth/{model['id']}" - if json.dumps(config, sort_keys = True) != before: - _write_private_json(path, config) - typer.echo(f"Updated {path}") - - -def hermes_config_path() -> Path: - return Path.home() / ".hermes" / "config.yaml" - - -def write_hermes_config(base: str, model: dict) -> None: - import yaml - - path = hermes_config_path() - config: dict = {} - if path.exists(): - try: - loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) - except (yaml.YAMLError, OSError): - typer.echo( - f"Warning: couldn't parse {path} — configure the custom endpoint " - "there yourself, or move the file aside and re-run.", - err = True, - ) - return - if isinstance(loaded, dict): - config = loaded - elif loaded is not None: - # Non-empty, non-mapping YAML is a user-managed file; leave it. - typer.echo( - f"Warning: couldn't parse {path} — configure the custom endpoint " - "there yourself, or move the file aside and re-run.", - err = True, - ) - return - # Hermes only reads the key for a *named* custom provider (a bare - # `provider: custom` ignores it), so register it under providers.*. - _subdict(config, "model").update( - provider = f"custom:{_HERMES_PROVIDER}", - default = model["id"], - api_mode = "openai", - ) - _subdict(config, "providers")[_HERMES_PROVIDER] = { - "base_url": f"{base}/v1", - "api_mode": "openai", - "key_env": _HERMES_ENV_KEY, - } - text = yaml.safe_dump(config, sort_keys = False) - if not path.exists() or path.read_text(encoding = "utf-8") != text: - path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(text, encoding = "utf-8") - typer.echo(f"Updated {path}") - - -@connect_app.command("claude", context_settings = _PASSTHROUGH) -def claude( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point Claude Code at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - model_id = entry["id"] - ensure_claude_attribution_header() - - env = { - "ANTHROPIC_BASE_URL": base, - "ANTHROPIC_AUTH_TOKEN": key, - "ANTHROPIC_MODEL": model_id, - # Update checks, beta features, and other background requests either - # stall against a local server or evict the conversation from - # llama-server's KV-cache slots, so turn off everything nonessential. - "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", - "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", - } - command = ["claude", "--model", model_id, *_claude_cache_flags(), *ctx.args] - install_hint = ( - "irm https://claude.ai/install.ps1 | iex" - if os.name == "nt" - else "curl -fsSL https://claude.ai/install.sh | bash" - ) - _run( - base, - entry, - env, - command, - launch = launch, - install_hint = install_hint, - unset_env = _CLAUDE_ENV_UNSET, - ) - - -@connect_app.command("codex", context_settings = _PASSTHROUGH) -def codex( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenAI Codex at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - _require_gguf_for_codex(base, key, entry["id"]) - write_codex_config(base, entry) - - env = {_CODEX_ENV_KEY: key} - command = ["codex", "--oss", "--profile", _CODEX_PROFILE, *ctx.args] - _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") - - -@connect_app.command("openclaw", context_settings = _PASSTHROUGH) -def openclaw( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenClaw at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_openclaw_config(base, key, entry) # key lives in the config, not the env - - command = ["openclaw", *ctx.args] - install_hint = ( - "iwr -useb https://openclaw.ai/install.ps1 | iex" - if os.name == "nt" - else "curl -fsSL https://openclaw.ai/install.sh | bash" - ) - _run(base, entry, {}, command, launch = launch, install_hint = install_hint) - - -@connect_app.command("opencode", context_settings = _PASSTHROUGH) -def opencode( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenCode at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_opencode_config(base, key, entry) # key lives in the config, not the env - - command = ["opencode", *ctx.args] - _run(base, entry, {}, command, launch = launch, install_hint = "npm install -g opencode-ai") - - -@connect_app.command("hermes", context_settings = _PASSTHROUGH) -def hermes( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point Hermes (Nous Research) at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_hermes_config(base, entry) - - env = {_HERMES_ENV_KEY: key} - command = ["hermes", *ctx.args] - install_hint = ( - "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" - "/main/scripts/install.sh | bash" - ) - _run(base, entry, env, command, launch = launch, install_hint = install_hint) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py new file mode 100644 index 0000000000..b188180188 --- /dev/null +++ b/unsloth_cli/commands/start.py @@ -0,0 +1,1497 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""`unsloth start` — launch a coding agent against a running Studio server.""" + +import atexit +import contextlib +import json +import os +import re +import shlex +import shutil +import signal +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import NamedTuple, NoReturn, Optional +from urllib.parse import urlparse + +import click +import typer + +from unsloth_cli._inference import ( + _USER_AGENT, + _studio_token, + ensure_studio_backend_path, + find_studio_server, + is_loopback_url, + urlopen_no_redirect, + verify_studio_identity, +) + +start_app = typer.Typer( + help = "Start a coding agent against a running Studio server.", + no_args_is_help = True, + context_settings = {"help_option_names": ["-h", "--help"]}, +) + +_CODEX_PROFILE = "unsloth_api" +_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" +_HERMES_ENV_KEY = "UNSLOTH_API_KEY" +_HERMES_PROVIDER = "unsloth" +# Hermes refuses to initialize when the model window is under 64,000 tokens; its +# error message points at the model.context_length / auxiliary.compression +# overrides in config.yaml. write_hermes_config claims this value for smaller +# windows and scales the compaction threshold back down to the real window. +_HERMES_MIN_CONTEXT = 65536 +_PI_PROVIDER = "unsloth" +_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" +_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} +_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") + +# Shared by every agent command; only the config/env/command differ. +_MODEL_OPTION = typer.Option( + None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." +) +_KEY_OPTION = typer.Option( + None, + "--api-key", + envvar = "UNSLOTH_API_KEY", + help = ( + "Studio API key. For a local Studio it is minted automatically and " + "remembered per server. For a remote server, pass one with --api-key " + "(or UNSLOTH_API_KEY); it is remembered for next time." + ), +) +_LAUNCH_OPTION = typer.Option( + True, + "--launch/--no-launch", + help = "--no-launch prints the env and command instead (remote shells, WSL).", +) +_SERVE_OPTION = typer.Option( + True, + "--serve/--no-serve", + help = ( + "If no Studio server is running, auto-start one for --model and stop it when the " + "agent exits. --no-serve keeps the old behavior of erroring out." + ), +) +# Model-load knobs mirrored from `unsloth run`; only used when --model triggers a +# load on the server. Server-startup flags (--host/--port/--cloudflare/...) do not +# apply here because `unsloth start` attaches to an already-running server. +_GGUF_VARIANT_OPTION = typer.Option( + None, "--gguf-variant", help = "GGUF quant variant to load (e.g. UD-Q4_K_XL)." +) +_CONTEXT_OPTION = typer.Option( + 0, + "--max-seq-length", + "--context-length", + help = "Context length in tokens for the load (0 = model default).", +) +_LOAD_4BIT_OPTION = typer.Option( + True, "--load-in-4bit/--no-load-in-4bit", help = "Load hub models in 4-bit (ignored for GGUF)." +) +_TENSOR_PARALLEL_OPTION = typer.Option( + False, + "--tensor-parallel/--no-tensor-parallel", + help = "Split a GGUF across GPUs by tensor instead of by layer (multi-GPU only).", +) +# One normalized "run tools without prompting" switch. Each agent spells this +# differently and it's easy to forget which is which, so accept every spelling and +# route to the agent's own mechanism in _yolo_command_flags / the config writers. +_YOLO_OPTION = typer.Option( + False, + "--yolo", + "--dangerously-skip-permissions", + "--dangerously-bypass-approvals-and-sandbox", + help = ( + "Auto-approve all tool actions for this session; routed to the agent's own " + "flag/config. Any of the three spellings works for any agent." + ), +) + +# Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no +# such flag (config only) and are handled in their config writers, so they are absent. +_YOLO_COMMAND_FLAGS = { + "claude": ["--dangerously-skip-permissions"], + "codex": ["--dangerously-bypass-approvals-and-sandbox"], + "hermes": ["--yolo"], + # Pi never prompts per tool call; its only approval gate is project trust, so -a + # (trust project resources) is the closest "don't ask me" equivalent. + "pi": ["--approve"], +} + + +def _yolo_command_flags(agent: str, yolo: bool) -> list: + # .get so a config-based agent (or a typo) yields no flag instead of a KeyError. + return _YOLO_COMMAND_FLAGS.get(agent, []) if yolo else [] + + +class LoadOptions(NamedTuple): + """Model-load knobs forwarded to /api/inference/load when --model triggers a load.""" + + gguf_variant: Optional[str] = None + max_seq_length: int = 0 + load_in_4bit: bool = True + tensor_parallel: bool = False + + +def _split_repo_variant(model: str) -> tuple: + """Split ``org/name:QUANT`` into ``(repo, variant)`` -> ``("org/name", "QUANT")``. + + ``unsloth run`` and llama.cpp accept ``--model org/name:QUANT`` as shorthand for + ``--model org/name --gguf-variant QUANT``. Mirror that here so a ``:variant`` suffix + resolves against the already-loaded ``org/name`` (which /v1/models lists without the + suffix) instead of trying to load a repo id containing ``:`` -- which Hugging Face + rejects, and which would evict a model another session is using. Local paths, Windows + drive letters, and ids without a ``:`` pass through unchanged. + """ + s = (model or "").strip() + if not s or s.startswith(("/", "./", "../", "~")) or s == ".": + return s, None + if len(s) >= 2 and s[1] == ":" and s[0].isalpha(): # Windows drive, e.g. C:\models\x + return s, None + if ":" not in s: + return s, None + repo, _, variant = s.rpartition(":") + if not repo or not variant or "/" in variant: + return s, None + return repo, variant + + +def _fail(message: str) -> NoReturn: + typer.echo(message, err = True) + raise typer.Exit(code = 1) + + +def _http_error_detail(exc: urllib.error.HTTPError) -> str: + try: + body = json.loads(exc.read().decode()) + return body.get("detail") or body["error"]["message"] + except Exception: + return str(exc) + + +def _http_json( + method: str, + url: str, + token: str, + payload = None, + timeout = 30, + error = None, +): + """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" + request = urllib.request.Request( + url, + data = None if payload is None else json.dumps(payload).encode(), + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": _USER_AGENT, + }, + method = method, + ) + try: + # No redirects: a 3xx would leak this bearer token to an unvetted base. + with urlopen_no_redirect(request, timeout = timeout) as response: + return json.loads(response.read().decode() or "{}") + except urllib.error.HTTPError as exc: + if error is None: + raise + _fail(f"{error}: {_http_error_detail(exc)}") + except (urllib.error.URLError, TimeoutError) as exc: + if error is None: + raise + _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") + + +# A server that WE auto-started (never one we merely found). Kept at module scope so +# _run's finally and the atexit backstop can tear it down without threading a handle +# through all six agent commands. Only one agent runs per process, so one slot is enough. +_auto_served_server: Optional[subprocess.Popen] = None +# Model download + load can be slow; give the auto-started server room before giving up. +_SERVER_START_TIMEOUT_S = 900 + + +def _studio_healthy(base: str, timeout: float = 3.0) -> bool: + request = urllib.request.Request(f"{base}/api/health", headers = {"User-Agent": _USER_AGENT}) + try: + with urllib.request.urlopen(request, timeout = timeout) as response: + return json.loads(response.read(65536).decode() or "{}").get("status") == "healthy" + except Exception: + return False + + +def _log_tail(path: Path, lines: int = 20) -> str: + try: + return "\n".join(path.read_text(encoding = "utf-8", errors = "replace").splitlines()[-lines:]) + except OSError: + return "(no server log)" + + +def _shutdown_server(server: Optional[subprocess.Popen]) -> None: + # Idempotent teardown of a server WE started, plus its own children (llama-server, + # cloudflared). A no-op once the process is already gone. + if server is None or server.poll() is not None: + return + if os.name == "nt": + # terminate()/kill() reach only the parent `unsloth run`; taskkill /T walks the + # whole tree so the llama-server child doesn't keep the port and GPU (matches the + # taskkill /T /F pattern already used in unsloth/dataprep/synthetic.py). + try: + subprocess.run( + ["taskkill", "/PID", str(server.pid), "/T", "/F"], + capture_output = True, + timeout = 15, + check = False, + ) + server.wait(timeout = 5) + except Exception: + with contextlib.suppress(Exception): + server.kill() + return + try: + os.killpg(os.getpgid(server.pid), signal.SIGTERM) + except OSError: + server.terminate() + try: + server.wait(timeout = 15) + except Exception: + try: + os.killpg(os.getpgid(server.pid), signal.SIGKILL) + except OSError: + server.kill() + + +def _shutdown_auto_served() -> None: + global _auto_served_server + server, _auto_served_server = _auto_served_server, None + if server is not None and server.poll() is None: + typer.echo("Stopping the auto-started Studio server…") + _shutdown_server(server) + + +def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess.Popen: + """Spawn `unsloth run` for `model`, wait until it is fully ready, and return it.""" + global _auto_served_server + unsloth = shutil.which("unsloth") or "unsloth" + parsed = urlparse(base) + # --disable-tools = passthrough mode (relay the agent's own tools); --no-cloudflare = + # loopback only, no tunnel. Mirrors .github/scripts/serve-unsloth-run.sh. + command = [ + unsloth, + "run", + "-H", + parsed.hostname or "127.0.0.1", + "-p", + str(parsed.port or 8888), + "--disable-tools", + "--no-cloudflare", + "--model", + model, + ] + if load.gguf_variant: + command += ["--gguf-variant", load.gguf_variant] + if load.max_seq_length: + command += ["--context-length", str(load.max_seq_length)] + if not load.load_in_4bit: + command += ["--no-load-in-4bit"] + if load.tensor_parallel: + command += ["--tensor-parallel"] + + log_path = Path(tempfile.gettempdir()) / f"unsloth-start-server-{os.getpid()}.log" + typer.echo( + f"No Studio server at {base}. Starting one for {model} (loading the model can take a while)…" + ) + typer.echo(f"Server log: {log_path}") + # 0600: the `unsloth run` banner in this log carries the minted sk-unsloth- key, and + # the tempdir is world-traversable. Unlink first so a stale looser-mode file (pid + # reuse) can't survive with its old permissions. + log_path.unlink(missing_ok = True) + log = os.fdopen(os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb") + # Own session/process group so a mid-session Ctrl+C (cancel a turn) doesn't reach the + # server; we tear it down explicitly when the agent exits. + kwargs: dict = {"stdout": log, "stderr": subprocess.STDOUT, "stdin": subprocess.DEVNULL} + if os.name == "nt": + kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP + else: + kwargs["start_new_session"] = True + try: + server = subprocess.Popen(command, **kwargs) + finally: + log.close() # Popen dup'd the fd; drop the parent's copy + _auto_served_server = server + atexit.register(_shutdown_auto_served) + + deadline = time.monotonic() + _SERVER_START_TIMEOUT_S + while time.monotonic() < deadline: + if server.poll() is not None: + tail = _log_tail(log_path) + _shutdown_auto_served() + _fail(f"The Studio server stopped before it was ready. Last log lines:\n{tail}") + # `unsloth run` prints the minted key only after the server is up AND the model is + # loaded, so it is the fully-ready signal (same contract serve-unsloth-run.sh uses). + if _studio_healthy(base) and "sk-unsloth-" in _log_tail(log_path, lines = 400): + typer.echo(f"Studio server ready at {base}.") + return server + time.sleep(2.0) + _shutdown_auto_served() + _fail( + f"The Studio server didn't become ready within {_SERVER_START_TIMEOUT_S}s. See {log_path}." + ) + + +def _effective_base(base: str) -> str: + # `unsloth run` binds to `parsed.port or 8888` and serves at the root, so normalize + # UNSLOTH_STUDIO_URL to plain scheme://host:port. A portless http://127.0.0.1 would + # otherwise launch on 8888 but poll port 80, and a path like /studio would poll + # /studio/api/health (404) -- either way hitting the startup timeout. IPv6 literals + # stay bracketed. + parsed = urlparse(base) + host = parsed.hostname or "127.0.0.1" + if ":" in host: # bare IPv6 literal (urlparse strips the brackets) + host = f"[{host}]" + return f"{parsed.scheme or 'http'}://{host}:{parsed.port or 8888}" + + +def _require_studio( + model: Optional[str] = None, + load: Optional[LoadOptions] = None, + *, + serve: bool = False, + launch: bool = True, +) -> tuple: + """Return (base, server). server is a Popen only when WE auto-started it.""" + base = find_studio_server() + if base is not None: + return base, None + expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/") + # Auto-start a local server only for an interactive launch with a model to serve, and + # only for a plain-HTTP loopback target: never stand in for an explicit remote + # UNSLOTH_STUDIO_URL, and never for an https:// one -- `unsloth run` serves plain + # HTTP, so the health poll against https would spin until the startup timeout. + if ( + serve + and launch + and model + and is_loopback_url(expected) + and urlparse(expected).scheme == "http" + ): + # Normalize to the port unsloth run actually binds, so the health poll and the + # returned base hit the same server we launch (not a portless :80). + expected = _effective_base(expected) + return expected, _start_studio_server(expected, model, load or LoadOptions()) + model_hint = "" if model else " Pass --model to have it start one for you, or" + _fail( + f"No running Studio server found at {expected}.{model_hint} start one with " + "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." + ) + + +def _key_cache_path() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agent_api_key.json" + + +def _read_cache(cache: Path) -> dict: + try: + data = json.loads(cache.read_text(encoding = "utf-8")) + except Exception: + return {} + return data if isinstance(data, dict) else {} + + +def _server_buckets(servers: dict, base: str) -> dict: + # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a + # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). + entry = servers.get(base) if isinstance(servers, dict) else None + if isinstance(entry, list): + return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} + if not isinstance(entry, dict): + return {"saved": [], "minted": []} + + def _strs(name: str) -> list: + value = entry.get(name) + return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] + + return {"saved": _strs("saved"), "minted": _strs("minted")} + + +def _cached_keys(cache: Path, base: str, source: str) -> list: + # Keys are scoped per server. `source` splits user-supplied --api-key keys + # ("saved", trusted for that base) from auto-minted ones ("minted", replayed + # only after the identity check). Legacy unscoped caches are ignored. + return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] + + +def _write_private_json(path: Path, data: dict) -> None: + # O_CREAT with 0o600 so a file holding an API key is never world-readable, + # even briefly (existing files keep whatever perms the user set). + path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as handle: + handle.write(json.dumps(data, indent = 2) + "\n") + + +def _read_json_object(path: Path) -> Optional[dict]: + # {} when missing, None when it can't be parsed as an object (so the caller + # leaves a user-managed file untouched rather than clobbering it). + if not path.exists(): + return {} + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (ValueError, OSError): + return None + return data if isinstance(data, dict) else None + + +def _subdict(parent: dict, key: str) -> dict: + child = parent.get(key) + if not isinstance(child, dict): + child = parent[key] = {} + return child + + +def _remember_key(cache: Path, base: str, key: str, source: str) -> None: + data = _read_cache(cache) + servers = data.get("servers") + if not isinstance(servers, dict): + servers = data["servers"] = {} + buckets = _server_buckets(servers, base) + other = "minted" if source == "saved" else "saved" + buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] + buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance + new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} + if servers.get(base) == new_entry: + return + servers[base] = new_entry + # Collapse legacy unscoped fields. + data.pop("keys", None) + data.pop("key", None) + try: + _write_private_json(cache, data) + except OSError: + pass # worst case the next launch mints another key + + +def _key_accepted(base: str, key: str) -> bool: + # Only a genuine auth rejection (401/403) means "this key is bad -- skip it and try + # the next cached key or mint a fresh one". A 5xx or a network blip is a server-side + # outage, not a bad key: fail with a clean message (never a traceback) instead of + # silently discarding a working key and minting extras against a struggling server. + try: + _http_json("GET", f"{base}/v1/models", key) + return True + except urllib.error.HTTPError as exc: + if exc.code in (401, 403): + return False + _fail( + f"Studio server error while checking an API key ({exc.code}). " + "The server may be starting up or unhealthy; try again shortly." + ) + except (urllib.error.URLError, TimeoutError) as exc: + _fail( + "Couldn't reach the Studio server while checking an API key: " + f"{getattr(exc, 'reason', None) or exc}" + ) + + +def _agent_api_key( + base: str, + explicit: Optional[str], + *, + auto_started: bool = False, +) -> str: + cache = _key_cache_path() + if explicit: + if not auto_started or _key_accepted(base, explicit): + _remember_key(cache, base, explicit, "saved") + return explicit + # The server was auto-started for this run, so an exported + # UNSLOTH_API_KEY meant for some other server must not fail the + # launch: the loopback mint path below is guaranteed to work. + # (An explicit key that the fresh server accepts, e.g. one persisted + # in this Studio home's auth db, is still honored above.) + + # Replay a key the user saved for *this exact* server first (scoped per base, + # so it only goes back there -- including a remote/SSH-tunnelled Studio whose + # secret the local handshake can't match). Skip ones the server rejects. + for key in _cached_keys(cache, base, "saved"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "saved") + return key + + # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() + # trusts a base after only a health check, so both are limited to a loopback + # server we can cryptographically confirm is ours. + if not is_loopback_url(base): + _fail( + f"No saved API key for {base} and automatic minting only runs against " + "a local Studio. Create an API key in Studio → Settings → API and " + "pass it with --api-key (it is remembered per server), or set " + "UNSLOTH_API_KEY." + ) + if not verify_studio_identity(base): + _fail( + f"Couldn't verify that {base} is your Studio (it may be running as a " + "different OS user, or another process took the port). Create an API " + "key in Studio → Settings → API and pass it with --api-key, or set " + "UNSLOTH_API_KEY." + ) + + # Identity verified: replay a previously auto-minted key, else mint a new one. + for key in _cached_keys(cache, base, "minted"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "minted") + return key + + # Self-issue a JWT (signed with the local secret) and mint a key. + token = _studio_token() + if token is None: + _fail( + "Couldn't authenticate with the Studio server automatically. Create " + "an API key in Studio → Settings → API and pass it with --api-key, " + "or set UNSLOTH_API_KEY." + ) + key = _http_json( + "POST", + f"{base}/api/auth/api-keys", + token, + {"name": "Coding agents (unsloth start)"}, + error = "Couldn't create an API key", + )["key"] + _remember_key(cache, base, key, "minted") + return key + + +def _loaded_models(base: str, key: str) -> list: + return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) + + +def _resolve_model( + base: str, + key: str, + requested: Optional[str], + load: LoadOptions = LoadOptions(), +) -> dict: + models = _loaded_models(base, key) + # /v1/models reports the model id but not the active GGUF variant or runtime load + # settings, so an id match alone can hide the wrong quant (Q8_0 serving while the + # user asked for UD-Q4_K_XL). When the user passed any explicit load knob, defer to + # /api/inference/load: the server's already-loaded dedup answers "already_loaded" + # without reloading when the variant AND settings match, so a second session running + # the same command still attaches without evicting the first. + load_has_overrides = bool( + load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel + ) + match = ( + None + if requested and load_has_overrides + else next((m for m in models if m["id"] == requested), None) + ) + if requested and match is None: + typer.echo( + f"Ensuring {requested} is loaded with the requested settings…" + if load_has_overrides + else f"Loading {requested} on the Studio server (this can take a while)…" + ) + # Mirror `unsloth run`'s load knobs; keep the default payload as just + # model_path so a bare `--model` load is unchanged. + payload = {"model_path": requested} + if load.gguf_variant: + payload["gguf_variant"] = load.gguf_variant + if load.max_seq_length: + payload["max_seq_length"] = load.max_seq_length + if not load.load_in_4bit: + payload["load_in_4bit"] = False + if load.tensor_parallel: + payload["tensor_parallel"] = True + loaded = _http_json( + "POST", + f"{base}/api/inference/load", + key, + payload, + timeout = 3600, + error = "Model load failed", + ) + # Studio registers the model under a canonical id (resolved identifier, + # casing) that /v1/models echoes but which may differ from the path we + # passed; match on the id the load reports so we don't silently fall + # through to models[0] and connect to a different loaded model. + wanted = {requested} + if isinstance(loaded, dict): + wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} + models = _loaded_models(base, key) + match = next((m for m in models if m["id"] in wanted), None) + if match is not None: + return match + if requested: + # We asked Studio to load it and it didn't surface in /v1/models; don't + # silently hand back an unrelated loaded model. + _fail( + f"Studio didn't report '{requested}' as loaded. Double-check the model " + "id, or load it from the model dropdown in the UI." + ) + if not models: + _fail( + "No model is loaded in Studio. Load one from the model dropdown in " + "the UI, or pass --model to load it from here." + ) + return models[0] + + +def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: + # Codex always streams, and Studio only streams /v1/responses from llama-server. + try: + status = _http_json("GET", f"{base}/api/inference/status", key) + except urllib.error.HTTPError as exc: + if exc.code == 404: + return # older server without the endpoint; don't block the launch + raise + if status.get("is_gguf"): + return + hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" + _fail( + f"Codex needs a GGUF model served by llama-server, but {model_id} is on " + f"the transformers backend. Try: unsloth start codex --model {hint}" + ) + + +_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" +# Session overlay applied via `claude --settings`; suppresses the attribution header +# for THIS run only (no ~/.claude write) so llama.cpp KV-cache reuse is preserved. It +# reinforces the CLAUDE_CODE_ATTRIBUTION_HEADER env var on builds that read the setting +# only from settings.json. +_CLAUDE_SETTINGS_OVERLAY = '{"env":{"CLAUDE_CODE_ATTRIBUTION_HEADER":"0"}}' + + +def _claude_version() -> Optional[tuple]: + # None = no local `claude` (a --no-launch printout for another machine; assume a + # current build). An unparseable version is treated as too old for the new flags. + executable = shutil.which("claude") + if executable is None: + return None + try: + result = subprocess.run( + [executable, "--version"], capture_output = True, text = True, timeout = 10 + ) + # Pull the X.Y.Z out of the output rather than assuming it is the first token. + # claude prints it first today ("2.1.98 (Claude Code)"), but a format change + # (e.g. "claude version 2.1.98") shouldn't silently drop the optimization flags; + # no match falls through to "too old", same as an unparseable version. + match = re.search(r"(\d+)\.(\d+)\.(\d+)", result.stdout) + return tuple(int(part) for part in match.groups()) if match else (0,) + except Exception: + return (0,) + + +def _claude_flags() -> list: + # Both knobs preserve llama.cpp KV-cache reuse: --exclude-dynamic-system-prompt-sections + # moves per-session context out of the system prompt, and --settings suppresses the + # attribution header for this session only (no persistent ~/.claude write; the env var + # sets it too). Claude Code < 2.1.98 aborts on unknown flags, so gate on the version; + # no local binary means a printout for another machine, so assume a current build. + version = _claude_version() + if version is not None and version < (2, 1, 98): + return [] + return [_DYNAMIC_SECTIONS_FLAG, "--settings", _CLAUDE_SETTINGS_OVERLAY] + + +def _merge_codex_config(existing: str, base: str) -> str: + chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table + if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): + if chunks[0] and not chunks[0].endswith("\n"): + chunks[0] += "\n" + chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' + # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. + stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") + text = "".join(c for c in chunks if not c.startswith(stale)) + if not text.endswith("\n"): + text += "\n" + if not text.endswith("\n\n"): + text += "\n" + return text + ( + f"{_PROVIDER_HEADER}\n" + 'name = "Unsloth Studio"\n' + f"base_url = {json.dumps(base + '/v1')}\n" + f'env_key = "{_CODEX_ENV_KEY}"\n' + 'wire_api = "responses"\n' + "requires_openai_auth = false\n" + ) + + +def write_codex_config(base: str, model: dict, home: Path) -> None: + home.mkdir(parents = True, exist_ok = True) + + config = home / "config.toml" + existing = config.read_text(encoding = "utf-8") if config.exists() else "" + merged = _merge_codex_config(existing, base) + if merged != existing: + config.write_text(merged, encoding = "utf-8") + typer.echo(f"Updated {config}") + + # oss_provider here too: codex --oss picks the provider from it, and the + # profile layer must beat a user-set value (e.g. "ollama") in config.toml. + profile_text = ( + f'oss_provider = "{_CODEX_PROFILE}"\n' + f'model_provider = "{_CODEX_PROFILE}"\n' + f"model = {json.dumps(model['id'])}\n" + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + profile_text += f"model_context_window = {int(window)}\n" + profile = home / f"{_CODEX_PROFILE}.config.toml" + if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: + profile.write_text(profile_text, encoding = "utf-8") + typer.echo(f"Updated {profile}") + + +def _wsl_windows_executable(command: list) -> Optional[str]: + if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): + return None + executable = shutil.which(command[0]) + if executable and executable.startswith("/mnt/"): + return executable + return None + + +def _looks_like_path(value: str) -> bool: + # A var only wants the WSLENV /p flag if its value is a filesystem path: an + # absolute POSIX path (/...), a UNC path (\\...), or a drive-qualified Windows + # path (C:...). Scalar knobs (e.g. a numeric context window) must pass through + # untranslated, so they get no flag. + return bool(value) and (value.startswith(("/", "\\")) or (len(value) >= 2 and value[1] == ":")) + + +def _wsl_bridge_names(env: dict, unset_env: tuple) -> tuple: + # Build the WSLENV share list for a Windows shim reached from WSL. Path-valued + # vars get /p so WSLENV translates them to the Windows path the /mnt shim can + # actually open; a cleared var carries no value to translate. + names = [name + ("/p" if _looks_like_path(value) else "") for name, value in env.items()] + names.extend(unset_env) + return tuple(dict.fromkeys(names)) + + +def _merge_wslenv(current: str, names: tuple) -> str: + # Index WSLENV entries by bare var name, preserving first-seen order. The vars we + # bridge are applied last so our entry wins: a user's pre-existing unflagged "HOME" + # is upgraded to "HOME/p" (rather than left as-is), since WSLENV ignores a duplicate + # name and a bare entry would leave the path untranslated for a Windows shim. + ordered = [] + by_name = {} + for entry in (*current.split(":"), *names): + if not entry: + continue + base = entry.split("/", 1)[0] + if base not in by_name: + ordered.append(base) + by_name[base] = entry + return ":".join(by_name[base] for base in ordered) + + +def _powershell_quote(arg: str) -> str: + # PowerShell reads single-quoted strings literally (an embedded ' is doubled), so + # JSON args such as `--settings {"env":...}` survive intact. list2cmdline's + # backslash-escaped double quotes are cmd.exe syntax and PowerShell mis-parses them. + if arg and re.fullmatch(r"[A-Za-z0-9_./:=+-]+", arg): + return arg + return "'" + arg.replace("'", "''") + "'" + + +def _print_env( + env: dict, + command: list, + unset_env: tuple = (), + wsl_env_bridge: tuple = (), +) -> None: + if os.name == "nt": + for name in unset_env: + typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") + for name, value in env.items(): + # PowerShell: ` is the escape char, and $ triggers expansion inside "". + escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") + typer.echo(f'$env:{name} = "{escaped}"') + typer.echo(" ".join(_powershell_quote(arg) for arg in command)) + return + for name in unset_env: + typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") + for name, value in env.items(): + typer.echo(f"export {name}={shlex.quote(value)}") + if wsl_env_bridge: + typer.echo( + f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + # The final line is a SELF-CONTAINED one-liner (inline env, VAR=... cmd) rather than a + # bare command. People copy just the last line, and a bare `codex`/`claude` would then + # run against their real ~/.codex or Anthropic credentials with zero isolation -- e.g. + # inheriting a pre-existing damaged ~/.codex state DB and blaming the recipe. Inline + # assignments scope every var (and empty-string the conflicting ones) to this single + # invocation, so a partial copy behaves the same as pasting the whole block. + inline = [f"{name}=" for name in unset_env] + inline += [f"{name}={shlex.quote(value)}" for name, value in env.items()] + if wsl_env_bridge: + inline.append( + f"WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + typer.echo(" ".join((*inline, shlex.join(command)))) + + +def _install_agent(name: str, install_hint: str) -> Optional[str]: + # Missing agent under --launch: offer to run its documented install command, then + # re-resolve it on PATH. Consent-based (we never auto-run a remote install script + # silently), and a non-interactive stdin cannot answer the prompt, so both the + # no-TTY and declined cases return None and let the caller print the hint and exit. + if not sys.stdin.isatty(): + return None + typer.echo(f"`{name}` is not installed.") + if not typer.confirm(f"Install it now with `{install_hint}`?", default = False): + return None + # Run each hint through the shell it is written for: PowerShell (irm | iex, or npm) + # on Windows, /bin/sh (curl | bash, or npm) everywhere else. + if os.name == "nt": + install_command = ["powershell", "-NoProfile", "-Command", install_hint] + else: + install_command = ["/bin/sh", "-c", install_hint] + if subprocess.run(install_command).returncode != 0: + _fail(f"Install command failed. Run it yourself, then re-run: {install_hint}") + executable = shutil.which(name) + if executable is None: + _fail( + f"`{name}` installed but isn't on PATH yet. Open a new shell (or add it to " + f"PATH), then re-run. Install command: {install_hint}" + ) + return executable + + +def _launch( + command: list, + env: dict, + install_hint: str, + unset_env: tuple = (), +) -> NoReturn: + executable = shutil.which(command[0]) or _install_agent(command[0], install_hint) + if executable is None: + _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") + wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else () + child_env = dict(os.environ) + if wsl_env_bridge: + child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) + for name in unset_env: + child_env[name] = "" + else: + for name in unset_env: + child_env.pop(name, None) + child_env.update(env) + # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. + previous = signal.signal(signal.SIGINT, signal.SIG_IGN) + try: + code = subprocess.run([executable, *command[1:]], env = child_env).returncode + finally: + signal.signal(signal.SIGINT, previous) + # Negative returncode means killed by signal N; shells expect 128+N. + raise typer.Exit(code = code if code >= 0 else 128 - code) + + +def _connect( + api_key: Optional[str], + model: Optional[str], + load: LoadOptions = LoadOptions(), + *, + serve: bool = False, + launch: bool = True, +) -> tuple: + # `--model org/name:QUANT` is shorthand for `--model org/name --gguf-variant QUANT`. + # Split it before we match/serve so the attach path resolves against the already-loaded + # `org/name` (listed without the suffix) instead of reloading a `:`-suffixed repo id -- + # which Studio rejects and which would evict a model another session is using. + if model: + repo, variant = _split_repo_variant(model) + if variant: + model = repo + if not load.gguf_variant: + load = load._replace(gguf_variant = variant) + base, server = _require_studio(model, load, serve = serve, launch = launch) + try: + key = _agent_api_key(base, api_key, auto_started = server is not None) + # A server we just started has exactly the requested model loaded, so resolve to + # whatever it is serving instead of re-matching the raw --model string. + entry = _resolve_model(base, key, None if server is not None else model, load) + except BaseException: + _shutdown_auto_served() + raise + return base, key, entry + + +def _run( + base: str, + entry: dict, + env: dict, + command: list, + *, + launch: bool, + install_hint: str, + unset_env: tuple = (), + clear_screen: bool = False, +) -> None: + # Some agents (Pi) render inline from wherever the cursor sits: their first + # paint assumes a clean screen rather than clearing or entering the + # alternate screen themselves. Hand them one so the session doesn't start + # mid-scroll under our connection output. click.clear() is cross-platform + # and a no-op when stdout is not a terminal (piped/CI), so transcripts and + # --no-launch recipes stay intact. + if launch and clear_screen: + click.clear() + typer.echo(f"Studio {base} · model {entry['id']}") + wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else () + if not launch: + _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) + return + try: + _launch(command, env, install_hint = install_hint, unset_env = unset_env) + finally: + # Tear down a server we auto-started once the agent session ends (no-op otherwise). + _shutdown_auto_served() + + +def _agents_config_root() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agents" + + +@contextlib.contextmanager +def _session_config(agent: str, launch: bool): + """Yield a private directory for an agent's session config (never the user's own). + + launch: an ephemeral temp dir removed after the agent process exits, so nothing + persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later + on this machine), reused across runs. Either way the user's real ~/. + config is left untouched. + """ + if launch: + path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-")) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors = True) + else: + # Never wipe this dir: a previously printed recipe may still be running + # an agent whose sessions/state live here, and every config writer + # merges idempotently into an existing home anyway. + path = _agents_config_root() / agent + path.mkdir(parents = True, exist_ok = True, mode = 0o700) + yield path + + +def write_openclaw_config( + base: str, + key: str, + model: dict, + path: Path, + yolo: bool = False, +) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). + provider_model = {"id": model["id"], "name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + provider_model["contextWindow"] = int(window) + models = _subdict(config, "models") + models.setdefault("mode", "merge") + _subdict(models, "providers")["unsloth"] = { + "baseUrl": f"{base}/v1", + "apiKey": key, + "api": "openai-completions", + "models": [provider_model], + } + # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). + defaults = _subdict(_subdict(config, "agents"), "defaults") + _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" + # Unauthenticated loopback gateway: without auth.mode=none the client won't open + # the websocket. The daemon must still be started separately (`openclaw gateway`). + gateway = _subdict(config, "gateway") + gateway.setdefault("mode", "local") + _subdict(gateway, "auth").setdefault("mode", "none") + if yolo: + # OpenClaw has no --yolo flag, and it gates tool execution on BOTH the + # tools.exec config AND a host-local approvals file (the stricter wins), so + # setting only the config still lets the agent prompt/deny. Set both, mirroring + # `openclaw exec-policy preset yolo`. + exec_policy = _subdict(_subdict(config, "tools"), "exec") + exec_policy["host"] = "gateway" + exec_policy["security"] = "full" + exec_policy["ask"] = "off" + # Approvals file in OPENCLAW_STATE_DIR (== this config's dir). ask=off means + # nothing is ever prompted, so the runtime socket block is unnecessary here. + approvals = path.parent / "exec-approvals.json" + _write_private_json( + approvals, + {"version": 1, "defaults": {"security": "full", "ask": "off", "askFallback": "full"}}, + ) + typer.echo(f"Updated {approvals}") + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def write_opencode_config( + base: str, + key: str, + model: dict, + path: Path, + yolo: bool = False, +) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + config.setdefault("$schema", "https://opencode.ai/config.json") + model_entry = {"name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # A custom-provider model with no limit defaults to context 0, which silently + # disables OpenCode's auto-compaction; declare the real window (and a sane + # output cap) so it compacts instead of overflowing the server. + model_entry["limit"] = {"context": window, "output": min(window // 4, 8192)} + _subdict(config, "provider")["unsloth"] = { + "npm": "@ai-sdk/openai-compatible", + "name": "Unsloth Studio", + "options": {"baseURL": f"{base}/v1", "apiKey": key}, + "models": {model["id"]: model_entry}, + } + # OpenCode selects a model by "/". + config["model"] = f"unsloth/{model['id']}" + if window: + # Compact with ~10% headroom (near 90% full). The fixed 20k-token default + # buffer over-compacts, or never settles, on a small local context. + compaction = _subdict(config, "compaction") + compaction["auto"] = True + compaction["reserved"] = max(1, window // 10) + if yolo: + # OpenCode has no --yolo flag; auto-approve is the config `permission` block + # (singular). Allow the prompting tools so tool calls don't block on the TUI. + config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def write_hermes_config(base: str, model: dict, path: Path) -> None: + import yaml + + config: dict = {} + if path.exists(): + try: + loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) + except (yaml.YAMLError, OSError): + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + if isinstance(loaded, dict): + config = loaded + elif loaded is not None: + # Non-empty, non-mapping YAML is a user-managed file; leave it. + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + # Hermes only reads the key for a *named* custom provider (a bare + # `provider: custom` ignores it), so register it under providers.*. + _subdict(config, "model").update( + provider = f"custom:{_HERMES_PROVIDER}", + default = model["id"], + api_mode = "openai", + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # Hermes auto-detects context from GET /v1/models, but OpenAI's schema has no + # context field, so it can fall back to a 256k default that overflows a small + # local model. Pin the real window (top-level model.context_length is the + # highest-priority override) and compact at 90% of it (Hermes defaults to 50%). + if window >= _HERMES_MIN_CONTEXT: + _subdict(config, "model")["context_length"] = window + _subdict(config, "compression").update(enabled = True, threshold = 0.9) + else: + # Below Hermes' 64,000-token floor it refuses to initialize, so claim + # the floor and shrink the threshold so compaction still fires at 90% + # of the REAL window (the threshold is a fraction of the claimed + # context_length). The auxiliary override keeps the same floor check + # from rejecting the compression model mid-session. + _subdict(config, "model")["context_length"] = _HERMES_MIN_CONTEXT + threshold = round(0.9 * window / _HERMES_MIN_CONTEXT, 4) + _subdict(config, "compression").update(enabled = True, threshold = threshold) + auxiliary = _subdict(_subdict(config, "auxiliary"), "compression") + auxiliary["context_length"] = _HERMES_MIN_CONTEXT + _subdict(config, "providers")[_HERMES_PROVIDER] = { + "base_url": f"{base}/v1", + "api_mode": "openai", + "key_env": _HERMES_ENV_KEY, + } + text = yaml.safe_dump(config, sort_keys = False) + if not path.exists() or path.read_text(encoding = "utf-8") != text: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(text, encoding = "utf-8") + typer.echo(f"Updated {path}") + + +def write_pi_config(base: str, key: str, model: dict, path: Path) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Pi reads custom providers from ~/.pi/agent/models.json (HOME-relocated for the + # session). Studio is a generic OpenAI-compatible /v1 endpoint, and the key lives + # in the config rather than the env (matching openclaw/opencode). + provider_model = {"id": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # An unspecified model defaults to contextWindow 128000 / maxTokens 16384, + # far larger than a small Studio context, so Pi compacts too late and overflows + # the server. Pin the real window and a sane output cap (mirrors OpenCode). + provider_model["contextWindow"] = window + provider_model["maxTokens"] = min(window // 4, 8192) + _subdict(config, "providers")[_PI_PROVIDER] = { + "api": "openai-completions", + "baseUrl": f"{base}/v1", + "apiKey": key, + "models": [provider_model], + } + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +@start_app.command("claude", context_settings = _PASSTHROUGH) +def claude( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Claude Code at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + model_id = entry["id"] + + env = { + "ANTHROPIC_BASE_URL": base, + "ANTHROPIC_AUTH_TOKEN": key, + "ANTHROPIC_MODEL": model_id, + # Session-only (no ~/.claude write): suppress the attribution header so + # llama.cpp KV-cache reuse is preserved; --settings below reinforces it. + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + # Update checks, beta features, and other background requests either + # stall against a local server or evict the conversation from + # llama-server's KV-cache slots, so turn off everything nonessential. + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + # A local server streams in bursts; disable the full-screen TUI redraw so the + # terminal doesn't flicker between tokens. + "CLAUDE_CODE_NO_FLICKER": "1", + } + # Claude Code auto-compacts against its native (~600k token) window; a local + # model's context is usually far smaller, so size the window to the loaded + # model's real context length. Otherwise the conversation overflows the + # server's window (silent truncation) long before Claude decides to compact. + # codex/openclaw get the same value through their config (model_context_window + # / contextWindow); Claude has no config file, so it rides on the env var. + window = entry.get("context_length") or entry.get("max_context_length") + if window: + env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] = str(int(window)) + # Compact at 90% of that window; the override only takes effect once the + # window is set, and it can only lower the threshold, so it just guarantees + # headroom before the server's context limit instead of relying on Claude's + # default (which is tuned for its native 200K/1M window). + env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] = "90" + # --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions. + # IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a + # sandbox is detected, and we don't want to falsely claim one on the user's host. + command = [ + "claude", + "--model", + model_id, + *_claude_flags(), + *_yolo_command_flags("claude", yolo), + *ctx.args, + ] + install_hint = ( + "irm https://claude.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://claude.ai/install.sh | bash" + ) + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + unset_env = _CLAUDE_ENV_UNSET, + ) + + +@start_app.command("codex", context_settings = _PASSTHROUGH) +def codex( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenAI Codex at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + # This preflight runs after _connect may have auto-started a server but before _run + # installs its teardown finally, so tear the server down here if it rejects the model + # (e.g. a transformers-backend model) rather than leaving it on the atexit backstop. + try: + _require_gguf_for_codex(base, key, entry["id"]) + except BaseException: + _shutdown_auto_served() + raise + command = [ + "codex", + "--oss", + "--profile", + _CODEX_PROFILE, + *_yolo_command_flags("codex", yolo), + *ctx.args, + ] + with _session_config("codex", launch) as home: + write_codex_config(base, entry, home) + env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)} + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") + + +@start_app.command("openclaw", context_settings = _PASSTHROUGH) +def openclaw( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenClaw at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["openclaw", *ctx.args] + install_hint = ( + "iwr -useb https://openclaw.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://openclaw.ai/install.sh | bash" + ) + with _session_config("openclaw", launch) as cfg: + config_path = cfg / "openclaw.json" + # key lives in the config, not the env; --yolo writes the exec policy here too. + write_openclaw_config(base, key, entry, config_path, yolo = yolo) + # Scope both config and state so OpenClaw never touches the user's ~/.openclaw. + env = {"OPENCLAW_CONFIG_PATH": str(config_path), "OPENCLAW_STATE_DIR": str(cfg)} + _run(base, entry, env, command, launch = launch, install_hint = install_hint) + + +@start_app.command("opencode", context_settings = _PASSTHROUGH) +def opencode( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenCode at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["opencode", *ctx.args] + with _session_config("opencode", launch) as cfg: + config_path = cfg / "opencode.json" + # OPENCODE_CONFIG is an overlay (loaded between the user's global and project + # configs), so this adds the Unsloth provider/model for the session without + # changing the user's default model. Key lives in the config, not the env. + write_opencode_config(base, key, entry, config_path, yolo = yolo) + # A project's own opencode.json outranks OPENCODE_CONFIG, so the session model + # pin (and --yolo permissions) would silently lose to a repo config. Carry the + # settings that must win in OPENCODE_CONFIG_CONTENT, which outranks project + # config; the API key stays in the private file, never in the printed env. + inline_config: dict = {"model": f"unsloth/{entry['id']}"} + if yolo: + inline_config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + env = { + "OPENCODE_CONFIG": str(config_path), + "OPENCODE_CONFIG_CONTENT": json.dumps(inline_config), + } + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g opencode-ai") + + +@start_app.command("hermes", context_settings = _PASSTHROUGH) +def hermes( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Hermes (Nous Research) at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] + install_hint = ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash" + ) + with _session_config("hermes", launch) as home: + # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) + # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session. + write_hermes_config(base, entry, home / "config.yaml") + env = {_HERMES_ENV_KEY: key, "HERMES_HOME": str(home)} + _run(base, entry, env, command, launch = launch, install_hint = install_hint) + + +@start_app.command("pi", context_settings = _PASSTHROUGH) +def pi( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Pi (coding agent) at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + # Pi defaults to the google provider, so pin our provider/model on the command + # line; the custom OpenAI-compatible endpoint itself is only configurable via + # ~/.pi/agent/models.json. + command = [ + "pi", + "--provider", + _PI_PROVIDER, + "--model", + entry["id"], + *_yolo_command_flags("pi", yolo), + *ctx.args, + ] + # --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs + # no install scripts), so accepting the prompt skips dependency lifecycle scripts. + install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent" + with _session_config("pi", launch) as home: + # Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers + # it over $HOME/.pi/agent), so pin it at the session dir: an inherited + # PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real + # config and skip our provider/key. HOME is relocated too so any other ~/.pi paths + # stay in the session. The key rides in the config rather than the env. + pi_agent_dir = home / ".pi" / "agent" + write_pi_config(base, key, entry, pi_agent_dir / "models.json") + env = {"HOME": str(home), "PI_CODING_AGENT_DIR": str(pi_agent_dir)} + if os.name == "nt" or os.environ.get("WSL_DISTRO_NAME"): + # Node resolves ~/.pi via USERPROFILE (then HOMEDRIVE + HOMEPATH) on Windows, + # not HOME. Set them whenever Pi may run as a Windows process: native Windows, + # or a /mnt Windows shim launched from WSL (the WSLENV bridge then translates + # the path). Otherwise the Windows process falls back to the user's real + # %USERPROFILE%\.pi. splitdrive yields no drive off a POSIX path, so + # HOMEDRIVE/HOMEPATH stay unset there. + env["USERPROFILE"] = str(home) + drive, tail = os.path.splitdrive(str(home)) + if drive: + env["HOMEDRIVE"], env["HOMEPATH"] = drive, tail + # Pi paints inline from the current cursor position (no alternate screen, + # no clear on first render), so give it the clean screen it assumes. + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + clear_screen = True, + ) diff --git a/unsloth_cli/tests/test_connect.py b/unsloth_cli/tests/test_connect.py deleted file mode 100644 index e76a892647..0000000000 --- a/unsloth_cli/tests/test_connect.py +++ /dev/null @@ -1,954 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -"""Tests for `unsloth connect` — config merging and launch env, no network.""" - -from __future__ import annotations - -import json -import os -import sys -import urllib.error -from pathlib import Path -from types import SimpleNamespace - -_REPO_ROOT = Path(__file__).resolve().parents[2] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - - -import pytest -from typer.testing import CliRunner - -import unsloth_cli.commands.connect as connect - -BASE = "http://127.0.0.1:8888" -MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} - - -# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and -# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. -def _assert_env_set(output: str, name: str, value: str) -> None: - needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" - assert needle in output, f"{needle!r} not found in:\n{output}" - - -def _assert_env_unset(output: str, name: str) -> None: - needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" - assert needle in output, f"{needle!r} not found in:\n{output}" - - -@pytest.fixture() -def claude_settings(tmp_path, monkeypatch): - path = tmp_path / "claude" / "settings.json" - monkeypatch.setattr(connect, "claude_settings_path", lambda: path) - return path - - -def test_claude_settings_created_when_missing(claude_settings): - connect.ensure_claude_attribution_header() - settings = json.loads(claude_settings.read_text()) - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_claude_settings_merge_preserves_existing(claude_settings): - claude_settings.parent.mkdir(parents = True) - claude_settings.write_text( - json.dumps({"effortLevel": "high", "env": {"CLAUDE_CODE_ENABLE_TELEMETRY": "0"}}) - ) - connect.ensure_claude_attribution_header() - settings = json.loads(claude_settings.read_text()) - assert settings["effortLevel"] == "high" - assert settings["env"]["CLAUDE_CODE_ENABLE_TELEMETRY"] == "0" - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_claude_settings_already_set_untouched(claude_settings): - claude_settings.parent.mkdir(parents = True) - original = json.dumps({"env": {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}}) - claude_settings.write_text(original) - connect.ensure_claude_attribution_header() - assert claude_settings.read_text() == original - - -def test_claude_settings_bad_json_left_alone(claude_settings, capsys): - claude_settings.parent.mkdir(parents = True) - claude_settings.write_text("{not json") - connect.ensure_claude_attribution_header() - assert claude_settings.read_text() == "{not json" - assert "couldn't parse" in capsys.readouterr().err - - -def _fake_claude(monkeypatch, version_output: str) -> None: - monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") - monkeypatch.setattr( - connect.subprocess, - "run", - lambda *args, **kwargs: SimpleNamespace(stdout = version_output), - ) - - -def test_cache_flags_passed_to_supported_claude(monkeypatch): - _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") - assert connect._claude_cache_flags() == ["--exclude-dynamic-system-prompt-sections"] - - -def test_cache_flags_skipped_on_old_claude(monkeypatch): - _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") - assert connect._claude_cache_flags() == [] - - -def test_cache_flags_skipped_on_unparseable_version(monkeypatch): - _fake_claude(monkeypatch, "weird build string\n") - assert connect._claude_cache_flags() == [] - - -def _parse_toml(text: str) -> dict: - tomllib = pytest.importorskip("tomllib") - return tomllib.loads(text) - - -def test_merge_codex_config_fresh(): - merged = connect._merge_codex_config("", BASE) - parsed = _parse_toml(merged) - assert parsed["oss_provider"] == "unsloth_api" - provider = parsed["model_providers"]["unsloth_api"] - assert provider["base_url"] == f"{BASE}/v1" - assert provider["wire_api"] == "responses" - assert provider["requires_openai_auth"] is False - - -def test_merge_codex_config_replaces_stale_block(): - existing = ( - 'model = "gpt-5"\n' - "\n" - "[model_providers.unsloth_api]\n" - 'base_url = "http://old-host:9999/v1"\n' - 'wire_api = "chat"\n' - "\n" - "[model_providers.unsloth_api.http_headers]\n" - 'x-old = "1"\n' - "\n" - "[model_providers.ollama]\n" - 'base_url = "http://localhost:11434/v1"\n' - ) - merged = connect._merge_codex_config(existing, BASE) - parsed = _parse_toml(merged) - assert parsed["model"] == "gpt-5" - assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" - assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" - assert "http_headers" not in parsed["model_providers"]["unsloth_api"] - assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" - assert connect._merge_codex_config(merged, BASE) == merged - - -def test_merge_codex_config_keeps_user_oss_provider(): - merged = connect._merge_codex_config('oss_provider = "ollama"\n', BASE) - assert _parse_toml(merged)["oss_provider"] == "ollama" - - -def test_write_codex_config_profile(tmp_path, monkeypatch): - monkeypatch.setenv("CODEX_HOME", str(tmp_path)) - connect.write_codex_config(BASE, MODEL) - profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) - assert profile["oss_provider"] == "unsloth_api" - assert profile["model_provider"] == "unsloth_api" - assert profile["model"] == MODEL["id"] - assert profile["model_context_window"] == 131072 - config = _parse_toml((tmp_path / "config.toml").read_text()) - assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" - - -@pytest.fixture() -def fake_studio(tmp_path, monkeypatch, claude_settings): - calls = [] - state = {"models": [MODEL]} - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - calls.append((method, url, payload)) - if url.endswith("/v1/models"): - return {"object": "list", "data": state["models"]} - if url.endswith("/api/inference/status"): - return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} - if url.endswith("/api/auth/api-keys"): - return {"key": "sk-unsloth-feedfacefeedface"} - if url.endswith("/api/inference/load"): - state["models"] = [{"id": payload["model_path"], "context_length": 4096}] - return {} - raise AssertionError(f"unexpected request: {method} {url}") - - monkeypatch.setattr(connect, "find_studio_server", lambda: BASE) - # Identity handshake has its own tests; trust the loopback server here. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: True) - # _studio_token / api-keys are faked so the mint flow stays offline. - monkeypatch.setattr(connect, "_studio_token", lambda: "jwt-token") - monkeypatch.setattr(connect, "_http_json", http_json) - monkeypatch.setattr(connect, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") - # No `claude` on PATH, so _claude_cache_flags never probes the real binary. - monkeypatch.setattr(connect.shutil, "which", lambda _: None) - monkeypatch.setenv("CODEX_HOME", str(tmp_path / "codex")) - monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) - return calls - - -def test_connect_claude_no_launch(fake_studio, claude_settings): - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_unset(result.output, "ANTHROPIC_API_KEY") - _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") - _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) - _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") - _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") - assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output - settings = json.loads(claude_settings.read_text()) - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): - captured = {} - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") - monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") - monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") - monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) - - def run(command, env): - captured["command"] = command - captured["env"] = env - return SimpleNamespace(returncode = 0) - - monkeypatch.setattr(connect.subprocess, "run", run) - result = CliRunner().invoke(connect.connect_app, ["claude"]) - - assert result.exit_code == 0, result.output - assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] - assert "ANTHROPIC_API_KEY" not in captured["env"] - assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] - assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" - assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE - assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] - - -@pytest.mark.skipif( - os.name == "nt", - reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " - "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", -) -def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): - captured = {} - monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") - monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") - monkeypatch.setattr( - connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" - ) - monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) - - def run(command, env): - captured["command"] = command - captured["env"] = env - return SimpleNamespace(returncode = 0) - - monkeypatch.setattr(connect.subprocess, "run", run) - result = CliRunner().invoke(connect.connect_app, ["claude"]) - - assert result.exit_code == 0, result.output - assert captured["command"] == [ - "/mnt/c/Users/samle/AppData/Roaming/npm/claude", - "--model", - MODEL["id"], - ] - assert captured["env"]["ANTHROPIC_API_KEY"] == "" - assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" - assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" - assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE - assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] - for name in ( - "ANTHROPIC_AUTH_TOKEN", - "ANTHROPIC_BASE_URL", - "ANTHROPIC_MODEL", - "ANTHROPIC_API_KEY", - "CLAUDE_CODE_OAUTH_TOKEN", - ): - assert name in captured["env"]["WSLENV"].split(":") - - -@pytest.mark.skipif( - os.name == "nt", - reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " - "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", -) -def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): - monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setattr( - connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" - ) - - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - - assert result.exit_code == 0, result.output - assert "export ANTHROPIC_API_KEY=" in result.output - assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output - assert "export WSLENV=" in result.output - assert "ANTHROPIC_AUTH_TOKEN" in result.output - assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output - - -def test_connect_codex_no_launch(fake_studio, tmp_path): - result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - assert "codex --oss --profile unsloth_api" in result.output - assert (tmp_path / "codex" / "config.toml").exists() - assert (tmp_path / "codex" / "unsloth_api.config.toml").exists() - - -def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): - CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - # First run mints; second reuses the minted key cached for this server. - mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] - assert len(mints) == 1 - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] - - -def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): - CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - # Reused, not re-minted (a mint would return the feedface stand-in). - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - # An explicit key is remembered as "saved" so it replays without the handshake. - assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] - - -def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): - cache = tmp_path / "agent_api_key.json" - cache.write_text( - json.dumps( - {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} - ) - ) - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/v1/models") and token == "sk-unsloth-stale": - raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - # The working key moves to the front so the next run tries it first. - cached = json.loads(cache.read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] - - -def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): - # Legacy unscoped caches have no server binding (could leak across servers), - # so they're ignored: a fresh key is minted and stored scoped to this server. - (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] - assert "key" not in cached # legacy field collapsed away - - -def test_connect_model_flag_loads_on_server(fake_studio): - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] - ) - assert result.exit_code == 0, result.output - loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] - assert loads == [ - ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) - ] - _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") - - -def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): - # Studio registers a loaded model under a canonical id (resolved identifier - # / casing) that can differ from the path we passed. The agent must connect - # to that model, not silently fall through to the first loaded one. - requested = "Unsloth/Qwen3.5-35B-A3B" - canonical = "unsloth/Qwen3.5-35B-A3B" - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/load"): - return {"model": canonical, "display_name": canonical} - if url.endswith("/v1/models"): - # Decoy sorts first, so models[0] is the wrong pick on the old code. - return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", requested] - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) - - -def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): - monkeypatch.setattr( - connect, - "_http_json", - lambda method, url, token, payload = None, timeout = 30, error = None: ( - {"key": "sk-unsloth-feedfacefeedface"} - if url.endswith("/api/auth/api-keys") - else {"object": "list", "data": []} - ), - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "No model is loaded" in result.output - - -def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): - # Studio never surfaces the requested model; fail loudly rather than - # silently connecting to whatever else happens to be loaded. - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/load"): - return {} - if url.endswith("/v1/models"): - return {"object": "list", "data": [MODEL]} # decoy; request never appears - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] - ) - assert result.exit_code == 1 - assert "unsloth/Missing-7B" in result.output - - -def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/status"): - return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) - assert result.exit_code == 1 - assert "GGUF" in result.output - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - - -def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): - # A server known only by URL + health check is unverified: keyless connect - # must refuse and make no request at all. - monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.evil.example:8888") - result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) - assert result.exit_code == 1 - assert "Settings → API" in result.output - assert "--api-key" in result.output - assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) - - -def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): - # User named both server and key, so it's their choice; only auto-send is blocked. - monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.example:8888") - result = CliRunner().invoke( - connect.connect_app, - ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - - -def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): - # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; - # auto-minting stays blocked for non-loopback. - remote = "http://studio.example:8888" - monkeypatch.setattr(connect, "find_studio_server", lambda: remote) - (tmp_path / "agent_api_key.json").write_text( - json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted - - -def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): - # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an - # error, not a silent local model load (which they did not ask for). - import typer - - import unsloth_cli._inference as inference - - monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") - monkeypatch.setattr( - inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" - ) - with pytest.raises(typer.Exit): - inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) - - -def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): - # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback - # server can't be verified, fall back to a local load rather than erroring. - import unsloth_cli._inference as inference - - monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) - monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") - monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) - assert ( - inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) - is None - ) - - -def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( - fake_studio, tmp_path, monkeypatch -): - # With no saved key, the next step would auto-mint; an unverified loopback - # server (port squatter) must be refused, with nothing sent. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "--api-key" in result.output - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted - - -def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): - # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) - # replays on keyless runs without the handshake, scoped to its own base. - cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted - - -def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): - # A "minted" key is NOT replayed to an unverified loopback server: minting and - # minted-key replay both sit behind the handshake, so a squatter can't grab it. - cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "--api-key" in result.output - assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent - - -def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): - # An explicit key is the user's deliberate choice, so it does not require - # the automatic identity handshake. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - - -def _serve_identity(proof_for): - """Start a localhost HTTP server answering /api/auth/identity with - proof_for(nonce_bytes). Returns (base_url, shutdown).""" - import base64 - import threading - from http.server import BaseHTTPRequestHandler, HTTPServer - from urllib.parse import parse_qs, urlparse - - class Handler(BaseHTTPRequestHandler): - def do_GET(self): - parsed = urlparse(self.path) - if parsed.path != "/api/auth/identity": - self.send_response(404) - self.end_headers() - return - nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) - host, port = self.server.server_address[0], self.server.server_address[1] - body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(body) - - def log_message(self, *a): - pass - - server = HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target = server.serve_forever, daemon = True).start() - base = f"http://127.0.0.1:{server.server_address[1]}" - return base, server.shutdown - - -def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): - # Real crypto end to end: verify_studio_identity reads the install secret from - # an isolated DB; a "good" server proves the same secret, a spoofing one can't. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: # backend not importable here (e.g. missing deps) - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - good = lambda nonce, host, port: storage.compute_identity_proof( - nonce, host, port - ) # real secret - bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret - base_ok, stop_ok = _serve_identity(good) - base_bad, stop_bad = _serve_identity(bad) - try: - assert inference.verify_studio_identity(base_ok) is True - assert inference.verify_studio_identity(base_bad) is False - finally: - stop_ok() - stop_bad() - - -def _serve_redirect(target): - """Start a localhost server that 302-redirects every GET to target+path.""" - import threading - from http.server import BaseHTTPRequestHandler, HTTPServer - - class Handler(BaseHTTPRequestHandler): - def do_GET(self): - self.send_response(302) - self.send_header("Location", target + self.path) - self.end_headers() - - def log_message(self, *a): - pass - - server = HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target = server.serve_forever, daemon = True).start() - base = f"http://127.0.0.1:{server.server_address[1]}" - return base, server.shutdown - - -def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): - # A squatter could 302 /api/auth/identity to the real Studio and relay its - # proof; redirects must be refused so the squatter's base isn't accepted. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - real_base, stop_real = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) - ) - squatter_base, stop_squatter = _serve_redirect(real_base) - try: - assert inference.verify_studio_identity(real_base) is True # direct: ok - assert inference.verify_studio_identity(squatter_base) is False # relayed: refused - finally: - stop_real() - stop_squatter() - - -def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): - # A squatter that proxies the nonce to the real Studio on another port gets a - # proof bound to *that* port; the client expects one bound to the port it - # connected to, so the relayed proof is rejected. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - real_base, stop_real = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) - ) - real_port = int(real_base.rsplit(":", 1)[1]) - # The squatter answers on its own port but returns the proof for the real port. - squatter_base, stop_squatter = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) - ) - try: - assert inference.verify_studio_identity(real_base) is True - assert inference.verify_studio_identity(squatter_base) is False - finally: - stop_real() - stop_squatter() - - -@pytest.mark.parametrize( - "url, loopback", - [ - ("http://127.0.0.1:8888", True), - ("http://localhost:8888", True), - ("http://[::1]:8888", True), - ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 - ("http://0.0.0.0:8888", False), - ("http://10.0.0.5:8888", False), - ("http://studio.evil.example:8888", False), - ("https://studio.example.com", False), - ], -) -def test_is_loopback_url(url, loopback): - assert connect.is_loopback_url(url) is loopback - - -def test_connect_no_studio_errors(fake_studio, monkeypatch): - monkeypatch.setattr(connect, "find_studio_server", lambda: None) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "No running Studio server" in result.output - - -def test_connect_explicit_api_key_skips_mint(fake_studio): - result = CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) - - -# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── - - -def test_write_openclaw_config_fresh(tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - provider = config["models"]["providers"]["unsloth"] - assert provider["baseUrl"] == f"{BASE}/v1" - assert provider["apiKey"] == "sk-unsloth-abc" - assert provider["api"] == "openai-completions" - assert provider["models"] == [ - {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} - ] - # The default model must be pinned or OpenClaw has nothing active. - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - assert config["gateway"]["mode"] == "local" - assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway - if os.name != "nt": # the file holds an API key - assert path.stat().st_mode & 0o777 == 0o600 - - -def test_write_openclaw_config_preserves_and_idempotent(tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - path.write_text( - json.dumps( - { - "theme": "dark", - "agents": {"defaults": {"temperature": 0.5}}, - "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, - } - ) - ) - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - assert config["theme"] == "dark" - assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - assert config["models"]["mode"] == "replace" # user's mode is left as-is - assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" - assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" - before = path.read_text() - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == before - - -def test_write_openclaw_config_corrupt_left_alone(tmp_path, monkeypatch, capsys): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - path.write_text("{not json") - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == "{not json" - assert "couldn't parse" in capsys.readouterr().err - - -def test_connect_openclaw_no_launch(fake_studio, tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - result = CliRunner().invoke(connect.connect_app, ["openclaw", "--no-launch"]) - assert result.exit_code == 0, result.output - assert "openclaw" in result.output - assert "export" not in result.output # key lives in the config, not the env - config = json.loads(path.read_text()) - assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - # OpenAI /v1/chat/completions works on either backend — no GGUF gate. - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) - - -# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── - - -def test_write_opencode_config_fresh(tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - provider = config["provider"]["unsloth"] - assert provider["npm"] == "@ai-sdk/openai-compatible" - assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} - assert provider["models"] == {MODEL["id"]: {"name": MODEL["id"]}} - assert config["model"] == f"unsloth/{MODEL['id']}" - - -def test_write_opencode_config_preserves_and_idempotent(tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - path.write_text( - json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) - ) - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - assert config["theme"] == "tokyonight" - assert config["provider"]["anthropic"]["name"] == "Anthropic" - assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" - before = path.read_text() - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == before - - -def test_connect_opencode_no_launch(fake_studio, tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) - assert result.exit_code == 0, result.output - assert "opencode" in result.output - config = json.loads(path.read_text()) - assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["model"] == f"unsloth/{MODEL['id']}" - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) - - -# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── - - -@pytest.fixture() -def hermes_config(tmp_path, monkeypatch): - path = tmp_path / "config.yaml" - monkeypatch.setattr(connect, "hermes_config_path", lambda: path) - return path - - -def test_write_hermes_config_fresh(hermes_config): - yaml = pytest.importorskip("yaml") - connect.write_hermes_config(BASE, MODEL) - config = yaml.safe_load(hermes_config.read_text()) - # Hermes only honors the key for a *named* custom provider, so the endpoint - # is registered under providers.* and model.provider points at it. - assert config["model"]["provider"] == "custom:unsloth" - assert config["model"]["default"] == MODEL["id"] - assert config["model"]["api_mode"] == "openai" - provider = config["providers"]["unsloth"] - assert provider["base_url"] == f"{BASE}/v1" - assert provider["api_mode"] == "openai" - assert provider["key_env"] == "UNSLOTH_API_KEY" - # The key is resolved from the launch env, never written to disk. - assert "sk-unsloth" not in hermes_config.read_text() - - -def test_write_hermes_config_preserves_and_idempotent(hermes_config): - yaml = pytest.importorskip("yaml") - hermes_config.write_text( - yaml.safe_dump( - { - "terminal": {"backend": "local"}, - "model": {"temperature": 0.7}, - "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, - } - ) - ) - connect.write_hermes_config(BASE, MODEL) - config = yaml.safe_load(hermes_config.read_text()) - assert config["terminal"] == {"backend": "local"} # unrelated sections kept - assert config["model"]["temperature"] == 0.7 # unrelated model keys kept - assert config["model"]["provider"] == "custom:unsloth" - assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" - assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" - before = hermes_config.read_text() - connect.write_hermes_config(BASE, MODEL) - assert hermes_config.read_text() == before - - -def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): - pytest.importorskip("yaml") - original = "- just\n- a\n- list\n" # valid YAML, but not a mapping - hermes_config.write_text(original) - connect.write_hermes_config(BASE, MODEL) - assert hermes_config.read_text() == original # user-managed file left untouched - assert "couldn't parse" in capsys.readouterr().err - - -def test_connect_hermes_no_launch(fake_studio, hermes_config): - yaml = pytest.importorskip("yaml") - result = CliRunner().invoke(connect.connect_app, ["hermes", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") - assert "hermes" in result.output - config = yaml.safe_load(hermes_config.read_text()) - assert config["model"]["provider"] == "custom:unsloth" - assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" - assert config["model"]["default"] == MODEL["id"] - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py new file mode 100644 index 0000000000..a6a092a17c --- /dev/null +++ b/unsloth_cli/tests/test_start.py @@ -0,0 +1,1848 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for `unsloth start` — config merging and launch env, no network.""" + +from __future__ import annotations + +import json +import os +import shlex +import sys +import urllib.error +from pathlib import Path +from types import SimpleNamespace + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +import pytest +from typer.testing import CliRunner + +import unsloth_cli.commands.start as start + +BASE = "http://127.0.0.1:8888" +MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} + + +# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and +# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. +def _assert_env_set(output: str, name: str, value: str) -> None: + needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _assert_env_unset(output: str, name: str) -> None: + needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _launch_command(output: str) -> list: + # The --no-launch recipe ends with a self-contained one-liner: inline NAME=value + # assignments, then the command. Return just the command argv. + last = [ln for ln in output.splitlines() if ln.strip()][-1] + parts = shlex.split(last) + for i, part in enumerate(parts): + name = part.partition("=")[0] + if "=" not in part or not name.replace("_", "").isalnum(): + return parts[i:] + return [] + + +def _fake_claude(monkeypatch, version_output: str) -> None: + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr( + start.subprocess, + "run", + lambda *args, **kwargs: SimpleNamespace(stdout = version_output), + ) + + +def test_claude_flags_passed_to_supported_claude(monkeypatch): + _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") + assert start._claude_flags() == [ + "--exclude-dynamic-system-prompt-sections", + "--settings", + start._CLAUDE_SETTINGS_OVERLAY, + ] + + +def test_claude_flags_skipped_on_old_claude(monkeypatch): + _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") + assert start._claude_flags() == [] + + +def test_claude_flags_skipped_on_unparseable_version(monkeypatch): + _fake_claude(monkeypatch, "weird build string\n") + assert start._claude_flags() == [] + + +def test_claude_flags_detected_when_version_not_first_token(monkeypatch): + # The X.Y.Z is pulled from anywhere in the output, so a format change (version not + # the first token) doesn't silently drop the optimization flags. + _fake_claude(monkeypatch, "claude version 2.1.98\n") + assert start._claude_flags() == [ + "--exclude-dynamic-system-prompt-sections", + "--settings", + start._CLAUDE_SETTINGS_OVERLAY, + ] + + +def test_install_agent_prompts_then_installs(monkeypatch): + # TTY + yes: run the documented install command, then re-resolve the now-present binary. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True) + ran = [] + monkeypatch.setattr( + start.subprocess, + "run", + lambda command, *a, **k: ran.append(command) or SimpleNamespace(returncode = 0), + ) + # _install_agent only re-resolves after installing (the pre-install check is the + # caller's job), so `which` reports the now-present binary. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + executable = start._install_agent("codex", "npm install -g @openai/codex") + assert executable == "/usr/local/bin/codex" + assert ran == [["/bin/sh", "-c", "npm install -g @openai/codex"]] + + +def test_install_agent_declined_returns_none(monkeypatch): + # TTY + no: never runs anything; caller falls back to the print-hint failure. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) + monkeypatch.setattr(start.shutil, "which", lambda _: None) + monkeypatch.setattr( + start.subprocess, "run", lambda *a, **k: pytest.fail("should not install when declined") + ) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + + +def test_install_agent_non_interactive_returns_none(monkeypatch): + # No TTY (piped stdin): cannot prompt, so don't install; return None silently. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: False)) + monkeypatch.setattr( + start.subprocess, "run", lambda *a, **k: pytest.fail("should not install without a TTY") + ) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + + +def _parse_toml(text: str) -> dict: + tomllib = pytest.importorskip("tomllib") + return tomllib.loads(text) + + +def test_merge_codex_config_fresh(): + merged = start._merge_codex_config("", BASE) + parsed = _parse_toml(merged) + assert parsed["oss_provider"] == "unsloth_api" + provider = parsed["model_providers"]["unsloth_api"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["wire_api"] == "responses" + assert provider["requires_openai_auth"] is False + + +def test_merge_codex_config_replaces_stale_block(): + existing = ( + 'model = "gpt-5"\n' + "\n" + "[model_providers.unsloth_api]\n" + 'base_url = "http://old-host:9999/v1"\n' + 'wire_api = "chat"\n' + "\n" + "[model_providers.unsloth_api.http_headers]\n" + 'x-old = "1"\n' + "\n" + "[model_providers.ollama]\n" + 'base_url = "http://localhost:11434/v1"\n' + ) + merged = start._merge_codex_config(existing, BASE) + parsed = _parse_toml(merged) + assert parsed["model"] == "gpt-5" + assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" + assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" + assert "http_headers" not in parsed["model_providers"]["unsloth_api"] + assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" + assert start._merge_codex_config(merged, BASE) == merged + + +def test_merge_codex_config_keeps_user_oss_provider(): + merged = start._merge_codex_config('oss_provider = "ollama"\n', BASE) + assert _parse_toml(merged)["oss_provider"] == "ollama" + + +def test_write_codex_config_profile(tmp_path): + start.write_codex_config(BASE, MODEL, tmp_path) + profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) + assert profile["oss_provider"] == "unsloth_api" + assert profile["model_provider"] == "unsloth_api" + assert profile["model"] == MODEL["id"] + assert profile["model_context_window"] == 131072 + config = _parse_toml((tmp_path / "config.toml").read_text()) + assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" + + +@pytest.fixture() +def fake_studio(tmp_path, monkeypatch): + calls = [] + state = {"models": [MODEL]} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/v1/models"): + return {"object": "list", "data": state["models"]} + if url.endswith("/api/inference/status"): + return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} + if url.endswith("/api/auth/api-keys"): + return {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/inference/load"): + state["models"] = [{"id": payload["model_path"], "context_length": 4096}] + return {} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "find_studio_server", lambda: BASE) + # Identity handshake has its own tests; trust the loopback server here. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: True) + # _studio_token / api-keys are faked so the mint flow stays offline. + monkeypatch.setattr(start, "_studio_token", lambda: "jwt-token") + monkeypatch.setattr(start, "_http_json", http_json) + monkeypatch.setattr(start, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") + # --no-launch session configs land under tmp instead of the real Unsloth dir. + monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents") + # No `claude` on PATH, so _claude_flags never probes the real binary. + monkeypatch.setattr(start.shutil, "which", lambda _: None) + monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) + return calls + + +def test_connect_claude_no_launch(fake_studio): + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_unset(result.output, "ANTHROPIC_API_KEY") + _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") + _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") + # Suppress the full-screen TUI redraw so a bursty local server doesn't flicker. + _assert_env_set(result.output, "CLAUDE_CODE_NO_FLICKER", "1") + # Attribution header is suppressed for the session via env + --settings, never + # by writing the user's ~/.claude/settings.json. + _assert_env_set(result.output, "CLAUDE_CODE_ATTRIBUTION_HEADER", "0") + # Auto-compact window is sized to the loaded model's real context length so the + # session compacts before it overflows the local server's (much smaller) window, + # and compaction is forced at 90% of it for headroom. + _assert_env_set(result.output, "CLAUDE_CODE_AUTO_COMPACT_WINDOW", str(MODEL["context_length"])) + _assert_env_set(result.output, "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE", "90") + assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output + # Overlay is passed inline (session-only), not a path into the user's ~/.claude. + assert "--settings" in result.output + assert ".claude/settings.json" not in result.output + + +def test_connect_claude_compact_window_omitted_without_context(fake_studio, monkeypatch): + # A model that doesn't report a context length -> leave Claude's default window + # rather than guessing one. + monkeypatch.setattr(start, "_resolve_model", lambda *a, **k: {"id": "local-model"}) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW" not in result.output + assert "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE" not in result.output + + +def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] + assert "ANTHROPIC_API_KEY" not in captured["env"] + assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + assert captured["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) +def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + monkeypatch.setattr(start, "_claude_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == [ + "/mnt/c/Users/samle/AppData/Roaming/npm/claude", + "--model", + MODEL["id"], + ] + assert captured["env"]["ANTHROPIC_API_KEY"] == "" + assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + for name in ( + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_BASE_URL", + "ANTHROPIC_MODEL", + "ANTHROPIC_API_KEY", + "CLAUDE_CODE_OAUTH_TOKEN", + ): + assert name in captured["env"]["WSLENV"].split(":") + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) +def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_API_KEY=" in result.output + assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output + assert "export WSLENV=" in result.output + assert "ANTHROPIC_AUTH_TOKEN" in result.output + assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output + + +def test_connect_codex_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + assert "codex --oss --profile unsloth_api" in result.output + # Config lands in the session-scoped CODEX_HOME, not the user's ~/.codex. + home = tmp_path / "agents" / "codex" + _assert_env_set(result.output, "CODEX_HOME", str(home)) + assert (home / "config.toml").exists() + assert (home / "unsloth_api.config.toml").exists() + + +def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch): + # Launch mode writes config to a throwaway temp CODEX_HOME and removes it after + # the agent exits; the user's real ~/.codex is never the target. + captured = {} + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + + def run(command, env): + captured["home"] = env["CODEX_HOME"] + captured["config_present"] = (Path(env["CODEX_HOME"]) / "config.toml").exists() + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["codex"]) + assert result.exit_code == 0, result.output + home = Path(captured["home"]) + assert captured["config_present"] # config existed while codex ran + assert "unsloth-codex-" in home.name # an ephemeral temp dir, not ~/.codex + assert not home.exists() # cleaned up after the agent exits + + +@pytest.mark.skipif( + os.name == "nt", + reason = "the #6547 CI parser is bash-only; on Windows --no-launch prints PowerShell", +) +def test_no_launch_output_is_parseable(fake_studio): + # Mirror the #6547 CI parser: status lines, then `export`/`unset`, then exactly + # one launch command on the last line (now an inline-env one-liner, so the parser + # matches by substring rather than prefix). + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + lines = [ln for ln in result.output.splitlines() if ln.strip()] + skip = ("export ", "unset ", "Studio ", "Updated ", "Disabled ", "Warning", "Loading") + body = [ln for ln in lines if not ln.startswith(skip)] + assert "codex --oss --profile unsloth_api" in body[-1] + assert any(ln.startswith("export CODEX_HOME=") for ln in lines) + + +def test_no_launch_last_line_is_self_contained(fake_studio, tmp_path): + # People copy just the last line. A bare `codex` there would run against the user's + # real ~/.codex (e.g. a pre-existing damaged state DB) with zero isolation, so the + # last line must inline every session env var ahead of the command. + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + last = [ln for ln in result.output.splitlines() if ln.strip()][-1] + parts = shlex.split(last) + assignments = {} + command = [] + for i, part in enumerate(parts): + if "=" not in part: + command = parts[i:] + break + name, _, value = part.partition("=") + assignments[name] = value + assert command and command[0] == "codex" + assert assignments["CODEX_HOME"] == str(tmp_path / "agents" / "codex") + assert assignments["UNSLOTH_STUDIO_AUTH_TOKEN"].startswith("sk-unsloth-") + + +def test_no_launch_claude_last_line_blanks_conflicting_auth(fake_studio): + # The unset vars must be neutralized inline too, or a partial copy would send the + # user's own ANTHROPIC_API_KEY to the Studio base. + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + last = [ln for ln in result.output.splitlines() if ln.strip()][-1] + assert "ANTHROPIC_API_KEY= " in last + assert "CLAUDE_CODE_OAUTH_TOKEN= " in last + assert "ANTHROPIC_AUTH_TOKEN=" in last # the real key still applied after the blanks + + +def test_opencode_inline_config_beats_project_config(fake_studio): + # A project's opencode.json outranks OPENCODE_CONFIG, so the model pin (and --yolo + # permissions) ride in OPENCODE_CONFIG_CONTENT, which outranks project config. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "--yolo"]) + assert result.exit_code == 0, result.output + content_line = next( + ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") + ) + inline = json.loads( + shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] + ) + assert inline["model"] == f"unsloth/{MODEL['id']}" + assert inline["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + assert "sk-unsloth" not in content_line # key stays in the private file + + +def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + content_line = next( + ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") + ) + inline = json.loads( + shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] + ) + assert inline == {"model": f"unsloth/{MODEL['id']}"} + + +def test_https_loopback_never_auto_serves(fake_studio, monkeypatch): + # `unsloth run` serves plain HTTP; auto-serving behind an https:// target would poll + # the wrong scheme until the startup timeout. Keep the plain "no server" error. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "https://127.0.0.1:8443") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_connect_alias_still_works(fake_studio): + # `unsloth connect` remains a compat alias for `unsloth start`. + from unsloth_cli import app + + result = CliRunner().invoke(app, ["connect", "claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): + CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + # First run mints; second reuses the minted key cached for this server. + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert len(mints) == 1 + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + + +def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): + CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + # Reused, not re-minted (a mint would return the feedface stand-in). + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + # An explicit key is remembered as "saved" so it replays without the handshake. + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] + + +def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): + cache = tmp_path / "agent_api_key.json" + cache.write_text( + json.dumps( + {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} + ) + ) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-stale": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + # The working key moves to the front so the next run tries it first. + cached = json.loads(cache.read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] + + +def test_connect_saved_key_server_outage_surfaces_not_reminted(fake_studio, tmp_path, monkeypatch): + # A 5xx/timeout while checking a saved key is a server outage, not a rejected key: + # surface it instead of discarding the key and minting a new one against a sick server. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-saved"]}}})) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-saved": + raise urllib.error.HTTPError(url, 503, "Service Unavailable", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code != 0, result.output + # The outage did not cause a fresh key to be minted. + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert mints == [] + + +def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): + # Legacy unscoped caches have no server binding (could leak across servers), + # so they're ignored: a fresh key is minted and stored scoped to this server. + (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + assert "key" not in cached # legacy field collapsed away + + +def test_connect_model_flag_loads_on_server(fake_studio): + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) + ] + _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") + + +def test_connect_model_flag_forwards_load_options(fake_studio): + # The model-load knobs mirrored from `unsloth run` reach /api/inference/load. + result = CliRunner().invoke( + start.start_app, + [ + "claude", + "--no-launch", + "--model", + "unsloth/Qwen3-4B-GGUF", + "--gguf-variant", + "UD-Q4_K_XL", + "--context-length", + "8192", + "--no-load-in-4bit", + "--tensor-parallel", + ], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + { + "model_path": "unsloth/Qwen3-4B-GGUF", + "gguf_variant": "UD-Q4_K_XL", + "max_seq_length": 8192, + "load_in_4bit": False, + "tensor_parallel": True, + }, + ) + ] + + +def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): + # Studio registers a loaded model under a canonical id (resolved identifier + # / casing) that can differ from the path we passed. The agent must connect + # to that model, not silently fall through to the first loaded one. + requested = "Unsloth/Qwen3.5-35B-A3B" + canonical = "unsloth/Qwen3.5-35B-A3B" + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {"model": canonical, "display_name": canonical} + if url.endswith("/v1/models"): + # Decoy sorts first, so models[0] is the wrong pick on the old code. + return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch", "--model", requested]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) + + +@pytest.mark.parametrize( + "model, expected", + [ + ("unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL", ("unsloth/Qwen3-1.7B-GGUF", "UD-Q4_K_XL")), + ("unsloth/gemma-4-E2B-it-GGUF:Q8_0", ("unsloth/gemma-4-E2B-it-GGUF", "Q8_0")), + ("unsloth/Qwen3-1.7B-GGUF", ("unsloth/Qwen3-1.7B-GGUF", None)), # no suffix + ("/models/local.gguf", ("/models/local.gguf", None)), # absolute path + ("./rel.gguf", ("./rel.gguf", None)), # relative path + ("C:\\models\\x.gguf", ("C:\\models\\x.gguf", None)), # Windows drive + ("repo:with/slash", ("repo:with/slash", None)), # slash in variant -> not a variant + ("", ("", None)), + ], +) +def test_split_repo_variant(model, expected): + assert start._split_repo_variant(model) == expected + + +def test_connect_model_bare_id_matches_loaded_without_reload(fake_studio): + # A bare `--model ` (no load knobs) attaches to the already-loaded model + # without touching /api/inference/load, so it can never evict another session. + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch", "--model", MODEL["id"]]) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [] + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_model_variant_suffix_defers_to_server_dedup(fake_studio): + # `--model repo:QUANT` splits into a VALID load payload (bare repo + gguf_variant), + # never the `:`-suffixed repo id Studio rejects. The variant knob defers to + # /api/inference/load, whose already-loaded dedup answers without reloading when the + # active variant+settings match -- so a second session running the same command + # attaches without evicting the first, while a genuinely different quant reloads. + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", MODEL["id"] + ":UD-Q4_K_XL"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": MODEL["id"], "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_load_knobs_reach_server_even_when_id_loaded(fake_studio): + # /v1/models can't reveal the active quant, so an id match alone would silently keep + # the wrong variant loaded. Explicit knobs must always consult the load endpoint. + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--model", MODEL["id"], "--gguf-variant", "Q8_0"], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": MODEL["id"], "gguf_variant": "Q8_0"}) + ] + + +def test_connect_model_variant_suffix_loads_split_repo(fake_studio): + # When the model is not already loaded, the `:QUANT` suffix becomes the gguf_variant + # and the load uses the bare (valid) repo id, mirroring `unsloth run repo --gguf-variant`. + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--model", "unsloth/Qwen3-4B-GGUF:UD-Q4_K_XL"], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": "unsloth/Qwen3-4B-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + + +def test_connect_explicit_gguf_variant_wins_over_suffix(fake_studio): + # An explicit --gguf-variant takes precedence; the suffix is still stripped so the + # repo id stays valid. + result = CliRunner().invoke( + start.start_app, + [ + "claude", + "--no-launch", + "--model", + "unsloth/Qwen3-4B-GGUF:Q8_0", + "--gguf-variant", + "UD-Q4_K_XL", + ], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": "unsloth/Qwen3-4B-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + + +def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): + monkeypatch.setattr( + start, + "_http_json", + lambda method, url, token, payload = None, timeout = 30, error = None: ( + {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/auth/api-keys") + else {"object": "list", "data": []} + ), + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No model is loaded" in result.output + + +def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): + # Studio never surfaces the requested model; fail loudly rather than + # silently connecting to whatever else happens to be loaded. + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {} + if url.endswith("/v1/models"): + return {"object": "list", "data": [MODEL]} # decoy; request never appears + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] + ) + assert result.exit_code == 1 + assert "unsloth/Missing-7B" in result.output + + +def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 1 + assert "GGUF" in result.output + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): + # A server known only by URL + health check is unverified: keyless connect + # must refuse and make no request at all. + monkeypatch.setattr(start, "find_studio_server", lambda: "http://studio.evil.example:8888") + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 1 + assert "Settings → API" in result.output + assert "--api-key" in result.output + assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) + + +def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): + # User named both server and key, so it's their choice; only auto-send is blocked. + monkeypatch.setattr(start, "find_studio_server", lambda: "http://studio.example:8888") + result = CliRunner().invoke( + start.start_app, + ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): + # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; + # auto-minting stays blocked for non-loopback. + remote = "http://studio.example:8888" + monkeypatch.setattr(start, "find_studio_server", lambda: remote) + (tmp_path / "agent_api_key.json").write_text( + json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): + # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an + # error, not a silent local model load (which they did not ask for). + import typer + + import unsloth_cli._inference as inference + + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") + monkeypatch.setattr( + inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" + ) + with pytest.raises(typer.Exit): + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + + +def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): + # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback + # server can't be verified, fall back to a local load rather than erroring. + import unsloth_cli._inference as inference + + monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) + monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") + monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) + assert ( + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + is None + ) + + +def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( + fake_studio, tmp_path, monkeypatch +): + # With no saved key, the next step would auto-mint; an unverified loopback + # server (port squatter) must be refused, with nothing sent. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): + # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) + # replays on keyless runs without the handshake, scoped to its own base. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted + + +def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): + # A "minted" key is NOT replayed to an unverified loopback server: minting and + # minted-key replay both sit behind the handshake, so a squatter can't grab it. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent + + +def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): + # An explicit key is the user's deliberate choice, so it does not require + # the automatic identity handshake. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + + +def _serve_identity(proof_for): + """Start a localhost HTTP server answering /api/auth/identity with + proof_for(nonce_bytes). Returns (base_url, shutdown).""" + import base64 + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + from urllib.parse import parse_qs, urlparse + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + parsed = urlparse(self.path) + if parsed.path != "/api/auth/identity": + self.send_response(404) + self.end_headers() + return + nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) + host, port = self.server.server_address[0], self.server.server_address[1] + body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): + # Real crypto end to end: verify_studio_identity reads the install secret from + # an isolated DB; a "good" server proves the same secret, a spoofing one can't. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: # backend not importable here (e.g. missing deps) + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + good = lambda nonce, host, port: storage.compute_identity_proof( + nonce, host, port + ) # real secret + bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret + base_ok, stop_ok = _serve_identity(good) + base_bad, stop_bad = _serve_identity(bad) + try: + assert inference.verify_studio_identity(base_ok) is True + assert inference.verify_studio_identity(base_bad) is False + finally: + stop_ok() + stop_bad() + + +def _serve_redirect(target): + """Start a localhost server that 302-redirects every GET to target+path.""" + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(302) + self.send_header("Location", target + self.path) + self.end_headers() + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): + # A squatter could 302 /api/auth/identity to the real Studio and relay its + # proof; redirects must be refused so the squatter's base isn't accepted. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + squatter_base, stop_squatter = _serve_redirect(real_base) + try: + assert inference.verify_studio_identity(real_base) is True # direct: ok + assert inference.verify_studio_identity(squatter_base) is False # relayed: refused + finally: + stop_real() + stop_squatter() + + +def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): + # A squatter that proxies the nonce to the real Studio on another port gets a + # proof bound to *that* port; the client expects one bound to the port it + # connected to, so the relayed proof is rejected. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + real_port = int(real_base.rsplit(":", 1)[1]) + # The squatter answers on its own port but returns the proof for the real port. + squatter_base, stop_squatter = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) + ) + try: + assert inference.verify_studio_identity(real_base) is True + assert inference.verify_studio_identity(squatter_base) is False + finally: + stop_real() + stop_squatter() + + +@pytest.mark.parametrize( + "url, loopback", + [ + ("http://127.0.0.1:8888", True), + ("http://localhost:8888", True), + ("http://[::1]:8888", True), + ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 + ("http://0.0.0.0:8888", False), + ("http://10.0.0.5:8888", False), + ("http://studio.evil.example:8888", False), + ("https://studio.example.com", False), + ], +) +def test_is_loopback_url(url, loopback): + assert start.is_loopback_url(url) is loopback + + +def test_connect_no_studio_errors(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + + +@pytest.fixture(autouse = True) +def _reset_auto_served(): + # Never let a test leave a fake server in the module slot (an atexit backstop would + # otherwise try to signal it at interpreter shutdown). + yield + start._auto_served_server = None + + +def test_start_studio_server_builds_command_and_waits(monkeypatch): + captured = {} + + class FakePopen: + def __init__(self, command, **kwargs): + captured["command"] = command + captured["kwargs"] = kwargs + self.pid = 4321 + + def poll(self): + return None + + monkeypatch.setattr(start.subprocess, "Popen", FakePopen) + monkeypatch.setattr(start, "_studio_healthy", lambda base, timeout = 3.0: True) + monkeypatch.setattr(start, "_log_tail", lambda path, lines = 20: "API Key: sk-unsloth-abc123") + monkeypatch.setattr(start.time, "sleep", lambda _s: None) + + server = start._start_studio_server( + "http://127.0.0.1:8888", + "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL", + start.LoadOptions( + gguf_variant = "UD-Q4_K_XL", max_seq_length = 8192, load_in_4bit = True, tensor_parallel = True + ), + ) + cmd = captured["command"] + assert cmd[1] == "run" + assert "--disable-tools" in cmd and "--no-cloudflare" in cmd + assert cmd[cmd.index("--model") + 1] == "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL" + assert cmd[cmd.index("--gguf-variant") + 1] == "UD-Q4_K_XL" + assert cmd[cmd.index("--context-length") + 1] == "8192" + assert "--tensor-parallel" in cmd + assert cmd[cmd.index("-p") + 1] == "8888" + assert start.LoadOptions().load_in_4bit is True and "--no-load-in-4bit" not in cmd + assert captured["kwargs"].get("start_new_session") is True # own process group + assert server.pid == 4321 + + +def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started.update(base = base, model = model, load = load) + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr( + start, "_shutdown_server", lambda server: started.__setitem__("down", server) + ) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL"] + ) + assert result.exit_code == 0, result.output + # The `:QUANT` suffix is split off into the gguf_variant so `unsloth run` gets a valid + # repo id plus `--gguf-variant`, mirroring how `unsloth run` accepts either form. + assert started["model"] == "unsloth/Qwen3-1.7B-GGUF" + assert started["load"].gguf_variant == "UD-Q4_K_XL" + assert started["base"] == BASE + # Torn down after the agent session ended. + assert started.get("down") is fake + + +def test_codex_preflight_failure_tears_down_auto_served(fake_studio, monkeypatch): + # The Codex GGUF preflight runs after _connect may have auto-started a server but + # before _run's teardown finally, so a preflight rejection must not leave the server + # holding the port/GPU (waiting on the atexit backstop). + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started.update(base = base, model = model) + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr( + start, "_shutdown_server", lambda server: started.__setitem__("down", server) + ) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "transformers-model"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke( + start.start_app, ["codex", "--model", "unsloth/Qwen3-1.7B", "--launch"] + ) + assert result.exit_code != 0, result.output + assert "GGUF" in result.output + # Torn down at the point the preflight rejected the model, not only via atexit. + assert started.get("down") is fake + + +def test_no_serve_preserves_error(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-serve"] + ) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_no_launch_never_serves(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-launch"] + ) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_no_server_no_model_hints_model_flag(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + result = CliRunner().invoke(start.start_app, ["claude"]) + assert result.exit_code == 1 + assert "--model" in result.output + + +@pytest.mark.parametrize( + "base, expected", + [ + ("http://127.0.0.1", "http://127.0.0.1:8888"), # portless -> unsloth run's :8888 + ("http://127.0.0.1:8888", "http://127.0.0.1:8888"), # explicit port kept + ("http://127.0.0.1:9000", "http://127.0.0.1:9000"), + ("http://localhost", "http://localhost:8888"), + ("http://[::1]", "http://[::1]:8888"), # IPv6 literal stays bracketed + ("http://[::1]:8888", "http://[::1]:8888"), + # Paths are stripped: unsloth run serves at the root, so /studio would make the + # health poll hit /studio/api/health (404) until the startup timeout. + ("http://127.0.0.1:8888/studio", "http://127.0.0.1:8888"), + ("http://127.0.0.1/studio", "http://127.0.0.1:8888"), + ], +) +def test_effective_base(base, expected): + assert start._effective_base(base) == expected + + +def test_auto_serve_normalizes_portless_url(fake_studio, monkeypatch): + # A portless UNSLOTH_STUDIO_URL must launch AND poll :8888 (what unsloth run binds), + # not port 80, or readiness never matches and we hit the startup timeout. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://127.0.0.1") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started["base"] = base + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr(start, "_shutdown_server", lambda server: None) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + + result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"]) + assert result.exit_code == 0, result.output + assert started["base"] == "http://127.0.0.1:8888" + + +def test_connect_explicit_api_key_skips_mint(fake_studio): + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) + + +# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── + + +def test_write_openclaw_config_fresh(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["models"]["providers"]["unsloth"] + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + assert provider["api"] == "openai-completions" + assert provider["models"] == [ + {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} + ] + # The default model must be pinned or OpenClaw has nothing active. + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["gateway"]["mode"] == "local" + assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway + if os.name != "nt": # the file holds an API key + assert path.stat().st_mode & 0o777 == 0o600 + + +def test_write_openclaw_config_preserves_and_idempotent(tmp_path): + path = tmp_path / "openclaw.json" + path.write_text( + json.dumps( + { + "theme": "dark", + "agents": {"defaults": {"temperature": 0.5}}, + "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, + } + ) + ) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["theme"] == "dark" + assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["models"]["mode"] == "replace" # user's mode is left as-is + assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" + assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_write_openclaw_config_corrupt_left_alone(tmp_path, capsys): + path = tmp_path / "openclaw.json" + path.write_text("{not json") + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == "{not json" + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_openclaw_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "openclaw" in result.output + config_path = tmp_path / "agents" / "openclaw" / "openclaw.json" + # Config + state are scoped to the session dir, not the user's ~/.openclaw. + _assert_env_set(result.output, "OPENCLAW_CONFIG_PATH", str(config_path)) + _assert_env_set(result.output, "OPENCLAW_STATE_DIR", str(tmp_path / "agents" / "openclaw")) + config = json.loads(config_path.read_text()) + assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + # OpenAI /v1/chat/completions works on either backend — no GGUF gate. + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── + + +def test_write_opencode_config_fresh(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["provider"]["unsloth"] + assert provider["npm"] == "@ai-sdk/openai-compatible" + assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} + # Context limit must be declared, or OpenCode treats it as 0 and disables compaction. + assert provider["models"] == { + MODEL["id"]: {"name": MODEL["id"], "limit": {"context": 131072, "output": 8192}} + } + assert config["model"] == f"unsloth/{MODEL['id']}" + # Compaction buffer scaled to ~10% of the window (compact near 90%). + assert config["compaction"] == {"auto": True, "reserved": 131072 // 10} + + +def test_write_opencode_config_preserves_and_idempotent(tmp_path): + path = tmp_path / "opencode.json" + path.write_text( + json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) + ) + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["theme"] == "tokyonight" + assert config["provider"]["anthropic"]["name"] == "Anthropic" + assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" + before = path.read_text() + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_connect_opencode_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "opencode" in result.output + config_path = tmp_path / "agents" / "opencode" / "opencode.json" + # OPENCODE_CONFIG overlay points at the session file, not the user's global config. + _assert_env_set(result.output, "OPENCODE_CONFIG", str(config_path)) + config = json.loads(config_path.read_text()) + assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["model"] == f"unsloth/{MODEL['id']}" + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── + + +@pytest.fixture() +def hermes_config(tmp_path): + return tmp_path / "config.yaml" + + +def test_write_hermes_config_fresh(hermes_config): + yaml = pytest.importorskip("yaml") + start.write_hermes_config(BASE, MODEL, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes only honors the key for a *named* custom provider, so the endpoint + # is registered under providers.* and model.provider points at it. + assert config["model"]["provider"] == "custom:unsloth" + assert config["model"]["default"] == MODEL["id"] + assert config["model"]["api_mode"] == "openai" + # Pin the real context window (top-level override) and compact at 90% of it. + assert config["model"]["context_length"] == MODEL["context_length"] + assert config["compression"] == {"enabled": True, "threshold": 0.9} + # Windows at or above Hermes' floor need no auxiliary compression override. + assert "auxiliary" not in config + provider = config["providers"]["unsloth"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["api_mode"] == "openai" + assert provider["key_env"] == "UNSLOTH_API_KEY" + # The key is resolved from the launch env, never written to disk. + assert "sk-unsloth" not in hermes_config.read_text() + + +def test_write_hermes_config_small_window_claims_floor(hermes_config): + yaml = pytest.importorskip("yaml") + small = {"id": "unsloth/Qwen3-1.7B-GGUF", "context_length": 40960} + start.write_hermes_config(BASE, small, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes refuses to initialize below its 64,000-token floor, so the recipe + # claims the floor and scales the compaction threshold so it still fires at + # 90% of the REAL window: 0.9 * 40960 / 65536. + assert config["model"]["context_length"] == 65536 + assert config["compression"] == {"enabled": True, "threshold": 0.5625} + # The same floor check runs against the compression model mid-session. + assert config["auxiliary"]["compression"]["context_length"] == 65536 + + +def test_write_hermes_config_preserves_and_idempotent(hermes_config): + yaml = pytest.importorskip("yaml") + hermes_config.write_text( + yaml.safe_dump( + { + "terminal": {"backend": "local"}, + "model": {"temperature": 0.7}, + "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, + } + ) + ) + start.write_hermes_config(BASE, MODEL, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + assert config["terminal"] == {"backend": "local"} # unrelated sections kept + assert config["model"]["temperature"] == 0.7 # unrelated model keys kept + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + before = hermes_config.read_text() + start.write_hermes_config(BASE, MODEL, hermes_config) + assert hermes_config.read_text() == before + + +def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): + pytest.importorskip("yaml") + original = "- just\n- a\n- list\n" # valid YAML, but not a mapping + hermes_config.write_text(original) + start.write_hermes_config(BASE, MODEL, hermes_config) + assert hermes_config.read_text() == original # user-managed file left untouched + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_hermes_no_launch(fake_studio, tmp_path): + yaml = pytest.importorskip("yaml") + result = CliRunner().invoke(start.start_app, ["hermes", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") + # HERMES_HOME relocates the whole hermes home, so the user's ~/.hermes is untouched. + home = tmp_path / "agents" / "hermes" + _assert_env_set(result.output, "HERMES_HOME", str(home)) + assert "hermes" in result.output + config = yaml.safe_load((home / "config.yaml").read_text()) + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + assert config["model"]["default"] == MODEL["id"] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Pi (OpenAI-compatible /v1, key in config, ~/.pi relocated via HOME) ── + + +def test_write_pi_config_fresh(tmp_path): + path = tmp_path / ".pi" / "agent" / "models.json" + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["providers"]["unsloth"] + assert provider["api"] == "openai-completions" + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + # Pin the loaded window (and a sane output cap) so Pi compacts instead of + # overflowing; without it Pi assumes its 128000 default. + assert provider["models"] == [ + {"id": MODEL["id"], "contextWindow": MODEL["context_length"], "maxTokens": 8192} + ] + + +def test_write_pi_config_preserves_and_idempotent(tmp_path): + path = tmp_path / ".pi" / "agent" / "models.json" + path.parent.mkdir(parents = True) + path.write_text(json.dumps({"providers": {"google": {"api": "gemini"}}})) + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["providers"]["google"] == {"api": "gemini"} # unrelated provider kept + assert config["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_connect_pi_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + # Pi resolves its config dir from PI_CODING_AGENT_DIR first, so pin it at the session + # dir (and relocate HOME) to keep the user's real ~/.pi untouched and their own + # PI_CODING_AGENT_DIR from redirecting Pi away from our provider/key. + home = tmp_path / "agents" / "pi" + _assert_env_set(result.output, "HOME", str(home)) + _assert_env_set(result.output, "PI_CODING_AGENT_DIR", str(home / ".pi" / "agent")) + # Provider/model pinned on the command (Pi defaults to google otherwise). + assert f"pi --provider unsloth --model {MODEL['id']}" in result.output + config = json.loads((home / ".pi" / "agent" / "models.json").read_text()) + assert config["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["providers"]["unsloth"]["models"] == [ + {"id": MODEL["id"], "contextWindow": MODEL["context_length"], "maxTokens": 8192} + ] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +def test_connect_pi_no_launch_windows_relocates_userprofile(fake_studio, tmp_path, monkeypatch): + # On native Windows Node resolves ~/.pi via USERPROFILE, not HOME, so the session + # must point USERPROFILE at the relocated home or Pi reads the user's real ~/.pi. + monkeypatch.setattr(start.os, "name", "nt") + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + home = tmp_path / "agents" / "pi" + assert f'$env:HOME = "{home}"' in result.output + assert f'$env:USERPROFILE = "{home}"' in result.output + + +# ── WSLENV path translation + PowerShell quoting (helper units) ── + + +def test_wsl_bridge_names_flags_paths_not_scalars(): + # WSLENV only translates a var to a Windows path when its entry carries /p. + # Path-valued vars must get it; scalar knobs and URLs must not, or WSLENV would + # mangle them when handing off to a Windows shim under /mnt. + env = { + "CODEX_HOME": "/tmp/sess/codex", + "HOME": "/tmp/sess/pi", + "CLAUDE_CODE_AUTO_COMPACT_WINDOW": "4096", + "ANTHROPIC_BASE_URL": "http://127.0.0.1:8888", + "USERPROFILE": r"C:\Users\x", + } + names = start._wsl_bridge_names(env, ("ANTHROPIC_API_KEY",)) + assert "CODEX_HOME/p" in names + assert "HOME/p" in names + assert "USERPROFILE/p" in names # drive-qualified Windows path + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW" in names # scalar: no /p + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW/p" not in names + assert "ANTHROPIC_BASE_URL" in names # URL is not a filesystem path + assert "ANTHROPIC_API_KEY" in names # cleared var carries no value to translate + + +def test_merge_wslenv_dedups_on_base_name(): + # An already-shared var must not be appended again just because the flag differs. + merged = start._merge_wslenv("CODEX_HOME/p:FOO", ("CODEX_HOME/p", "BAR/p")) + parts = merged.split(":") + assert parts.count("CODEX_HOME/p") == 1 + assert "FOO" in parts and "BAR/p" in parts + + +def test_merge_wslenv_upgrades_existing_unflagged_entry(): + # A user's pre-existing bare "HOME" must be upgraded to "HOME/p" (not left bare or + # duplicated), or the Windows shim gets the path without WSL translation. + merged = start._merge_wslenv("HOME:FOO", ("HOME/p", "CODEX_HOME/p")) + parts = merged.split(":") + assert "HOME/p" in parts and "HOME" not in parts # upgraded in place + assert parts.count("HOME/p") == 1 + assert "FOO" in parts # untouched user var preserved + assert "CODEX_HOME/p" in parts + + +def test_powershell_quote_single_quotes_json(): + # Bare flags/paths pass through; JSON payloads get single-quoted so PowerShell + # keeps the embedded double quotes literal (list2cmdline's backslashes would not). + assert start._powershell_quote("--settings") == "--settings" + assert start._powershell_quote("unsloth/gemma-4-26B") == "unsloth/gemma-4-26B" + quoted = start._powershell_quote(start._CLAUDE_SETTINGS_OVERLAY) + assert quoted == "'" + start._CLAUDE_SETTINGS_OVERLAY + "'" + assert "\\" not in quoted # no cmd.exe backslash escaping + assert start._powershell_quote("a'b") == "'a''b'" # embedded quote doubled + + +# ── --yolo: one switch routed to each agent's own auto-approve form ── + +# The native "run tools without prompting" CLI flag each agent should receive. +_NATIVE_YOLO = { + "claude": "--dangerously-skip-permissions", + "codex": "--dangerously-bypass-approvals-and-sandbox", + "hermes": "--yolo", + "pi": "--approve", +} + + +@pytest.mark.parametrize("agent, native", sorted(_NATIVE_YOLO.items())) +def test_yolo_routes_to_native_flag(fake_studio, agent, native): + result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + assert native in result.output + + +@pytest.mark.parametrize("agent, native", sorted(_NATIVE_YOLO.items())) +def test_no_yolo_omits_native_flag(fake_studio, agent, native): + result = CliRunner().invoke(start.start_app, [agent, "--no-launch"]) + assert result.exit_code == 0, result.output + # pi's --approve is a real flag only added under --yolo; assert it's absent here. + command = _launch_command(result.output) + assert command and command[0] == agent, result.output + assert native not in command + + +@pytest.mark.parametrize( + "alias", + ["--yolo", "--dangerously-skip-permissions", "--dangerously-bypass-approvals-and-sandbox"], +) +def test_yolo_aliases_are_interchangeable(fake_studio, alias): + # Any spelling on any agent routes to that agent's own flag, even the "wrong" one. + claude = CliRunner().invoke(start.start_app, ["claude", alias, "--no-launch"]) + assert claude.exit_code == 0, claude.output + assert "--dangerously-skip-permissions" in claude.output + # The codex spelling must not leak through to Claude's command line. + assert "--dangerously-bypass-approvals-and-sandbox" not in claude.output + + codex = CliRunner().invoke(start.start_app, ["codex", alias, "--no-launch"]) + assert codex.exit_code == 0, codex.output + assert "--dangerously-bypass-approvals-and-sandbox" in codex.output + assert "--dangerously-skip-permissions" not in codex.output + + +def test_yolo_opencode_writes_permission_block(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + + +def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + assert "permission" not in config + + +def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + state = tmp_path / "agents" / "openclaw" + config = json.loads((state / "openclaw.json").read_text()) + assert config["tools"]["exec"] == {"host": "gateway", "security": "full", "ask": "off"} + # Both layers: the host approvals file in OPENCLAW_STATE_DIR must also be set, or + # OpenClaw can still prompt/deny despite the config. + approvals = json.loads((state / "exec-approvals.json").read_text()) + assert approvals["defaults"] == {"security": "full", "ask": "off", "askFallback": "full"} + + +def test_no_yolo_openclaw_has_no_exec_policy(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + state = tmp_path / "agents" / "openclaw" + config = json.loads((state / "openclaw.json").read_text()) + assert "exec" not in config.get("tools", {}) # no auto-approve policy without --yolo + assert not (state / "exec-approvals.json").exists() + + +def test_write_opencode_config_yolo_unit(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + config = json.loads(path.read_text()) + assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + + +def test_write_openclaw_config_yolo_unit(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"host": "gateway", "security": "full", "ask": "off"} + approvals = json.loads((path.parent / "exec-approvals.json").read_text()) + assert approvals == { + "version": 1, + "defaults": {"security": "full", "ask": "off", "askFallback": "full"}, + } + + +def test_yolo_command_flags_unmapped_agent_is_empty(): + # Config-based agents (and any typo) must yield no flag, not a KeyError. + assert start._yolo_command_flags("opencode", True) == [] + assert start._yolo_command_flags("openclaw", True) == [] + assert start._yolo_command_flags("claude", True) == ["--dangerously-skip-permissions"] + assert start._yolo_command_flags("claude", False) == [] + + +def test_yolo_config_agents_add_no_command_flag(fake_studio): + # opencode/openclaw auto-approve is config-only; nothing should leak onto argv. + for agent in ("opencode", "openclaw"): + result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + command = _launch_command(result.output) + assert command and command[0] == agent, result.output + assert not any("--yolo" in arg or "--dangerous" in arg for arg in command) + + +def test_pi_launch_clears_screen_first(fake_studio, monkeypatch): + # Pi paints inline from the current cursor position (no alternate screen, no + # clear on its first render), so the launcher hands it a clean screen. The + # clear must come BEFORE the exec, and only on the launch path. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/pi") + + def run(command, env): + calls.append("exec") + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["pi"]) + assert result.exit_code == 0, result.output + assert calls == ["clear", "exec"] + + +def test_pi_no_launch_does_not_clear(fake_studio, monkeypatch): + # The --no-launch recipe is meant to be read (and piped); never wipe it. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + assert calls == [] + + +def test_claude_launch_does_not_clear(fake_studio, monkeypatch): + # Alternate-screen agents manage the terminal themselves; leave it alone. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + result = CliRunner().invoke(start.start_app, ["claude"]) + assert result.exit_code == 0, result.output + assert calls == [] + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario: a Windows pi shim under /mnt called from WSL " + "(os.name is 'posix' under WSL), so this can't run on a native Windows runner.", +) +def test_connect_pi_wsl_windows_shim_relocates_userprofile(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr(start.shutil, "which", lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/pi") + + def run(command, env): + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["pi"]) + assert result.exit_code == 0, result.output + home = captured["env"]["HOME"] + # A Windows pi shim resolves ~/.pi via USERPROFILE, so it must match the session + # HOME and ride the WSLENV bridge (with /p) so the path is translated for Windows. + assert captured["env"]["USERPROFILE"] == home + wslenv = captured["env"]["WSLENV"].split(":") + assert "HOME/p" in wslenv + assert "USERPROFILE/p" in wslenv + + +def test_agent_api_key_auto_started_rejected_env_key_falls_back(fake_studio, tmp_path, monkeypatch): + # UNSLOTH_API_KEY exported for some OTHER server must not fail the launch + # against a server this run just auto-started: validate, then fall back to + # the local mint path, and never remember the foreign key for this base. + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-other-server": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + key = start._agent_api_key(BASE, "sk-unsloth-other-server", auto_started = True) + assert key == "sk-unsloth-feedfacefeedface" # minted for the fresh server + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert "sk-unsloth-other-server" not in json.dumps(cached["servers"].get(BASE, {})) + + +def test_agent_api_key_auto_started_accepted_key_is_honored(fake_studio, tmp_path): + # An explicit key the fresh server accepts (e.g. persisted in this Studio + # home's auth db across restarts) keeps working exactly as before. + key = start._agent_api_key(BASE, "sk-unsloth-deadbeefdeadbeef", auto_started = True) + assert key == "sk-unsloth-deadbeefdeadbeef" + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] + + +def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path): + # A previously printed recipe may still be running an agent whose sessions + # or sqlite state live in the stable home; a re-run must not wipe it. + with start._session_config("codex", launch = False) as home: + marker = home / "sessions" / "live.sqlite" + marker.parent.mkdir(parents = True) + marker.write_text("state") + with start._session_config("codex", launch = False) as home2: + assert home2 == home + assert (home2 / "sessions" / "live.sqlite").read_text() == "state" From 308ea5a93cd07b72e0e3ed1afa0cd987300769fa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 08:22:42 -0700 Subject: [PATCH 15/23] Tool-call healing (default on) and opt-in nudging for the client-tool passthrough (#6801) * inference: add passthrough tool-call healing core (heal_gate, heal_openai_message, StreamToolCallHealer, nudge helpers) Small GGUF models often emit tool calls as text ({...}, Gemma <|tool_call>, XML) instead of structured tool_calls. Studio's enable-tools loop already heals these, but the client-tool passthrough (unsloth run --disable-tools, unsloth start agents) relays them verbatim, so the agent sees prose and the turn dies. This module is the shared response-side repair layer the passthrough routes will call: promote parsed text-form calls to structured calls, but only for function names the client actually declared; coerce arguments through the same canonical-key healing as the tool loop; never touch the upstream request body (llama-server KV/slot reuse stays byte-identical). StreamToolCallHealer is the streaming buffer-and-repair state machine: prose forwards immediately, only a partial-signal tail or a suspected tool block is held, false alarms flush verbatim, and a 64 KiB bound caps memory. nudge_should_retry/nudge_messages support an opt-in single-retry nudge for non-streaming routes (wired later). Kill-switch: UNSLOTH_DISABLE_TOOL_CALL_HEALING=1. Reuses core/tool_healing.parse_tool_calls_from_text, strip_tool_call_markup, and tool_loop_controller.coerce_tool_arguments unchanged. * inference: heal text-form tool calls on the OpenAI and Responses passthrough Wire the passthrough healing core into /v1/chat/completions and /v1/responses, default ON whenever the request declares client tools: Non-streaming: heal_openai_message runs inside the existing response-mutation loop; a promoted call flips finish_reason to tool_calls and nulls the content, and the verbatim-bytes fast path still applies when nothing was healed. /v1/responses non-streaming inherits this through openai_chat_completions. Streaming: a StreamToolCallHealer per stream. Ordinary prose relays byte-for-byte (a fast path keeps upstream bytes when the healer passes a chunk through whole); once a tool signal appears, content is held, and at the finish/[DONE] boundary either synthetic delta.tool_calls chunks replace the markup (finish_reason rewritten to tool_calls, including the synthetic-finish path) or a false alarm flushes the held text verbatim. Structured upstream deltas put the healer to sleep after flushing anything held, so grammar-mode responses stay byte-identical. The Responses stream feeds healed calls through the same per-call state machinery as structured deltas (indexes live in a disjoint range so a healed call can never merge into a structured call's state), and the visible/reasoning split runs first so reasoning text is never promoted. parallel_tool_calls=false caps healed calls on every path. The upstream request body is never touched and healing issues no extra generation, so llama-server slot/KV-cache reuse is unchanged. Opt-out per request with auto_heal_tool_calls=false (Responses reads it from the extra-body); requests without tools relay verbatim. * inference: heal text-form tool calls on the Anthropic /v1/messages passthrough Streaming: AnthropicPassthroughEmitter.enable_healing(allowed_tools) routes content deltas through the shared StreamToolCallHealer. A promoted call closes any open text block (only the safe prose prefix ever streamed into it), opens a synthetic tool_use block with a fresh toolu_* id, carries one input_json_delta, and closes; finish() then forces stop_reason to tool_use unless a truncation (max_tokens) wins. Structured upstream deltas flush anything held and put the healer to sleep, so grammar-mode responses are untouched, as is every stream where enable_healing is never called (Studio's own loop, no-tools requests). disable_parallel_tool_use caps healed calls too. Non-streaming: the OpenAI message dict is healed BEFORE block building, so the existing tool_use promotion loop and stop_reason line treat promoted calls exactly like native ones (finish_reason length still maps to max_tokens). The legacy tool-XML strip still runs on remaining text, so opted-out requests keep today's cleanup behavior byte-for-byte. auto_heal_tool_calls is now a typed field on AnthropicMessagesRequest (default True, mirroring Chat Completions) and threads into both passthrough calls. Healing never touches the upstream request body. * inference: opt-in single-retry tool-call nudge on the non-streaming passthrough When the model clearly tried to call a tool (a tool signal in the text) but healing produced nothing usable, re-ask once: the retry body is the original body plus an assistant turn (the model's own failed text) and a short user nudge naming the declared tools. The prompt prefix stays byte-identical, so llama-server reuses the slot's KV cache and only the two-message suffix is prefilled. The retry replaces the original response only when it actually yields a promotable or structured call; on any error or still-garbage output the original response is returned unchanged. Exactly one retry, non-streaming OpenAI and Anthropic passthroughs only (a stream has already emitted bytes). OPT-IN per user decision: nudge_tool_calls=true per request (typed on both ChatCompletionRequest and AnthropicMessagesRequest, lifted from the Responses extra-body), or UNSLOTH_TOOL_CALL_NUDGE=1 to flip the process default. auto_heal_tool_calls=false disables healing AND the nudge. Also align the non-streaming heal on allow_incomplete=True: the response is final, so a trailing unclosed tool block is a model failure worth repairing, matching the enable-tools loop's drain semantics. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * inference: never assume the upstream response shape in the nudge helpers llama-server error bodies can carry message: null (or no choices at all), and _last_assistant_text / response_has_promotable_calls / nudge_should_retry called .get() on the message without a dict check, so a malformed upstream response raised an AttributeError the surrounding except tuples did not catch, failing the request instead of degrading to 'nothing to heal'. Route the shape probing through one _first_choice_message helper that returns None for any non-dict message, and add a parametrized test over the malformed shapes. * inference: constrain healing by tool_choice, preserve length finish_reason, keep healed event order in Responses streams Three review findings on the passthrough healer: - heal_gate now honors the request's tool_choice: "none" disables healing outright and a forced function narrows the promotion allowlist to that one function, so healing can never contradict the request's tool-choice constraint. Wired through the OpenAI chat (stream and non-stream), Responses, and Anthropic (converted shape) passthroughs. - The OpenAI non-streaming heal only upgrades finish_reason "stop" to "tool_calls"; a truncated generation keeps "length" (the healed call stays attached) matching the streaming and Anthropic paths. - The Responses stream emits healer events in order instead of collapsing all text ahead of the healed calls, so text after a healed call no longer jumps ahead of the function_call item and output indexes are claimed in the order the model produced them. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * inference: all-or-nothing promotion when a response mixes declared and undeclared text-form calls Promoting a subset used to strip ALL tool markup from the content, which silently deleted the text of any call naming an undeclared tool. The heal now declines entirely when any parsed call is unpromotable, so the whole message relays verbatim (pre-PR behavior) and no bytes are ever lost. In streaming, a declared call that completed before an undeclared one arrived is already emitted; the late undeclared markup still flushes as raw text. The nudge helpers mirror the same contract via a shared predicate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests: wrap long lines in the Responses healing tests to the project style * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * inference: span-exact healing, disjoint healed stream indexes, per-call Responses message items, allowlisted nudge acceptance Four review findings on the passthrough healer: - parse_tool_calls_from_text gains an optional with_spans return so healing removes EXACTLY the promoted calls' markup. This supersedes the previous all-or-nothing rule: declared calls promote and every unpromoted byte (undeclared calls, unparseable closed blocks, suppressed alternate formats such as a block after a JSON call) relays as text. The stream healer also processes one block per pass, so text between two healed calls keeps its document position instead of trailing them. - The OpenAI chat stream shifts native tool-call delta indexes past any already-emitted healed calls; clients merge deltas by index, so a healed call and a later native call can no longer merge into one. - A healed call in the Responses stream closes the open message item and trailing text opens a fresh one with a later output index, matching the native stream shape; response.completed snapshots every message item with its own text. - The nudge retry only replaces the original response when the retry's structured call names a DECLARED tool; a hallucinated undeclared call is not an improvement. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: stop the heal path folding trailing prose into a closed function call parse_tool_calls_from_text(allow_incomplete=True) cut a body only at an end-anchored , so a fully closed call followed by trailing prose (.. words) folded and the prose into the tool argument and deleted the prose from visible content. The strict path (allow_incomplete=False) already cut at the real via rfind. Do the same in both modes: trim the body at the real when present and end the removal span there, falling back to the end-anchored strip and body_end only when the call is genuinely truncated. Add a regression test. * inference: one shared single-call budget for healed and native calls Codex round 5: the parallel-call caps counted healed and native calls separately, so a healed text-form call followed by a native structured delta double-emitted on all three streaming surfaces when the client disabled parallel calls. - OpenAI SSE: once a healed call went out with parallel_tool_calls false, native tool_call deltas are dropped instead of index-shifted. - Anthropic emitter: native deltas skip block allocation when the healed-plus-native count already filled the single slot, and healed emission counts open native states too. - Responses stream: native deltas that survived the chunk-level cap are skipped once a healed call claimed the slot. Also adds a span assertion for the closed- trailing-prose parse fixed in the previous commit. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: relay undeclared text-form calls as text on Anthropic non-streaming heal_openai_message promotes only declared text-form tool calls and span-trims just their markup, deliberately leaving every unpromoted byte (undeclared text-form calls included) in the content to relay as text. The Anthropic non-streaming builder then ran a blanket _TOOL_XML_RE strip over that content unconditionally, deleting the undeclared block before building the text part, so Anthropic clients silently lost a call the OpenAI non-streaming path preserves. The strip was harmless when healing was all-or-nothing but became data loss once healing turned span-exact. Gate the legacy strip on whether healing promoted a call, matching the OpenAI passthrough and the intent already stated in the comment above. Add a route-level regression test for the mixed declared+undeclared case. * inference: require fully declared nudge retries; keep unpromoted Anthropic text Codex round 6, two findings: - response_has_promotable_calls accepted a nudge retry when any one structured call named a declared tool, so a mixed retry (hallucinated undeclared call plus a declared one) replaced the original and the caller forwarded the undeclared call, or with parallel_tool_calls false could keep only it. All structured retry calls must be declared. - The Anthropic non-streaming builder still ran the legacy _TOOL_XML_RE strip after span-exact healing, deleting undeclared or malformed call text that healing deliberately preserved. The legacy strip now runs only when healing is off (no declared tools, or opted out), matching the OpenAI passthrough. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * inference: keep unpromoted Anthropic text whenever healing is active The previous commit skipped the legacy strip only when a call was actually promoted, so an undeclared-only (or malformed-only) response was still silently emptied: exactly the dead-turn shape this path exists to fix, and inconsistent with the OpenAI passthrough, which relays those bytes verbatim. Gate the strip on healing being active instead; opt-out and no-tools requests keep the legacy strip. * Fix schema-aware tool healing for PR #6801 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix passthrough healing ordering for PR #6801 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix stream finish ordering for PR #6801 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: wasimysaid Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> --- .../core/inference/anthropic_compat.py | 147 +- .../core/inference/passthrough_healing.py | 535 +++++++ studio/backend/core/tool_healing.py | 186 +-- studio/backend/models/inference.py | 18 + studio/backend/routes/inference.py | 666 ++++++-- .../backend/tests/test_passthrough_healing.py | 1358 +++++++++++++++++ .../tests/test_responses_tool_passthrough.py | 174 +++ .../tests/test_tool_call_parser_strict.py | 37 + 8 files changed, 2919 insertions(+), 202 deletions(-) create mode 100644 studio/backend/core/inference/passthrough_healing.py create mode 100644 studio/backend/tests/test_passthrough_healing.py diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 0307336dde..7b572a28ff 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -494,6 +494,29 @@ class AnthropicPassthroughEmitter: self._usage: dict = {} self._stop_reason: str = "end_turn" self._stop_sequence: Optional[str] = None + # Optional text-form tool-call healing (client-tool passthrough only). + self._healer = None + self._healed_tool_use = False + self._healed_call_count = 0 + self._heal_disable_parallel = False + + def enable_healing( + self, + allowed_tools: set, + tools: Optional[list] = None, + *, + disable_parallel_tool_use: bool = False, + ) -> None: + """Promote text-form tool calls in streamed content to tool_use blocks. + + Only calls naming a tool in ``allowed_tools`` (the client's declared + tools) are promoted; everything else streams as text exactly as before. + Never enabled for Studio's own tool loop. + """ + from core.inference.passthrough_healing import StreamToolCallHealer + + self._healer = StreamToolCallHealer(allowed_tools, tools) + self._heal_disable_parallel = disable_parallel_tool_use def start( self, @@ -542,29 +565,42 @@ class AnthropicPassthroughEmitter: delta = choice.get("delta") or {} finish_reason = choice.get("finish_reason") + # ── Structured tool calls take precedence over healing ── + # Grammar mode worked: flush anything the healer held (it preceded the + # call in the model's output) and relay verbatim from here on. + if delta.get("tool_calls") and self._healer is not None and not self._healer.dormant: + for kind, value in self._healer.structured_tool_call_seen(): + if kind == "text" and value: + events.extend(self._emit_text_delta(value)) + # ── Text content ── content = delta.get("content") - if content: - if self._current_block_type != "text": - if self._current_block_type is not None: - events.append(self._close_current_block()) - events.extend(self._open_text_block()) - events.append( - build_anthropic_sse_event( - "content_block_delta", - { - "type": "content_block_delta", - "index": self.block_index, - "delta": {"type": "text_delta", "text": content}, - }, - ) - ) + if content and self._healer is not None and not self._healer.dormant: + # Route text through the healer: held/promoted portions become + # synthetic tool_use blocks, the rest streams as text unchanged. + for kind, value in self._healer.feed(content): + if kind == "text": + events.extend(self._emit_text_delta(value)) + else: + events.extend(self._emit_healed_tool_use(value)) + elif content: + events.extend(self._emit_text_delta(content)) # ── Tool calls (streaming deltas) ── tool_calls = delta.get("tool_calls") or [] for tc in tool_calls: tc_idx = tc.get("index", 0) fn = tc.get("function") or {} + if ( + self._heal_disable_parallel + and tc_idx not in self._tool_call_states + and (self._healed_call_count + len(self._tool_call_states)) >= 1 + ): + # disable_parallel_tool_use: a healed call already consumed the + # single allowed slot. The caller's chunk-level cap only sees + # native indexes, so drop this native call (and its later + # argument deltas, which never allocate a state either). + continue if tc_idx not in self._tool_call_states: # New tool call — close prior block, open tool_use block if self._current_block_type is not None: @@ -618,6 +654,17 @@ class AnthropicPassthroughEmitter: def finish(self) -> list[str]: events: list[str] = [] + if self._healer is not None: + # Last-chance heal of any held residue (e.g. an unclosed tool block). + for kind, value in self._healer.finalize(): + if kind == "text" and value: + events.extend(self._emit_text_delta(value)) + elif kind == "tool_call": + events.extend(self._emit_healed_tool_use(value)) + if self._healed_tool_use and self._stop_reason != "max_tokens": + # A promoted call must stop for tool use; a truncation still wins + # (its arguments may be incomplete). + self._stop_reason = "tool_use" if self._current_block_type is not None: events.append(self._close_current_block()) events.append( @@ -641,6 +688,76 @@ class AnthropicPassthroughEmitter: ) return events + def _emit_text_delta(self, content: str) -> list[str]: + events: list[str] = [] + if self._current_block_type != "text": + if self._current_block_type is not None: + events.append(self._close_current_block()) + events.extend(self._open_text_block()) + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": content}, + }, + ) + ) + return events + + def _emit_healed_tool_use(self, call: dict) -> list[str]: + # A healed call arrives complete, so its tool_use block opens, carries + # one input_json_delta, and closes immediately; an open text block is + # closed first (only the safe prefix ever streamed into it). + if ( + self._heal_disable_parallel + and (self._healed_call_count + len(self._tool_call_states)) >= 1 + ): + # Healed and native calls share the single allowed slot. + return [] + events: list[str] = [] + if self._current_block_type is not None: + events.append(self._close_current_block()) + function = call.get("function") or {} + tool_id = anthropic_tool_use_id("") + self.block_index += 1 + self._current_block_type = "tool_use" + events.append( + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": { + "type": "tool_use", + "id": tool_id, + "name": function.get("name", ""), + "input": {}, + }, + }, + ) + ) + arguments = function.get("arguments") or "" + if arguments: + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": { + "type": "input_json_delta", + "partial_json": arguments, + }, + }, + ) + ) + events.append(self._close_current_block()) + self._healed_tool_use = True + self._healed_call_count += 1 + return events + def _open_text_block(self) -> list[str]: self.block_index += 1 self._current_block_type = "text" diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py new file mode 100644 index 0000000000..c73134b4a2 --- /dev/null +++ b/studio/backend/core/inference/passthrough_healing.py @@ -0,0 +1,535 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tool-call healing for the client-tool passthrough. + +With server-side tools disabled (``unsloth run --disable-tools``, every +``unsloth start`` coding agent), requests carrying the client's own ``tools`` +bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small +GGUF models often emit their tool calls as TEXT (``{...}``, +Gemma ``<|tool_call>...``, ```` XML) instead of structured +``tool_calls`` -- on the passthrough that text reaches the agent as prose and +the turn dies. This module promotes such text back into structured calls on the +RESPONSE side only: the upstream request body is never touched, no extra +generation is issued, so llama-server slot/KV-cache reuse is byte-identical. + +Healing only ever fires when the request declared client tools, and only +promotes calls whose function name exactly matches a declared tool. Promotion +removes EXACTLY the promoted calls' markup spans (the parser reports them): +undeclared calls, unparseable blocks, and suppressed alternate formats keep +every byte and relay as text, so healing can never silently delete model +output. Responses without a tool signal, requests without tools, and Studio's +own enable-tools loop are untouched. Per-request opt-out: +``auto_heal_tool_calls: false``. Process kill-switch: +``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``. +""" + +import json +import os +from collections.abc import Mapping +from typing import Any, Optional + +from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal +from core.inference.tool_loop_controller import coerce_tool_arguments +from core.tool_healing import parse_tool_calls_from_text + +# Read once at import (same convention as the other UNSLOTH_* switches). +_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1" +# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process +# default with UNSLOTH_TOOL_CALL_NUDGE=1 (e.g. an `unsloth run` operator). +_NUDGE_DEFAULT = os.environ.get("UNSLOTH_TOOL_CALL_NUDGE", "0") == "1" + + +def nudge_enabled(request_flag: Optional[bool]) -> bool: + return _NUDGE_DEFAULT if request_flag is None else bool(request_flag) + + +_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS) +# A suspected-but-unclosed tool block larger than this is declared a false +# alarm and flushed, bounding memory on a model rambling XML-lookalike text. +_MAX_HOLD_CHARS = 64 * 1024 + + +def heal_gate( + auto_heal: Optional[bool], + tools: Optional[list], + tool_choice: Any = None, +) -> Optional[set]: + """Return the declared client-tool name set when healing applies, else None. + + ``tools`` is the OpenAI-shaped list forwarded to llama-server + (``[{"type": "function", "function": {"name": ...}}, ...]``). The name set + doubles as the promotion allowlist so healed calls can never invent a tool + the client did not declare. + + ``tool_choice`` (OpenAI shape) constrains the allowlist so healing never + contradicts the request: ``"none"`` forbids tool calls outright (text-form + markup stays text), and a forced ``{"type": "function", "function": + {"name": N}}`` narrows promotion to that one function. ``"auto"`` / + ``"required"`` / absent keep the full declared set. + """ + if _HEALING_DISABLED or auto_heal is False: + return None + if tool_choice == "none": + return None + names = set() + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if isinstance(function, dict) and isinstance(function.get("name"), str): + names.add(function["name"]) + if isinstance(tool_choice, dict): + function = tool_choice.get("function") + forced = function.get("name") if isinstance(function, dict) else None + if isinstance(forced, str): + names &= {forced} + return names or None + + +def _tool_schemas_by_name(tools: Optional[list]) -> dict[str, Any]: + schemas: dict[str, Any] = {} + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if not isinstance(function, dict): + continue + name = function.get("name") + if isinstance(name, str): + schemas[name] = function.get("parameters") + return schemas + + +def _string_arg_key_from_schema(schema: Any) -> Optional[str]: + if not isinstance(schema, dict): + return None + properties = schema.get("properties") + required = schema.get("required") + if not isinstance(properties, dict) or not isinstance(required, list): + return None + required_names = [name for name in required if isinstance(name, str)] + if len(required_names) != 1: + return None + key = required_names[0] + + if key not in properties: + return None + prop_schema = properties.get(key) + if isinstance(prop_schema, dict): + prop_type = prop_schema.get("type") + if isinstance(prop_type, list): + if "string" not in prop_type: + return None + elif prop_type is not None and prop_type != "string": + return None + return key + + +def _coerce_promoted_arguments( + raw_args: Any, tool_name: str, tool_schemas: Optional[dict] +) -> Optional[dict]: + if isinstance(raw_args, Mapping): + return dict(raw_args) + if isinstance(raw_args, str): + try: + parsed = json.loads(raw_args) + if isinstance(parsed, Mapping): + return dict(parsed) + except (json.JSONDecodeError, ValueError): + pass + if tool_schemas is not None: + key = _string_arg_key_from_schema(tool_schemas.get(tool_name)) + return {key: raw_args} if key else None + coerced = coerce_tool_arguments(raw_args, heal = True, tool_name = tool_name) + return coerced.arguments + + +def _promote( + calls: list, + allowed_tools: set, + id_offset: int = 0, + tool_schemas: Optional[dict] = None, +) -> list: + """Filter parsed calls to declared tools and normalize their arguments. + + Bare string arguments on the client-tool passthrough use the declared + schema's single required string property. If the schema is ambiguous, the + call stays text instead of inventing a generic key. + """ + promoted = [] + for call in calls: + function = call.get("function") if isinstance(call, dict) else None + name = function.get("name") if isinstance(function, dict) else None + if name not in allowed_tools: + continue + arguments = _coerce_promoted_arguments(function.get("arguments"), name, tool_schemas) + if arguments is None: + continue + promoted.append( + { + "id": f"call_{id_offset + len(promoted)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments, ensure_ascii = False), + }, + } + ) + return promoted + + +def _remove_spans(text: str, spans: list) -> str: + """Text with the given non-overlapping, sorted (start, end) ranges removed.""" + pieces = [] + pos = 0 + for start, end in spans: + pieces.append(text[pos:start]) + pos = end + pieces.append(text[pos:]) + return "".join(pieces) + + +def heal_openai_message_events( + msg: dict, + allowed_tools: set, + tools: Optional[list] = None, +) -> Optional[list]: + if not isinstance(msg, dict) or msg.get("tool_calls"): + return None + content = msg.get("content") + if not isinstance(content, str) or not has_tool_signal(content): + return None + parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True) + tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + events: list = [] + pos = 0 + call_count = 0 + for call, (start, end) in zip(parsed, spans): + promoted = _promote([call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas) + if promoted: + if content[pos:start]: + events.append(("text", content[pos:start])) + events.append(("tool_call", promoted[0])) + call_count += 1 + else: + events.append(("text", content[pos:end])) + pos = end + if not call_count: + return None + if content[pos:]: + events.append(("text", content[pos:])) + return events + + +def heal_openai_message( + msg: dict, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """Promote text-form tool calls in a non-streaming OpenAI message. In place. + + No-op (returns False) unless the message has NO structured ``tool_calls`` + (grammar mode already worked when it does) and its content carries a tool + signal that parses into at least one declared call. Only the promoted + calls' markup spans are removed from the content; undeclared calls and + anything the parser did not consume stay in the text byte-intact. + """ + events = heal_openai_message_events(msg, allowed_tools, tools) + if not events: + return False + calls = [value for kind, value in events if kind == "tool_call"] + content = "".join(value for kind, value in events if kind == "text").strip() + msg["tool_calls"] = calls + # OpenAI requires content = null on a pure tool-call turn. + msg["content"] = content or None + return True + + +def _earliest_signal(buffer: str) -> int: + best = -1 + for signal in TOOL_XML_SIGNALS: + index = buffer.find(signal) + if index >= 0 and (best < 0 or index < best): + best = index + return best + + +def _closed_signal_span(buffer: str) -> Optional[tuple[int, int]]: + spans = [] + for open_tag, close_tag in ( + ("", ""), + ("<|tool_call>", ""), + (""), + ): + start = buffer.find(open_tag) + if start < 0: + continue + end = buffer.find(close_tag, start) + if end >= 0: + spans.append((start, end + len(close_tag))) + return min(spans, key = lambda span: span[0]) if spans else None + + +def _partial_signal_suffix(buffer: str) -> int: + """Length of the longest buffer suffix that is a proper prefix of a signal.""" + for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1): + tail = buffer[-length:] + if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS): + return length + return 0 + + +class StreamToolCallHealer: + """Buffer-and-repair state machine for streamed passthrough content. + + ``feed(text)`` / ``finalize()`` yield ``("text", str)`` events for content + to relay and ``("tool_call", dict)`` events carrying an OpenAI-shaped call + (string ``function.arguments``). Normal prose is forwarded immediately; only + a trailing partial-signal window (< max signal length) or a suspected tool + block is ever withheld, so streaming latency stays bounded. A false alarm + (the buffer can no longer become a parseable declared call) flushes the held + text verbatim. + """ + + def __init__( + self, + allowed_tools: set, + tools: Optional[list] = None, + ) -> None: + self._allowed = set(allowed_tools) + + self._tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + self._buffer = "" + self._holding = False + self._id_offset = 0 + # Structured delta.tool_calls seen upstream: grammar mode already + # worked, so healing goes dormant and text relays verbatim. + self.dormant = False + + @property + def healed(self) -> bool: + return self._id_offset > 0 + + def structured_tool_call_seen(self) -> list: + """Go dormant; flush anything held so no text is swallowed.""" + self.dormant = True + held, self._buffer, self._holding = self._buffer, "", False + return [("text", held)] if held else [] + + def feed(self, text: str) -> list: + if self.dormant: + return [("text", text)] if text else [] + self._buffer += text + return self._drain() + + def _drain(self) -> list: + events: list = [] + while True: + if not self._holding: + start = _earliest_signal(self._buffer) + if start >= 0: + if start: + events.append(("text", self._buffer[:start])) + self._buffer = self._buffer[start:] + self._holding = True + else: + keep = _partial_signal_suffix(self._buffer) + emit = self._buffer[: len(self._buffer) - keep] + if emit: + events.append(("text", emit)) + self._buffer = self._buffer[len(self._buffer) - keep :] + return events + # HOLD: handle the FIRST complete block per pass so events keep + # document order (a later declared call must not overtake an + # earlier undeclared one flushing as text). + parsed, spans = parse_tool_calls_from_text( + self._buffer, + id_offset = self._id_offset, + allow_incomplete = False, + with_spans = True, + ) + if not parsed: + closed_span = _closed_signal_span(self._buffer) + if closed_span: + _start, end = closed_span + events.append(("text", self._buffer[:end])) + self._buffer = self._buffer[end:] + self._holding = False + continue + if len(self._buffer) > _MAX_HOLD_CHARS: + events.append(("text", self._buffer)) + self._buffer = "" + self._holding = False + continue + return events + start, end = spans[0] + promoted = _promote( + [parsed[0]], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + if start: + events.append(("text", self._buffer[:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + # Drop exactly the promoted markup span; everything else + # (leading text, later blocks) stays and is rescanned. + self._buffer = self._buffer[end:] + else: + # Undeclared or unusable name: its markup is DATA, flush it + # (and anything before it) verbatim, then rescan the rest. + events.append(("text", self._buffer[:end])) + self._buffer = self._buffer[end:] + self._holding = False + + def finalize(self) -> list: + """End of stream: last-chance heal of the residue, else flush it. + + Events keep document order; only the promoted calls' markup spans are + dropped, every other residue byte flushes as text. + """ + if not self._buffer: + return [] + residue, self._buffer = self._buffer, "" + holding, self._holding = self._holding, False + if self.dormant or not holding: + return [("text", residue)] + parsed, spans = parse_tool_calls_from_text( + residue, + id_offset = self._id_offset, + allow_incomplete = True, + with_spans = True, + ) + events: list = [] + pos = 0 + any_promoted = False + for call, (start, end) in zip(parsed, spans): + promoted = _promote( + [call], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + if residue[pos:start]: + events.append(("text", residue[pos:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + any_promoted = True + else: + events.append(("text", residue[pos:end])) + pos = end + if not any_promoted: + return [("text", residue)] + tail = residue[pos:].strip() + if tail: + events.append(("text", tail)) + return events + + +def _first_choice_message(data: Any) -> Optional[dict]: + """First-choice message dict of a non-streaming chat response, else None. + + Upstream error bodies can carry ``"message": null`` (or no choices at all), + so never assume the shape: a non-dict message means "nothing to heal". + """ + try: + message = data["choices"][0]["message"] + except (KeyError, IndexError, TypeError): + return None + return message if isinstance(message, dict) else None + + +def _last_assistant_text(data: Any) -> str: + """First-choice assistant content of a non-streaming chat response, or ''.""" + message = _first_choice_message(data) + content = message.get("content") if message else None + return content if isinstance(content, str) else "" + + +def _heal_would_promote( + text: str, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """Whether ``heal_openai_message`` would promote at least one call.""" + parsed = parse_tool_calls_from_text(text, allow_incomplete = True) + tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + return bool(_promote(parsed, allowed_tools, tool_schemas = tool_schemas)) + + +def response_has_promotable_calls( + data: Any, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """True when a non-streaming chat response carries a usable tool call + (structured naming a DECLARED tool, or text-form that healing would + promote). Used to decide whether a nudge retry actually improved on the + original response; a hallucinated undeclared call is not an improvement.""" + message = _first_choice_message(data) + if not message: + return False + tool_calls = message.get("tool_calls") + if tool_calls: + # ALL structured calls must be declared: the caller forwards the whole + # list (and a parallel cap could keep only the FIRST one), so a mixed + # response with a single hallucinated name could still hand the client + # an undeclared tool. + return all( + isinstance(tc, dict) + and isinstance(tc.get("function"), dict) + and tc["function"].get("name") in allowed_tools + for tc in tool_calls + ) + text = message.get("content") + if not isinstance(text, str): + return False + return _heal_would_promote(text, allowed_tools, tools) + + +def nudge_should_retry( + data: Any, + allowed_tools: Optional[set], + tools: Optional[list] = None, +) -> bool: + """True when the first response tried to call a tool but nothing healed. + + Trigger only on: healing enabled (allowed_tools set), zero structured + calls, a tool signal present in the text, and zero promotable calls -- the + exact failure a single re-ask can fix. Clean prose never retries. + """ + if not allowed_tools: + return False + message = _first_choice_message(data) + if not message or message.get("tool_calls"): + return False + text = message.get("content") + if not isinstance(text, str) or not has_tool_signal(text): + return False + return not _heal_would_promote(text, allowed_tools, tools) + + +def nudge_messages(data: Any, allowed_tools: set) -> list: + """The two-message suffix appended for the single nudge retry. + + The retry body is the original body plus this suffix, so the prompt prefix + is byte-identical and llama-server's slot/prefix cache is reused (same + shape as the enable-tools loop's reprompt). + """ + tool_hint = " or ".join(f"`{name}`" for name in sorted(allowed_tools)) or "an available tool" + return [ + {"role": "assistant", "content": _last_assistant_text(data)}, + { + "role": "user", + "content": ( + "You have access to the declared tools. If a tool is needed to " + f"complete the action you described, call {tool_hint} now using the " + "native tool-call format with valid JSON arguments, not prose. If no " + "tool is needed, provide the final answer directly." + ), + }, + ] diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index fe26d48c7f..e8367ad08c 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -301,28 +301,28 @@ def parse_tool_calls_from_text( *, id_offset: int = 0, allow_incomplete: bool = True, -) -> list[dict]: + with_spans: bool = False, +): """Parse OpenAI-format tool calls from model text. Handles formats like: {"name":"web_search","arguments":{"query":"..."}} <|tool_call>call:web_search{query:"..."} ... + + With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]`` + is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup + in ``content`` (including its close tag when present), so a caller can + remove exactly the parsed markup and keep every other byte intact. """ tool_calls: list[dict] = [] - # Collect JSON- and Gemma-format candidates with their byte spans, then - # accept them in document order. Both order and spans matter: - # * tools execute in returned order, so a call appearing earlier in the - # text must be emitted first even across the two formats; - # * a tool-call marker INSIDE another call's argument string is data, not a - # call, so a candidate starting within an already accepted span is - # skipped (covers a JSON marker nested in a Gemma arg and a Gemma marker - # nested in a JSON arg alike, regardless of which format is outer). + call_spans: list[tuple] = [] + # Collect every supported call format with spans, then emit in document + # order. A marker inside another call's argument string is data, not a + # separate executable call. + parsed_items = [] # (start, span_end, name, arguments) candidates = [] # (start, brace_end, kind, match) for m in _TC_JSON_START_RE.finditer(content): - # A marker that begins inside an open value - # is that parameter's data, not its own call; skip it (same guard the - # XML-style parser below applies to nested = 0: + span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) + body = body[:close_idx] + elif not allow_incomplete: + continue + else: + body = _TC_FUNC_CLOSE_RE.sub("", body) + span_end = body_end + + arguments: dict = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + pm = param_starts[0] + val = body[pm.end() :] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + continue + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[pm.group(1)] = val.strip() + else: + valid_params = True + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + ) + val = body[val_start:next_param] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + valid_params = False + break + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[param_name] = val.strip() + if not valid_params: + continue + + span_start = fm.start() + wrap_open = re.search(r"\s*$", content[:span_start]) + wrap_close = re.match(r"\s*", content[span_end:]) + if wrap_open and wrap_close: + span_start = wrap_open.start() + span_end += wrap_close.end() + parsed_items.append((span_start, span_end, func_name, json.dumps(arguments))) + + parsed_items.sort(key = lambda item: item[0]) + for start, span_end, name, arguments in parsed_items: tool_calls.append( { "id": f"call_{id_offset + len(tool_calls)}", @@ -369,77 +443,9 @@ def parse_tool_calls_from_text( "function": {"name": name, "arguments": arguments}, } ) - - if not tool_calls: - func_starts = [ - fm - for fm in _TC_FUNC_START_RE.finditer(content) - if not _inside_open_parameter(content, fm.start()) - ] - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - if not allow_incomplete: - close_idx = body.rfind(_FUNC_CLOSE_TAG) - if close_idx < 0: - continue - body = body[:close_idx] - else: - body = _TC_FUNC_CLOSE_RE.sub("", body) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - pm = param_starts[0] - val = body[pm.end() :] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - continue - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - valid_params = True - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - valid_params = False - break - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() - if not valid_params: - continue - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) + call_spans.append((start, span_end)) + if with_spans: + return tool_calls, call_spans return tool_calls diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 4a3162b09e..31c100dbec 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -780,6 +780,16 @@ class ChatCompletionRequest(BaseModel): True, description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Opt-in, non-streaming client-tool passthrough only: when the " + "model emitted a tool signal that healing could not repair, retry ONCE with " + "a short nudge appended (the retry shares the full prompt prefix, so the " + "server's KV cache is reused). Default off; UNSLOTH_TOOL_CALL_NUDGE=1 flips " + "the process default." + ), + ) context_overflow: Optional[Literal["error", "truncate_middle"]] = Field( None, description = ( @@ -1612,6 +1622,14 @@ class AnthropicMessagesRequest(BaseModel): False, description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.", ) + auto_heal_tool_calls: Optional[bool] = Field( + True, + description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).", + ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] Opt-in, non-streaming only: retry once with a nudge when the model emitted a tool signal healing could not repair (mirrors the Chat Completions field).", + ) model_config = {"extra": "allow"} @model_validator(mode = "before") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a948a6eaf5..ccf36e8f71 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1135,6 +1135,16 @@ from core.inference.key_exchange import decrypt_api_key from core.inference.model_ids import public_model_id from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client +from core.inference.passthrough_healing import ( + StreamToolCallHealer, + heal_gate, + heal_openai_message, + heal_openai_message_events, + nudge_enabled, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override @@ -8065,6 +8075,13 @@ def _build_chat_request( if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw: chat_kwargs["enable_thinking"] = bool(_tpl_kw["enable_thinking"]) explicit_enable_thinking = True + # auto_heal_tool_calls / nudge_tool_calls are not typed on + # ResponsesRequest; lift them from the extra-body so passthrough + # healing (and the opt-in nudge) honor them on both paths. + if isinstance(_extra.get("auto_heal_tool_calls"), bool): + chat_kwargs["auto_heal_tool_calls"] = _extra["auto_heal_tool_calls"] + if isinstance(_extra.get("nudge_tool_calls"), bool): + chat_kwargs["nudge_tool_calls"] = _extra["nudge_tool_calls"] if isinstance(payload.reasoning, dict): effort = payload.reasoning.get("effort") @@ -8299,16 +8316,112 @@ async def _responses_stream( parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend) ) reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} - message_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} + message_state: dict[str, Any] = { + "output_index": None, + "item_id": None, + "opened": False, + "text": "", + } + # Message items already closed mid-stream (a healed tool call splits + # the assistant text into separate message items, as native Responses + # streams do). Kept for the final response.completed snapshot. + closed_message_states: list[dict] = [] # Per-tool-call state keyed by Chat Completions `tool_calls[].index`, # stable across chunks for the same call. Values: # {output_index, item_id, call_id, name, arguments, opened} tool_call_state: dict[int, dict] = {} next_output_index = 0 + # Text-form tool calls promoted back to structured calls (declared + # client tools only); dormant once grammar-mode structured deltas appear. + _allowed_tools = heal_gate( + getattr(chat_req, "auto_heal_tool_calls", None), + body.get("tools"), + body.get("tool_choice"), + ) + healer = StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + healed_tc_index = 0 + + def _healed_tc(call: dict): + # Chat-delta shape for a healed call. Indexes live in a disjoint + # range so a healed call can never merge into a structured call's + # state slot; parallel_tool_calls=false caps healed calls too (the + # upstream cap ran before injection). + nonlocal healed_tc_index + if payload.parallel_tool_calls is False and healed_tc_index >= 1: + return None + tc = { + "index": 1_000_000 + healed_tc_index, + "id": call["id"], + "type": "function", + "function": call["function"], + } + healed_tc_index += 1 + return tc def _sse(event_name: str, payload: dict) -> str: return f"event: {event_name}\ndata: {json.dumps(payload)}\n\n" + def _tool_call_delta_events(tc: dict) -> list: + # One Chat Completions tool_calls delta -> Responses SSE events, + # allocating/merging per-call state (shared by the structured loop + # and the healer's promoted calls). + events = [] + idx = tc.get("index", 0) + st = tool_call_state.get(idx) + fn = tc.get("function") or {} + if st is None: + # First chunk for this tool call -- allocate an + # output_index and emit output_item.added. + st = { + "output_index": _claim_output_index(), + "item_id": f"fc_{uuid.uuid4().hex[:12]}", + "call_id": tc.get("id") or "", + "name": fn.get("name") or "", + "arguments": "", + "opened": False, + } + tool_call_state[idx] = st + else: + # Later chunks sometimes carry id/name only once; merge + # when present. + if tc.get("id") and not st["call_id"]: + st["call_id"] = tc["id"] + if fn.get("name") and not st["name"]: + st["name"] = fn["name"] + + if not st["opened"] and st["call_id"] and st["name"]: + item_added = { + "type": "response.output_item.added", + "output_index": st["output_index"], + "item": { + "type": "function_call", + "id": st["item_id"], + "status": "in_progress", + "call_id": st["call_id"], + "name": st["name"], + "arguments": "", + }, + } + events.append(_sse("response.output_item.added", item_added)) + st["opened"] = True + + arg_delta = fn.get("arguments") or "" + if arg_delta and st["opened"]: + st["arguments"] += arg_delta + args_delta_event = { + "type": "response.function_call_arguments.delta", + "item_id": st["item_id"], + "output_index": st["output_index"], + "delta": arg_delta, + } + events.append(_sse("response.function_call_arguments.delta", args_delta_event)) + elif arg_delta: + # Buffer args until we can open the item (some models + # send id/name in the same chunk as the first arg delta; + # if not, stash). + st["arguments"] += arg_delta + return events + def _claim_output_index() -> int: nonlocal next_output_index output_index = next_output_index @@ -8393,6 +8506,98 @@ async def _responses_stream( ), ] + def _close_message_item() -> list[str]: + """Close the open message item so later text opens a fresh one. + + Emits the same done-event triplet the end-of-stream close loop + would, records the item for the final snapshot, and resets the + state in place. No-op when no message item is open. + """ + if not message_state["opened"]: + return [] + text = message_state["text"] + events = [ + _sse( + "response.output_text.done", + { + "type": "response.output_text.done", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "text": text, + }, + ), + _sse( + "response.content_part.done", + { + "type": "response.content_part.done", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "part": {"type": "output_text", "text": text, "annotations": []}, + }, + ), + _sse( + "response.output_item.done", + { + "type": "response.output_item.done", + "output_index": message_state["output_index"], + "item": { + "type": "message", + "id": message_state["item_id"], + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + }, + }, + ), + ] + closed_message_states.append(dict(message_state)) + message_state.update( + {"output_index": None, "item_id": None, "opened": False, "text": ""} + ) + return events + + def _healed_event_sse(events) -> list[str]: + """Serialize healer events preserving their order. + + Text around a healed call must keep its position relative to the + function_call item (output indexes are claimed in emission order), + so never split an event list into all-text-then-all-calls. A healed + call also CLOSES any open message item, so trailing text opens a + fresh message with a later output index, exactly like a native + Responses stream that interleaves messages and calls. + """ + nonlocal full_text + out: list[str] = [] + for kind, value in events: + if kind == "text": + if not value: + continue + out.extend(_ensure_message_open()) + full_text += value + message_state["text"] += value + api_monitor.append_reply(monitor_id, value) + out.append( + _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": value, + }, + ) + ) + else: + tc = _healed_tc(value) + if tc is None: + continue + out.extend(_close_message_item()) + out.extend(_tool_call_delta_events(tc)) + return out + def _snapshot_output() -> list[dict]: """Snapshot of all completed output items for response.completed.""" indexed_items: list[tuple[int, dict]] = [] @@ -8409,19 +8614,23 @@ async def _responses_stream( }, ) ) - if message_state["opened"]: + # Closed copies keep opened=True (snapshotted before reset); the + # live state contributes only when a message is currently open. + for msg_st in [*closed_message_states, message_state]: + if not msg_st["opened"]: + continue indexed_items.append( ( - message_state["output_index"], + msg_st["output_index"], { "type": "message", - "id": message_state["item_id"], + "id": msg_st["item_id"], "status": "completed", "role": "assistant", "content": [ { "type": "output_text", - "text": full_text, + "text": msg_st["text"], "annotations": [], } ], @@ -8605,10 +8814,30 @@ async def _responses_stream( "delta": reasoning_delta, }, ) + # Heal text-form tool calls in the visible stream (never in + # reasoning text): promoted calls join the structured tc loop + # below through the same state machinery, and healer events are + # emitted IN ORDER so text after a healed call never jumps ahead + # of the function_call item. Once a structured delta arrives, + # grammar mode worked and the healer goes dormant. + if healer is not None and not healer.dormant: + healed_events = [] + if delta.get("tool_calls"): + # Held text preceded the structured call; the call's own + # deltas follow in the structured loop below. + healed_events = healer.structured_tool_call_seen() + if visible_delta: + healed_events.append(("text", visible_delta)) + elif visible_delta: + healed_events = healer.feed(visible_delta) + visible_delta = "" + for event in _healed_event_sse(healed_events): + yield event if visible_delta: for event in _ensure_message_open(): yield event full_text += visible_delta + message_state["text"] += visible_delta api_monitor.append_reply(monitor_id, visible_delta) yield _sse( "response.output_text.delta", @@ -8622,60 +8851,19 @@ async def _responses_stream( ) for tc in delta.get("tool_calls") or []: - idx = tc.get("index", 0) - st = tool_call_state.get(idx) - fn = tc.get("function") or {} - if st is None: - # First chunk for this tool call -- allocate an - # output_index and emit output_item.added. - st = { - "output_index": _claim_output_index(), - "item_id": f"fc_{uuid.uuid4().hex[:12]}", - "call_id": tc.get("id") or "", - "name": fn.get("name") or "", - "arguments": "", - "opened": False, - } - tool_call_state[idx] = st - else: - # Later chunks sometimes carry id/name only once; merge - # when present. - if tc.get("id") and not st["call_id"]: - st["call_id"] = tc["id"] - if fn.get("name") and not st["name"]: - st["name"] = fn["name"] - - if not st["opened"] and st["call_id"] and st["name"]: - item_added = { - "type": "response.output_item.added", - "output_index": st["output_index"], - "item": { - "type": "function_call", - "id": st["item_id"], - "status": "in_progress", - "call_id": st["call_id"], - "name": st["name"], - "arguments": "", - }, - } - yield _sse("response.output_item.added", item_added) - st["opened"] = True - - arg_delta = fn.get("arguments") or "" - if arg_delta and st["opened"]: - st["arguments"] += arg_delta - args_delta_event = { - "type": "response.function_call_arguments.delta", - "item_id": st["item_id"], - "output_index": st["output_index"], - "delta": arg_delta, - } - yield _sse("response.function_call_arguments.delta", args_delta_event) - elif arg_delta: - # Buffer args until we can open the item (some models - # send id/name in the same chunk as the first arg delta; - # if not, stash). - st["arguments"] += arg_delta + if ( + payload.parallel_tool_calls is False + and healed_tc_index >= 1 + and tc.get("index", 0) not in tool_call_state + ): + # A healed call already consumed the single allowed slot; + # _drop_parallel_tool_call_deltas only sees native indexes, + # so a native index-0 call would still open a second + # function_call item. Skip it (and its later argument + # deltas, which never allocate a state either). + continue + for event in _tool_call_delta_events(tc): + yield event _apply_usage(chunk_data.get("usage")) except asyncio.CancelledError: @@ -8731,10 +8919,19 @@ async def _responses_stream( "delta": final_reasoning, }, ) + # Last-chance heal of any held residue (e.g. a tool block the model + # never closed) before the trailing visible text is flushed; events + # keep healer order so trailing text stays behind a healed call. + if healer is not None: + events = (healer.feed(final_visible) if final_visible else []) + healer.finalize() + final_visible = "" + for event in _healed_event_sse(events): + yield event if final_visible: for event in _ensure_message_open(): yield event full_text += final_visible + message_state["text"] += final_visible api_monitor.append_reply(monitor_id, final_visible) yield _sse( "response.output_text.delta", @@ -8793,6 +8990,10 @@ async def _responses_stream( continue if kind == "message": + # Per-item text: message items closed mid-stream (healed-call + # rotation) already emitted their done events, so this state + # carries only its own text, not the whole stream's. + _msg_text = st["text"] yield _sse( "response.output_text.done", { @@ -8800,7 +9001,7 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "text": full_text, + "text": _msg_text, }, ) yield _sse( @@ -8810,7 +9011,7 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "part": {"type": "output_text", "text": full_text, "annotations": []}, + "part": {"type": "output_text", "text": _msg_text, "annotations": []}, }, ) yield _sse( @@ -8824,7 +9025,7 @@ async def _responses_stream( "status": "completed", "role": "assistant", "content": [ - {"type": "output_text", "text": full_text, "annotations": []} + {"type": "output_text", "text": _msg_text, "annotations": []} ], }, }, @@ -9430,6 +9631,7 @@ async def anthropic_messages( session_id = payload.session_id, cancel_id = payload.cancel_id, disable_parallel_tool_use = _disable_parallel, + auto_heal_tool_calls = payload.auto_heal_tool_calls, ) ) return await _monitored_anthropic( @@ -9449,6 +9651,8 @@ async def anthropic_messages( presence_penalty = presence_penalty, tool_choice = openai_tool_choice, disable_parallel_tool_use = _disable_parallel, + auto_heal_tool_calls = payload.auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, ) ) @@ -10019,6 +10223,7 @@ async def _anthropic_passthrough_stream( session_id = None, cancel_id = None, disable_parallel_tool_use = False, + auto_heal_tool_calls = None, ): """Streaming client-side pass-through: forward tools to llama-server and translate its stream to Anthropic SSE without executing anything.""" @@ -10055,6 +10260,16 @@ async def _anthropic_passthrough_stream( async def _stream(): emitter = AnthropicPassthroughEmitter() + # Promote text-form tool calls (declared client tools only) into + # tool_use blocks; verbatim behavior when healing is off or no tools. + # tool_choice arrives here already converted to the OpenAI shape. + _allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice) + if _allowed_tools: + emitter.enable_healing( + _allowed_tools, + openai_tools, + disable_parallel_tool_use = disable_parallel_tool_use, + ) for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line @@ -10191,6 +10406,8 @@ async def _anthropic_passthrough_non_streaming( presence_penalty = None, tool_choice = "auto", disable_parallel_tool_use = False, + auto_heal_tool_calls = None, + nudge_tool_calls = None, ): """Non-streaming client-side pass-through.""" target_url = f"{llama_backend.base_url}/v1/chat/completions" @@ -10223,34 +10440,98 @@ async def _anthropic_passthrough_non_streaming( ) data = resp.json() + # tool_choice arrives here already converted to the OpenAI shape. + _allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice) + + # Opt-in single-retry nudge (mirrors the OpenAI passthrough): the model + # tried to call a tool but nothing usable came out; re-ask once with the + # prompt prefix intact so llama-server's KV cache is reused. + if ( + _allowed_tools + and nudge_enabled(nudge_tool_calls) + and nudge_should_retry(data, _allowed_tools, openai_tools) + ): + retry_body = { + **body, + "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + } + try: + retry_resp = await nonstreaming_client().post( + target_url, + json = retry_body, + timeout = _llama_non_streaming_generation_timeout(), + ) + if retry_resp.status_code == 200: + retry_data = retry_resp.json() + if response_has_promotable_calls(retry_data, _allowed_tools, openai_tools): + data = retry_data + except (httpx.RequestError, ValueError) as exc: + logger.warning("tool-call nudge retry failed; keeping original: %s", exc) + choice = (data.get("choices") or [{}])[0] message = choice.get("message") or {} finish_reason = choice.get("finish_reason") - content_blocks = [] - text = message.get("content") or "" - if text: - text = _TOOL_XML_RE.sub("", text).strip() - if text: - content_blocks.append(AnthropicResponseTextBlock(text = text)) + healing_active = bool(_allowed_tools) + healed_events = ( + heal_openai_message_events(message, _allowed_tools, openai_tools) + if healing_active + else None + ) - tool_calls = message.get("tool_calls") or [] - # disable_parallel_tool_use: keep only the first tool_use block. - if disable_parallel_tool_use and len(tool_calls) > 1: - tool_calls = tool_calls[:1] - for tc in tool_calls: - fn = tc.get("function") or {} - try: - args = json.loads(fn.get("arguments", "{}")) - except json.JSONDecodeError: - args = {} - content_blocks.append( - AnthropicResponseToolUseBlock( - id = anthropic_tool_use_id(tc.get("id")), - name = fn.get("name", ""), - input = args, + content_blocks = [] + tool_calls = [] + if healed_events: + emitted_tool_uses = 0 + for kind, value in healed_events: + if kind == "text": + text = str(value).strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + continue + if disable_parallel_tool_use and emitted_tool_uses >= 1: + continue + fn = value.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + tool_calls.append(value) + emitted_tool_uses += 1 + content_blocks.append( + AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(value.get("id")), + name = fn.get("name", ""), + input = args, + ) + ) + else: + text = message.get("content") or "" + if text: + # Keep unpromoted bytes when healing is active; legacy stripping is + # only for opted-out or no-client-tool requests. + if not healing_active: + text = _TOOL_XML_RE.sub("", text) + text = text.strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + + tool_calls = message.get("tool_calls") or [] + if disable_parallel_tool_use and len(tool_calls) > 1: + tool_calls = tool_calls[:1] + for tc in tool_calls: + fn = tc.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + content_blocks.append( + AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(tc.get("id")), + name = fn.get("name", ""), + input = args, + ) ) - ) stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls)) @@ -10598,6 +10879,13 @@ async def _openai_passthrough_stream( body = _build_openai_passthrough_body( payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) + # Text-form tool calls from small models get promoted to structured calls on + # the way back (declared client tools only); requests without tools or with + # auto_heal_tool_calls=false keep the verbatim relay. tool_choice constrains + # the allowlist ("none" disables, a forced function narrows to it). + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) @@ -10700,9 +10988,14 @@ async def _openai_passthrough_stream( last_chunk_id = completion_id last_chunk_model = model_name last_chunk_created = int(time.time()) + healer = ( + StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + ) + healed_call_index = 0 def _synthetic_finish_line() -> str: - finish_reason = "tool_calls" if saw_tool_call_delta else "stop" + healed = healer is not None and healer.healed + finish_reason = "tool_calls" if (saw_tool_call_delta or healed) else "stop" chunk = ChatCompletionChunk( id = last_chunk_id, created = last_chunk_created, @@ -10716,6 +11009,108 @@ async def _openai_passthrough_stream( ) return f"data: {chunk.model_dump_json(exclude_none = True)}" + def _healer_sse_lines(events) -> list: + # Serialize healer events as chunks matching the upstream stream's + # id/model/created so clients see one coherent completion. + nonlocal healed_call_index + lines = [] + for kind, value in events: + if kind == "text": + if not value: + continue + delta = {"content": value} + else: + # parallel_tool_calls=false caps healed calls too (the SSE + # line cap only sees structured upstream deltas). + if payload.parallel_tool_calls is False and healed_call_index >= 1: + continue + delta = { + "tool_calls": [ + { + "index": healed_call_index, + "id": value["id"], + "type": "function", + "function": value["function"], + } + ] + } + healed_call_index += 1 + chunk = { + "id": last_chunk_id, + "object": "chat.completion.chunk", + "created": last_chunk_created, + "model": last_chunk_model, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], + } + lines.append("data: " + json.dumps(chunk, ensure_ascii = False)) + return lines + + def _heal_transform(chunk_data: dict, raw_line: str) -> list: + """SSE lines to emit in place of one upstream line (healing on).""" + choices = chunk_data.get("choices") + if not (isinstance(choices, list) and choices and isinstance(choices[0], dict)): + return [raw_line] + choice = choices[0] + delta = choice.get("delta") + delta = delta if isinstance(delta, dict) else {} + if delta.get("tool_calls"): + # Structured call streamed: grammar mode worked. Flush any held + # text (it preceded the call) and relay verbatim from here on. + lines = _healer_sse_lines(healer.structured_tool_call_seen()) + if healed_call_index: + if payload.parallel_tool_calls is False: + # A healed call already consumed the single allowed + # slot; the upstream SSE cap keeps native index 0, so + # drop the native call here or the client gets two. + del delta["tool_calls"] + if delta or choice.get("finish_reason") or chunk_data.get("usage"): + lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + return lines + # A healed call already went out on index 0..n-1; OpenAI + # clients merge tool-call deltas by index, so shift the + # native calls into the next indexes or they would merge + # into the healed call. + for tc in delta["tool_calls"]: + if isinstance(tc, dict) and isinstance(tc.get("index"), int): + tc["index"] += healed_call_index + return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [raw_line] + content = delta.get("content") + finish = choice.get("finish_reason") + if not isinstance(content, str) or not content: + if not finish: + return [raw_line] + # Finish chunk: last-chance heal of the residue, and rewrite a + # "stop" into "tool_calls" when text-form calls were promoted. + lines = _healer_sse_lines(healer.finalize()) + if healer.healed and finish == "stop": + choice["finish_reason"] = "tool_calls" + return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [raw_line] + events = healer.feed(content) + if finish: + events += healer.finalize() + if not finish and events == [("text", content)]: + # Nothing held or promoted: the healer passed the chunk + # through whole, so keep the verbatim upstream bytes. + return [raw_line] + del delta["content"] + prefix_lines = [] + if delta: + prefix_chunk = {k: v for k, v in chunk_data.items() if k != "usage"} + prefix_choice = dict(choice) + prefix_choice["delta"] = dict(delta) + prefix_choice["finish_reason"] = None + prefix_chunk["choices"] = [prefix_choice] + prefix_lines.append("data: " + json.dumps(prefix_chunk, ensure_ascii = False)) + delta.clear() + lines = prefix_lines + _healer_sse_lines(events) + if delta or finish or chunk_data.get("usage"): + if healer.healed and finish == "stop": + choice["finish_reason"] = "tool_calls" + lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + return lines + try: lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( @@ -10732,6 +11127,14 @@ async def _openai_passthrough_stream( data_text = raw_line[6:].strip() if data_text == "[DONE]": saw_done = True + # Upstream ended without a finish chunk: heal the residue + # first so the synthetic finish sees healer.healed. + if healer is not None and not saw_stream_error: + for held_line in _healer_sse_lines(healer.finalize()): + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" if ( not saw_finish_reason and not saw_stream_error @@ -10787,13 +11190,18 @@ async def _openai_passthrough_stream( # emit a successful finish_reason after a failed stream. if _monitor_openai_error_message(chunk_data): saw_stream_error = True - monitor_event = _monitor_openai_sse_line( - monitor_id, - raw_line, - llama_backend.context_length, - ) - if monitor_event == "error": - saw_stream_error = True + # With healing active, a content-bearing line may be replaced by + # held/promoted chunks; otherwise the single upstream line + # relays verbatim (monitored exactly as emitted either way). + if ( + healer is not None + and not healer.dormant + and isinstance(chunk_data, dict) + and not saw_stream_error + ): + out_lines = _heal_transform(chunk_data, raw_line) + else: + out_lines = [raw_line] # If a trailing usage-only chunk (include_usage) arrives before # any finish chunk, emit the synthetic finish first so the order # stays finish -> usage -> [DONE], matching the other streams. @@ -10807,23 +11215,46 @@ async def _openai_passthrough_stream( and not saw_stream_error and not cancel_event.is_set() ): + if healer is not None: + # Residue must precede the finish it may upgrade. + held = _healer_sse_lines(healer.finalize()) + for held_line in held: + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" finish_line = _synthetic_finish_line() _monitor_openai_sse_line( monitor_id, finish_line, llama_backend.context_length ) yield finish_line + "\n\n" saw_finish_reason = True - # Relay verbatim to preserve llama-server's native id, - # finish_reason, delta.tool_calls, and usage chunks. - yield raw_line + "\n\n" - if monitor_event == "done": - monitor_done = True + for out_line in out_lines: + monitor_event = _monitor_openai_sse_line( + monitor_id, + out_line, + llama_backend.context_length, + ) + if monitor_event == "error": + saw_stream_error = True + # Relay to preserve llama-server's native id, + # finish_reason, delta.tool_calls, and usage chunks. + yield out_line + "\n\n" + if monitor_event == "done": + monitor_done = True + if monitor_done: break if not saw_done and not saw_stream_error and not cancel_event.is_set(): # Synthesize a finish chunk only if one was not already # emitted (e.g. before a trailing usage-only chunk), but # always close with [DONE] whenever the upstream omitted it, # so the stream ends on the [DONE] sentinel either way. + if healer is not None: + for held_line in _healer_sse_lines(healer.finalize()): + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" if not saw_finish_reason: finish_line = _synthetic_finish_line() _monitor_openai_sse_line( @@ -10962,6 +11393,9 @@ async def _openai_passthrough_non_streaming( _guided_fence = bool((payload.model_extra or {}).get("_unsloth_guided_fence")) _do_fence = _guided_fence and _extract_response_format(payload) is not None _cap_parallel = payload.parallel_tool_calls is False + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) try: data = resp.json() @@ -10974,6 +11408,33 @@ async def _openai_passthrough_non_streaming( api_monitor.finish(monitor_id) return Response(content = resp.content, media_type = "application/json") + # Opt-in single-retry nudge: the model clearly tried to call a tool (signal + # present) but nothing parseable/declared came out, so re-ask once with the + # original prompt prefix intact (llama-server reuses the slot's KV cache) + # plus a two-message nudge suffix. The retry replaces the original response + # only when it actually yields a usable call. + if ( + _allowed_tools + and nudge_enabled(payload.nudge_tool_calls) + and nudge_should_retry(data, _allowed_tools, body.get("tools")) + ): + retry_body = { + **body, + "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + } + try: + retry_resp = await nonstreaming_client().post( + target_url, + json = retry_body, + timeout = _llama_non_streaming_generation_timeout(), + ) + if retry_resp.status_code == 200: + retry_data = retry_resp.json() + if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")): + resp, data = retry_resp, retry_data + except (httpx.RequestError, ValueError) as exc: + logger.warning("tool-call nudge retry failed; keeping original: %s", exc) + changed = False for choice in data.get("choices", []): if not isinstance(choice, dict): @@ -10982,6 +11443,17 @@ async def _openai_passthrough_non_streaming( if not isinstance(msg, dict): continue + # Small models emit tool calls as text instead of structured tool_calls; + # promote them (declared client tools only) so the agent sees a real call. + # Truncation wins over the upgrade (same rule as the streaming and + # Anthropic paths): a call cut off at max_tokens keeps + # finish_reason="length" so the client knows the arguments may be + # incomplete, while the healed call itself stays attached. + if _allowed_tools and heal_openai_message(msg, _allowed_tools, body.get("tools")): + if choice.get("finish_reason") == "stop": + choice["finish_reason"] = "tool_calls" + changed = True + # OpenAI requires content=null on a pure tool-call turn; llama-server # emits content="". if msg.get("tool_calls") and msg.get("content") == "": diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py new file mode 100644 index 0000000000..06316e2243 --- /dev/null +++ b/studio/backend/tests/test_passthrough_healing.py @@ -0,0 +1,1358 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for core/inference/passthrough_healing.py: promoting text-form +tool calls back into structured calls on the client-tool passthrough. The +route-level wiring (OpenAI / Anthropic / Responses endpoints) is covered in +their own endpoint test files; this file exercises the shared state machine +and helpers directly. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.passthrough_healing import ( # noqa: E402 + StreamToolCallHealer, + heal_gate, + heal_openai_message, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) + +TOOLS = [ + {"type": "function", "function": {"name": "Bash", "parameters": {}}}, + {"type": "function", "function": {"name": "Read", "parameters": {}}}, +] + +BASH_COMMAND_TOOL = { + "type": "function", + "function": { + "name": "Bash", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, +} +XML_BASH = '{"name":"Bash","arguments":{"cmd":"ls"}}' +XML_UNDECLARED = '{"name":"Nuke","arguments":{}}' + + +def _events_text(events): + return "".join(text for kind, text in events if kind == "text") + + +def _events_calls(events): + return [call for kind, call in events if kind == "tool_call"] + + +class TestHealGate: + def test_returns_declared_names(self): + assert heal_gate(None, TOOLS) == {"Bash", "Read"} + assert heal_gate(True, TOOLS) == {"Bash", "Read"} + + def test_opt_out_and_no_tools(self): + assert heal_gate(False, TOOLS) is None + assert heal_gate(None, []) is None + assert heal_gate(None, None) is None + + def test_malformed_tool_entries_ignored(self): + assert heal_gate(None, ["nonsense", {"function": "x"}, {}]) is None + + def test_tool_choice_none_disables(self): + assert heal_gate(None, TOOLS, "none") is None + + def test_tool_choice_forced_function_narrows_allowlist(self): + forced = {"type": "function", "function": {"name": "Bash"}} + assert heal_gate(None, TOOLS, forced) == {"Bash"} + + def test_tool_choice_forced_undeclared_function_disables(self): + forced = {"type": "function", "function": {"name": "Nuke"}} + assert heal_gate(None, TOOLS, forced) is None + + def test_tool_choice_auto_and_required_keep_full_set(self): + assert heal_gate(None, TOOLS, "auto") == {"Bash", "Read"} + assert heal_gate(None, TOOLS, "required") == {"Bash", "Read"} + + def test_tool_choice_unrecognized_dict_keeps_full_set(self): + assert heal_gate(None, TOOLS, {"type": "function"}) == {"Bash", "Read"} + + +class TestHealOpenaiMessage: + def test_promotes_xml_and_strips_content(self): + msg = {"role": "assistant", "content": XML_BASH} + assert heal_openai_message(msg, {"Bash"}) is True + assert msg["content"] is None + (call,) = msg["tool_calls"] + assert call["function"]["name"] == "Bash" + assert json.loads(call["function"]["arguments"]) == {"cmd": "ls"} + + def test_keeps_surrounding_prose(self): + msg = {"role": "assistant", "content": f"Let me check.\n{XML_BASH}"} + assert heal_openai_message(msg, {"Bash"}) is True + assert msg["content"] == "Let me check." + + def test_undeclared_name_not_promoted(self): + msg = {"role": "assistant", "content": XML_UNDECLARED} + assert heal_openai_message(msg, {"Bash"}) is False + assert msg["content"] == XML_UNDECLARED + assert "tool_calls" not in msg + + def test_structured_calls_untouched(self): + msg = {"role": "assistant", "content": XML_BASH, "tool_calls": [{"id": "x"}]} + assert heal_openai_message(msg, {"Bash"}) is False + assert msg["content"] == XML_BASH + + def test_prose_only_untouched(self): + msg = {"role": "assistant", "content": "just an answer"} + assert heal_openai_message(msg, {"Bash"}) is False + + def test_bare_string_arguments_use_schema_key(self): + msg = { + "role": "assistant", + "content": '{"name":"Bash","arguments":"echo hi"}', + } + assert heal_openai_message(msg, {"Bash"}, [BASH_COMMAND_TOOL]) is True + args = json.loads(msg["tool_calls"][0]["function"]["arguments"]) + assert args == {"command": "echo hi"} + + def test_bare_string_arguments_decline_ambiguous_schema(self): + msg = { + "role": "assistant", + "content": '{"name":"Bash","arguments":"echo hi"}', + } + assert heal_openai_message(msg, {"Bash"}, TOOLS) is False + assert "tool_calls" not in msg + + def test_mixed_declared_and_undeclared_promotes_declared_keeps_undeclared_text(self): + # Span-exact removal: only the promoted Bash markup is dropped; the + # undeclared Nuke call's text stays in the content byte-intact. + content = f"pre {XML_BASH} mid {XML_UNDECLARED} post" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + (call,) = msg["tool_calls"] + assert call["function"]["name"] == "Bash" + assert XML_UNDECLARED in msg["content"] + assert "pre" in msg["content"] and "post" in msg["content"] + assert XML_BASH not in msg["content"] + + def test_multiple_declared_calls_all_promoted(self): + content = f"{XML_BASH} and {XML_BASH}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + assert len(msg["tool_calls"]) == 2 + + def test_mixed_formats_promote_in_document_order(self): + func_read = "a.txt" + content = f"{func_read} then {XML_BASH}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash", "Read"}) is True + assert [call["function"]["name"] for call in msg["tool_calls"]] == ["Read", "Bash"] + assert msg["content"] == "then" + + def test_unparseable_closed_block_not_deleted(self): + # A closed block whose body never parses is model output, + # not a promotable call; it must survive promotion of its neighbor. + garbage = "not json at all" + content = f"{XML_BASH} {garbage}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + assert garbage in msg["content"] + + +class TestStreamHealer: + def test_plain_text_passes_through(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("hello ") + healer.feed("world") + healer.finalize() + assert _events_text(events) == "hello world" + assert not _events_calls(events) + + def test_complete_call_in_one_chunk(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"On it. {XML_BASH}") + healer.finalize() + assert _events_text(events) == "On it. " + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + assert healer.healed + + def test_signal_split_across_chunks(self): + healer = StreamToolCallHealer({"Bash"}) + events = [] + for piece in ["{"name":"Bash",', '"arguments":{}}']: + events += healer.feed(piece) + events += healer.finalize() + assert _events_text(events) == "" + assert len(_events_calls(events)) == 1 + + def test_closed_malformed_tool_block_flushes_immediately(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("not json after") + assert _events_text(events) == "not json after" + assert not _events_calls(events) + + def test_mixed_formats_stream_in_document_order(self): + healer = StreamToolCallHealer({"Bash", "Read"}) + func_read = "a.txt" + events = healer.feed(f"{func_read} then {XML_BASH}") + healer.finalize() + assert [call["function"]["name"] for call in _events_calls(events)] == ["Read", "Bash"] + assert _events_text(events).strip() == "then" + + def test_false_alarm_html_flushes(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("use the
tag") + healer.finalize() + assert _events_text(events) == "use the
tag" + assert not _events_calls(events) + + def test_partial_signal_tail_held_then_flushed_at_end(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("trailing text -> call B, never both calls then the text. + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_BASH} middle {XML_BASH}") + healer.finalize() + kinds = [k for k, _ in events] + assert kinds == ["tool_call", "text", "tool_call"] + assert events[1][1] == " middle " + + def test_undeclared_then_declared_keeps_document_order(self): + # The undeclared block precedes the declared call; its raw text must + # be emitted BEFORE the promoted call event, never after. + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_UNDECLARED} then {XML_BASH}") + healer.finalize() + kinds = [k for k, _ in events] + assert kinds.index("tool_call") == len(kinds) - 1 + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + assert XML_UNDECLARED in _events_text(events) + + def test_declared_promoted_then_late_undeclared_flushes_raw(self): + # Streaming causality: the declared call completed and was already + # emitted before the undeclared one arrived. The undeclared markup + # must still reach the client as raw text (no data loss). + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_BASH} then ") + assert len(_events_calls(events)) == 1 + events += healer.feed(XML_UNDECLARED) + healer.finalize() + assert XML_UNDECLARED in _events_text(events) + assert len(_events_calls(events)) == 1 + + def test_undeclared_tool_flushes_raw(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(XML_UNDECLARED) + healer.finalize() + assert _events_text(events) == XML_UNDECLARED + assert not _events_calls(events) + + def test_two_calls_and_text_between(self): + healer = StreamToolCallHealer({"Bash", "Read"}) + xml_read = '{"name":"Read","arguments":{"path":"f"}}' + events = healer.feed(f"{XML_BASH} then {xml_read}") + healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["Bash", "Read"] + assert [c["id"] for c in calls] == ["call_0", "call_1"] + assert _events_text(events).strip() == "then" + + def test_incomplete_call_healed_at_finalize(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed('{"name":"Bash","arguments":{"cmd":"ls"}}') + assert events == [] # held + events = healer.finalize() + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + + def test_teaching_text_flushes_at_finalize(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(" is the marker syntax") + healer.finalize() + assert _events_text(events) == " is the marker syntax" + assert not _events_calls(events) + + def test_hold_bound_flushes(self): + healer = StreamToolCallHealer({"Bash"}) + blob = "" + "x" * (64 * 1024 + 10) + events = healer.feed(blob) + healer.finalize() + assert _events_text(events) == blob + assert not _events_calls(events) + + def test_dormant_after_structured_delta(self): + healer = StreamToolCallHealer({"Bash"}) + held = healer.feed("prefix call Bash somehow???") + assert nudge_should_retry(data, {"Read"}) is True + + def test_no_retry_on_clean_prose(self): + assert nudge_should_retry(self._resp("all done"), {"Bash"}) is False + + def test_no_retry_when_heal_would_succeed(self): + assert nudge_should_retry(self._resp(XML_BASH), {"Bash"}) is False + + def test_no_retry_with_structured_calls(self): + data = self._resp("", tool_calls = [{"id": "x"}]) + assert nudge_should_retry(data, {"Bash"}) is False + + def test_no_retry_when_healing_disabled(self): + assert nudge_should_retry(self._resp("???"), None) is False + + def test_nudge_messages_shape(self): + data = self._resp("garbage") + suffix = nudge_messages(data, {"Bash", "Read"}) + assert [m["role"] for m in suffix] == ["assistant", "user"] + assert suffix[0]["content"] == "garbage" + assert "`Bash` or `Read`" in suffix[1]["content"] + + def test_retry_with_undeclared_structured_call_is_not_an_improvement(self): + # The retry replaces the original only when it carries a USABLE call: + # a structured call naming an undeclared tool must not count. + undeclared = [ + {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}} + ] + declared = [ + {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}} + ] + assert response_has_promotable_calls(self._resp("", undeclared), {"Bash"}) is False + assert response_has_promotable_calls(self._resp("", declared), {"Bash"}) is True + + def test_retry_with_mixed_structured_calls_is_not_an_improvement(self): + # ALL structured calls must be declared: the caller forwards the whole + # list (and a parallel cap could keep only the FIRST), so a mixed retry + # could still hand the client an undeclared tool. + mixed = [ + {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}}, + {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}, + ] + assert response_has_promotable_calls(self._resp("", mixed), {"Bash"}) is False + assert ( + response_has_promotable_calls(self._resp("", list(reversed(mixed))), {"Bash"}) is False + ) + + @pytest.mark.parametrize( + "data", + [ + None, + "not a dict", + {}, + {"choices": []}, + {"choices": [{}]}, + {"choices": [{"message": None}]}, # llama-server error bodies do this + {"choices": [{"message": "not a dict"}]}, + {"choices": [{"message": {"content": None}}]}, + {"error": {"message": "boom"}}, + ], + ) + def test_malformed_response_shapes_never_raise(self, data): + # A malformed upstream body must degrade to "nothing to heal/nudge", + # never crash the request with an AttributeError. + assert nudge_should_retry(data, {"Bash"}) is False + assert response_has_promotable_calls(data, {"Bash"}) is False + suffix = nudge_messages(data, {"Bash"}) + assert suffix[0] == {"role": "assistant", "content": ""} + + +# ── Route-level wiring (OpenAI passthrough) ───────────────────────────── +# Mirrors the fake-llama-server patterns in test_openai_tool_passthrough.py. + +import asyncio # noqa: E402 +import threading # noqa: E402 +from types import SimpleNamespace # noqa: E402 + +import httpx # noqa: E402 + +from core.inference.api_monitor import ApiMonitor # noqa: E402 +from models.inference import ChatCompletionRequest, ChatMessage # noqa: E402 +from routes.inference import ( # noqa: E402 + _openai_passthrough_non_streaming, + _openai_passthrough_stream, +) + +LOOKUP_TOOL = { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object", "properties": {}}}, +} +LOOKUP_XML = '{"name":"lookup","arguments":{"q":"x"}}' + + +def _payload(**kwargs): + defaults = dict( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [LOOKUP_TOOL], + ) + defaults.update(kwargs) + return ChatCompletionRequest(**defaults) + + +def _llama_backend(): + return SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + + +def _upstream_message( + content, + tool_calls = None, + finish_reason = "stop", +): + message = {"role": "assistant", "content": content} + if tool_calls is not None: + message["tool_calls"] = tool_calls + return { + "id": "chatcmpl-up", + "object": "chat.completion", + "created": 1, + "model": "gguf", + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + +class ScriptedClient: + """Fake nonstreaming_client() returning scripted JSON bodies, counting POSTs.""" + + def __init__(self, bodies): + self.bodies = list(bodies) + self.posts = [] + + async def post( + self, + _url, + json = None, + timeout = None, + ): + self.posts.append(json) + return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)]) + + +async def _drive_non_streaming(monkeypatch, payload, bodies): + import routes.inference as inf_mod + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _openai_passthrough_non_streaming( + _llama_backend(), payload, "gguf", monitor_id = None + ) + return client, json.loads(response.body) + + +async def _drive_stream(monkeypatch, payload, lines): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in lines: + yield line + + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monkeypatch.setattr(inf_mod, "api_monitor", ApiMonitor(max_entries = 3)) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + _llama_backend(), + payload, + "gguf", + "chatcmpl-test", + monitor_id = None, + ) + return [chunk async for chunk in response.body_iterator] + + +def _stream_payloads(chunks): + out = [] + for chunk in chunks: + for line in chunk.splitlines(): + if line.startswith("data: ") and line[6:] != "[DONE]": + out.append(json.loads(line[6:])) + return out + + +class TestOpenaiNonStreamingRoute: + def test_heals_xml_to_tool_calls(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(LOOKUP_XML)] + ) + choice = data["choices"][0] + assert choice["finish_reason"] == "tool_calls" + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert json.loads(call["function"]["arguments"]) == {"q": "x"} + assert choice["message"]["content"] is None + assert data["usage"]["total_tokens"] == 3 # usage preserved + assert len(client.posts) == 1 # healing never re-requests + + asyncio.run(_run()) + + def test_bare_string_uses_client_schema_key(self, monkeypatch): + async def _run(): + content = '{"name":"Bash","arguments":"echo hi"}' + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tools = [BASH_COMMAND_TOOL]), + [_upstream_message(content)], + ) + (call,) = data["choices"][0]["message"]["tool_calls"] + assert json.loads(call["function"]["arguments"]) == {"command": "echo hi"} + + asyncio.run(_run()) + + def test_opt_out_relays_verbatim(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(auto_heal_tool_calls = False), + [_upstream_message(LOOKUP_XML)], + ) + choice = data["choices"][0] + assert choice["message"]["content"] == LOOKUP_XML + assert "tool_calls" not in choice["message"] + assert choice["finish_reason"] == "stop" + + asyncio.run(_run()) + + def test_no_tools_untouched(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, _payload(tools = None), [_upstream_message(LOOKUP_XML)] + ) + assert data["choices"][0]["message"]["content"] == LOOKUP_XML + + asyncio.run(_run()) + + def test_undeclared_tool_not_promoted(self, monkeypatch): + async def _run(): + xml = '{"name":"rogue","arguments":{}}' + _, data = await _drive_non_streaming(monkeypatch, _payload(), [_upstream_message(xml)]) + assert data["choices"][0]["message"]["content"] == xml + assert "tool_calls" not in data["choices"][0]["message"] + + asyncio.run(_run()) + + def test_structured_calls_untouched(self, monkeypatch): + async def _run(): + native = [ + { + "id": "call_up", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + _, data = await _drive_non_streaming( + monkeypatch, + _payload(), + [_upstream_message("", tool_calls = native, finish_reason = "tool_calls")], + ) + assert data["choices"][0]["message"]["tool_calls"] == native + + asyncio.run(_run()) + + def test_length_finish_reason_preserved(self, monkeypatch): + async def _run(): + # Truncated generation: the healed call stays attached but the + # client must still see the truncation, so length is never + # upgraded to tool_calls. + _, data = await _drive_non_streaming( + monkeypatch, + _payload(), + [_upstream_message(LOOKUP_XML, finish_reason = "length")], + ) + choice = data["choices"][0] + assert choice["finish_reason"] == "length" + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + + asyncio.run(_run()) + + def test_tool_choice_none_relays_verbatim(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tool_choice = "none"), + [_upstream_message(LOOKUP_XML)], + ) + message = data["choices"][0]["message"] + assert message["content"] == LOOKUP_XML + assert "tool_calls" not in message + + asyncio.run(_run()) + + def test_tool_choice_forcing_other_function_not_promoted(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tool_choice = {"type": "function", "function": {"name": "other"}}), + [_upstream_message(LOOKUP_XML)], + ) + message = data["choices"][0]["message"] + assert message["content"] == LOOKUP_XML + assert "tool_calls" not in message + + asyncio.run(_run()) + + def test_mixed_declared_and_undeclared_promotes_and_keeps_text(self, monkeypatch): + async def _run(): + rogue = '{"name":"rogue","arguments":{}}' + mixed = f"{LOOKUP_XML} also {rogue}" + _, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(mixed)] + ) + choice = data["choices"][0] + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert rogue in choice["message"]["content"] + assert choice["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + def test_healed_then_native_stream_indexes_disjoint(self, monkeypatch): + async def _run(): + # A healed text-form call goes out first (index 0); a native + # structured delta follows. Clients merge deltas by index, so the + # native call must be shifted off index 0 or the two would merge. + native_line = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_native","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":' + + json.dumps(LOOKUP_XML) + + "}}]}", + native_line, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + indexes = {} + for payload_data in _stream_payloads(chunks): + for ch in payload_data.get("choices", []): + for tc in (ch.get("delta") or {}).get("tool_calls") or []: + indexes.setdefault(tc["index"], tc.get("id")) + assert indexes.get(0, "").startswith("call_") and indexes[0] != "call_native" + assert indexes.get(1) == "call_native" + + asyncio.run(_run()) + + def test_role_delta_precedes_healed_stream_content(self, monkeypatch): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":' + + json.dumps(LOOKUP_XML) + + "}}]}", + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + payloads = _stream_payloads(chunks) + first_delta = payloads[0]["choices"][0]["delta"] + assert first_delta == {"role": "assistant"} + assert "tool_calls" in payloads[1]["choices"][0]["delta"] + + asyncio.run(_run()) + + def test_same_chunk_role_content_finish_delays_finish_until_after_healed_tool( + self, monkeypatch + ): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":' + + json.dumps(LOOKUP_XML) + + '},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + payloads = _stream_payloads(chunks) + assert payloads[0]["choices"][0]["finish_reason"] is None + assert payloads[0]["choices"][0]["delta"] == {"role": "assistant"} + assert "tool_calls" in payloads[1]["choices"][0]["delta"] + assert payloads[-1]["choices"][0]["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + +GARBAGE_SIGNAL = "call lookup somehow???" + + +class TestNudgeRetryOpenai: + def test_retry_recovers_call(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(LOOKUP_XML)], + ) + assert len(client.posts) == 2 # exactly one retry + # Prefix byte-identical, nudge suffix appended (KV-cache reuse guard). + original, retry = client.posts + assert retry["messages"][: len(original["messages"])] == original["messages"] + suffix = retry["messages"][len(original["messages"]) :] + assert [m["role"] for m in suffix] == ["assistant", "user"] + assert suffix[0]["content"] == GARBAGE_SIGNAL + # The healed retry response is returned. + (call,) = data["choices"][0]["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert data["choices"][0]["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + def test_retry_still_garbage_returns_original(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(GARBAGE_SIGNAL + "2")], + ) + assert len(client.posts) == 2 + assert data["choices"][0]["message"]["content"] == GARBAGE_SIGNAL + assert "tool_calls" not in data["choices"][0]["message"] + + asyncio.run(_run()) + + def test_default_off_single_post(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(GARBAGE_SIGNAL)] + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + def test_no_retry_on_clean_prose(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message("all done")], + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + def test_no_retry_when_heal_succeeds(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(LOOKUP_XML)], + ) + assert len(client.posts) == 1 + assert data["choices"][0]["message"]["tool_calls"] + + asyncio.run(_run()) + + def test_heal_opt_out_disables_nudge_too(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, + _payload(auto_heal_tool_calls = False, nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL)], + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + +class TestNudgeRetryAnthropic: + async def _drive( + self, + monkeypatch, + bodies, + nudge = None, + ): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + nudge_tool_calls = nudge, + ) + return client, json.loads(response.body) + + def test_retry_recovers_tool_use(self, monkeypatch): + async def _run(): + client, data = await self._drive( + monkeypatch, + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(LOOKUP_XML)], + nudge = True, + ) + assert len(client.posts) == 2 + (block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert block["name"] == "lookup" + assert data["stop_reason"] == "tool_use" + + asyncio.run(_run()) + + def test_healed_tool_use_precedes_trailing_text(self, monkeypatch): + async def _run(): + _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} done")]) + assert [block["type"] for block in data["content"]] == ["tool_use", "text"] + assert data["content"][1]["text"] == "done" + + asyncio.run(_run()) + + def test_default_off(self, monkeypatch): + async def _run(): + client, _ = await self._drive(monkeypatch, [_upstream_message(GARBAGE_SIGNAL)]) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + +class TestAnthropicPassthroughHealingText: + """Non-streaming Anthropic passthrough must relay unpromoted (undeclared) + text-form calls as text, matching the OpenAI passthrough contract. Once + heal_openai_message promotes the declared call it span-trims only that + markup and deliberately leaves the undeclared bytes in the content; the + legacy blanket _TOOL_XML_RE strip must not delete them. + """ + + async def _drive(self, monkeypatch, upstream): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient([upstream]) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + ) + return json.loads(response.body) + + def test_mixed_declared_and_undeclared_relays_undeclared_as_text(self, monkeypatch): + async def _run(): + content = f"Running now. {LOOKUP_XML} then {XML_UNDECLARED} done." + data = await self._drive(monkeypatch, _upstream_message(content)) + # Declared lookup call is promoted into a structured tool_use block. + (tool_use,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert tool_use["name"] == "lookup" + text = " ".join(b["text"] for b in data["content"] if b["type"] == "text") + assert XML_UNDECLARED in text + assert "Running now." in text and "done." in text + assert LOOKUP_XML not in text + + asyncio.run(_run()) + + +class TestAnthropicEmitterHealing: + def _events( + self, + emitter, + chunks, + finish = True, + ): + lines = [] + for chunk in chunks: + lines += emitter.feed_chunk(chunk) + if finish: + lines += emitter.finish() + return [json.loads(ln.split("data: ", 1)[1]) for ln in lines if "data: " in ln] + + def _emitter( + self, + allowed = ("lookup",), + **kwargs, + ): + from core.inference.anthropic_compat import AnthropicPassthroughEmitter + + emitter = AnthropicPassthroughEmitter() + emitter.enable_healing(set(allowed), **kwargs) + return emitter + + def _chunk( + self, + content = None, + tool_calls = None, + finish_reason = None, + ): + delta = {} + if content is not None: + delta["content"] = content + if tool_calls is not None: + delta["tool_calls"] = tool_calls + return {"choices": [{"delta": delta, "finish_reason": finish_reason}]} + + def test_xml_becomes_tool_use_block_and_stop_reason(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = LOOKUP_XML), + self._chunk(finish_reason = "stop"), + ], + ) + starts = [e for e in events if e.get("type") == "content_block_start"] + (tool_start,) = [e for e in starts if e["content_block"]["type"] == "tool_use"] + assert tool_start["content_block"]["name"] == "lookup" + assert tool_start["content_block"]["id"].startswith("toolu_") + (args,) = [ + e["delta"]["partial_json"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "input_json_delta" + ] + assert json.loads(args) == {"q": "x"} + (message_delta,) = [e for e in events if e.get("type") == "message_delta"] + assert message_delta["delta"]["stop_reason"] == "tool_use" + + def test_mid_block_signal_closes_text_block_first(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = f"Let me check {LOOKUP_XML}"), + self._chunk(finish_reason = "stop"), + ], + ) + kinds = [ + (e["type"], (e.get("content_block") or e.get("delta") or {}).get("type")) + for e in events + if e["type"].startswith("content_block") + ] + # text opens, streams the safe prefix, closes; then the tool_use block. + assert kinds[0] == ("content_block_start", "text") + assert kinds[1] == ("content_block_delta", "text_delta") + assert kinds[2] == ("content_block_stop", None) + assert kinds[3] == ("content_block_start", "tool_use") + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == "Let me check " + + def test_false_alarm_streams_as_text(self): + events = self._events( + self._emitter(), + [self._chunk(content = "use the
tag"), self._chunk(finish_reason = "stop")], + ) + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == "use the
tag" + (message_delta,) = [e for e in events if e.get("type") == "message_delta"] + assert message_delta["delta"]["stop_reason"] == "end_turn" + + def test_signal_split_across_chunks(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = "{"name":"lookup","arguments":{"q":"y"}}' + events = self._events( + self._emitter(disable_parallel_tool_use = True), + [self._chunk(content = two), self._chunk(finish_reason = "stop")], + ) + starts = [ + e + for e in events + if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(starts) == 1 + + def test_disable_parallel_drops_native_after_healed(self): + # A healed call consumed the single allowed slot; a later native + # structured call (index 0, so it survives the caller's chunk-level + # cap) must not open a second tool_use block. + structured = [ + { + "index": 0, + "id": "call_up", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + events = self._events( + self._emitter(disable_parallel_tool_use = True), + [ + self._chunk(content = LOOKUP_XML), + self._chunk(tool_calls = structured), + self._chunk(finish_reason = "tool_calls"), + ], + ) + starts = [ + e + for e in events + if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(starts) == 1 + + def test_no_healing_means_verbatim_text(self): + from core.inference.anthropic_compat import AnthropicPassthroughEmitter + + emitter = AnthropicPassthroughEmitter() # enable_healing never called + events = self._events( + emitter, + [self._chunk(content = LOOKUP_XML), self._chunk(finish_reason = "stop")], + ) + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == LOOKUP_XML + + +class TestAnthropicNonStreamingRoute: + async def _drive( + self, + monkeypatch, + bodies, + auto_heal = None, + tools = None, + tool_choice = "auto", + ): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + tools if tools is not None else [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + tool_choice = tool_choice, + auto_heal_tool_calls = auto_heal, + ) + return client, json.loads(response.body) + + def test_promotes_xml_to_tool_use(self, monkeypatch): + async def _run(): + _, data = await self._drive(monkeypatch, [_upstream_message(LOOKUP_XML)]) + (block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert block["name"] == "lookup" + assert block["input"] == {"q": "x"} + assert data["stop_reason"] == "tool_use" + assert not any(b["type"] == "text" for b in data["content"]) + + asyncio.run(_run()) + + def test_opt_out_keeps_legacy_strip(self, monkeypatch): + async def _run(): + _, data = await self._drive( + monkeypatch, [_upstream_message(f"plan {LOOKUP_XML}")], auto_heal = False + ) + assert data["stop_reason"] == "end_turn" + (block,) = data["content"] + assert block["type"] == "text" + assert block["text"] == "plan" # XML stripped, nothing promoted + + asyncio.run(_run()) + + def test_undeclared_tool_not_promoted(self, monkeypatch): + async def _run(): + xml = '{"name":"rogue","arguments":{}}' + _, data = await self._drive(monkeypatch, [_upstream_message(xml)]) + assert data["stop_reason"] == "end_turn" + assert not any(b["type"] == "tool_use" for b in data["content"]) + # Healing preserves what it does not promote: the undeclared call + # reaches the client as text instead of being silently stripped. + (text_block,) = [b for b in data["content"] if b["type"] == "text"] + assert text_block["text"] == xml + + asyncio.run(_run()) + + def test_mixed_undeclared_text_preserved_after_heal(self, monkeypatch): + async def _run(): + # Declared call promoted to tool_use; the undeclared call's markup + # stays in the text block (the legacy strip must not run after a + # span-exact heal), matching the OpenAI passthrough. + rogue = '{"name":"rogue","arguments":{}}' + _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} {rogue}")]) + (tool_block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert tool_block["name"] == "lookup" + (text_block,) = [b for b in data["content"] if b["type"] == "text"] + assert rogue in text_block["text"] + assert data["stop_reason"] == "tool_use" + + asyncio.run(_run()) + + def test_length_beats_tool_use(self, monkeypatch): + async def _run(): + _, data = await self._drive( + monkeypatch, [_upstream_message(LOOKUP_XML, finish_reason = "length")] + ) + assert data["stop_reason"] == "max_tokens" + assert any(b["type"] == "tool_use" for b in data["content"]) + + asyncio.run(_run()) + + def test_tool_choice_none_keeps_legacy_strip(self, monkeypatch): + async def _run(): + # Anthropic {"type": "none"} arrives here converted to "none": + # the request forbade tool calls, so nothing is promoted and the + # legacy XML strip applies as before healing existed. + _, data = await self._drive( + monkeypatch, + [_upstream_message(f"plan {LOOKUP_XML}")], + tool_choice = "none", + ) + assert data["stop_reason"] == "end_turn" + (block,) = data["content"] + assert block["type"] == "text" + assert block["text"] == "plan" + + asyncio.run(_run()) + + +class TestOpenaiStreamingRoute: + def test_heals_streamed_xml(self, monkeypatch): + async def _run(): + pieces = ["", '{"name":"lookup",', '"arguments":{"q":"x"}}', ""] + lines = [ + 'data: {"id":"c1","model":"gguf","created":1,"choices":[{"index":0,"delta":{"content":%s}}]}' + % json.dumps(p) + for p in pieces + ] + lines += [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + (call,) = tool_deltas + assert call["function"]["name"] == "lookup" + assert json.loads(call["function"]["arguments"]) == {"q": "x"} + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["tool_calls"] + # None of the XML leaked as visible content. + text = "".join( + (c.get("delta") or {}).get("content") or "" + for p in payloads + for c in p.get("choices", []) + ) + assert "" not in text + assert chunks[-1] == "data: [DONE]\n\n" + + asyncio.run(_run()) + + def test_parallel_cap_drops_native_after_healed(self, monkeypatch): + async def _run(): + # parallel_tool_calls=false: a healed call consumed the single + # allowed slot, and the upstream SSE cap keeps native index 0, so + # the route must drop the later native call itself. + xml = '{"name":"lookup","arguments":{"q":"x"}}' + native = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_up","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + 'data: {"id":"c1","model":"gguf","created":1,"choices":' + '[{"index":0,"delta":{"content":%s}}]}' % json.dumps(xml), + native, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(parallel_tool_calls = False), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + (call,) = tool_deltas + assert call["id"] == "call_0" # the healed call; native was dropped + + asyncio.run(_run()) + + def test_false_alarm_text_flushes(self, monkeypatch): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":"use the
tag"}}]}', + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + text = "".join( + (c.get("delta") or {}).get("content") or "" + for p in payloads + for c in p.get("choices", []) + ) + assert text == "use the
tag" + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["stop"] + + asyncio.run(_run()) + + def test_incomplete_xml_healed_at_done(self, monkeypatch): + async def _run(): + # No close tag and no finish chunk: healed at the [DONE] boundary, + # synthetic finish must say tool_calls. + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":"{\\"name\\":\\"lookup\\",\\"arguments\\":{}}"}}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + assert len(tool_deltas) == 1 + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["tool_calls"] + + asyncio.run(_run()) + + def test_structured_upstream_calls_relay_verbatim(self, monkeypatch): + async def _run(): + line = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_up","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + line, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + assert chunks[0] == line + "\n\n" # byte-for-byte relay + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 4147746b54..a7ceb49ed9 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -1986,3 +1986,177 @@ class TestTranslatedMessagesValidate: msgs = _normalise_responses_input(payload) for m in msgs: ChatMessage(**m.model_dump(exclude_none = True)) + + +# ===================================================================== +# Streaming passthrough healing — text-form calls promoted in order +# ===================================================================== + + +class TestResponsesStreamHealing: + """Route-level healing on the /v1/responses stream: text-form tool calls + are promoted through the same per-call item state machinery as structured + deltas, and healer events keep their order (text around a healed call must + not move relative to the function_call item).""" + + _XML = '{"name":"lookup","arguments":{"q":"x"}}' + _TOOL = {"type": "function", "name": "lookup", "parameters": {"type": "object"}} + + @staticmethod + def _ordered_events(lines): + events = [] + for line in lines: + if not line.startswith("event: "): + continue + name, _, rest = line.partition("\n") + payload = json.loads(rest.split("data: ", 1)[1].strip()) + events.append((name[len("event: ") :], payload)) + return events + + def _run_stream(self, monkeypatch, content, **payload_kwargs): + TestResponsesStreamAdapter._install_stream_mock( + monkeypatch, [{"choices": [{"delta": {"content": content}}]}] + ) + payload = ResponsesRequest(input = "hi", stream = True, tools = [self._TOOL], **payload_kwargs) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, messages, TestResponsesStreamAdapter._Request() + ) + return await TestResponsesStreamAdapter._collect(response) + + return self._ordered_events(asyncio.run(run())) + + def test_text_around_healed_call_keeps_order(self, monkeypatch): + events = self._run_stream(monkeypatch, f"before {self._XML} after.") + pos_before = pos_item = pos_after = None + for i, (name, payload) in enumerate(events): + if name == "response.output_text.delta": + if "before" in payload["delta"] and pos_before is None: + pos_before = i + if "after" in payload["delta"]: + pos_after = i + if ( + name == "response.output_item.added" + and payload["item"]["type"] == "function_call" + and pos_item is None + ): + pos_item = i + assert payload["item"]["name"] == "lookup" + assert pos_before is not None and pos_item is not None and pos_after is not None + assert pos_before < pos_item < pos_after + + def test_call_before_trailing_text_claims_lower_output_index(self, monkeypatch): + events = self._run_stream(monkeypatch, f"{self._XML} done.") + item_added = [ + (name, payload) for name, payload in events if name == "response.output_item.added" + ] + # The call came first in the model output, so its item is added first + # and claims the lower output_index; the trailing text's message item + # follows. + assert [payload["item"]["type"] for _, payload in item_added] == [ + "function_call", + "message", + ] + call_idx = item_added[0][1]["output_index"] + msg_idx = item_added[1][1]["output_index"] + assert call_idx < msg_idx + text = "".join( + payload["delta"] for name, payload in events if name == "response.output_text.delta" + ) + assert "done." in text + assert "" not in text + + def test_tool_choice_none_streams_raw_text(self, monkeypatch): + events = self._run_stream(monkeypatch, self._XML, tool_choice = "none") + assert not any( + payload["item"]["type"] == "function_call" + for name, payload in events + if name == "response.output_item.added" + ) + text = "".join( + payload["delta"] for name, payload in events if name == "response.output_text.delta" + ) + assert text == self._XML + + def test_healed_call_splits_message_items(self, monkeypatch): + # Text on both sides of a healed call becomes TWO message items: the + # healed function_call closes the first, trailing text opens a fresh + # one with a later output index (native Responses stream shape). + events = self._run_stream(monkeypatch, f"before {self._XML} after.") + added = [ + (payload["output_index"], payload["item"]["type"], payload["item"].get("id")) + for name, payload in events + if name == "response.output_item.added" + ] + assert [item_type for _, item_type, _ in added] == [ + "message", + "function_call", + "message", + ] + assert [idx for idx, _, _ in added] == sorted(idx for idx, _, _ in added) + assert added[0][2] != added[2][2] # distinct message item ids + # Text deltas attribute to their OWN message item. + deltas = [ + (payload["item_id"], payload["delta"]) + for name, payload in events + if name == "response.output_text.delta" + ] + assert [d for i, d in deltas if i == added[0][2]] == ["before "] + assert [d for i, d in deltas if i == added[2][2]] == [" after."] + # The completed snapshot lists all three items with per-item text. + completed = [payload for name, payload in events if name == "response.completed"] + output = completed[0]["response"]["output"] + assert [item["type"] for item in output] == ["message", "function_call", "message"] + assert output[0]["content"][0]["text"] == "before " + assert output[2]["content"][0]["text"] == " after." + + def test_parallel_cap_drops_native_after_healed(self, monkeypatch): + # parallel_tool_calls=false: a healed call consumed the single allowed + # slot; a later native structured call (index 0, so it survives + # _drop_parallel_tool_call_deltas) must not open a second + # function_call item. + TestResponsesStreamAdapter._install_stream_mock( + monkeypatch, + [ + {"choices": [{"delta": {"content": self._XML}}]}, + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_up", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + } + } + ] + }, + ], + ) + payload = ResponsesRequest( + input = "hi", + stream = True, + tools = [self._TOOL], + parallel_tool_calls = False, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, messages, TestResponsesStreamAdapter._Request() + ) + return await TestResponsesStreamAdapter._collect(response) + + events = self._ordered_events(asyncio.run(run())) + calls = [ + payload + for name, payload in events + if name == "response.output_item.added" and payload["item"]["type"] == "function_call" + ] + assert len(calls) == 1 + assert calls[0]["item"]["name"] == "lookup" diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 931d8a705d..39fdd151be 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -71,6 +71,28 @@ class TestFunctionStyleTrailingText: call = _only(text) assert call == {"name": "python", "arguments": {"code": 'print("")'}} + def test_closed_function_with_trailing_prose_heal_path(self): + # Regression: the heal / finalize path (allow_incomplete=True) used to fold + # and the trailing prose into the argument and drop + # the prose from visible content. It must now match the strict path -- keep a + # clean argument and leave the trailing prose outside the call span. + text = "cats trailing words" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + fn = calls[0]["function"] + assert fn["name"] == "web_search" + assert json.loads(fn["arguments"]) == {"query": "cats"} + # The trailing prose sits outside the removed span, so it stays visible. + from core.tool_healing import ( + parse_tool_calls_from_text as _parse_with_spans, + ) + + _calls, spans = _parse_with_spans(text, allow_incomplete = True, with_spans = True) + out = text + for s, e in sorted(spans, reverse = True): + out = out[:s] + out[e:] + assert out == " trailing words" + def test_incomplete_function_without_close_is_still_rejected(self): text = "weather london" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] @@ -160,3 +182,18 @@ class TestHealingPathUnaffected: calls = parse_tool_calls_from_text(text, allow_incomplete = True) assert len(calls) == 1 assert calls[0]["function"]["name"] == "web_search" + + def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self): + # allow_incomplete exists for truncated output; a call that DID close + # must parse identically to strict mode, leaving prose after + # out of the last parameter and out of the removal span. + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + text = "cats trailing" + calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "cats"} + (span,) = spans + assert text[span[0] : span[1]] == ( + "cats" + ) From 026141a4a4026aa4a9bdf6d8752d8d3515f388b9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 08:25:10 -0700 Subject: [PATCH 16/23] Studio: multi-select export formats, portable FP8/INT8, GGUF LoRA, and source parity (#6767) * Studio: expose full compressed-tensors scheme set in an export formats dropdown * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: multi-select export formats, portable torchao FP8/INT8, GGUF LoRA, source parity Export page overhaul on top of the formats dropdown: - Unify merged precision into one sorted multi-select list (16-bit first, then 8-bit, then 4-bit). Drop "vLLM" from labels, add INT8 (W8A8), INT8 (W8A16), INT4 (W4A16), MXFP4, MXFP8. Quick formats render as toggle pills; the rest live in a multi-select "More formats" dropdown, so several formats export in one run. - Add a portable torchao FP8/INT8 save path (Float8WeightOnlyConfig / Int8WeightOnlyConfig) that needs no NVIDIA GPU to produce and loads in vLLM. FP8 serializes to safetensors, INT8 to .bin. Wired into save_pretrained_merged and push_to_hub_merged via a TORCHAO_EXPORT_SCHEMES registry and _unsloth_save_torchao, parallel to the compressed-tensors path. - Hide NVIDIA-only compressed-tensors formats when no NVIDIA GPU is present; keep 16-bit and portable FP8/INT8. The backend also rejects a compressed request on non-NVIDIA hardware so it stays authoritative. - Relax merged export to non-PEFT models so Local Model and Hugging Face sources get the same 16-bit / compressed / portable options. - GGUF: send the whole quant list in one call (merge once, quantize many). - LoRA: add a GGUF adapter option (convert_lora_to_gguf.py) with an outtype select (f16/bf16/f32/q8_0/auto), alongside the safetensors adapter. - Thread the new fields through models, routes, orchestrator, and worker; extend the export tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: gate export by accelerator with a torch-aware reason; fix export save dir naming Export runs through Unsloth, which requires a compute accelerator (NVIDIA/AMD/Intel GPU or Apple MLX) and has no CPU code path, so a bare-CPU host cannot export even with PyTorch installed. Add export_capability() in utils/hardware that reports export_supported plus a precise reason so the UI stops showing a generic "no GPU": - pytorch_not_installed: a --no-torch install (even a physical GPU is unusable) - no_accelerator: PyTorch present but no supported accelerator (bare CPU) - mlx_unavailable: Apple Silicon where the MLX stack is missing or too old Expose the fields on /api/system/hardware and /api/system, and guard the mutating export routes (load-checkpoint, export/merged|base|gguf|lora) with HTTP 400 and the reason, leaving read-only endpoints usable so the Export page still renders. Make core/export/export.py import without PyTorch and without a usable accelerator (the Unsloth import is caught) so the export worker degrades to a clear message instead of crashing at import. Frontend: keep /export reachable on chat-only hosts and gray out the method and format options with the backend reason (Alert plus disabled MethodPicker) instead of silently redirecting to /chat, so users see why export is unavailable. Also fix the export save directory producing "model/null" for Local Model and Hugging Face sources that have no run/checkpoint, naming the folder from the model id. * CI: validate Studio export capability gating on Linux, Windows and macOS Add a small pytest matrix that runs studio/backend/tests/test_export_capability.py on ubuntu-latest, windows-latest and macos-latest. It confirms, on each real OS, that hardware.export_capability() reports the right decision and reason (pytorch_not_installed, no_accelerator, or mlx_unavailable) and that the export backend imports without PyTorch and degrades to a clear message instead of crashing. Hosted runners have no GPU/MLX, so this covers the "export unavailable, here is why" path a Mac/Windows user without an accelerator sees; a real accelerator export is validated separately. The job installs only a CPU PyTorch plus the backend import deps (no unsloth, triton, or llama.cpp), so it runs in seconds with no GPU. * Studio export: address Codex review (source-aware gating, GGUF LoRA token/MLX/guard) Frontend (export-page): - Gate LoRA and quantized-model restrictions on the active source. isAdapter / isQuantized come from the selected checkpoint; in Local Model / Hugging Face ("model") source mode they were stale, so LoRA stayed wrongly enabled for a direct base model (backend then rejects "No adapter to export") and a stale "quantized" flag disabled every method for an unrelated, exportable model. Add effectiveIsAdapter / effectiveIsQuantized (false outside checkpoint mode) and use them in the method-reset effect and the MethodPicker disabled state. - Hide the GGUF LoRA option on a macOS/MLX host (the backend rejects GGUF LoRA on MLX), so users no longer pick it, wait through the load, and always fail. Disable the "GGUF adapter" button on a Mac host and never send loraGguf there. Backend (core/export/export.py): - Pass the HF token into the GGUF LoRA conversion (save_pretrained_gguf), so a gated/private base model's config fetch in convert_lora_to_gguf.py is authenticated; without it the load can succeed but the conversion fails. - Guard the save_pretrained_gguf capability check with getattr so an older Unsloth model that lacks the method returns the clean "not supported" message instead of an AttributeError that surfaces as a generic 500. * Studio export: address 2nd Codex review (CI index, empty merged, test import) - studio-export-capability-ci.yml: add --extra-index-url https://pypi.org/simple to the torch install so torch's transitive deps still resolve; --index-url alone replaces PyPI with only the CPU wheel index, which does not serve all of them. - export-page handleStart: reject an empty merged selection (mirrors canExport), so clicking the panel's Start button with every precision pill deselected no longer submits mergedSelections: [] and launches an unintended default 16-bit export. - test_export_imatrix_compressed: the torchao-registry test now reads unsloth/save.py as text (like the other ast/string checks) instead of `import unsloth.save`, which raised ModuleNotFoundError in the CPU studio-backend suite that has no unsloth installed. * Studio export: make comments succinct across the export changes * Studio export: use load token for local GGUF LoRA export of gated bases * Studio export: harden portable torchao path and gate multi-format Hub push torchao (_unsloth_save_torchao): - merge to an isolated temp staging dir so a co-selected 16-bit output at save_directory is not deleted - narrow VLM detection to vision_config / ForVisionText2Text so T5/BART/Whisper are not misrouted - forward trust_remote_code (from auto_map) to the reload so custom-code models export Export UI: - hide portable torchao formats on macOS/MLX (backend rejects quantized export there) - restrict a Hub merged export to a single format (each writes to the repo root) * Studio export: torchao tokenizer remote-code + XPU offload, scale GGUF timeout torchao (_unsloth_save_torchao): - honor auto_map in the staged tokenizer/processor configs (not just model.config) when deriving trust_remote_code, so custom-code tokenizers reload after the merge - offload single-device XPU models to CPU (and empty the XPU cache) before the reload, matching the CUDA path, so an Intel GPU that fits the model once does not OOM on the second copy Export orchestrator: - scale the GGUF wait timeout by the number of requested quants so a multi-quant list export of a large model does not time out at a flat 3600s * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio export: show portable torchao formats only on non-NVIDIA (CPU) hosts Portable torchao FP8/INT8 is the fallback for hosts without the NVIDIA compressed-tensors path. On an NVIDIA GPU the compressed-tensors FP8/FP4/INT formats are the intended path (llm-compressor auto-installs), so hide the portable duplicates there; keep them on CPU / non-NVIDIA hosts and continue hiding them on macOS/MLX. * Studio export: report all output folders and the exported formats - Multi-format merged export now collects every sibling output directory (one per selected precision) instead of only the last; the success banner lists them all. - Show the selected precision formats in the run summary (a Formats row, like GGUF Quantizations), so the panel says what is being exported rather than just 'Merged Model'. - Persist the selected formats in the run summary and seed them on mount, so navigating away and back (or toggling the export method) restores the selection instead of resetting to 16-bit. * Studio export: list all output formats, add GGUF LoRA target, default Q8_0, auto-select newest checkpoint - Progress/summary panel now shows a Formats row with the selected merged formats, and the success banner lists every output folder a multi-format merged run creates (one line per format) instead of only the last one. - Merged format selection is seeded from the active run, so navigating away and back (or switching method cards) no longer resets it to 16-bit. - GGUF / Llama.cpp now offers an Export target toggle (Full model or LoRA adapter) for adapter checkpoints, reusing the LoRA GGUF export path. - Removed the Auto GGUF LoRA output type and defaulted to Q8_0 in the UI, the request model, and the backend defaults; the outtype list is now Q8_0/F16/BF16/F32. Core save.py still accepts auto for external callers. - When a finetune has no checkpoint selected, auto-select the newest one. * Studio torchao export: robust reload class + optional VLM import Two fixes to the portable torchao FP8/INT8 export reload, from review of the narrowed VLM detection: - Encoder-decoder seq2seq checkpoints (T5/BART/Whisper) are not causal LMs. With the narrowed is_vlm test they now correctly skip the image-text class, but fell through to AutoModelForCausalLM and failed to reload after the merge. Reload them with their own architecture class from the config instead. - AutoModelForImageTextToText was imported unconditionally at the top of the torchao path, so on Transformers builds without that class the import aborted every torchao export (even text-only). Import it lazily only for a VLM, with the AutoModelForVision2Seq fallback used elsewhere in Unsloth. * Studio: enable FP8/FP4 compressed export for newer-transformers models The shipped llm-compressor 0.10.x pins transformers<=4.57.6, so FP8/FP4 export failed for models needing a transformers 5.x sidecar (Qwen3.5, Gemma-4, Qwen3-Next): the quantization subprocess crashed importing the removed TORCH_INIT_FUNCTIONS. Run the quantization against a dedicated llm-compressor-main "shadow": a --target package dir (transformers 5.10.2 + llm-compressor main + compressed-tensors) layered over the existing torch. It installs --no-deps so torch is never touched (works on any Studio torch build), is provisioned lazily and fingerprint-cached, and can be turned off with UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN. - transformers_version.py: provision + validate .venv_llmcompressor. - export.py: route all compressed exports through the shadow when available; else keep the workspace 0.10.x path and fail fast past its transformers ceiling. - save.py: launch _compressed_quantize.py with a clean PYTHONPATH = shadow. - _compressed_quantize.py: skip linear_attn / vision tower / MTP modules (matches the RedHatAI and NVIDIA reference quants, and is required by the grouped schemes). Verified all four schemes (fp8, w8a8, w4a16, mxfp4) on Qwen3.5-9B and Llama-3.2-1B, and fp8 on Gemma-4, end to end through Studio. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix GGUF LoRA export tests * Fix export CI expectations * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- .../workflows/studio-export-capability-ci.yml | 76 +++ studio/backend/core/export/export.py | 305 +++++++++-- studio/backend/core/export/orchestrator.py | 18 +- studio/backend/core/export/worker.py | 3 + studio/backend/main.py | 8 +- studio/backend/models/export.py | 27 +- studio/backend/routes/export.py | 25 + .../backend/tests/test_export_capability.py | 156 ++++++ .../tests/test_export_imatrix_compressed.py | 140 ++++- studio/backend/utils/hardware/__init__.py | 7 + studio/backend/utils/hardware/hardware.py | 43 ++ studio/backend/utils/transformers_version.py | 151 ++++++ studio/frontend/src/app/routes/__root.tsx | 3 + .../frontend/src/components/app-sidebar.tsx | 16 +- .../src/features/export/api/export-api.ts | 9 +- .../export/components/export-run-panel.tsx | 51 +- .../frontend/src/features/export/constants.ts | 168 +++++- .../src/features/export/export-page.tsx | 501 +++++++++++++++--- .../export/stores/export-runtime-store.ts | 116 ++-- .../frontend/src/hooks/use-hardware-info.ts | 15 + tests/studio/playwright_extra_ui.py | 14 +- unsloth/_compressed_quantize.py | 4 + unsloth/save.py | 494 ++++++++++++++++- 23 files changed, 2120 insertions(+), 230 deletions(-) create mode 100644 .github/workflows/studio-export-capability-ci.yml create mode 100644 studio/backend/tests/test_export_capability.py diff --git a/.github/workflows/studio-export-capability-ci.yml b/.github/workflows/studio-export-capability-ci.yml new file mode 100644 index 0000000000..1ee6489209 --- /dev/null +++ b/.github/workflows/studio-export-capability-ci.yml @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs studio/backend/tests/test_export_capability.py on Linux, Windows and macOS. +# +# export_capability() is per-OS (is_apple_silicon() and the PyTorch-import probe differ per +# platform) and the export backend must import without PyTorch, so this confirms the gating and +# import-safety on hosted Windows/macOS. Hosted runners have no GPU/MLX, so a real accelerator +# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block +# torch/unsloth, so the job installs only a CPU PyTorch plus import deps. + +name: Studio export capability + +on: + pull_request: + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + push: + branches: [main] + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + capability: + name: capability (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + env: + # No accelerator on hosted runners; keep detection on the CPU path. + CUDA_VISIBLE_DEVICES: "" + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Upgrade pip + run: python -m pip install --upgrade pip + - name: Install CPU PyTorch + # CPU wheel index so every OS gets a CPU build; keep PyPI as an extra index so torch's + # transitive deps still resolve (matching the other workflows in this repo). + run: python -m pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple "torch>=2.4,<2.13" + - name: Install backend import deps + # Enough to import utils.hardware and core.export.export; NOT unsloth (needs a GPU, and + # the import-safety test blocks it) or triton/llama.cpp (Linux-only / native builds). + run: python -m pip install + transformers peft accelerate safetensors huggingface_hub datasets + sentencepiece protobuf fastapi starlette structlog psutil + python-multipart pydantic httpx "numpy<3" pytest + - name: Export capability + import-safety tests + working-directory: studio/backend + run: python -m pytest tests/test_export_capability.py -q diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index d0461dae95..c8be50b08b 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -13,7 +13,18 @@ import shutil import contextlib from pathlib import Path from typing import Optional, Tuple, List -from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + +# unsloth imports torch on non-MLX hosts, so a --no-torch install raises here. Stay importable +# (null the classes) so exports return a clean "PyTorch is not installed" error, not an import crash. +try: + from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX + _UNSLOTH_IMPORT_ERROR = None +except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load + FastLanguageModel = None + FastVisionModel = None + _IS_MLX = False + _UNSLOTH_IMPORT_ERROR = _unsloth_exc + from huggingface_hub import HfApi, ModelCard from utils.hardware import clear_gpu_cache @@ -27,14 +38,46 @@ from utils.paths import ( ) from core.inference import get_inference_backend -# GPU-only imports — guarded for Apple Silicon where these aren't needed +# GPU/PyTorch-only imports, skipped on MLX and on a --no-torch install so the module stays +# importable; export then degrades to a clear "PyTorch is not installed" error. +torch = None +_TORCH_IMPORT_ERROR: Optional[BaseException] = None if not _IS_MLX: - from peft import PeftModel, PeftModelForCausalLM - from transformers.modeling_utils import PushToHubMixin - import torch + try: + from peft import PeftModel, PeftModelForCausalLM + from transformers.modeling_utils import PushToHubMixin + import torch + except Exception as _torch_exc: # ImportError, or a broken native torch load + _TORCH_IMPORT_ERROR = _torch_exc logger = get_logger(__name__) + +def _export_runtime_available() -> bool: + """True if export can run: MLX active, or Unsloth imported (only succeeds on a GPU host).""" + return bool(_IS_MLX) or (FastLanguageModel is not None) + + +def _export_runtime_message() -> str: + """Precise reason the export runtime is unavailable, mirroring hardware.export_capability().""" + if torch is None: + return ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + return ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export on " + "CPU only.)" + ) + + +# Kept for call sites / tests referencing the PyTorch-missing text. +_PYTORCH_MISSING_MESSAGE = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." +) + _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False @@ -58,6 +101,28 @@ def _compressed_export_supported(): return False +def _torchao_export_supported(): + """True if the installed unsloth build has the portable torchao FP8/INT8 export path.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_torchao_method") + except Exception: + return False + + +def _has_nvidia_gpu(): + """True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it.""" + try: + from utils.hardware import hardware as _hw + return _hw.DEVICE == _hw.DeviceType.CUDA and not _hw.IS_ROCM + except Exception: + try: + import torch + return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None + except Exception: + return False + + def _hf_offline(timeout = 3): """True if export should avoid the Hub: honors the HF offline env vars, else does one cheap TCP reachability probe so a network-down load uses local files / the HF cache @@ -394,13 +459,17 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """ Export merged model (for PEFT models). Args: save_directory: Local directory to save model - format_type: "16-bit (FP16)" or "4-bit (FP4)" + format_type: "16-bit (FP16)", "4-bit (FP4)", or a compressed-tensors label + compressed_method: Optional compressed-tensors scheme alias (e.g. "fp8", + "fp8_static", "w8a8", "w4a16", "mxfp4", "mxfp8", "nvfp4"). Overrides + format_type and is resolved against unsloth.save COMPRESSED_EXPORT_SCHEMES. push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID (username/model-name) hf_token: Hugging Face token @@ -409,38 +478,108 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - if not self.is_peft: - return ( - False, - "This is not a PEFT model. Use 'Export Base Model' instead.", - None, - ) + # Merged export works for PEFT adapters and non-PEFT Local/HF base models alike + # (save_pretrained_merged is a no-op merge that just saves the base). output_path: Optional[str] = None - # compressed-tensors formats run save_pretrained_merged with an FP8/FP4 save_method and - # write to a sibling "-" directory (for vLLM). - _COMPRESSED = { - "FP8 (compressed-tensors)": ("fp8", "fp8"), - "NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"), + # Quantized formats save to a sibling "-". Two backends: compressed-tensors + # (llm-compressor, NVIDIA-only) and portable torchao FP8/INT8 (device-agnostic). The alias + # comes from `compressed_method` (the "all formats" dropdown) or the `format_type` label. + _LABEL_TO_ALIAS = { + "FP8 (compressed-tensors)": "fp8", + "NVFP4 (compressed-tensors)": "nvfp4", } - is_compressed = format_type in _COMPRESSED + compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type) + compressed_suffix: Optional[str] = None + # Classify the alias: torchao-portable vs compressed-tensors. + torchao_info = None + if compressed_alias and _torchao_export_supported(): + try: + import unsloth.save as _us_t + torchao_info = _us_t._normalize_torchao_method(compressed_alias) + except Exception: + torchao_info = None + is_torchao = torchao_info is not None + is_compressed = compressed_alias is not None and not is_torchao try: - if _IS_MLX: - if is_compressed: - return False, "Compressed-tensors export is not supported on macOS/MLX.", None - mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" - elif is_compressed: + if _IS_MLX and (is_compressed or is_torchao): + return ( + False, + "Quantized (FP8/FP4/INT) export is not supported on macOS/MLX. " + "Use 16-bit or GGUF.", + None, + ) + + if is_torchao: + # Portable torchao: no NVIDIA GPU, no calibration. + compressed_suffix = torchao_info[1] + + if is_compressed: + # compressed-tensors needs CUDA; enforce in the backend even if the UI gate is bypassed. + if not _has_nvidia_gpu(): + return ( + False, + "Compressed-tensors (FP8/FP4) export requires an NVIDIA GPU. On other " + "hardware use the portable FP8/INT8 (torchao) formats or 16-bit.", + None, + ) if not _compressed_export_supported(): return ( False, - "Compressed-tensors (FP8/NVFP4) export requires an Unsloth build with " + "Compressed-tensors (FP8/FP4) export requires an Unsloth build with " "compressed-tensors support. Upgrade unsloth, or choose 16-bit.", None, ) - save_method = _COMPRESSED[format_type][0] + import unsloth.save as _us + + # Prefer the llm-compressor-main shadow (transformers 5.x): it quantizes newer models + # (Qwen3.5, Gemma-4, ...) the shipped 0.10.x cannot. Route all compressed exports + # through it when available; else fall back to the workspace 0.10.x path below. + _shadow_pp = None + try: + from utils.transformers_version import llmcompressor_shadow_pythonpath + _shadow_pp = llmcompressor_shadow_pythonpath() + except Exception as e: + logger.warning(f"llm-compressor-main shadow unavailable: {e}") + if _shadow_pp: + os.environ[_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV] = _shadow_pp + else: + # No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its + # transformers ceiling, so fail fast for sidecar models; default-tier still works. + os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None) + _exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + return ( + False, + "FP8/FP4 compressed-tensors export is not available for this model: it " + f"runs under transformers {_tf_ver}, but the installed llm-compressor " + f"supports transformers <= {_us._LLM_COMPRESSOR_MAX_TRANSFORMERS} and the " + "llm-compressor-main runtime could not be provisioned (offline or " + "UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN). Export to GGUF or 16-bit instead.", + None, + ) + + try: + info = _us._normalize_compressed_method(compressed_alias) + except Exception as e: + return False, f"Unsupported compressed export '{compressed_alias}': {e}", None + if info is None: + return ( + False, + f"'{compressed_alias}' is not a recognized compressed-tensors export.", + None, + ) + compressed_suffix = info[2] + + if _IS_MLX: + mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + elif is_compressed or is_torchao: + save_method = compressed_alias elif format_type == "4-bit (FP4)": save_method = "merged_4bit_forced" elif self._audio_type == "whisper": @@ -464,10 +603,10 @@ class ExportBackend: save_directory, self.current_tokenizer, save_method = save_method ) - # Compressed export writes to the "-" sibling; report that as output. + # Compressed / torchao writes to the "-" sibling; report that as output. final_dir = ( - f"{save_directory}-{_COMPRESSED[format_type][1]}" - if is_compressed + f"{save_directory}-{compressed_suffix}" + if (is_compressed or is_torchao) else save_directory ) self._write_export_metadata(final_dir) @@ -507,10 +646,9 @@ class ExportBackend: token = hf_token, private = private, ) - elif is_compressed and output_path and Path(output_path).is_dir(): - # The compressed model was already built locally in output_path; upload it - # directly so we do not re-run the (expensive, OOM-prone) compression that - # push_to_hub_merged(save_method=fp8/nvfp4) would otherwise do a second time. + elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir(): + # Already built in output_path; upload it directly instead of re-running the + # expensive quantization that push_to_hub_merged(save_method=...) would redo. hf_api = HfApi(token = hf_token) repo_id = PushToHubMixin._create_repo( PushToHubMixin, @@ -522,7 +660,7 @@ class ExportBackend: username = repo_id.split("/")[0], base_model = getattr(self.current_model.config, "_name_or_path", "unknown"), model_type = getattr(self.current_model.config, "model_type", "llm"), - method = format_type, + method = compressed_alias or format_type, extra = "unsloth", ) ModelCard(content).push_to_hub( @@ -568,6 +706,8 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -686,7 +826,7 @@ class ExportBackend: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, @@ -697,7 +837,9 @@ class ExportBackend: Args: save_directory: Local directory to save model - quantization_method: GGUF quantization method (e.g., "Q4_K_M") + quantization_method: A single GGUF quant method (e.g., "Q4_K_M") or a list of them + (e.g., ["Q4_K_M", "Q8_0"]). A list produces one GGUF per quant from a single + model load (unsloth save_to_gguf loops internally). push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID hf_token: Hugging Face token @@ -705,11 +847,13 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - # Only forward imatrix_file to an unsloth build that accepts it; otherwise even a plain - # no-imatrix export would fail with an unexpected-keyword error against an older unsloth. + # Only forward imatrix_file to an unsloth build that accepts it, else older builds raise + # an unexpected-keyword error even for a plain no-imatrix export. if imatrix_file is not None and not _supports_kwarg( self.current_model.save_pretrained_gguf, "imatrix_file" ): @@ -724,8 +868,14 @@ class ExportBackend: output_path: Optional[str] = None model_tmp_to_cleanup: Optional[str] = None try: - # unsloth expects lowercase quant method - quant_method = quantization_method.lower() + # Normalize to a lowercased list so multiple quants come from one model load. + if isinstance(quantization_method, (list, tuple)): + quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()] + else: + quant_methods = [str(quantization_method).lower()] + if not quant_methods: + quant_methods = ["q4_k_m"] + quant_method = quant_methods if len(quant_methods) > 1 else quant_methods[0] # Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it # can't drift past the pinned llama-quantize binary's gguf API. @@ -847,7 +997,7 @@ class ExportBackend: return ( True, - f"GGUF model exported successfully ({quantization_method})", + f"GGUF model exported successfully ({', '.join(quant_methods)})", output_path, ) @@ -867,19 +1017,56 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: """ Export LoRA adapter only (not merged). + Args: + gguf: If True, also convert the adapter to a GGUF LoRA file (llama.cpp + convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`. + gguf_outtype: GGUF LoRA output float type; one of q8_0/f16/bf16/f32. + Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: return False, "This is not a PEFT model. No adapter to export.", None + _GGUF_LORA_OUTTYPES = ("q8_0", "f16", "bf16", "f32") + if gguf: + if _IS_MLX: + return ( + False, + "GGUF LoRA adapter export is not supported on macOS/MLX. " + "Use the safetensors adapter instead.", + None, + ) + outtype = str(gguf_outtype).lower() + if outtype not in _GGUF_LORA_OUTTYPES: + return ( + False, + f"Invalid GGUF LoRA outtype '{gguf_outtype}'. " + f"Choose one of {', '.join(_GGUF_LORA_OUTTYPES)}.", + None, + ) + # getattr so an older build without save_pretrained_gguf returns a clean message + # instead of an AttributeError (a generic 500). + _save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None) + if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"): + return ( + False, + "This Unsloth build does not support GGUF LoRA adapter export. " + "Upgrade unsloth and unsloth_zoo, or export the safetensors adapter.", + None, + ) + output_path: Optional[str] = None try: if save_directory: @@ -887,7 +1074,24 @@ class ExportBackend: logger.info(f"Saving LoRA adapter locally to: {save_directory}") ensure_dir(Path(save_directory)) - if _IS_MLX: + if gguf: + # Writes the adapter files plus "-lora-.gguf". + _apply_wsl_sudo_patch() + self.current_model.save_pretrained_gguf( + save_directory, + self.current_tokenizer, + save_method = "lora", + quantization_method = outtype, + # Forward the token so convert_lora_to_gguf.py can fetch a gated base's config. + token = hf_token or None, + ) + final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf"))) + logger.info( + "LoRA GGUF export complete. Files in %s:\n %s", + save_directory, + "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", + ) + elif _IS_MLX: # MLX: save adapters.safetensors + tokenizer files self.current_model.save_lora_adapters(save_directory) self.current_tokenizer.save_pretrained(save_directory) @@ -907,7 +1111,24 @@ class ExportBackend: logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") - if _IS_MLX: + if gguf: + # Upload the locally-built GGUF folder; needs a local save_directory so the + # conversion is not re-run. + if not (output_path and Path(output_path).is_dir()): + return ( + False, + "GGUF LoRA Hub upload requires a local save directory; set one and " + "retry.", + None, + ) + hf_api = HfApi(token = hf_token) + hf_api.create_repo(repo_id, private = private, exist_ok = True) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) + elif _IS_MLX: with tempfile.TemporaryDirectory() as tmp_dir: self.current_model.save_lora_adapters(tmp_dir) self.current_tokenizer.save_pretrained(tmp_dir) diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 052a47dd80..671ef363f5 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -456,6 +456,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """Export merged PEFT model.""" return self._run_export( @@ -467,6 +468,7 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "compressed_method": compressed_method, }, ) @@ -495,13 +497,13 @@ class ExportOrchestrator: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, imatrix_file = None, ) -> Tuple[bool, str, Optional[str]]: - """Export model in GGUF format.""" + """Export model in GGUF format. `quantization_method` may be a single method or a list.""" return self._run_export( "gguf", { @@ -521,8 +523,10 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: - """Export LoRA adapter only.""" + """Export LoRA adapter only (optionally also as a GGUF LoRA file).""" return self._run_export( "lora", { @@ -531,6 +535,8 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "gguf": gguf, + "gguf_outtype": gguf_outtype, }, ) @@ -557,9 +563,13 @@ class ExportOrchestrator: cmd = {"type": "export", "export_type": export_type, **params} try: self._send_cmd(cmd) + # GGUF for 30B+ models can take 30+ min per quant; a multi-quant list runs them + # all in one op off a single merge, so scale the timeout by the quant count. + _qm = params.get("quantization_method") + _n = len(_qm) if isinstance(_qm, (list, tuple)) and _qm else 1 resp = self._wait_response( f"export_{export_type}_done", - timeout = 3600, # GGUF for 30B+ models can take 30+ min + timeout = 3600 * max(1, _n), ) op_success = resp.get("success", False) op_message = resp.get("message", "") diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index d473dcb54f..7828116236 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -397,6 +397,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + compressed_method = cmd.get("compressed_method"), ) elif export_type == "base": success, message, output_path = backend.export_base_model( @@ -423,6 +424,8 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + gguf = cmd.get("gguf", False), + gguf_outtype = cmd.get("gguf_outtype", "q8_0"), ) else: success, message = False, f"Unknown export type: {export_type}" diff --git a/studio/backend/main.py b/studio/backend/main.py index 0613a5ae53..8762c43195 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1145,7 +1145,7 @@ def get_system_info(current_subject: str = Depends(get_current_subject)): import os import time import logging - from utils.hardware import get_device + from utils.hardware import get_device, export_capability from utils.hardware.hardware import _backend_label logger = logging.getLogger(__name__) @@ -1218,6 +1218,8 @@ def get_system_info(current_subject: str = Depends(get_current_subject)): }, "gpu": gpu_info, "ml_packages": ml_packages, + # Export capability + torch-aware reason. See /api/system/hardware. + **export_capability(), } @@ -1240,11 +1242,13 @@ def get_hardware_info( method auto-selection. Sync def (not async): hardware/detail probes can shell out, and FastAPI runs sync endpoints in a threadpool. """ - from utils.hardware import get_gpu_summary, get_package_versions + from utils.hardware import get_gpu_summary, get_package_versions, export_capability body = { "gpu": get_gpu_summary(), "versions": get_package_versions(), + # Export capability + torch-aware reason; the Export UI grays out with the message. + **export_capability(), } if include_details: from utils.llama_cpp_update import get_installed_llama_version diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 7e05373f11..9dc4d9451a 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -6,7 +6,7 @@ from pathlib import Path, PureWindowsPath from pydantic import BaseModel, Field, field_validator -from typing import List, Optional, Literal, Dict, Any +from typing import List, Optional, Literal, Dict, Any, Union def _validate_save_directory(value: str) -> str: @@ -168,6 +168,15 @@ class ExportMergedModelRequest(ExportCommonOptions): description = "Export precision / format for the merged model. The compressed-tensors " "options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).", ) + compressed_method: Optional[str] = Field( + None, + description = "Optional quantized-export alias. Either a compressed-tensors scheme " + "(e.g. 'fp8', 'fp8_static', 'w8a8', 'w4a16', 'mxfp4', 'mxfp8', 'nvfp4' - NVIDIA only) " + "from unsloth.save COMPRESSED_EXPORT_SCHEMES, or a portable torchao alias " + "('torchao_fp8', 'torchao_int8') from TORCHAO_EXPORT_SCHEMES that needs no NVIDIA GPU. " + "When set, it overrides format_type. Lets the export UI expose the full set of formats " + "beyond the quick buttons.", + ) class ExportBaseModelRequest(ExportCommonOptions): @@ -189,9 +198,10 @@ class ExportGGUFRequest(BaseModel): def _check_save_directory(cls, v): return _validate_save_directory(v) - quantization_method: str = Field( + quantization_method: Union[str, List[str]] = Field( "Q4_K_M", - description = 'GGUF quantization method (e.g. "Q4_K_M")', + description = 'GGUF quantization method(s). A single method (e.g. "Q4_K_M") or a list ' + '(e.g. ["Q4_K_M", "Q8_0"]) to produce multiple GGUFs from one model load.', ) push_to_hub: bool = Field( False, @@ -219,4 +229,13 @@ class ExportGGUFRequest(BaseModel): class ExportLoRAAdapterRequest(ExportCommonOptions): """Request for exporting only the LoRA adapter (not merged).""" - # Uses fields from ExportCommonOptions only + gguf: bool = Field( + False, + description = "If True, also convert the adapter to a GGUF LoRA file " + "(llama.cpp convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`.", + ) + gguf_outtype: Literal["q8_0", "f16", "bf16", "f32"] = Field( + "q8_0", + description = "GGUF LoRA output float type (only used when gguf=True). " + "Q8_0 falls back to F16 per tensor for dims not divisible by the block size (32).", + ) diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index cf2cb2fa70..a7fd7cbec7 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -46,6 +46,23 @@ router = APIRouter() logger = get_logger(__name__) +def _ensure_export_supported() -> None: + """Reject a mutating export request up front (HTTP 400) when the host can't export. + + Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints + (scan/status/logs) are intentionally NOT gated so the Export page can still render the reason. + """ + from utils.hardware import export_capability + + cap = export_capability() + if not cap.get("export_supported", True): + raise HTTPException( + status_code = 400, + detail = cap.get("export_unsupported_message") + or "Export is not supported on this platform.", + ) + + @router.post("/load-checkpoint", response_model = ExportOperationResponse) async def load_checkpoint( request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject) @@ -58,6 +75,7 @@ async def load_checkpoint( a clear error instead of tearing down the user's other running workloads. """ try: + _ensure_export_supported() backend = get_export_backend() # Run in a worker thread (spawns and waits on a subprocess, can take # minutes) so the event loop stays free to serve the live log SSE stream. @@ -266,6 +284,7 @@ async def export_merged_model( Wraps ExportBackend.export_merged_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_merged_model, @@ -275,6 +294,7 @@ async def export_merged_model( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + compressed_method = request.compressed_method, ) if not success: @@ -304,6 +324,7 @@ async def export_base_model( Wraps ExportBackend.export_base_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_base_model, @@ -342,6 +363,7 @@ async def export_gguf( Wraps ExportBackend.export_gguf. """ try: + _ensure_export_supported() backend = get_export_backend() # A custom path wins; otherwise the imatrix toggle requests the upstream auto-download. imatrix_file = request.imatrix_path or (True if request.imatrix else None) @@ -382,6 +404,7 @@ async def export_lora_adapter( Wraps ExportBackend.export_lora_adapter. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_lora_adapter, @@ -390,6 +413,8 @@ async def export_lora_adapter( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + gguf = request.gguf, + gguf_outtype = request.gguf_outtype, ) if not success: diff --git a/studio/backend/tests/test_export_capability.py b/studio/backend/tests/test_export_capability.py new file mode 100644 index 0000000000..e04417f933 --- /dev/null +++ b/studio/backend/tests/test_export_capability.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for export capability gating. + +Export is supported iff ``get_device() in {CUDA, XPU, MLX}``, with a torch-aware reason otherwise +(pytorch_not_installed / no_accelerator / mlx_unavailable), and the backend must import without +PyTorch. The matrix mocks the hardware probes; wiring is checked with ast so it runs on CPU. +""" + +import ast +import builtins +from pathlib import Path + +import pytest + +import utils.hardware.hardware as hw + +_BACKEND = Path(__file__).resolve().parent.parent + + +def _src(rel): + return (_BACKEND / rel).read_text(encoding = "utf-8") + + +def _func_src(rel, name): + src = _src(rel) + node = next( + n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(src, node) + + +# -- capability matrix -------------------------------------------------------------------------- + + +def _patch(monkeypatch, *, torch: bool, device, apple: bool): + monkeypatch.setattr(hw, "_has_torch", lambda: torch) + monkeypatch.setattr(hw, "get_device", lambda: device) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: apple) + + +def test_cpu_with_torch_unsupported_no_accelerator(monkeypatch): + # PyTorch present but no accelerator: unsupported with no_accelerator, not "PyTorch missing". + _patch(monkeypatch, torch = True, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "no_accelerator" + assert "accelerator" in cap["export_unsupported_message"].lower() + # Must NOT tell a user with PyTorch installed to install PyTorch. + assert "PyTorch is not installed" not in cap["export_unsupported_message"] + + +def test_cuda_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.CUDA, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is True + assert cap["export_unsupported_reason"] is None + assert cap["export_unsupported_message"] is None + + +def test_xpu_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.XPU, apple = False) + assert hw.export_capability()["export_supported"] is True + + +def test_mlx_without_torch_supports_export(monkeypatch): + # Apple Silicon MLX exports without PyTorch. + _patch(monkeypatch, torch = False, device = hw.DeviceType.MLX, apple = True) + assert hw.export_capability()["export_supported"] is True + + +def test_no_torch_non_apple_reports_pytorch_missing(monkeypatch): + _patch(monkeypatch, torch = False, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "pytorch_not_installed" + assert "PyTorch is not installed" in cap["export_unsupported_message"] + + +def test_apple_without_mlx_reports_mlx_unavailable(monkeypatch): + # Apple + CPU means the MLX stack is missing; reason is mlx_unavailable regardless of torch. + for has_torch in (False, True): + _patch(monkeypatch, torch = has_torch, device = hw.DeviceType.CPU, apple = True) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "mlx_unavailable" + assert "MLX" in cap["export_unsupported_message"] + + +# -- import safety without PyTorch -------------------------------------------------------------- + + +def test_export_backend_imports_without_torch(monkeypatch): + """core/export/export.py must import on a --no-torch host (unsloth/torch blocked) and return a + clean 'PyTorch is not installed' message from an export attempt, not crash at import.""" + import importlib + import sys + + real_import = builtins.__import__ + + def blocking_import(name, *args, **kwargs): + top = name.split(".")[0] + if top in {"torch", "unsloth"}: + raise ImportError(f"simulated: {top} not installed") + return real_import(name, *args, **kwargs) + + # Drop any preloaded copies so the guarded import paths re-run under the block. + for m in [k for k in sys.modules if k.split(".")[0] in {"torch", "unsloth"}]: + monkeypatch.delitem(sys.modules, m, raising = False) + monkeypatch.delitem(sys.modules, "core.export.export", raising = False) + monkeypatch.setattr(builtins, "__import__", blocking_import) + + mod = importlib.import_module("core.export.export") + assert mod._IS_MLX is False + assert mod.torch is None + assert mod._export_runtime_available() is False + + be = mod.ExportBackend.__new__(mod.ExportBackend) + be.current_model = None + be.current_tokenizer = None + be.is_peft = False + be._audio_type = None + ok, message, out = be.export_merged_model("/tmp/does-not-matter") + assert ok is False + assert "PyTorch is not installed" in message + + +# -- endpoint / backend wiring (ast) ------------------------------------------------------------ + + +def test_main_endpoints_expose_export_capability(): + m = _src("main.py") + # Both system endpoints spread export_capability() into their response. + assert m.count("**export_capability()") >= 2 + assert '"/api/system/hardware"' in m and '"/api/system"' in m + + +def test_routes_guard_mutating_endpoints(): + r = _src("routes/export.py") + assert "def _ensure_export_supported()" in r + # load + all four export endpoints call the guard. + assert r.count("_ensure_export_supported()") >= 6 + + +def test_export_methods_check_runtime(): + e = _src("core/export/export.py") + assert "def _export_runtime_available()" in e + # Each export method returns the clear message when the runtime is missing. + assert e.count("_export_runtime_available()") >= 5 + assert "_PYTORCH_MISSING_MESSAGE" in e + + +def test_export_capability_reads_no_torch_helper(): + cap = _func_src("utils/hardware/hardware.py", "export_capability") + assert "_has_torch()" in cap and "DeviceType.MLX" in cap and "is_apple_silicon()" in cap diff --git a/studio/backend/tests/test_export_imatrix_compressed.py b/studio/backend/tests/test_export_imatrix_compressed.py index d914ff8651..f499390add 100644 --- a/studio/backend/tests/test_export_imatrix_compressed.py +++ b/studio/backend/tests/test_export_imatrix_compressed.py @@ -54,8 +54,7 @@ def test_merged_request_rejects_unknown_format(): def test_export_gguf_threads_imatrix_to_save_and_push(): - # imatrix_file must reach both save_pretrained_gguf and push_to_hub_gguf, but only via the - # conditional **imatrix_kw so a no-imatrix export never sends an unsupported keyword. + # imatrix_file must reach both save paths, but only via the conditional **imatrix_kw. g = _func_src("core/export/export.py", "export_gguf") assert g.count("**imatrix_kw") >= 2 assert 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' in g @@ -109,8 +108,139 @@ def test_export_merged_maps_compressed_to_save_method(): def test_compressed_hub_push_uploads_local_dir_without_recompressing(): - # A compressed Hub push must upload the already-built output_path, not re-run compression - # via push_to_hub_merged (which would compress a second time). + # A compressed / torchao Hub push must upload the built output_path, not re-quantize. m = _func_src("core/export/export.py", "export_merged_model") - assert "elif is_compressed and output_path and Path(output_path).is_dir():" in m + assert "elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():" in m assert "hf_api.upload_folder(" in m and "folder_path = output_path" in m + + +# -- torchao portable FP8/INT8 (device-agnostic, no NVIDIA GPU) --------------------------------- + + +def test_merged_request_accepts_torchao_aliases(): + # Portable torchao aliases pass through compressed_method (validated in the backend registry). + for alias in ("torchao_fp8", "torchao_int8"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_routes_torchao_and_skips_nvidia_guard(): + m = _func_src("core/export/export.py", "export_merged_model") + # torchao is classified separately and its suffix comes from the torchao normalizer. + assert "_normalize_torchao_method(compressed_alias)" in m + assert "is_torchao = torchao_info is not None" in m + assert "is_compressed = compressed_alias is not None and not is_torchao" in m + # The NVIDIA guard applies to compressed-tensors only, not torchao. + assert "_has_nvidia_gpu()" in m + # torchao routes through save_method just like compressed. + assert "elif is_compressed or is_torchao:" in m + + +def test_export_merged_nvidia_guard_present(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "requires an NVIDIA GPU" in m + + +def test_has_nvidia_gpu_helper_reads_hardware_module(): + h = _func_src("core/export/export.py", "_has_nvidia_gpu") + assert "DeviceType.CUDA" in h and "IS_ROCM" in h + + +def test_export_merged_relaxes_is_peft_guard(): + # Non-PEFT (Local/HF base) models can now export merged; the old hard block must be gone. + m = _func_src("core/export/export.py", "export_merged_model") + assert "Use 'Export Base Model' instead." not in m + + +def test_unsloth_save_has_torchao_registry_and_path(): + # Read unsloth/save.py as text (not import) so this runs in the CPU suite without unsloth. + save_py = (_BACKEND.parent.parent / "unsloth" / "save.py").read_text(encoding = "utf-8") + assert "def _normalize_torchao_method" in save_py + assert "def _unsloth_save_torchao" in save_py + assert "TORCHAO_EXPORT_SCHEMES = {" in save_py + # torchao aliases must map to (scheme, suffix) so the backend routes to the torchao path. + assert '"torchao_fp8": ("fp8", "torchao-fp8")' in save_py + assert '"torchao_int8": ("int8", "torchao-int8")' in save_py + + +# -- GGUF multi-quant list ---------------------------------------------------------------------- + + +def test_gguf_request_accepts_list_of_quants(): + r = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = ["Q4_K_M", "Q8_0"]) + assert r.quantization_method == ["Q4_K_M", "Q8_0"] + r2 = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = "Q4_K_M") + assert r2.quantization_method == "Q4_K_M" + + +def test_export_gguf_normalizes_quant_list(): + g = _func_src("core/export/export.py", "export_gguf") + assert "isinstance(quantization_method, (list, tuple))" in g + assert "quant_methods" in g + + +# -- GGUF LoRA adapter export ------------------------------------------------------------------- + + +def test_lora_request_has_gguf_fields(): + from models.export import ExportLoRAAdapterRequest + + r = ExportLoRAAdapterRequest(save_directory = "/tmp/x") + assert r.gguf is False and r.gguf_outtype == "q8_0" + r2 = ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf = True, gguf_outtype = "q8_0") + assert r2.gguf is True and r2.gguf_outtype == "q8_0" + + +def test_lora_request_rejects_bad_outtype(): + from models.export import ExportLoRAAdapterRequest + with pytest.raises(ValidationError): + ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf_outtype = "q3_k") + + +def test_export_lora_wires_gguf_save_method(): + la = _func_src("core/export/export.py", "export_lora_adapter") + assert 'save_method = "lora"' in la + assert "quantization_method = outtype" in la + + +def test_orchestrator_and_worker_pass_lora_gguf(): + o = _func_src("core/export/orchestrator.py", "export_lora_adapter") + assert '"gguf": gguf' in o and '"gguf_outtype": gguf_outtype' in o + w = _src("core/export/worker.py") + assert 'gguf = cmd.get("gguf", False)' in w + assert 'gguf_outtype = cmd.get("gguf_outtype", "q8_0")' in w + + +def test_route_passes_lora_gguf(): + r = _src("routes/export.py") + assert "gguf = request.gguf" in r and "gguf_outtype = request.gguf_outtype" in r + + +# -- compressed_method ("all formats" dropdown) ------------------------------------------------- + + +def test_merged_request_accepts_compressed_method(): + # Defaults to None; any scheme alias is accepted (validation happens in the backend registry). + assert ExportMergedModelRequest(save_directory = "/tmp/x").compressed_method is None + for alias in ("fp8", "fp8_static", "w8a8", "w8a16", "w4a16", "mxfp4", "mxfp8", "nvfp4"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_resolves_alias_via_registry(): + # The scheme + suffix must come from unsloth.save's registry normalizer, not a hardcoded dict. + m = _func_src("core/export/export.py", "export_merged_model") + assert "compressed_method" in m + assert "_normalize_compressed_method(compressed_alias)" in m + assert "compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)" in m + assert "compressed_suffix" in m and 'f"{save_directory}-{compressed_suffix}"' in m + + +def test_orchestrator_and_worker_pass_compressed_method(): + o = _func_src("core/export/orchestrator.py", "export_merged_model") + assert "compressed_method" in o and '"compressed_method": compressed_method' in o + assert 'compressed_method = cmd.get("compressed_method")' in _src("core/export/worker.py") + + +def test_route_passes_compressed_method(): + assert "compressed_method = request.compressed_method" in _src("routes/export.py") diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 5f2b2abbcf..62b537fbac 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -44,6 +44,12 @@ from .vram_estimation import ( estimate_training_vram, ) + +def export_capability() -> dict: + """Return live export capability from the hardware module.""" + return _hardware.export_capability() + + __all__ = [ "DeviceType", "DEVICE", @@ -51,6 +57,7 @@ __all__ = [ "IS_ROCM", "detect_hardware", "get_device", + "export_capability", "is_apple_silicon", "clear_gpu_cache", "get_gpu_memory_info", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index cde7070075..8d6c919ebd 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -263,6 +263,49 @@ def get_device() -> DeviceType: return DEVICE +def export_capability() -> dict: + """Whether model export can run here, with a torch-aware reason when it cannot. + + Export runs through Unsloth, which hard-requires an accelerator (it calls ``torch.cuda`` at + import and has no CPU path), so it is supported iff ``get_device() in {CUDA, XPU, MLX}``. The + reason distinguishes a --no-torch install from a bare-CPU host. Safe to call without torch. + + Returns {export_supported, export_unsupported_reason, export_unsupported_message}. + """ + if get_device() in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX): + return { + "export_supported": True, + "export_unsupported_reason": None, + "export_unsupported_message": None, + } + # No accelerator: name the blocker. Apple Silicon first -- its path is MLX, so "install PyTorch" + # would be wrong advice on a Mac even when torch is also absent. + if is_apple_silicon(): + reason = "mlx_unavailable" + message = ( + "Export on Apple Silicon requires the MLX stack, which is unavailable or too old. Run " + "`unsloth studio update` to restore MLX and enable export." + ) + elif not _has_torch(): + reason = "pytorch_not_installed" + message = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + else: + reason = "no_accelerator" + message = ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export " + "on CPU only.)" + ) + return { + "export_supported": False, + "export_unsupported_reason": reason, + "export_unsupported_message": message, + } + + def clear_gpu_cache(): """ Clear GPU memory cache for the current device. diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index ac0ecfbfcd..2a63caa5b2 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -226,6 +226,11 @@ _VENV_T5_510_DIR = str(_studio_root() / ".venv_t5_510") # Backwards-compat alias _VENV_T5_DIR = _VENV_T5_550_DIR +# llm-compressor-main shadow for FP8/FP4 export of newer-transformers models. Like the .venv_t5_* +# sidecars but also shadows llm-compressor main + compressed-tensors; installed --no-deps so it +# reuses the workspace torch (torch-agnostic). +_VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor") + # Tier precedence: higher rank wins in _higher_tier. _TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3} @@ -1518,6 +1523,152 @@ def _ensure_venv_t5_exists() -> bool: return _ensure_venv_t5_550_exists() +# --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) --------------------- +# Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize +# Qwen3.5 / Gemma-4 / Llama. +_LLMC_MAIN_TRANSFORMERS = "5.10.2" +_LLMC_MAIN_SHA = "973c9c539a84dd9efaf74e115ede5ca419704c18" +_LLMC_MAIN_COMPRESSED_TENSORS = "0.17.2a20260702" +# Installed --no-deps (torch untouched); the full runtime set llm-compressor main needs, pinned. +_VENV_LLMCOMPRESSOR_SPECS = ( + f"transformers=={_LLMC_MAIN_TRANSFORMERS}", + f"llmcompressor @ git+https://github.com/vllm-project/llm-compressor@{_LLMC_MAIN_SHA}", + f"compressed-tensors=={_LLMC_MAIN_COMPRESSED_TENSORS}", + "huggingface-hub==1.21.0", + "hf-xet==1.5.1", + "tokenizers==0.22.2", + "safetensors==0.8.0", + "accelerate==1.14.0", + "datasets==5.0.0", + "pydantic==2.13.4", + "pydantic-core==2.46.4", + "typing-inspection==0.4.2", + "loguru==0.7.3", + "pyyaml==6.0.3", + "nvidia-ml-py==13.610.43", + "pillow==12.3.0", + "auto-round==0.13.1", + "regex==2026.6.28", +) +# Fingerprint of the pin set; bump the trailing schema version to force a rebuild on layout changes. +_LLMC_SHADOW_FINGERPRINT = ( + f"{_LLMC_MAIN_SHA}|{_LLMC_MAIN_TRANSFORMERS}|{_LLMC_MAIN_COMPRESSED_TENSORS}|schema=1" +) +_LLMC_SHADOW_MARKER = ".unsloth_llmc_fingerprint" + + +def _llmcompressor_main_disabled() -> bool: + """True if the operator forbids the llm-compressor-main shadow (air-gapped / locked-down).""" + return os.environ.get("UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +def _llmcompressor_shadow_is_valid() -> bool: + """True if the shadow dir exists with a marker matching the current pin fingerprint.""" + marker = Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER + try: + return marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT + except Exception: + return False + + +def _ensure_venv_llmcompressor_exists() -> bool: + """Ensure .venv_llmcompressor/ has the pinned llm-compressor-main stack. Install if missing. + + All specs are installed with --no-deps into a --target dir (mirrors the transformers sidecars), + so the workspace torch is never touched. Returns True on success. + """ + if _llmcompressor_shadow_is_valid(): + return True + if _llmcompressor_main_disabled(): + logger.warning( + "llm-compressor-main shadow needed but UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN is set; " + "compressed export of newer-transformers models will fail fast." + ) + return False + if _env_offline(): + logger.warning( + "llm-compressor-main shadow missing and HF/offline mode is set; cannot provision it." + ) + return False + + logger.warning( + "Provisioning llm-compressor-main shadow at %s (one-time, ~a few hundred MB, no torch) ...", + _VENV_LLMCOMPRESSOR_DIR, + ) + shutil.rmtree(_VENV_LLMCOMPRESSOR_DIR, ignore_errors = True) + os.makedirs(_VENV_LLMCOMPRESSOR_DIR, exist_ok = True) + + # Prefer uv (faster) then pip; install every spec at once, --no-deps, prereleases allowed + # (compressed-tensors ships as a pre-release). + base = [ + "--target", + _VENV_LLMCOMPRESSOR_DIR, + "--no-deps", + "--prerelease=allow", + *_VENV_LLMCOMPRESSOR_SPECS, + ] + cmds = [] + if shutil.which("uv"): + cmds.append(["uv", "pip", "install", "--python", sys.executable, *base]) + cmds.append( + [ + sys.executable, + "-m", + "pip", + "install", + *[a for a in base if a != "--prerelease=allow"], + "--pre", + ] + ) + + last_out = "" + for cmd in cmds: + result = subprocess.run( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = child_env_without_native_path_secret(), + **_windows_hidden_subprocess_kwargs(), + ) + last_out = result.stdout or "" + if result.returncode == 0: + try: + (Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER).write_text( + _LLMC_SHADOW_FINGERPRINT + ) + except Exception: + pass + logger.info("Provisioned llm-compressor-main shadow at %s", _VENV_LLMCOMPRESSOR_DIR) + return True + logger.warning("llm-compressor-main shadow install failed with %s; trying next", cmd[0]) + + logger.error( + "Failed to provision llm-compressor-main shadow (spec: llmcompressor@%s). Output:\n%s", + _LLMC_MAIN_SHA, + last_out[-4000:], + ) + return False + + +def llmcompressor_shadow_pythonpath() -> str | None: + """Provision (lazily) the llm-compressor-main shadow and return its sys.path entry, or None. + + Returns None when the shadow is disabled (UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN), offline, or + provisioning failed - callers then fall back to the fail-fast path. + """ + if _llmcompressor_main_disabled(): + return None + if _ensure_venv_llmcompressor_exists(): + return _VENV_LLMCOMPRESSOR_DIR + return None + + def _activate_venv(venv_dir: str, label: str) -> None: """Prepend *venv_dir* to sys.path, purge stale modules, reimport.""" if venv_dir not in sys.path: diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index e5fa6f0191..8c6ddd197a 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -70,6 +70,9 @@ const CHAT_ONLY_ALLOWED = new Set([ "/login", "/signup", "/change-password", + // Export stays reachable on chat-only hosts so the page can show its own grayed-out reason + // instead of a silent redirect; it self-gates via export capability, so nothing runs. + "/export", ]); function isChatOnlyAllowed(pathname: string): boolean { diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index ef4178ea42..fb1a1fc9c7 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -290,13 +290,12 @@ export function AppSidebar() { const chatOnly = usePlatformStore((s) => s.isChatOnly()); const chatOnlyReason = usePlatformStore((s) => s.chatOnlyReason); - // When Train/Export are greyed out (chat-only host), explain why on hover - // instead of disabling them silently. mlx_unavailable is the common macOS case - // after a reinstall/update dropped MLX and is recoverable via `unsloth studio update`. - const trainExportDisabledHint: string | undefined = !chatOnly + // Explain a greyed-out Train (chat-only host) on hover instead of disabling silently. Export is + // no longer disabled here: it stays navigable so its page can show a precise grayed-out reason. + const trainDisabledHint: string | undefined = !chatOnly ? undefined : chatOnlyReason === "mlx_unavailable" - ? "Training needs MLX. Run `unsloth studio update` to enable Train and Export." + ? "Training needs MLX. Run `unsloth studio update` to enable Train." : chatOnlyReason === "intel_mac" ? "Training needs Apple Silicon or a GPU. Intel Macs are chat-only." : chatOnlyReason === "no_gpu" @@ -1206,7 +1205,7 @@ export function AppSidebar() { pathname === "/studio" || pathname.startsWith("/studio/") } disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1235,7 +1234,7 @@ export function AppSidebar() { label={t("shell.navigation.train")} active={pathname === "/studio" || pathname.startsWith("/studio/")} disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1256,11 +1255,8 @@ export function AppSidebar() { icon={DownloadSquare01Icon} label={t("shell.navigation.export")} active={pathname === "/export" || pathname.startsWith("/export/")} - disabled={chatOnly} - tooltip={trainExportDisabledHint} spinner={exportInProgress} onClick={() => { - if (chatOnly) return; navigate({ to: "/export" }); closeMobileIfOpen(); }} diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index d1b4e88a6d..be9767a24f 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -127,6 +127,8 @@ export async function loadCheckpoint(params: { export async function exportMerged(params: { save_directory: string; format_type?: string; + /** Compressed-tensors scheme alias (e.g. "fp8", "w4a16", "mxfp4"); overrides format_type. */ + compressed_method?: string | null; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -158,7 +160,8 @@ export async function exportBase(params: { export async function exportGGUF(params: { save_directory: string; - quantization_method: string; + /** A single GGUF quant method or a list (list produces multiple GGUFs from one model load). */ + quantization_method: string | string[]; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -179,6 +182,10 @@ export async function exportLoRA(params: { repo_id?: string | null; hf_token?: string | null; private?: boolean; + /** Also convert the adapter to a GGUF LoRA file (llama.cpp `--lora`). */ + gguf?: boolean; + /** GGUF LoRA output float type (f32/f16/bf16/q8_0/auto); only used when gguf=true. */ + gguf_outtype?: string; }): Promise { const response = await authFetch("/api/export/export/lora", { method: "POST", diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx index 9e0edf7303..29ba5a703b 100644 --- a/studio/frontend/src/features/export/components/export-run-panel.tsx +++ b/studio/frontend/src/features/export/components/export-run-panel.tsx @@ -28,7 +28,11 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; -import { EXPORT_METHODS, type ExportMethod } from "../constants"; +import { + EXPORT_METHODS, + type ExportMethod, + findMergedFormat, +} from "../constants"; import type { ExportLogEntry } from "../api/export-api"; import { getExportLogLineClass } from "../lib/log-style"; import { @@ -200,6 +204,9 @@ export function ExportRunPanel(props: ExportRunPanelProps) { const summaryMethodLabel = summary?.methodLabel ?? methodTitle; const summaryQuants = summary?.quantLevels ?? quantLevels; const summaryMethod = summary?.method ?? exportMethod; + const summaryFormats = (summary?.mergedFormats ?? []).map( + (v) => findMergedFormat(v)?.label ?? v, + ); const showProgress = isExporting || isTerminal; return ( @@ -392,14 +399,32 @@ export function ExportRunPanel(props: ExportRunPanelProps) { ? "Export finished and pushed to Hugging Face Hub." : "Export finished successfully."} - {run.result?.outputPath ? ( - - {run.result.outputPath} - - ) : null} + {(() => { + // List every folder written; a multi-format merged run created one per format. + const paths = run.result?.outputPaths ?? []; + const items = + paths.length > 0 + ? paths + : run.result?.outputPath + ? [{ label: "", path: run.result.outputPath }] + : []; + const showLabels = items.length > 1; + return items.map((o, i) => ( +
+ {showLabels && o.label ? ( + + {o.label} + + ) : null} + + {o.path} + +
+ )); + })()}
)} @@ -432,6 +457,14 @@ export function ExportRunPanel(props: ExportRunPanelProps) { Export Method {summaryMethodLabel}
+ {summaryMethod === "merged" && summaryFormats.length > 0 && ( +
+ Formats + + {summaryFormats.join(", ")} + +
+ )} {summaryMethod === "gguf" && summaryQuants.length > 0 && (
Quantizations diff --git a/studio/frontend/src/features/export/constants.ts b/studio/frontend/src/features/export/constants.ts index 058de1edc4..be9080ba14 100644 --- a/studio/frontend/src/features/export/constants.ts +++ b/studio/frontend/src/features/export/constants.ts @@ -55,34 +55,172 @@ export const QUANT_OPTIONS: { { value: "f16", label: "F16" }, ]; -/** Merged-export precision formats. The compressed-tensors ones run llm-compressor for vLLM. */ -export type MergedFormat = - | "16-bit (FP16)" - | "FP8 (compressed-tensors)" - | "NVFP4 (compressed-tensors)"; +/** + * Merged-export precision formats, sorted by bit width. Three backends: + * - "plain": standard save (16-bit); `formatType` is the backend `format_type`. + * - "compressed": llm-compressor compressed-tensors (vLLM), NVIDIA-only; `value` is the alias. + * - "torchao": portable FP8/INT8, no NVIDIA GPU needed; `value` is the alias. + * `common` entries are quick pills, the rest the "More formats" dropdown; `needsNvidia` entries + * are hidden on non-NVIDIA hardware. + */ +export type MergedBackend = "plain" | "compressed" | "torchao"; -export const MERGED_FORMATS: { - value: MergedFormat; +export type MergedFormatOption = { + value: string; label: string; + bits: number; + backend: MergedBackend; + group: string; + common: boolean; + needsNvidia: boolean; + needsCalibration?: boolean; hint: string; -}[] = [ + /** Backend `format_type` for a "plain" save (unused for compressed/torchao). */ + formatType?: string; +}; + +/** Kept as a string alias for back-compat with callers that typed the old union. */ +export type MergedFormat = string; + +export const MERGED_FORMATS: MergedFormatOption[] = [ + // 16-bit { - value: "16-bit (FP16)", + value: "16-bit", label: "16-bit", + bits: 16, + backend: "plain", + group: "16-bit", + common: true, + needsNvidia: false, hint: "Full precision, runs anywhere.", + formatType: "16-bit (FP16)", + }, + // 8-bit + { + value: "fp8", + label: "FP8", + bits: 8, + backend: "compressed", + group: "FP8", + common: true, + needsNvidia: true, + hint: "Dynamic per-token FP8 (W8A8) for vLLM. Data-free.", }, { - value: "FP8 (compressed-tensors)", - label: "FP8 (vLLM)", - hint: "compressed-tensors FP8 for vLLM. Needs an NVIDIA GPU.", + value: "torchao_fp8", + label: "FP8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic FP8 (torchao). Produces on any hardware; loads in vLLM.", }, { - value: "NVFP4 (compressed-tensors)", - label: "NVFP4 (vLLM)", - hint: "compressed-tensors NVFP4 for vLLM. Needs an NVIDIA GPU; calibrates.", + value: "w8a8", + label: "INT8 (W8A8)", + bits: 8, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "8-bit weights and 8-bit activations for vLLM. Data-free.", + }, + { + value: "torchao_int8", + label: "INT8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic INT8 (torchao). Produces on any hardware; loads in vLLM.", + }, + { + value: "fp8_static", + label: "FP8 Static", + bits: 8, + backend: "compressed", + group: "FP8", + common: false, + needsNvidia: true, + needsCalibration: true, + hint: "Static per-tensor FP8. Calibrates on data.", + }, + { + value: "w8a16", + label: "INT8 (W8A16)", + bits: 8, + backend: "compressed", + group: "INT", + common: false, + needsNvidia: true, + hint: "8-bit weight-only. Data-free.", + }, + { + value: "mxfp8", + label: "MXFP8", + bits: 8, + backend: "compressed", + group: "MXFP", + common: false, + needsNvidia: true, + hint: "Microscaling FP8. Needs a newer compressed-tensors stack.", + }, + // 4-bit + { + value: "w4a16", + label: "INT4 (W4A16)", + bits: 4, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "4-bit weight-only (GPTQ-style) for vLLM. Data-free.", + }, + { + value: "mxfp4", + label: "MXFP4", + bits: 4, + backend: "compressed", + group: "MXFP", + common: true, + needsNvidia: true, + hint: "Microscaling FP4 (W4A4) for vLLM. Data-free.", + }, + { + value: "nvfp4", + label: "NVFP4", + bits: 4, + backend: "compressed", + group: "FP4", + common: true, + needsNvidia: true, + needsCalibration: true, + hint: "NVIDIA FP4 (W4A4) for vLLM. Calibrates on data.", }, ]; +/** Look up a merged format option by its stable value. */ +export function findMergedFormat(value: string): MergedFormatOption | undefined { + return MERGED_FORMATS.find((f) => f.value === value); +} + +/** Backend payload for one merged format: plain -> formatType, compressed/torchao -> the alias. */ +export function mergedFormatPayload(value: string): { + formatType: string; + compressedMethod: string | null; +} { + const opt = findMergedFormat(value); + if (!opt || opt.backend === "plain") { + return { + formatType: opt?.formatType ?? "16-bit (FP16)", + compressedMethod: null, + }; + } + return { formatType: "16-bit (FP16)", compressedMethod: opt.value }; +} + /** * llama.cpp effective bits-per-weight per quant; GGUF size ~= fp16_bytes * bpw / 16. * K-quant values are published average bit-rates (Q2_K_L = Unsloth Q2_K + Q8_0 diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 7b5c6ec6d4..07606a26ed 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -24,6 +24,19 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { + DropdownMenu, + DropdownMenuCheckboxItem, + DropdownMenuContent, + DropdownMenuLabel, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { + Alert, + AlertDescription, + AlertTitle, +} from "@/components/ui/alert"; import { Separator } from "@/components/ui/separator"; import { Spinner } from "@/components/ui/spinner"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -63,11 +76,14 @@ import { type ExportMethod, GUIDE_STEPS, MERGED_FORMATS, - type MergedFormat, + type MergedFormatOption, + mergedFormatPayload, QUANT_OPTIONS, buildQuantSizeLabels, getEstimatedSize, } from "./constants"; +import { useHardwareInfo } from "@/hooks/use-hardware-info"; +import { usePlatformStore } from "@/config/env"; import { isExportPanelActive, useExportRuntimeStore, @@ -78,6 +94,10 @@ import { exportTourSteps } from "./tour"; const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]); +// GGUF LoRA output float types (Q8_0 first / default). Q8_0 falls back to F16 per tensor for dims +// not divisible by the block size (32); no "auto" - the choice is explicit. +const LORA_GGUF_OUTTYPES = ["q8_0", "f16", "bf16", "f32"] as const; + type SourceTab = "local" | "checkpoint" | "hf"; type SourceMode = "checkpoint" | "model"; @@ -109,7 +129,16 @@ function buildRelativeSaveDirectory( : sourceBaseModelName; return `${safePathSegment(rawName)}-GGUF`; } - return `${selectedModelIdx ?? "model"}/${checkpoint}`; + // Merged / LoRA: a checkpoint keeps the "/" layout under outputs. + if (sourceMode === "checkpoint" && selectedModelIdx && checkpoint) { + return `${selectedModelIdx}/${checkpoint}`; + } + // Local / HF source (no checkpoint): name from the model id to avoid "model/null". + const rawName = + sourceMode === "checkpoint" + ? checkpoint ?? selectedModelIdx ?? sourceBaseModelName + : sourceBaseModelName; + return `${safePathSegment(rawName)}-${exportMethod === "lora" ? "adapter" : "merged"}`; } function siblingGgufDirectory(sourcePath: string): string | null { @@ -177,9 +206,57 @@ export function ExportPage() { }); // GGUF importance matrix (required for the IQ quants) and merged-export precision. const [useImatrix, setUseImatrix] = useState(false); - const [mergedFormat, setMergedFormat] = useState("16-bit (FP16)"); - // IQ quants are imatrix-only, so force it on when one is selected; otherwise we would submit - // an IQ quant with no imatrix and llama.cpp would reject it. + // Merged precision: one or more MERGED_FORMATS values, exported in one run. Seed from a live run + // so navigating away and back (which remounts this page) keeps the selection, like exportMethod. + const [selectedFormats, setSelectedFormats] = useState(() => { + const s = useExportRuntimeStore.getState(); + return isExportPanelActive(s) && + s.summary?.method === "merged" && + s.summary.mergedFormats.length > 0 + ? s.summary.mergedFormats + : ["16-bit"]; + }); + // LoRA-only export: optionally also emit a GGUF LoRA adapter, and its output float type. + const [loraAsGguf, setLoraAsGguf] = useState(false); + const [loraGgufOuttype, setLoraGgufOuttype] = useState("q8_0"); + // GGUF method: export the full model as GGUF quants, or (for an adapter checkpoint) a GGUF LoRA. + const [ggufTarget, setGgufTarget] = useState<"model" | "lora">("model"); + + const hardware = useHardwareInfo(); + // GGUF LoRA conversion is rejected on the macOS / MLX path, so gate it out on a Mac host. + const isMacHost = usePlatformStore((s) => s.deviceType) === "mac"; + // Real CUDA (not ROCm); gates the NVIDIA-only compressed-tensors formats. + const hasNvidia = hardware.cuda != null && hardware.rocm == null; + // Only gray out on an authoritative unsupported response; while unloaded the backend route guard + // stays authoritative. The backend supplies the precise reason; the fallback below is a backstop. + const exportUnsupported = + hardware.loaded && hardware.exportSupported === false; + const exportUnsupportedMessage = + hardware.exportUnsupportedMessage ?? + "Export requires a supported accelerator (NVIDIA, AMD, or Intel GPU, or Apple Silicon) with PyTorch or MLX installed."; + const availableFormats = useMemo( + () => + MERGED_FORMATS.filter((f) => { + // compressed-tensors (llm-compressor) is the NVIDIA path; shown only on an NVIDIA GPU. + if (f.backend === "compressed") return hasNvidia; + // Portable torchao is the fallback for hosts without the NVIDIA compressed path, i.e. a + // CPU / non-NVIDIA box. Hidden on NVIDIA (use compressed-tensors) and on macOS/MLX (the + // backend rejects quantized export there). + if (f.backend === "torchao") return !hasNvidia && !isMacHost; + // Plain 16-bit is available everywhere. + return true; + }), + [hasNvidia, isMacHost], + ); + const toggleFormat = useCallback((value: string) => { + setSelectedFormats((prev) => + prev.includes(value) + ? prev.filter((v) => v !== value) + : [...prev, value], + ); + }, []); + // availableFormats already drops NVIDIA-only formats on other hardware, so no pruning needed. + // IQ quants are imatrix-only: force imatrix on when one is selected, else llama.cpp rejects it. const requiresImatrix = quantLevels.some( (q) => QUANT_OPTIONS.find((o) => o.value === q)?.imatrix, ); @@ -304,6 +381,11 @@ export function ExportPage() { const baseModelName = selectedModelData?.base_model ?? "—"; const isAdapter = !!selectedModelData?.peft_type; const isQuantized = !!selectedModelData?.is_quantized; + // isAdapter / isQuantized come from the checkpoint's metadata and are stale in "model" source + // mode (a direct base export), so treat both as false outside checkpoint mode to avoid wrongly + // gating the methods. + const effectiveIsAdapter = sourceMode === "checkpoint" && isAdapter; + const effectiveIsQuantized = sourceMode === "checkpoint" && isQuantized; const loraRank = selectedModelData?.lora_rank ?? null; const trainingMethodLabel = selectedModelData?.peft_type ? "LoRA / QLoRA" @@ -416,25 +498,30 @@ export function ExportPage() { setCheckpoint(null); }, [selectedModelIdx]); - // For a ?run= deep link, default to the run's main checkpoint. Declared after - // the reset effect above so it runs last and isn't clobbered back to null. + // Default to the newest checkpoint when none is chosen (checkpoints are sorted newest-first). + // Declared after the reset effect above so it runs last and isn't clobbered back to null. Covers + // both a ?run= deep link and a plain finetune opened without an explicit checkpoint pick. useEffect(() => { - if (appliedRunRef.current == null) return; - if (appliedRunRef.current !== selectedModelIdx) return; + if (sourceMode !== "checkpoint") return; if (checkpoint != null || checkpointsForModel.length === 0) return; setCheckpoint(checkpointsForModel[0].display_name); - }, [selectedModelIdx, checkpoint, checkpointsForModel]); + }, [sourceMode, selectedModelIdx, checkpoint, checkpointsForModel]); // Auto-reset export method if incompatible with the selected model type useEffect(() => { - if (!isAdapter && (exportMethod === "merged" || exportMethod === "lora")) { + // Only LoRA needs a real adapter; Merged and GGUF work for non-PEFT base models too. + if (!effectiveIsAdapter && exportMethod === "lora") { setExportMethod(null); } // Quantized non-PEFT models can't export to any format - if (!isAdapter && isQuantized && exportMethod !== null) { + if (!effectiveIsAdapter && effectiveIsQuantized && exportMethod !== null) { setExportMethod(null); } - }, [isAdapter, isQuantized, exportMethod]); + // The GGUF LoRA target only applies to an adapter checkpoint on a non-Mac host. + if ((!effectiveIsAdapter || isMacHost) && ggufTarget !== "model") { + setGgufTarget("model"); + } + }, [effectiveIsAdapter, effectiveIsQuantized, exportMethod, isMacHost, ggufTarget]); const handleSourceTabChange = useCallback((next: string) => { if (next === "checkpoint") { @@ -442,7 +529,7 @@ export function ExportPage() { } else if (next === "hf" || next === "local") { setSourceMode("model"); setModelSource(next); - setExportMethod("gguf"); + // Don't force GGUF: Local / HF sources can export Merged too; a stale LoRA pick auto-resets. } else { return; } @@ -508,10 +595,26 @@ export function ExportPage() { sourceMode, ]); const saveDirectory = customSaveDirectory?.trim() || defaultSaveDirectory; + // Each merged format uploads a full model to the repo root, so several to one repo would collide. + // GGUF method exporting an adapter checkpoint as a GGUF LoRA (vs full-model quants). Reuses the + // LoRA-adapter export path; no quant list needed. + const ggufAsLora = + exportMethod === "gguf" && + ggufTarget === "lora" && + effectiveIsAdapter && + !isMacHost; + + // Restrict a Hub merged export to a single format; multi-format stays available for local export. + const hubMultiFormat = + destination === "hub" && exportMethod === "merged" && selectedFormats.length > 1; + const canExport = !!( selectedExportSource && exportMethod && - (exportMethod !== "gguf" || quantLevels.length > 0) + !exportUnsupported && + !hubMultiFormat && + (exportMethod !== "gguf" || ggufAsLora || quantLevels.length > 0) && + (exportMethod !== "merged" || selectedFormats.length > 0) ); const applyHfSourceModel = useCallback((value: string) => { @@ -576,9 +679,14 @@ export function ExportPage() { const handleStart = useCallback(async () => { const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel; if (!source || !exportMethod) return; - // A GGUF export with no quant selected runs zero exports yet would still - // settle as success with no file; require at least one (mirrors canExport). - if (exportMethod === "gguf" && quantLevels.length === 0) return; + // No supported accelerator (or PyTorch/MLX missing): the backend would reject anyway; don't submit. + if (exportUnsupported) return; + // GGUF with no quant, or merged with no format, would run an unintended/empty export; require + // at least one (mirrors canExport, in case the panel's Start button bypasses the outer one). + if (exportMethod === "gguf" && !ggufAsLora && quantLevels.length === 0) return; + if (exportMethod === "merged" && selectedFormats.length === 0) return; + // A Hub merged push writes each format to the repo root; several would collide (mirrors canExport). + if (hubMultiFormat) return; const selectedCp = sourceMode === "checkpoint" ? checkpointsForModel.find((cp) => cp.display_name === checkpoint) @@ -591,8 +699,13 @@ export function ExportPage() { ? `${hfUsername}/${modelName}` : undefined; const token = pushToHub && hfToken ? hfToken : undefined; - const methodLabel = - EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod; + // The GGUF method with the LoRA target reuses the LoRA-adapter export path. + const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod; + const emitLoraGguf = + ggufAsLora || (effectiveMethod === "lora" && loraAsGguf && !isMacHost); + const methodLabel = ggufAsLora + ? "GGUF LoRA adapter" + : (EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod); const adapterExport = sourceMode === "checkpoint" && isAdapter; // Consent gate for an HF source's custom (auto_map) code, run before we hand @@ -624,11 +737,16 @@ export function ExportPage() { trustRemoteCode, approvedRemoteCodeFingerprint, loadToken: hfToken || null, - exportMethod, + exportMethod: effectiveMethod, isAdapter: adapterExport, quantLevels, useImatrix: effectiveImatrix, - mergedFormat, + mergedSelections: selectedFormats.map((v) => ({ + ...mergedFormatPayload(v), + label: MERGED_FORMATS.find((f) => f.value === v)?.label ?? v, + })), + loraGguf: emitLoraGguf, + loraGgufOuttype, saveDirectory, destination, repoId, @@ -639,8 +757,9 @@ export function ExportPage() { baseModelName: sourceBaseModelName, checkpointLabel: selectedExportSource, methodLabel, - method: exportMethod, + method: effectiveMethod, quantLevels, + mergedFormats: exportMethod === "merged" ? selectedFormats : [], destination, }, }); @@ -656,7 +775,13 @@ export function ExportPage() { isAdapter, quantLevels, effectiveImatrix, - mergedFormat, + selectedFormats, + hubMultiFormat, + ggufAsLora, + loraAsGguf, + isMacHost, + loraGgufOuttype, + exportUnsupported, destination, saveDirectory, hfUsername, @@ -1163,75 +1288,303 @@ export function ExportPage() {
+ {exportUnsupported && ( + + + Export unavailable + {exportUnsupportedMessage} + + )} + - {exportMethod === "merged" && isAdapter && ( -
-
Precision
-
- {MERGED_FORMATS.map((f) => ( - - ))} -
-
- {MERGED_FORMATS.find((f) => f.value === mergedFormat)?.hint} + {exportMethod === "merged" && !exportUnsupported && ( +
+
+
+
Precision
+ + — select one or more + +
+
+ {availableFormats + .filter((f) => f.common) + .map((f) => { + const active = selectedFormats.includes(f.value); + return ( + + ); + })} + + {availableFormats.some((f) => !f.common) && ( + + + + + + + Additional formats + + + {availableFormats + .filter((f) => !f.common) + .map((f) => ( + toggleFormat(f.value)} + onSelect={(e) => e.preventDefault()} + > + + + {f.label} + {f.needsCalibration ? " *" : ""} + + + {f.hint} + + + + ))} + + + )} +
+ + {selectedFormats.length > 0 && ( +
+ + {selectedFormats.length} selected:{" "} + {selectedFormats + .map( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.label ?? v, + ) + .join(", ")} + + {selectedFormats.length > 1 && ( + + )} +
+ )} + + {hubMultiFormat && ( +
+ Hub export supports one format at a time (each writes to + the repository root). Select a single format, or export + locally to produce several at once. +
+ )} + + {selectedFormats.some( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.needsCalibration, + ) && ( +
+ * calibrates on data (uses a small calibration set). +
+ )} + + {!hasNvidia && ( +
+ No NVIDIA GPU detected: compressed-tensors formats are + hidden. 16-bit and portable FP8/INT8 (torchao) still + work here and load in vLLM. +
+ )}
)} - {exportMethod === "gguf" && ( - <> - -
-
-
- Importance matrix (imatrix) + {exportMethod === "lora" && effectiveIsAdapter && !exportUnsupported && ( +
+
+
Adapter format
+
+ + +
+
+ {isMacHost + ? "GGUF LoRA is not available on macOS/MLX; exporting the safetensors adapter." + : loraAsGguf + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Standard PEFT adapter files. Pair with the base model at inference."} +
+
+ + {loraAsGguf && ( +
+
Output type
+ +
+ )} +
+ )} + + {exportMethod === "gguf" && !exportUnsupported && ( +
+ {effectiveIsAdapter && !isMacHost && ( +
+
Export target
+
+ +
- {requiresImatrix - ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." - : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} + {ggufTarget === "lora" + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Merges the adapter into the base model, then quantizes the full model to GGUF."}
- -
- + )} + + {ggufAsLora ? ( +
+
Output type
+ +
+ ) : ( + <> + +
+
+
+ Importance matrix (imatrix) +
+
+ {requiresImatrix + ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." + : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} +
+
+ +
+ + )} +
)} {estimatedSize && (
diff --git a/studio/frontend/src/features/export/stores/export-runtime-store.ts b/studio/frontend/src/features/export/stores/export-runtime-store.ts index a87e7d93e9..6ee5c78994 100644 --- a/studio/frontend/src/features/export/stores/export-runtime-store.ts +++ b/studio/frontend/src/features/export/stores/export-runtime-store.ts @@ -5,7 +5,6 @@ import { create } from "zustand"; import { cancelExport, cleanupExport, - exportBase, exportGGUF, exportLoRA, exportMerged, @@ -119,6 +118,8 @@ export interface ExportRunSummary { methodLabel: string; method: ExportMethod; quantLevels: string[]; + /** Merged: the selected format values (for the summary "Formats" row and to reseed the picker). */ + mergedFormats: string[]; destination: ExportDestination; } @@ -140,8 +141,16 @@ export interface RunExportParams { quantLevels: string[]; /** GGUF: use an importance matrix (auto-download); required for the IQ quants. */ useImatrix?: boolean; - /** Merged: precision/format ("16-bit (FP16)" or a compressed-tensors option). */ - mergedFormat?: string; + /** Merged: precision formats, each exported to its own sibling directory. Defaults to 16-bit. + * `label` is the display name for the success banner's per-format output line. */ + mergedSelections?: { + formatType: string; + compressedMethod: string | null; + label: string; + }[]; + /** LoRA: also emit a GGUF LoRA adapter (llama.cpp `--lora`), and its output float type. */ + loraGguf?: boolean; + loraGgufOuttype?: string; saveDirectory: string; destination: ExportDestination; repoId?: string; @@ -172,7 +181,13 @@ interface ExportRuntimeState { * settling the run by polling /api/export/status instead. Logs keep streaming. */ reconnecting: boolean; startedAt: number | null; - result: { outputPath: string | null; destination: ExportDestination } | null; + /** `outputPath` is the first path (back-compat); `outputPaths` is one entry per written folder + * so a multi-format merged run can list every sibling directory it created. */ + result: { + outputPath: string | null; + outputPaths: { label: string; path: string }[]; + destination: ExportDestination; + } | null; error: string | null; cancelRequested: boolean; hasHydrated: boolean; @@ -317,6 +332,10 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success" as const, result: { outputPath: status.last_op_output_path ?? null, + // A run recovered from the backend only knows the last output path. + outputPaths: status.last_op_output_path + ? [{ label: "", path: status.last_op_output_path }] + : [], destination: state.result?.destination ?? "local", }, }; @@ -351,7 +370,9 @@ export const useExportRuntimeStore = create()((set, get) => const quantTotal = params.exportMethod === "gguf" ? Math.max(1, params.quantLevels.length) - : 1; + : params.exportMethod === "merged" + ? Math.max(1, params.mergedSelections?.length ?? 1) + : 1; set({ runId, @@ -431,67 +452,72 @@ export const useExportRuntimeStore = create()((set, get) => } if (!isCurrent()) return; - // 2. Run the export. Capture the resolved output_path for the success - // banner; multi-quant GGUF shares one directory, so keep the last. + // 2. Run the export. Collect every resolved output_path so the success + // banner can list each sibling directory a multi-format run created. set({ phase: "exporting" }); - let lastOutputPath: string | null = null; + const outputs: { label: string; path: string }[] = []; if (params.exportMethod === "merged") { - if (params.isAdapter) { + // Each selected format writes its own sibling directory (PEFT or non-PEFT base alike). + const selections = + params.mergedSelections && params.mergedSelections.length > 0 + ? params.mergedSelections + : [{ formatType: "16-bit (FP16)", compressedMethod: null, label: "16-bit" }]; + for (let i = 0; i < selections.length; i += 1) { + if (!isCurrent()) return; + set({ quantIndex: i }); + const sel = selections[i]; const { outputPath } = await runRecoverableOp(() => exportMerged({ save_directory: params.saveDirectory, - format_type: params.mergedFormat, + format_type: sel.formatType, + compressed_method: sel.compressedMethod, push_to_hub: pushToHub, repo_id: params.repoId, hf_token: params.token, private: params.privateRepo, }), ); - lastOutputPath = outputPath; - } else { - const { outputPath } = await runRecoverableOp(() => - exportBase({ - save_directory: params.saveDirectory, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - private: params.privateRepo, - base_model_id: params.baseModelId, - }), - ); - lastOutputPath = outputPath; - } - } else if (params.exportMethod === "gguf") { - for (let i = 0; i < params.quantLevels.length; i += 1) { - if (!isCurrent()) return; - set({ quantIndex: i }); - const quant = params.quantLevels[i]; - const { outputPath } = await runRecoverableOp(() => - exportGGUF({ - save_directory: params.saveDirectory, - quantization_method: quant, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - imatrix: params.useImatrix, - }), - ); - lastOutputPath = outputPath ?? lastOutputPath; + if (outputPath) outputs.push({ label: sel.label, path: outputPath }); if (!isCurrent()) return; set({ quantIndex: i + 1 }); } + } else if (params.exportMethod === "gguf") { + // Send the whole quant list in ONE call: the model is merged once and every GGUF comes + // from that single merge (unsloth save_to_gguf loops internally). + const { outputPath } = await runRecoverableOp(() => + exportGGUF({ + save_directory: params.saveDirectory, + quantization_method: params.quantLevels, + push_to_hub: pushToHub, + repo_id: params.repoId, + hf_token: params.token, + imatrix: params.useImatrix, + }), + ); + if (outputPath) outputs.push({ label: "GGUF", path: outputPath }); + if (!isCurrent()) return; + set({ quantIndex: get().quantTotal }); } else if (params.exportMethod === "lora") { const { outputPath } = await runRecoverableOp(() => exportLoRA({ save_directory: params.saveDirectory, push_to_hub: pushToHub, repo_id: params.repoId, - hf_token: params.token, + // A local GGUF LoRA export still reloads a possibly-gated base config, so fall back to + // the load token when there is no hub-upload token (both are the same HF token). + hf_token: params.token ?? params.loadToken ?? null, private: params.privateRepo, + gguf: params.loraGguf ?? false, + gguf_outtype: params.loraGgufOuttype ?? "q8_0", }), ); - lastOutputPath = outputPath; + if (outputPath) { + outputs.push({ + label: params.loraGguf ? "GGUF LoRA adapter" : "LoRA adapter", + path: outputPath, + }); + } } if (!isCurrent()) return; @@ -499,7 +525,11 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success", isExporting: false, reconnecting: false, - result: { outputPath: lastOutputPath, destination: params.destination }, + result: { + outputPath: outputs[0]?.path ?? null, + outputPaths: outputs, + destination: params.destination, + }, }); } catch (err) { if (!isCurrent()) return; diff --git a/studio/frontend/src/hooks/use-hardware-info.ts b/studio/frontend/src/hooks/use-hardware-info.ts index 36e3c532f3..4d63d4d6af 100644 --- a/studio/frontend/src/hooks/use-hardware-info.ts +++ b/studio/frontend/src/hooks/use-hardware-info.ts @@ -25,6 +25,13 @@ export interface HardwareInfo { transformers: string | null; unsloth: string | null; llamaCpp: string | null; + // Whether export can run here (true only on a supported accelerator), with a torch-aware + // reason. `null` until the authoritative response lands, so callers don't briefly enable + // export; `loaded` flips true once a real (non-error) response arrives. + exportSupported: boolean | null; + exportUnsupportedReason: string | null; + exportUnsupportedMessage: string | null; + loaded: boolean; } const DEFAULT: HardwareInfo = { @@ -38,6 +45,10 @@ const DEFAULT: HardwareInfo = { transformers: null, unsloth: null, llamaCpp: null, + exportSupported: null, + exportUnsupportedReason: null, + exportUnsupportedMessage: null, + loaded: false, }; // Module-level cache so multiple components share one fetch. @@ -87,6 +98,10 @@ async function fetchOnce(): Promise { transformers: data?.versions?.transformers ?? null, unsloth: data?.versions?.unsloth ?? null, llamaCpp: data?.llama_cpp ?? null, + exportSupported: data?.export_supported ?? null, + exportUnsupportedReason: data?.export_unsupported_reason ?? null, + exportUnsupportedMessage: data?.export_unsupported_message ?? null, + loaded: true, }; if (generation === cacheGeneration) { cached = info; diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 304de7a9af..209a8a06f1 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -256,7 +256,7 @@ with sync_playwright() as p: composer = page.locator('textarea[aria-label="Message input"]') composer.wait_for(state = "visible", timeout = 60_000) - # Detect chat-only mode (/api/health.chat_only): in chat-only mode /studio + /export redirect to /chat. + # Detect chat-only mode (/api/health.chat_only): /studio redirects to /chat while /export stays reachable and self-gated. health_resp = evaluate_fetch( page, f"{BASE}/api/health", @@ -404,15 +404,19 @@ with sync_playwright() as p: # ───────────────────────────────────────────────────── # 3. Export route. # ───────────────────────────────────────────────────── - step(f"Export route ({'chat-only redirect' if chat_only else 'form fields'})") + step(f"Export route ({'chat-only self-gated' if chat_only else 'form fields'})") page.goto(f"{BASE}/export") page.wait_for_timeout(1500) shoot("07-export") if chat_only: - if "/export" in page.url: - soft_fail(f"chat-only mode should redirect /export -> /chat; url={page.url}") + if "/export" not in page.url: + soft_fail(f"chat-only mode should keep /export reachable; url={page.url}") else: - info(f"OK chat-only redirected /export -> {page.url}") + unavailable = page.get_by_text(re.compile(r"Export unavailable", re.I)).first + if unavailable.count() == 0: + soft_fail("chat-only /export did not show the export unavailable gate") + else: + info("OK chat-only /export rendered the unavailable gate") else: # Non-chat-only: verify the export-cta button + HF token field. cta = page.locator('[data-tour="export-cta"]').first diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index 66ebdd75da..f0a843c380 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -245,6 +245,10 @@ def main(): # expert even if the sample set does not route tokens to all of them. is_moe = _is_moe(getattr(model, "config", None)) ignore = ["lm_head"] + # Skip the same modules RedHatAI/NVIDIA skip for the Qwen3.5 / Qwen3-Next family (these also have + # shapes not divisible by the grouped-scheme group_size, which would otherwise error). No-ops + # elsewhere. Hybrid linear attention, VLM vision tower, and the MTP/speculative head. + ignore += ["re:.*\\.linear_attn\\..*", "re:.*\\.visual\\..*", "re:.*mtp.*"] if is_moe: # Keep MoE routing layers unquantized: the router gate and (Qwen) shared-expert gate. ignore += ["re:.*\\.gate$", "re:.*\\.shared_expert_gate$"] diff --git a/unsloth/save.py b/unsloth/save.py index 226ca5fed8..50ae4119bd 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -205,6 +205,24 @@ COMPRESSED_EXPORT_SCHEMES = { } +# torchao "portable" quant export: device-agnostic FP8 / INT8, no NVIDIA GPU needed. +# alias -> (kind, sibling suffix). FP8 saves to safetensors, INT8 to .bin; both load in vLLM. +TORCHAO_EXPORT_SCHEMES = { + "torchao_fp8": ("fp8", "torchao-fp8"), + "torchao_int8": ("int8", "torchao-int8"), + "portable_fp8": ("fp8", "torchao-fp8"), + "portable_int8": ("int8", "torchao-int8"), +} + + +def _normalize_torchao_method(save_method): + """Return (kind, suffix) if `save_method` is a torchao portable FP8/INT8 export, else None.""" + if not isinstance(save_method, str): + return None + key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + return TORCHAO_EXPORT_SCHEMES.get(key) + + def _normalize_compressed_method(save_method): """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). @@ -215,6 +233,9 @@ def _normalize_compressed_method(save_method): if not isinstance(save_method, str): return None key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + # torchao aliases route to the torchao path, so skip them before the "fp8" near-miss check. + if key in TORCHAO_EXPORT_SCHEMES: + return None if key in COMPRESSED_EXPORT_SCHEMES: return COMPRESSED_EXPORT_SCHEMES[key] if any(tag in key for tag in ("fp8", "fp4", "mxfp", "nvfp", "w4a", "w8a", "int4", "int8")): @@ -1368,6 +1389,53 @@ def install_python_non_blocking(packages = []): # bump deliberately. Floor 0.6.0 keeps torch>=2.4 resolvable (0.7+ need torch>=2.7; torch pinned below). _LLM_COMPRESSOR_SPEC = "llmcompressor>=0.6.0,<=0.12.0" +# Highest transformers release llm-compressor 0.10.x/0.12.x can run against (its metadata pins +# transformers<=4.57.6). Models that require a newer-transformers sidecar (e.g. Qwen3.5 needs +# transformers 5.3.0) cannot be quantized by llm-compressor at all: it imports +# transformers.modeling_utils.TORCH_INIT_FUNCTIONS, which was removed in transformers 5.x, so the +# compressed-export subprocess dies with a cryptic ImportError AFTER the expensive 16bit merge. +# Detect that up front and fail fast with an actionable message. Bump this in lockstep with a +# llm-compressor release that supports newer transformers. +_LLM_COMPRESSOR_MAX_TRANSFORMERS = "4.57.6" + + +def _transformers_exceeds_llm_compressor_ceiling(transformers_version = None): + """Return (exceeds, active_version) comparing the active transformers to the llm-compressor ceiling. + + `exceeds` is True only when we can parse both versions and the active transformers is strictly + newer than `_LLM_COMPRESSOR_MAX_TRANSFORMERS`. Any parse failure returns False (fail open) so a + real quantization attempt still surfaces the underlying error rather than a false positive. + """ + if transformers_version is None: + try: + import transformers as _tf + transformers_version = _tf.__version__ + except Exception: + return False, "unknown" + try: + from packaging.version import parse as _parse + + # Drop any local build suffix ("4.57.6+abc") so it does not skew the comparison. + active = _parse(str(transformers_version).split("+", 1)[0]) + ceiling = _parse(_LLM_COMPRESSOR_MAX_TRANSFORMERS) + return active > ceiling, str(transformers_version) + except Exception: + return False, str(transformers_version) + + +# A caller (e.g. Unsloth Studio) can enable FP8/FP4 export of newer-transformers models (Qwen3.5, +# Gemma-4, ...) by provisioning a dedicated llm-compressor-main "shadow" (transformers>=5.9 layered +# over the existing torch) and pointing us at its sys.path entry via this env var. When set, the +# quantization subprocess uses it instead of the workspace llm-compressor and the ceiling fail-fast +# is bypassed. +_COMPRESSED_QUANTIZE_PYTHONPATH_ENV = "UNSLOTH_COMPRESSED_QUANTIZE_PYTHONPATH" + + +def _compressed_quantize_pythonpath(): + """Return the llm-compressor-main shadow PYTHONPATH, or None if not set.""" + pp = os.environ.get(_COMPRESSED_QUANTIZE_PYTHONPATH_ENV, "").strip() + return pp or None + def install_llm_compressor(): """Import llm-compressor, installing it on first use for FP8/FP4 export. @@ -2023,10 +2091,40 @@ def unsloth_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -2107,6 +2205,37 @@ def unsloth_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -2114,6 +2243,7 @@ def unsloth_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3812,10 +3942,40 @@ def unsloth_generic_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3896,6 +4056,37 @@ def unsloth_generic_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -3903,6 +4094,7 @@ def unsloth_generic_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -4114,22 +4306,36 @@ def _unsloth_save_compressed_tensors( if not is_main_process: return None - # 1) Install llm-compressor and gate on scheme availability BEFORE merging, so an unsupported - # scheme (e.g. mxfp8) fails fast instead of writing a full 16bit checkpoint first. - install_llm_compressor() - if not _scheme_is_available(scheme): - try: - import transformers as _tf - tf_ver = _tf.__version__ - except Exception: - tf_ver = "unknown" - raise RuntimeError( - f"Unsloth: scheme '{scheme}' is not available in your installed " - f"compressed-tensors / llm-compressor.\n" - f"It requires a newer llm-compressor that needs transformers>=5.9 " - f"(you have transformers {tf_ver}).\n" - "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." - ) + # 1) Prepare the quantization runtime BEFORE merging, so an unusable config fails fast instead of + # writing a full 16bit checkpoint first. With the llm-compressor-main shadow the subprocess + # validates everything itself, so skip the workspace install / ceiling / scheme checks; without + # it, install the workspace llm-compressor and fail fast past its transformers ceiling. + _shadow_pythonpath = _compressed_quantize_pythonpath() + if _shadow_pythonpath is None: + install_llm_compressor() + # llm-compressor cannot run under a newer transformers than its ceiling: the quantization + # subprocess would die with a cryptic ImportError (TORCH_INIT_FUNCTIONS) only AFTER the costly + # 16bit merge. Detect and fail fast with an actionable message instead. + _exceeds, _tf_ver = _transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + raise RuntimeError( + f"Unsloth: FP8/FP4 compressed-tensors export is not available for this model. It runs " + f"under transformers {_tf_ver}, but llm-compressor supports transformers " + f"<= {_LLM_COMPRESSOR_MAX_TRANSFORMERS}. Export to GGUF or 16-bit instead." + ) + if not _scheme_is_available(scheme): + try: + import transformers as _tf + tf_ver = _tf.__version__ + except Exception: + tf_ver = "unknown" + raise RuntimeError( + f"Unsloth: scheme '{scheme}' is not available in your installed " + f"compressed-tensors / llm-compressor.\n" + f"It requires a newer llm-compressor that needs transformers>=5.9 " + f"(you have transformers {tf_ver}).\n" + "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." + ) # 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and # quantize inside an isolated temp dir instead of writing ./ into the cwd. @@ -4307,8 +4513,15 @@ def _unsloth_save_compressed_tensors( env["HF_TOKEN"] = token env["HUGGING_FACE_HUB_TOKEN"] = token + # Clean PYTHONPATH = shadow only. torch still comes from the interpreter's site-packages; + # transformers 5.x + llm-compressor main come from the shadow. Dropping the inherited + # PYTHONPATH removes any parent transformers sidecar so the shadow's is authoritative. + if _shadow_pythonpath is not None: + env["PYTHONPATH"] = _shadow_pythonpath + print( f"Unsloth: Quantizing the merged model to {scheme} with llm-compressor " + f"{'(llm-compressor-main shadow) ' if _shadow_pythonpath is not None else ''}" "(in a separate process)..." ) try: @@ -4377,6 +4590,255 @@ def _unsloth_save_compressed_tensors( torch.cuda.empty_cache() +def _unsloth_save_torchao( + model, + save_directory: Union[str, os.PathLike], + tokenizer, + kind: str, + suffix: str, + push_to_hub: bool = False, + token: Optional[Union[str, bool]] = None, + is_main_process: bool = True, + **merge_kwargs, +): + """Export a device-agnostic torchao FP8 / INT8 "portable" checkpoint (no NVIDIA GPU needed). + + Merges LoRA to 16bit in a staging dir, then applies torchao weight-only quantization via + `TorchAoConfig` into `save_directory + "-" + suffix`. No calibration, subprocess, or CUDA. + `kind` is "fp8" (safetensors) or "int8" (.bin; torchao only whitelists float8 for safetensors). + """ + import tempfile + + if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): + tokenizer = patch_saving_functions(tokenizer) + if token is None: + token = get_token() + + # Only the main process merges, quantizes, and uploads; other ranks return at once. + if not is_main_process: + return None + + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + AutoProcessor, + TorchAoConfig, + ) + from torchao.quantization import Float8WeightOnlyConfig, Int8WeightOnlyConfig + + if kind == "fp8": + quant_type = Float8WeightOnlyConfig() + safe_serialization = True + elif kind == "int8": + quant_type = Int8WeightOnlyConfig() + safe_serialization = False # torchao only supports safetensors for float8 configs + else: + raise RuntimeError(f"Unsloth: unknown torchao export kind '{kind}' (expected fp8/int8).") + + # Always merge into an isolated temp staging dir (never save_directory itself), so a co-selected + # 16-bit export written to save_directory is not overwritten or deleted; the torchao output is + # the sibling "-" (or the repo id on a hub push). + repo_id, work_tmp, model_dev = None, None, None + work_tmp = tempfile.mkdtemp(prefix = "unsloth-torchao-") + if push_to_hub: + repo_id = os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(repo_id.rstrip("/")) or "model") + out_dir = staging + "-" + suffix + else: + base = os.fspath(save_directory).rstrip("/\\") or os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(base) or "model") + out_dir = base + "-" + suffix + + api = None + try: + if push_to_hub: + from huggingface_hub import HfApi + api = HfApi(token = token) + api.create_repo( + repo_id = repo_id, + repo_type = "model", + private = merge_kwargs.get("private", None), + exist_ok = True, + ) + + # 1) Merge to 16bit at a staging dir (LoRA and base alike). The reload reads default + # weight filenames, so never write variant-named shards here. + merge_kwargs.pop("variant", None) + print(f"Unsloth: Merging to 16bit before torchao {kind} quantization...") + merge_args = dict(merge_kwargs) + merge_args.update( + dict( + model = model, + tokenizer = tokenizer, + save_directory = staging, + save_method = "merged_16bit", + push_to_hub = False, + token = token, + is_main_process = is_main_process, + ) + ) + unsloth_generic_save(**merge_args) + + # 2) Detect VLM + trust_remote_code so the right auto class reloads the staged checkpoint. + # A bare *ForConditionalGeneration also matches text seq2seq (T5/BART/Whisper), so key off + # vision_config / a vision-named architecture only, like the compressed path. + is_vlm = False + trust_remote_code = False + if hasattr(model, "config"): + archs = getattr(model.config, "architectures", None) or [] + is_vlm = hasattr(model.config, "vision_config") or any( + x.endswith("ForVisionText2Text") for x in archs + ) + trust_remote_code = bool(getattr(model.config, "auto_map", None)) + # Custom code can be declared only in the tokenizer/processor config, so also honor an + # auto_map in any staged config (the original load already had the user's consent). + if not trust_remote_code: + for _cfg in ( + "config.json", + "tokenizer_config.json", + "processor_config.json", + "preprocessor_config.json", + ): + try: + _p = os.path.join(staging, _cfg) + if os.path.exists(_p): + with open(_p, "r", encoding = "utf-8") as _f: + if "auto_map" in json.load(_f): + trust_remote_code = True + break + except Exception: + pass + # Reload with the class that matches the checkpoint: an image-text VLM class (with a + # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own + # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and + # AutoModelForCausalLM would fail to load them); otherwise causal-LM. + if is_vlm: + try: + from transformers import AutoModelForImageTextToText as _reload_model + except ImportError: + from transformers import AutoModelForVision2Seq as _reload_model + auto_model = _reload_model + elif getattr(getattr(model, "config", None), "is_encoder_decoder", False): + import transformers as _tf + auto_model = next( + ( + getattr(_tf, _arch) + for _arch in (getattr(model.config, "architectures", None) or []) + if getattr(_tf, _arch, None) is not None + ), + AutoModelForCausalLM, + ) + else: + auto_model = AutoModelForCausalLM + auto_processor = AutoProcessor if is_vlm else AutoTokenizer + + # 3) Free the in-memory model's accelerator memory before reloading a fresh copy from disk. + # Covers CUDA and XPU (torchao runs on Intel GPUs too), so the original doesn't sit + # resident alongside the reloaded copy and OOM a device that fit the model once. + _has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available() + try: + if ( + (torch.cuda.is_available() or _has_xpu) + and hasattr(model, "parameters") + and not getattr(model, "is_loaded_in_4bit", False) + and not getattr(model, "is_loaded_in_8bit", False) + and not getattr(model, "is_quantized", False) + ): + _devs = {str(p.device) for p in model.parameters()} + if len(_devs) == 1 and next(iter(_devs)).startswith(("cuda", "xpu")): + _dev = next(model.parameters()).device + model.to("cpu") + model_dev = _dev + except Exception: + model_dev = None + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if _has_xpu: + torch.xpu.empty_cache() + + # 4) Reload the staged 16bit checkpoint with torchao applied. bfloat16 is required; + # device_map="auto" falls back to CPU, so this works on any hardware. + print(f"Unsloth: Quantizing the merged model to torchao {kind}...") + dtype_kw = {"torch_dtype": torch.bfloat16} if HAS_TORCH_DTYPE else {"dtype": torch.bfloat16} + quantized_model = auto_model.from_pretrained( + staging, + device_map = "auto", + quantization_config = TorchAoConfig(quant_type = quant_type), + trust_remote_code = trust_remote_code, + **dtype_kw, + ) + staged_tokenizer = auto_processor.from_pretrained( + staging, trust_remote_code = trust_remote_code + ) + + quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization) + staged_tokenizer.save_pretrained(out_dir) + del quantized_model + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # 5) Validate the artifact. + cfg_path = os.path.join(out_dir, "config.json") + cfg = {} + if os.path.exists(cfg_path): + with open(cfg_path, "r", encoding = "utf-8") as f: + cfg = json.load(f) + if "quantization_config" not in cfg: + raise RuntimeError( + f"Unsloth: torchao {kind} export failed - no quantization_config written to " + f"{cfg_path}" + ) + + # 6) Optional hub upload of the quantized artifact (the temp staging is cleaned in finally). + if push_to_hub: + print(f"Unsloth: Uploading torchao {kind} checkpoint to '{repo_id}' ...") + api.upload_folder( + folder_path = out_dir, + repo_id = repo_id, + repo_type = "model", + commit_message = merge_kwargs.get("commit_message", None), + commit_description = merge_kwargs.get("commit_description", None), + create_pr = merge_kwargs.get("create_pr", False), + revision = merge_kwargs.get("revision", None), + ) + datasets = merge_kwargs.get("datasets", None) + if datasets: + try: + from huggingface_hub import metadata_update + metadata_update(repo_id, {"datasets": datasets}, overwrite = True, token = token) + except Exception as meta_err: + logger.warning_once( + f"Unsloth: could not update datasets metadata for {repo_id}: {meta_err}" + ) + + result = repo_id if push_to_hub else out_dir + print( + f"Unsloth: Saved torchao {kind} checkpoint to '{result}'.\n" + f"Unsloth: This is portable (produced on any device, no NVIDIA GPU required). Load it " + f"with vLLM or transformers; FP8/INT8 acceleration is available on supported GPUs." + ) + return result + finally: + if model_dev is not None: + try: + model.to(model_dev) + except Exception: + logger.warning_once( + "Unsloth: could not restore the model to its original device after torchao " + "export; it may remain on CPU." + ) + if work_tmp is not None: + shutil.rmtree(work_tmp, ignore_errors = True) + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def unsloth_save_pretrained_torchao( self, save_directory: Union[str, os.PathLike], From fbb5b0968cd4318e483fc9918f0c592bf325c178 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Fri, 3 Jul 2026 11:26:38 -0400 Subject: [PATCH 17/23] Studio: flush passthrough stream headers before upstream prefill stalls (#6835) * Studio: flush passthrough stream headers before upstream prefill stalls * Studio: clean up delayed passthrough send failures * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: close passthrough preheader cleanup gaps * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: retry delayed passthrough overflow truncation * Studio: close completed passthrough send responses * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/routes/inference.py | 132 ++++- .../tests/test_openai_tool_passthrough.py | 547 ++++++++++++++++++ 2 files changed, 669 insertions(+), 10 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ccf36e8f71..17be222d93 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -784,6 +784,7 @@ def _set_stream_response_read_timeout( _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 +_OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 class _CompatSameTaskTimeout: @@ -10890,41 +10891,73 @@ async def _openai_passthrough_stream( _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + client = None + resp = None + send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None + + async def _aclose_send_task(task: Optional[asyncio.Task[Optional[httpx.Response]]]) -> None: + if task is None: + return + if not task.done(): + task.cancel() + try: + task_resp = await task + if task_resp is not None: + try: + await task_resp.aclose() + except Exception: + pass + except (asyncio.CancelledError, Exception): + pass # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: - # Dispatch BEFORE returning StreamingResponse so transport errors and - # non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs - # rely on status codes to raise APIError/BadRequestError. + # Keep the pre-header window short so accepted SSE clients receive + # immediate headers in the common timeout-reduced stall. client = httpx.AsyncClient( timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), trust_env = False, ) - resp = None _truncate_budget = ( _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 ) + while True: try: req = client.build_request( "POST", target_url, json = body, headers = {"Connection": "close"} ) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - resp = await _send_stream_with_preheader_cancel( - client, req, cancel_event, request = request + send_task = asyncio.create_task( + _send_stream_with_preheader_cancel(client, req, cancel_event, request = request) ) + done, _ = await asyncio.wait( + {send_task}, + timeout = _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task not in done: + break + + # Dispatch returned quickly enough to preserve pre-header status. + resp = await send_task + send_task = None except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) + await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) raise HTTPException( status_code = 502, detail = _friendly_error(e), ) + if resp is None and send_task is not None and not send_task.done(): + break if resp is None: api_monitor.finish(monitor_id, "cancelled") + await _aclose_send_task(send_task) try: await client.aclose() except Exception: @@ -10969,6 +11002,8 @@ async def _openai_passthrough_stream( api_monitor.fail(monitor_id, err_text[:500]) raise _openai_passthrough_error(upstream_status, err_text) + # Keep tracker cleanup paired if pre-header dispatch is cancelled after we + # have already committed headers. async def _stream(): # Same httpx lifecycle pattern as _anthropic_passthrough_stream: # save resp.aiter_lines() so the finally block can aclose() it on @@ -10976,10 +11011,10 @@ async def _openai_passthrough_stream( lines_iter = None # Watchers unblock aiter_lines() during prefill, before in-loop # cancel/disconnect checks can run. - cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) - disconnect_watcher = asyncio.create_task( - _await_disconnect_then_close(request, resp, cancel_event) - ) + cancel_watcher = None + disconnect_watcher = None + + nonlocal resp, send_task, first_token_deadline, _truncate_budget monitor_done = False saw_finish_reason = False saw_done = False @@ -11112,6 +11147,79 @@ async def _openai_passthrough_stream( return lines try: + while True: + if send_task is not None and not send_task.done(): + try: + resp = await send_task + except httpx.RequestError as e: + logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" + return + send_task = None + elif send_task is not None: + try: + resp = send_task.result() + except httpx.RequestError as e: + logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" + return + send_task = None + + if resp is None: + api_monitor.finish(monitor_id, "cancelled") + return + if resp.status_code == 200: + break + + err_bytes = await resp.aread() + err_text = err_bytes.decode("utf-8", errors = "replace") + logger.error( + "openai passthrough upstream error: status=%s body=%s", + resp.status_code, + err_text[:500], + ) + upstream_status = resp.status_code + try: + await resp.aclose() + except Exception: + pass + resp = None + if ( + _truncate_budget > 0 + and _classify_llama_generation_error(Exception(err_text)) + and _apply_overflow_truncation(body, err_text) + ): + _truncate_budget -= 1 + req = client.build_request( + "POST", target_url, json = body, headers = {"Connection": "close"} + ) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + send_task = asyncio.create_task( + _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) + ) + continue + + upstream_error = _openai_passthrough_error(upstream_status, err_text) + error_payload = ( + upstream_error.detail + if isinstance(upstream_error.detail, dict) + else openai_error_body( + str(upstream_error.detail), + status = upstream_status, + ) + ) + api_monitor.fail(monitor_id, err_text[:500]) + yield f"data: {json.dumps(error_payload)}\n\n" + return + + cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_close(request, resp, cancel_event) + ) lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( lines_iter, @@ -11298,6 +11406,7 @@ async def _openai_passthrough_stream( err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: + await _aclose_send_task(send_task) await _aclose_stream_resources( watchers = (cancel_watcher, disconnect_watcher), iterator = lines_iter, @@ -11311,6 +11420,7 @@ async def _openai_passthrough_stream( # finally never ran. Release the eagerly-opened upstream resp/client # and the cancel-registry entry here; the watchers and line iterator # are created inside _stream(), so there is nothing else to close. + await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) @@ -11325,6 +11435,8 @@ async def _openai_passthrough_stream( unstarted_cleanup = _unstarted_cleanup, ) except BaseException: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) raise diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 1d725acd45..ccbd78e2b1 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1856,6 +1856,553 @@ class TestApiMonitorProviderAndCompletionStreams: chunks = [chunk async for chunk in response.body_iterator] return SimpleNamespace(chunks = chunks, body = "".join(chunks), monitor = monitor) + def test_passthrough_stream_preheader_dispatched_with_timeout(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + assert "data: [DONE]\n\n" in "".join(chunks) + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + return httpx.Response(400, content = b'{"error":"bad"}') + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + assert exc.value.status_code == 400 + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_request_error_in_window(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + raise httpx.ConnectError("connectivity issue") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + assert exc.value.status_code == 502 + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_non_200_returns_sse_error(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(400, content = b'{"error":"bad"}') + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data:" in body + assert '"error"' in body + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "bad" in entry["error"] + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_context_error_keeps_error_envelope( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)" + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(400, content = ctx_msg.encode("utf-8")) + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + payload = json.loads(body.removeprefix("data: ").strip()) + assert payload["error"]["code"] == "context_length_exceeded" + assert payload["error"]["param"] == "messages" + assert isinstance(payload["error"], dict) + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_context_error_retries_truncation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + calls = [] + err_body = json.dumps( + { + "error": { + "message": "request (10000 tokens) exceeds the available context size (2048 tokens)", + "n_prompt_tokens": 10000, + "n_ctx": 2048, + } + } + ).encode("utf-8") + + async def fake_send(_client, req, *_args, **_kwargs): + calls.append(json.loads(req.content.decode("utf-8"))) + if len(calls) == 1: + await gate.wait() + return httpx.Response(400, content = err_body) + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + messages = [ + ChatMessage(role = "system", content = "system"), + *[ + ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000)) + for idx in range(8) + ], + ] + payload = ChatCompletionRequest( + model = "default", + messages = messages, + stream = True, + context_overflow = "truncate_middle", + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + assert "data: [DONE]\n\n" in "".join(chunks) + assert len(calls) == 2 + assert len(calls[1]["messages"]) < len(calls[0]["messages"]) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + cancel_id = "delayed-request-error-cancel" + + async def fake_send(*_args, **_kwargs): + await gate.wait() + raise httpx.ConnectError("delayed connectivity issue") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data:" in body + assert '"error"' in body + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "Lost connection" in entry["error"] + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_cancel_cleans_pending_send(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + entered = asyncio.Event() + cancelled = asyncio.Event() + cancel_id = "preheader-cancel-cleanup" + + async def fake_send(*_args, **_kwargs): + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + task = asyncio.create_task( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + ) + await asyncio.wait_for(entered.wait(), timeout = 0.2) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(cancelled.wait(), timeout = 0.2) + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_unstarted_cleanup_closes_completed_send_response(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + returned = asyncio.Event() + cancel_id = "unstarted-completed-send-cleanup" + + class Stream(httpx.AsyncByteStream): + async def __aiter__(self): + if False: + yield b"" + + stream = Stream() + upstream_response = httpx.Response(200, stream = stream) + + async def fake_send(*_args, **_kwargs): + await gate.wait() + returned.set() + return upstream_response + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + await asyncio.wait_for(returned.wait(), timeout = 0.2) + await asyncio.sleep(0) + await response._unstarted_cleanup() + assert upstream_response.is_closed + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + def test_external_non_streaming_json_updates_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod From 2b06616a7eebe84535b9afc7ce4f37f45fe456e7 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 3 Jul 2026 12:35:02 -0300 Subject: [PATCH 18/23] Fix TrainingArguments silently disabling unsloth gradient checkpointing (#6829) * Fix TrainingArguments silently disabling unsloth gradient checkpointing * Cover loaded adapters and preserve explicit None in GC restore * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_gradient_checkpointing_restore.py | 185 +++++++++++++++++++ unsloth/models/llama.py | 8 + unsloth/models/rl.py | 16 +- unsloth/models/rl_replacements.py | 3 +- unsloth/models/vision.py | 5 + 5 files changed, 212 insertions(+), 5 deletions(-) create mode 100644 tests/test_gradient_checkpointing_restore.py diff --git a/tests/test_gradient_checkpointing_restore.py b/tests/test_gradient_checkpointing_restore.py new file mode 100644 index 0000000000..4f9f3faccc --- /dev/null +++ b/tests/test_gradient_checkpointing_restore.py @@ -0,0 +1,185 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Regression for #4735: a plain ``TrainingArguments`` silently disabling the +gradient-checkpointing (GC) mode the model was configured with at setup. + +Setup records the effective GC mode as ``_unsloth_gradient_checkpointing``; the +trainer restores *that* value, falling back to ``args.gradient_checkpointing`` +only when nothing was recorded. The restore lines live inside exec'd template +strings, which ``py_compile`` never sees, so these tests pull the real snippets +out of the source and execute them against fakes. GPU-free. +""" + +from __future__ import annotations + +import ast +import re +from pathlib import Path + +_ROOT = Path(__file__).resolve().parent.parent / "unsloth" / "models" +_RL = (_ROOT / "rl.py").read_text() +_RL_REPLACEMENTS = (_ROOT / "rl_replacements.py").read_text() + +# The single-line ternary form used at the trainer call sites: +# ._unsloth_gradient_checkpointing if hasattr(, '...') else getattr(, 'gradient_checkpointing', True) +_TERNARY = re.compile( + r"(?P[\w.]+)\._unsloth_gradient_checkpointing " + r"if hasattr\((?P=model), '_unsloth_gradient_checkpointing'\) " + r"else getattr\((?P[\w.]+), 'gradient_checkpointing', True\)" +) + +_MISSING = object() + + +class _Obj: + """Bare attribute bag; ``_unsloth_gradient_checkpointing`` present only when recorded.""" + + def __init__( + self, + recorded = _MISSING, + gradient_checkpointing = _MISSING, + ): + if recorded is not _MISSING: + self._unsloth_gradient_checkpointing = recorded + if gradient_checkpointing is not _MISSING: + self.gradient_checkpointing = gradient_checkpointing + + +class _Self: + def __init__( + self, + model = None, + args = None, + ): + if model is not None: + self.model = model + self.args = args + + +# (recorded on model, args.gradient_checkpointing, expected restored value) +# The point of the fix: a recorded mode wins over args, and a recorded ``None`` +# (a valid setup value) is restored verbatim rather than collapsing to the +# args fallback the way a ``None`` sentinel would. +_MATRIX = [ + ("unsloth", False, "unsloth"), # the #4735 case: args=False must NOT win + (True, False, True), + (False, True, False), # user turned GC off; args=True must NOT re-enable it + (None, True, None), # explicit None is restored, not treated as "unrecorded" + (_MISSING, True, True), # nothing recorded -> fall back to args + (_MISSING, False, False), +] + + +def _eval_ternary(expr, recorded, args_gc): + """Eval a restore expression that references either ``model``/``args`` or ``self.model``/``self.args``.""" + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + self = _Self(model = model, args = args) + return eval( + expr, {"hasattr": hasattr, "getattr": getattr}, {"model": model, "args": args, "self": self} + ) + + +def test_ternary_restore_semantics(): + exprs = [m.group(0) for m in _TERNARY.finditer(_RL)] + exprs += [m.group(0) for m in _TERNARY.finditer(_RL_REPLACEMENTS)] + # Also guards against the lines being deleted/renamed (which reinstates the bug). + assert len(exprs) >= 3, f"expected the 3 trainer-call restore sites, found {len(exprs)}" + for expr in exprs: + for recorded, args_gc, expected in _MATRIX: + got = _eval_ternary(expr, recorded, args_gc) + assert got == expected and type(got) is type( + expected + ), f"{expr!r}: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def _extract_prepare_restore_block(): + """Pull the multi-line restore block out of ``prepare_for_training_mode``'s wrapper. + + It lives inside an exec'd template string, so grab it textually: from the + ``_model = getattr(self, 'model', None)`` line through the closing + ``else:``/``use_gc = ...`` pair. + """ + lines = _RL.splitlines() + start = next( + i for i, l in enumerate(lines) if l.strip() == "_model = getattr(self, 'model', None)" + ) + # End at the fallback assignment rather than a fixed line count, so inserting + # lines into the block can't silently truncate what gets exec'd. + end = next( + i + for i, l in enumerate(lines) + if i > start and "use_gc = getattr(self.args, 'gradient_checkpointing', True)" in l + ) + block = lines[start : end + 1] + # dedent to column 0 so it execs as a top-level block + indent = len(block[0]) - len(block[0].lstrip()) + return "\n".join(l[indent:] for l in block) + + +def test_prepare_for_training_mode_block_semantics(): + block = _extract_prepare_restore_block() + # Must be valid Python (it's never seen by py_compile in the outer file). + ast.parse(block) + + for recorded, args_gc, expected in _MATRIX: + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + ns = {"self": _Self(model = model, args = args), "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + got = ns["use_gc"] + assert ( + got == expected and type(got) is type(expected) + ), f"prepare block: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def test_prepare_block_tolerates_missing_model(): + # gemini flagged the unguarded self.model access: the block reads self.model via + # getattr(self, 'model', None), so a trainer without a .model attribute must fall + # back to args rather than raising AttributeError. + block = _extract_prepare_restore_block() + args = _Obj(gradient_checkpointing = True) + self_no_model = _Self(model = None, args = args) # _Self leaves .model unset when model is None + assert not hasattr(self_no_model, "model") + ns = {"self": self_no_model, "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + assert ns["use_gc"] is True + + +def test_recording_sites_are_real_module_code(): + # The recording side (unlike the restore side) is real module code, not a template + # string. Assert it's present at the choke point (patch_peft_model, so loaded adapters + # are covered) and at the pre-wrapped pass-through, both of which bypass the old + # get_peft_model-only recording. + llama = (_ROOT / "llama.py").read_text() + tree = ast.parse(llama) + + def assigns_marker(node): + return any( + isinstance(n, ast.Assign) + and any( + isinstance(t, ast.Attribute) and t.attr == "_unsloth_gradient_checkpointing" + for t in n.targets + ) + for n in ast.walk(node) + ) + + fns = {n.name: n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + assert "patch_peft_model" in fns and assigns_marker( + fns["patch_peft_model"] + ), "patch_peft_model must record _unsloth_gradient_checkpointing so loaded adapters are covered" + # The pass-through branch lives in get_peft_model. + assert assigns_marker( + fns["get_peft_model"] + ), "get_peft_model pass-through must record _unsloth_gradient_checkpointing" diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 14ee5ee24e..bb7289dfa8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3047,6 +3047,9 @@ class FastLlamaModel: # Pre-wrapped PEFT model passes through here; still arm the detector so an RL # trainer can reset a compile cache poisoned by a pre-train forward. _unsloth_install_pretrain_detector(model) + # This branch returns before patch_peft_model, so record here too; + # apply_unsloth_gradient_checkpointing above already re-patched global state to match (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing model = _exclude_rope_inv_freq_from_ddp(model) return model else: @@ -3406,6 +3409,11 @@ class FastLlamaModel: @staticmethod def patch_peft_model(model, use_gradient_checkpointing = "unsloth"): + # Persist the effective GC mode so the trainer restores it verbatim: for_inference() + # clears the module flags every GRPO step, and a plain TrainingArguments defaults it to + # False, which would otherwise silently disable it at train time (#4735). Recorded here, + # not in get_peft_model, so adapters loaded via loader.py's from_pretrained path are covered. + model._unsloth_gradient_checkpointing = use_gradient_checkpointing if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1": return FastBaseModel.patch_peft_model( model = model, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 602de69d3f..62ef9e916a 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -423,8 +423,14 @@ def prepare_for_training_mode(f): pass # Enable training mode _was_training = None - # Get gradient checkpointing setting from training arguments - use_gc = getattr(self.args, 'gradient_checkpointing', True) + # Restore the GC mode the model was configured with at setup; fall back to + # the training args only when it wasn't recorded (issue #4735). Use hasattr, + # not a None sentinel, so a deliberately-recorded None is restored verbatim. + _model = getattr(self, 'model', None) + if hasattr(_model, '_unsloth_gradient_checkpointing'): + use_gc = _model._unsloth_gradient_checkpointing + else: + use_gc = getattr(self.args, 'gradient_checkpointing', True) if hasattr(self, 'model') and hasattr(self.model, "training"): _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): @@ -532,7 +538,8 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): if getattr(args, "_n_gpu", 1) != 1: args._n_gpu = 1 if "model" in locals() and hasattr(model, "for_training"): - model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) + model.for_training(use_gradient_checkpointing=_use_gc) super().__init__({RLTrainer_call_args}{RLTrainer_kwargs}) if "model" in locals() and hasattr(model, "for_inference"): model.for_inference() @@ -1165,7 +1172,8 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): if "model" in call_args: training_check = ( "if model is not None and hasattr(model, 'for_training'):\n" - " model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))\n" + " _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True)\n" + " model.for_training(use_gradient_checkpointing=_use_gc)\n" "if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'\n" "if 'processing_class' in locals():\n" " if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'\n" diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index d3ada23cf9..3be614cf4a 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -761,7 +761,8 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # Left pad prompt before calculation old and ref hidden states left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() - self.model.for_training(use_gradient_checkpointing=getattr(self.args, 'gradient_checkpointing', True))""" + _use_gc = self.model._unsloth_gradient_checkpointing if hasattr(self.model, '_unsloth_gradient_checkpointing') else getattr(self.args, 'gradient_checkpointing', True) + self.model.for_training(use_gradient_checkpointing=_use_gc)""" function = function.replace(line_to_replace, replacement_lines) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 5ab55152db..689e362f95 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1873,6 +1873,11 @@ class FastBaseModel: float32_mixed_precision = float32_mixed_precision, patch_modules_to_save = True, ) + # Persist the configured GC mode so the trainer restores it verbatim. + # for_inference() clears the module flags (GRPO does this every generation + # step), and a plain TrainingArguments defaults gradient_checkpointing=False, + # which would otherwise silently disable this setting at train time (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing # Gemma3N audio conformer processes variable-length audio tensors # that cause stride mismatches in AOT autograd compiled backward From 9c2eacc35e3f5f3f33af3b2c3cad42a8dd4c5ec2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 09:07:30 -0700 Subject: [PATCH 19/23] Studio: reserve CUDA context and mmproj/MTP soft overhead in the GGUF fit budget (#6718) --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 337 ++++++++++++++++-- studio/backend/tests/test_compute_buffer.py | 142 +++++++- studio/backend/tests/test_slot_offload_fit.py | 115 ++++++ studio/backend/tests/test_tensor_parallel.py | 79 +++- 4 files changed, 638 insertions(+), 35 deletions(-) create mode 100644 studio/backend/tests/test_slot_offload_fit.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 035e5d12c7..3ccfc5cdfe 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -801,7 +801,11 @@ _MTP_MIN_SIZE_B = 3.0 # Cap total GPU occupancy at this fraction of the card. The fit reserves an # absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction # of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve. -_CTX_FIT_VRAM_FRACTION = 0.95 +# 3%: the context-linear compute buffer is now modelled (_compute_buffer_ctx_bytes), +# so this cushion no longer covers it - only fragmentation, the per-device CUDA +# context on a multi-GPU split, and MoE routing, which measure ~2-3% (Qwen3.5-397B on +# 3 GPUs under-predicts by 2.7%). Below 3% one fragmentation spike overflows to CPU. +_CTX_FIT_VRAM_FRACTION = 0.97 # Apple unified memory is shared with the OS, so tighter than VRAM. Matches the # 0.85 MLX uses in mlx_inference.py (_configure_memory_limits); not kept in sync. @@ -2464,9 +2468,10 @@ class LlamaCppBackend: prev = curr # Free-VRAM fraction at which Studio pins the GPU directly instead of - # deferring to ``--fit on``. 5% headroom covers CUDA context + compute - # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). - _GPU_PIN_VRAM_FRACTION = 0.95 + # deferring to ``--fit on``. 3% headroom: the compute buffer is now modelled in + # the fit, so this only guards fragmentation + multi-GPU per-device CUDA context + # (~2-3%); kept >= 3% as a floor (0.90 dropped 91-94% fits to CPU offload, #5106). + _GPU_PIN_VRAM_FRACTION = 0.97 # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived @@ -3022,6 +3027,27 @@ class LlamaCppBackend: _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate + # Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682). + _CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB) + _MMPROJ_VRAM_SAFETY = 1.4 # mmproj worst-case buffer vs file size (runtime ~1.3x) + _MTP_DRAFT_COMPUTE_BYTES = 224 * 1024 * 1024 # MTP draft decode graph beyond its KV + # The flash-attn KQ mask + attention scratch grow ~linearly with context; the flat + # _estimate_compute_buffer_bytes term only covers ctx -> 0. The per-token rate + # depends on the KV cache type: a QUANTIZED cache (q8_0/q5/q4/iq4) needs a + # context-sized dequant scratch that scales with n_embd, measured at 0.74-2.02 x + # n_embd across Qwen3.5/3.6 (2B/4B/9B/27B) and Gemma-4 (12B/31B) at q8_0; an + # f16/bf16/f32 cache skips the dequant and pays only the KQ mask, a flat n_ubatch*2 + # bytes per context token regardless of n_embd (measured 1024 B/tok on Qwen-9B and + # Gemma-31B alike). So Qwen3.5-4B at 256k is 1.30 GiB at q8_0 vs 0.31 GiB at f16. + # 2.25 covers the worst quantized case (Qwen3.5-4B, ~2.0x) plus the under-modeled + # flat base; the mask safety covers the f16 base gap. Without this term, tight tiers + # at extreme context over-pin and spill to CPU (the 3% cushion is only ~0.25 GiB on + # an 8 GB card, far below the ~1-2.4 GiB quantized buffer at 256k): e.g. Qwen3.5-4B + # Q4 at 256k needs ~8.5 GiB on a real 8 GB card (weights 2.4 + KV 4.3 + compute 1.3 + # + CUDA ctx) -> CPU spill; with this reserve the auto context caps to ~210k, fits. + _CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch) + _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) + _CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) def _estimate_compute_buffer_bytes( self, @@ -3052,6 +3078,85 @@ class LlamaCppBackend: compute = act_scratch + out_buffer * max(0, par - 1) return int(compute * self._COMPUTE_BUFFER_SAFETY) + def _compute_buffer_ctx_bytes( + self, + n_ctx: int, + n_ubatch: Optional[int] = None, + cache_type_kv: Optional[str] = None, + ) -> int: + """Context-linear growth of the per-device compute buffer (bytes), charged + on top of the flat ``_estimate_compute_buffer_bytes``. The flash-attn KQ + mask + attention scratch scale ~linearly with context and with the micro- + batch; the flat term only covers ctx -> 0. A quantized KV cache adds a + context-sized dequant scratch that scales with n_embd; f16/bf16/f32 pays only + the KQ mask, a flat n_ubatch*2 bytes per context token. ``cache_type_kv`` None + -> f16 (llama.cpp's default; an env-set quantized cache is budgeted as f16 on + the KV side, whose over-reservation absorbs the dequant scratch). Returns 0 + when dims are missing or ``n_ctx`` <= 0.""" + n_embd = self._embedding_length or 0 + if n_embd <= 0 or n_ctx <= 0: + return 0 + ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + if _kv_bytes_per_elem(cache_type_kv) < 2.0: + # Quantized cache: the dequant scratch dominates and scales with n_embd. + # MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on + # GLM-5.2 and Kimi-K2.7 vs up to 2.02x on regular attention. + ub_scale = ub / self._DEFAULT_N_UBATCH + rate = ( + self._CTX_COMPUTE_BYTES_PER_EMBD_MLA + if self._key_length_mla + else self._CTX_COMPUTE_BYTES_PER_EMBD + ) + per_tok = rate * n_embd * ub_scale + else: + # f16/bf16/f32: only the KQ mask ([n_kv, n_ubatch] f16), n_embd-independent. + per_tok = ub * 2 * self._CTX_COMPUTE_F16_MASK_SAFETY + return int(per_tok * n_ctx) + + def _slots_that_fit_on_gpu( + self, + n_parallel: int, + effective_ctx: int, + gpus: list[tuple[int, int]], + total_by_idx: Optional[dict[int, int]], + base_footprint_bytes: int, + cache_type_kv: Optional[str], + pin_fraction: float, + per_device_overhead_bytes: int, + min_gpus: int, + n_ubatch: Optional[int] = None, + ) -> tuple[Optional[list[int]], bool, int]: + """Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits, + so Studio keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers + to host and collapses decode ~3x (oobabooga #6718). ``base_footprint_bytes`` is the + slot-independent footprint (weights + soft overhead + MTP + context-linear compute, + minus the folded compute buffer); each candidate re-adds the slot-sized compute buffer + and KV, then re-selects GPUs like the explicit-context path. Returns (gpu_indices, + use_fit=False, slots) for the largest fitting count, else (None, True, n_parallel). + Only ever reduces; deterministic and unit-testable with synthetic VRAM maps.""" + for slots in range(n_parallel - 1, 0, -1): + cb = self._estimate_compute_buffer_bytes( + n_ubatch = n_ubatch, n_parallel = slots, per_device_tensor = False + ) + if cb <= 0: + cb = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 + total = ( + base_footprint_bytes + + cb + + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots) + ) + gpu_indices, use_fit = self._select_gpus( + total, + gpus, + usable_fraction = pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = per_device_overhead_bytes, + min_gpus = min_gpus, + ) + if not use_fit: + return gpu_indices, False, slots + return None, True, n_parallel + def _fit_context_to_vram( self, requested_ctx: int, @@ -3067,6 +3172,7 @@ class LlamaCppBackend: kv_on_gpu: bool = True, mtp_engaged: bool = False, mtp_overhead_fn: Optional[Callable[[int], int]] = None, + compute_ctx_bytes_fn: Optional[Callable[[int], int]] = None, budget_frac: Optional[float] = None, total_mib: Optional[int] = None, ) -> int: @@ -3118,9 +3224,14 @@ class LlamaCppBackend: def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + def _cc_at(ctx: int) -> int: + # Context-linear compute-buffer growth (flash-attn KQ mask + scratch); + # the flat term in model_footprint only covers ctx -> 0. + return compute_ctx_bytes_fn(ctx) if compute_ctx_bytes_fn is not None else 0 + # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes: + if model_footprint + kv + _mtp_at(requested_ctx) + _cc_at(requested_ctx) <= budget_bytes: return requested_ctx # Weights + compute buffer alone exceed budget -- reducing ctx can't help. @@ -3141,7 +3252,7 @@ class LlamaCppBackend: while lo <= hi: mid = (lo + hi) // 2 kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) - if kv + _mtp_at(mid) <= remaining: + if kv + _mtp_at(mid) + _cc_at(mid) <= remaining: best = mid lo = mid + 1 else: @@ -4288,6 +4399,7 @@ class LlamaCppBackend: max_target_ctx: Optional[int] = None, total_by_idx: Optional[dict[int, int]] = None, n_ubatch: Optional[int] = None, + soft_overhead_bytes: int = 0, ) -> tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -4299,9 +4411,11 @@ class LlamaCppBackend: ``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``. Policy (assumes >= 2 GPUs; the caller drops the toggle below that): - - Cap context to the KV that fits the pooled VRAM after the weights and - one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``, - deterministic from dims; flat fallback when dims are unavailable). + - Cap context to the KV that fits the pooled VRAM after the weights, one + per-device flat compute-graph buffer (``_estimate_compute_buffer_bytes``, + deterministic from dims; flat fallback when dims are unavailable), and the + per-device context-linear compute growth (``_compute_buffer_ctx_bytes``, + replicated on every device in tensor mode, so summed over the split). llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only cap, honored even for an explicit ``-c``. It is more accurate than the 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. @@ -4310,7 +4424,9 @@ class LlamaCppBackend: share fits the smallest GPU; otherwise it is weighted by usable budget so the roomier GPU absorbs more weight and the smallest keeps room for KV. ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes - the compute buffer. + the compute buffer. ``soft_overhead_bytes`` is the CUDA-context / mmproj / + MTP-draft-graph reserve the layer path folds into ``model_size_fit``; + charged against the pooled budget so tensor mode reserves the same overhead. """ # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a @@ -4356,16 +4472,40 @@ class LlamaCppBackend: flat_mtp_bytes = max(0, mtp_flat_reserve_bytes) if mtp_engaged and mtp_overhead_fn is None: flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3) + # soft_overhead_bytes is the CUDA-context / mmproj / MTP-draft-graph reserve + # the layer path folds into model_size_fit. Tensor mode has no --fit valve, so + # an unreserved overshoot OOMs at startup rather than offloading; charge it here + # too. Once (pooled), mirroring the layer path -- the per-device CUDA context is + # a known slight under-charge, left for real multi-GPU data. kv_budget_b = ( - (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes + (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 + - model_size + - flat_mtp_bytes + - max(0, soft_overhead_bytes) ) def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + # Context-linear compute buffer, summed over the split. Tensor mode + # replicates the compute graph on EVERY device (measured: the per-device + # buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at + # f16, independent of n_embd), so the growth is n_dev x the per-device + # term. cache_type_kv here is always non-quantized (tensor forces f16), so + # _compute_buffer_ctx_bytes returns the light KQ-mask term, not the heavy + # quantized dequant scratch. The flat reserve_mib above only covers ctx->0; + # without this the fit over-pins and OOMs at high context on a tight pool + # (0.5-4 GiB unreserved at 262k-1M across 2-4 GPUs), the tensor-mode analog + # of the layer-split compute bug. + n_dev = len(gpu_indices) + + def _cc_ctx(ctx: int) -> int: + return n_dev * self._compute_buffer_ctx_bytes(ctx, n_ubatch, cache_type_kv) + def _fit_ctx(ctx: int) -> int: - # Largest context whose KV (+ MTP draft reserve) fits the pooled - # budget. Floors small, but never raises an explicit ctx above asked. + # Largest context whose KV (+ MTP draft reserve + context-linear + # compute) fits the pooled budget. Floors small, but never raises an + # explicit ctx above asked. if self._can_estimate_kv() and ctx > 0: ctx_floor = min(2048, ctx) if kv_budget_b <= 0: @@ -4373,11 +4513,13 @@ class LlamaCppBackend: # falls back to layer split. return ctx_floor if mtp_overhead_fn is not None: - # kv(ctx)+mtp(ctx) is not single-linear, so binary search. + # kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search. def _consumer(c: int) -> int: - return self._estimate_kv_cache_bytes( - c, cache_type_kv, n_parallel = n_parallel - ) + _mtp_at(c) + return ( + self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel) + + _mtp_at(c) + + _cc_ctx(c) + ) if _consumer(ctx) <= kv_budget_b: return ctx @@ -4391,9 +4533,10 @@ class LlamaCppBackend: hi = mid - 1 return best kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) - if kv_at <= kv_budget_b: + total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin + if total_at <= kv_budget_b: return ctx - return max(ctx_floor, int(ctx * kv_budget_b / kv_at)) + return max(ctx_floor, int(ctx * kv_budget_b / total_at)) # KV size unknown -> can't prove a safe cap; floor. return min(4096, ctx) if ctx > 0 else 4096 @@ -4413,10 +4556,23 @@ class LlamaCppBackend: # The MTP reserve also has to fit the even split (mirror the pooled budget): # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes - even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024) + # Context-linear compute is replicated per device; charge the whole split so + # the weighted ratio reflects it (mirrors kv_budget_b's per-device reserve). + cc_bytes = _cc_ctx(effective_ctx) if effective_ctx > 0 else 0 + even_share_mib = ( + (model_size + kv_bytes + mtp_bytes + cc_bytes) / len(gpu_indices) / (1024 * 1024) + ) tensor_split: Optional[list[int]] = None if even_share_mib > (min_usable_mib - reserve_mib): - adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices] + # Each device also holds its replicated share of the context-linear + # compute (cc_bytes/n_dev) on top of the flat reserve. The even-share + # gate above charges cc_bytes; the split weights must subtract it too, or + # the smaller card is weighted above its real usable budget and OOMs (the + # per-device analog of the layer path's per-GPU overhead in _select_gpus). + cc_per_dev_mib = (cc_bytes // len(gpu_indices)) // (1024 * 1024) if cc_bytes else 0 + adj = [ + max(0, int(usable_by_idx[i] - reserve_mib - cc_per_dev_mib)) for i in gpu_indices + ] if sum(adj) > 0: tensor_split = adj return effective_ctx, max_available_ctx, gpu_indices, tensor_split @@ -5179,6 +5335,20 @@ class LlamaCppBackend: # compute buffer); None -> the 512 default in the estimate. _effective_ubatch = _extra_args_n_ubatch(extra_args) + def _cc_bytes(ctx: int, n_gpus: int = 1) -> int: + # Context-linear compute-buffer growth (flash-attn KQ mask + + # attention scratch); the flat _compute_buffer_pipeline folded + # into model_size_fit only covers ctx -> 0. Charged per + # candidate context so the fit can't over-pin and spill. The + # rate depends on the KV cache type (quantized adds a dequant + # scratch), so pass it through. In a layer split this buffer is + # replicated on EVERY device (measured ~equal per GPU), so scale + # by the device count; a large model at high context otherwise + # under-reserves ~(n-1)x it (e.g. Qwen3.5-397B on 3 GPUs). + return max(1, n_gpus) * self._compute_buffer_ctx_bytes( + ctx, _effective_ubatch, cache_type_kv + ) + # Layer-split compute buffer (one lump; tensor mode reserves it # per device in _plan_tensor_parallel). Context-independent, so # fold it into the model footprint for the branches below. Falls @@ -5193,7 +5363,6 @@ class LlamaCppBackend: _compute_buffer_pipeline = ( self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 ) - model_size_fit = model_size + _compute_buffer_pipeline # Layer split adds a fixed per-device overhead on every GPU. The # folded buffer covers one device; reserve the extra devices' @@ -5201,9 +5370,6 @@ class LlamaCppBackend: # (k=1 adds nothing). _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 - def _subset_model_size(n_gpus: int) -> int: - return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes - # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: # honor it, cap only if it fits no combination. Auto (native): # prefer fewer GPUs with reduced context (multi-GPU is slower). @@ -5230,6 +5396,21 @@ class LlamaCppBackend: else 0.0 ) _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve + + # Charge the soft overhead _CTX_FIT_VRAM_FRACTION under-covers on tight + # tiers, gated so plain dense loads (#5106) only pay the CUDA-ctx base. + # CUDA/cuBLAS context is discrete-GPU only (not Metal); the mmproj and + # MTP draft-graph buffers exist on every backend. + _soft_overhead = self._CUDA_CONTEXT_RESERVE_BYTES if gpus else 0 + if effective_is_vision and mmproj_size > 0: + _soft_overhead += int(mmproj_size * (self._MMPROJ_VRAM_SAFETY - 1.0)) + if _mtp_reserves_gpu: + _soft_overhead += self._MTP_DRAFT_COMPUTE_BYTES + model_size_fit = model_size + _compute_buffer_pipeline + _soft_overhead + + def _subset_model_size(n_gpus: int) -> int: + return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + # Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below. _apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024) @@ -5334,7 +5515,9 @@ class LlamaCppBackend: _tp_flat_mtp, _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), ) - _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024) + _tp_required_mib = (model_size + _tp_mtp_floor + _soft_overhead) / ( + 1024 * 1024 + ) if _tp_weight_budget_mib <= _tp_required_mib: logger.info( "Tensor parallelism requested but the pooled VRAM " @@ -5383,6 +5566,7 @@ class LlamaCppBackend: max_target_ctx = self._context_length or target_ctx, total_by_idx = total_by_idx, n_ubatch = _effective_ubatch, + soft_overhead_bytes = _soft_overhead, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -5407,6 +5591,9 @@ class LlamaCppBackend: # budget so the fit and the check below agree. pool_budget = _pool_budget_mib(subset, _cap_fraction) _ms = _subset_model_size(n_gpus) + # Compute buffer is replicated per device in a layer + # split, so scale the context term by the subset size. + _cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n) capped = self._fit_context_to_vram( native_ctx_for_cap, pool_budget, @@ -5415,13 +5602,16 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + footprint_mib = ( + _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) + ) / (1024 * 1024) if footprint_mib <= pool_budget: best_cap = max(best_cap, capped) if best_cap > 0: @@ -5442,13 +5632,18 @@ class LlamaCppBackend: effective_ctx, cache_type_kv, n_parallel = n_parallel ) + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx) ) + # The compute buffer is replicated on every device in a + # layer split; fold it into the per-device reserve so a + # multi-GPU pin sizes each card for its own copy. gpu_indices, use_fit = self._select_gpus( requested_total, gpus, usable_fraction = _pin_fraction, total_by_idx = total_by_idx, - per_device_overhead_bytes = _pipeline_overhead_bytes, + per_device_overhead_bytes = _pipeline_overhead_bytes + + _cc_bytes(effective_ctx), min_gpus = _layer_min_gpus, ) # No silent shrink: effective_ctx stays == requested_ctx. @@ -5479,6 +5674,9 @@ class LlamaCppBackend: subset = ranked[:n_gpus] pool_budget = _pool_budget_mib(subset, pin_fraction) _ms = _subset_model_size(n_gpus) + # Compute buffer is replicated per device in a layer + # split, so scale the context term by the subset size. + _cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n) capped = self._fit_context_to_vram( effective_ctx, pool_budget, @@ -5487,13 +5685,16 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + footprint_mib = ( + _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) + ) / (1024 * 1024) if footprint_mib <= pool_budget: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) @@ -5516,6 +5717,7 @@ class LlamaCppBackend: _subset_model_size(n_gpus) + kv + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx, n_gpus) ) / (1024 * 1024) if footprint_mib <= _pool_budget_mib(subset, pin_fraction): gpu_indices = sorted(idx for idx, _ in subset) @@ -5570,6 +5772,7 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_bytes, budget_frac = 1.0, total_mib = None, ) @@ -5579,6 +5782,7 @@ class LlamaCppBackend: cap, cache_type_kv, n_parallel = n_parallel ) + _mtp_bytes(cap) + + _cc_bytes(cap) ) / (1024 * 1024) # Fit returns the request unchanged when it fits OR weights # exceed budget; only the latter over-commits, so floor to 4096. @@ -5594,6 +5798,48 @@ class LlamaCppBackend: if not explicit_ctx: effective_ctx = max_available_ctx + # Prefer fewer serving slots on GPU over --fit on offload: when the extra + # --parallel slots push the footprint past the pin budget, llama-server + # offloads layers to host and decode collapses ~3x (#6718). Retry the fit + # at fewer slots, keeping the largest count that stays fully on GPU and the + # chosen context. Skips tensor mode / Metal / KV-inestimable paths. + if ( + use_fit + and n_parallel > 1 + and gpus + and self._can_estimate_kv() + and effective_ctx > 0 + ): + # Slot-independent footprint (folded compute buffer swapped out so the + # helper re-adds a slot-sized one per candidate). + _base_footprint = ( + model_size_fit + - _compute_buffer_pipeline + + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx) + ) + _gi_slots, _uf_slots, _slots = self._slots_that_fit_on_gpu( + n_parallel, + effective_ctx, + gpus, + total_by_idx, + _base_footprint, + cache_type_kv, + _pin_fraction, + _pipeline_overhead_bytes + _cc_bytes(effective_ctx), + _layer_min_gpus, + _effective_ubatch, + ) + if not _uf_slots: + logger.info( + "Serving slots reduced %d -> %d to keep the model on GPU " + "(avoid --fit offload) at context %d.", + n_parallel, + _slots, + effective_ctx, + ) + gpu_indices, use_fit, n_parallel = _gi_slots, False, _slots + # MTP reserve at the final context, for the logs below. _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 if _mtp_will_engage: @@ -5692,8 +5938,10 @@ class LlamaCppBackend: if use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: - # Fits on selected GPU(s) -- offload all layers - cmd.extend(["-ngl", "-1"]) + # Fits on selected GPU(s) -- force all layers on GPU. --fit off is + # required: without it llama.cpp's default --fit on second-guesses + # and offloads ~1 GB at --parallel 4 even though the model fits. + cmd.extend(["-ngl", "-1", "--fit", "off"]) fully_gpu_offloaded = True server_caps = self.probe_server_capabilities(binary) @@ -6078,6 +6326,33 @@ class LlamaCppBackend: _split_axis_crash = self._is_tensor_split_assert( "\n".join(self._stdout_lines[-50:]) ) + if ( + _spawn_attempt == 0 + and fully_gpu_offloaded + and _startup_crashed + and not _split_axis_crash + ): + # We forced --fit off because Studio's (conservative) VRAM + # math placed the model fully on GPU. A startup crash here + # means that estimate was optimistic, so fall back to --fit + # on and let llama.cpp offload rather than fail the load. + logger.warning( + "llama-server crashed during startup (exit code %s) " + "with forced --fit off; the fit estimate was optimistic, " + "retrying once with --fit on so it can offload. " + "Crash log: %s", + self._process.returncode, + self._llama_log_path, + ) + # Flip Studio's own --fit off (added first, before any + # user extra args) to on; a user's later --fit still wins + # by last-arg. Defensive: if absent, the default is already + # --fit on, so leave it. + _run = list(run_cmd) + if "--fit" in _run: + _run[_run.index("--fit") + 1] = "on" + run_cmd = _run + continue if ( _spawn_attempt == 0 and _fit_retry_allowed diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 42c400383e..5d14c5c5bd 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -61,11 +61,16 @@ from core.inference.llama_cpp import LlamaCppBackend MIB = 1024 * 1024 -def _backend(vocab = 248320, embd = 5120): +def _backend( + vocab = 248320, + embd = 5120, + mla = None, +): """Backend with just the dims the compute-buffer estimate reads.""" b = LlamaCppBackend.__new__(LlamaCppBackend) b._vocab_size = vocab b._embedding_length = embd + b._key_length_mla = mla # non-None -> MLA (compressed attention) return b @@ -150,3 +155,138 @@ class TestParallel1Default: def test_default_n_parallel(self): est = _backend()._estimate_compute_buffer_bytes() / MIB assert est < 128 + + +class TestContextLinearBuffer: + """``_compute_buffer_ctx_bytes``: the flash-attn KQ-mask + attention scratch + grow ~linearly with context; the flat estimate above only covers ctx -> 0. + Measured slope (q8_0 KV, ubatch 512) was 0.74-2.02 x n_embd; 2 x n_embd is the + worst-case upper bound the term must hold to.""" + + # (model, n_embd, ctx, measured CUDA0 compute buffer MiB at that ctx, q8_0/ub512) + _MEASURED = [ + ("Qwen3.5-2B", 2048, 262144, 796), + ("Qwen3.5-4B", 2560, 262144, 1330), # worst slope, 2.02 x n_embd + ("Qwen3.5-9B", 4096, 262144, 1336), + ("Qwen3.6-27B", 5120, 262144, 1360), + ("Gemma-4-31B", 5376, 262144, 2392), + ] + + def test_zero_by_default(self): + # Omitted/zero ctx -> no term (keeps the flat callers unchanged). + assert _backend()._compute_buffer_ctx_bytes(0) == 0 + + def test_zero_when_embd_missing(self): + assert _backend(embd = None)._compute_buffer_ctx_bytes(262144) == 0 + + def test_grows_linearly_with_context(self): + b = _backend(embd = 4096) + a = b._compute_buffer_ctx_bytes(65536) + d = b._compute_buffer_ctx_bytes(131072) + assert d == pytest.approx(2 * a, rel = 1e-6) + + def test_scales_with_embd(self): + # The quantized (dequant-scratch) rate scales with n_embd; f16 (mask) does not. + small = _backend(embd = 2048)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + big = _backend(embd = 5120)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + assert big > small + + def test_scales_with_ubatch(self): + b = _backend(embd = 4096) + lo = b._compute_buffer_ctx_bytes(131072, n_ubatch = 256) + hi = b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) + assert hi > lo + + @pytest.mark.parametrize("name,embd,ctx,measured", _MEASURED) + def test_upper_bounds_measured_compute_growth(self, name, embd, ctx, measured): + # flat term + context-linear term must cover the real (q8_0) buffer at full ctx. + b = _backend(embd = embd) + flat = b._estimate_compute_buffer_bytes(n_parallel = 1) + total = (flat + b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0")) / MIB + assert total >= measured, f"{name}: under-reserved {total:.0f} < {measured}" + + def test_worst_case_rate_covers_two_x_embd(self): + # >= 2 x n_embd bytes per context token at the default micro-batch (the worst + # measured quantized slope, Qwen3.5-4B), so flat + term upper-bounds the buffer. + embd = 4096 + b = _backend(embd = embd) + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "q8_0") / 100000 + assert per_tok >= 2 * embd + + +class TestContextBufferKVQuant: + """The context-linear rate depends on the KV cache type: a quantized cache adds a + context-sized dequant scratch (heavy); f16/bf16/f32 only pays the KQ mask (light). + Measured Qwen3.5-4B at 256k: 1.30 GiB (q8_0) vs 0.31 GiB (f16).""" + + def test_quantized_heavier_than_f16(self): + b = _backend(embd = 4096) + q = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + f = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16") + assert q > f + + def test_none_cache_type_is_f16(self): + # None -> f16 (llama.cpp's default); the env-quantized case is covered by the + # KV budget's f16 over-reservation, so we take the lighter mask-only rate. + b = _backend(embd = 4096) + assert b._compute_buffer_ctx_bytes( + 131072, cache_type_kv = None + ) == b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16") + + @pytest.mark.parametrize("ct", ["f16", "bf16", "f32"]) + def test_unquantized_uses_mask_only_rate(self, ct): + # f16/bf16/f32: KQ mask only, n_ubatch*2 B/tok, independent of n_embd. + b_small = _backend(embd = 2048) + b_big = _backend(embd = 8192) + per_small = b_small._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + per_big = b_big._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + assert per_small == per_big # no n_embd scaling on the f16 path + expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY # ubatch 512 + assert per_small == pytest.approx(expected, rel = 1e-6) + + @pytest.mark.parametrize("ct", ["q8_0", "q5_1", "q4_0", "iq4_nl"]) + def test_quantized_types_use_heavy_rate(self, ct): + embd = 4096 + b = _backend(embd = embd) + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + assert per_tok == pytest.approx( + LlamaCppBackend._CTX_COMPUTE_BYTES_PER_EMBD * embd, rel = 1e-6 + ) + + def test_f16_covers_measured_mask(self): + # f16 buffer is ~mask only (~n_ubatch*2 B/tok); 0.5 x n_embd must cover the + # measured Qwen3.5-4B f16 slope (~0.4 x n_embd = 0.31 GiB at 256k). + b = _backend(embd = 2560) # Qwen3.5-4B + est = b._compute_buffer_ctx_bytes(262144, cache_type_kv = "f16") / MIB + assert est >= 320 # measured 0.31 GiB growth + + +class TestContextBufferMLA: + """MLA (compressed attention) needs a smaller quantized dequant scratch than + regular attention: measured 0.94 x n_embd on GLM-5.2 and Kimi-K2.7 vs up to + 2.02x on Qwen/Gemma. Charging the regular rate would badly over-reserve a tight + multi-GPU MLA pin (per-device scaling multiplies the error).""" + + def test_mla_lighter_than_regular(self): + reg = _backend(embd = 6144, mla = None)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + mla = _backend(embd = 6144, mla = 256)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + assert mla < reg + + @pytest.mark.parametrize( + "name,embd,ctx,measured", + [ + ("GLM-5.2", 6144, 754688, 4141), # per-device compute MiB at q8_0 + ("Kimi-K2.7", 7168, 262144, 1690), + ], + ) + def test_mla_rate_covers_measured(self, name, embd, ctx, measured): + b = _backend(embd = embd, mla = 256) + est = b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0") / MIB + assert est >= measured, f"{name}: MLA under-reserved {est:.0f} < {measured}" + + def test_mla_not_wildly_over(self): + # 1.25 x n_embd should stay within ~1.6x of the measured 0.94x (not 2.4x like + # the regular 2.25 rate would), so a multi-GPU MLA pin keeps its context. + b = _backend(embd = 6144, mla = 256) + est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB + assert est <= 4141 * 1.7 diff --git a/studio/backend/tests/test_slot_offload_fit.py b/studio/backend/tests/test_slot_offload_fit.py new file mode 100644 index 0000000000..ac606e4627 --- /dev/null +++ b/studio/backend/tests/test_slot_offload_fit.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the offload-avoidance serving-slot reduction (`_slots_that_fit_on_gpu`). + +When a pinned context does not fit at the requested `--parallel` slot count, Studio would +flip to `--fit on` and llama-server offloads layers to host RAM, collapsing decode ~3x +(oobabooga #6718). Instead the loader retries the on-GPU fit at fewer slots and keeps the +largest count that stays fully on GPU (`-ngl -1`). These tests drive the real helper with +synthetic VRAM maps; the KV term is mocked so totals are controlled and the reduction logic +is asserted directly (no GPU, network, or subprocess). +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from core.inference.llama_cpp import LlamaCppBackend + +MIB = 1024 * 1024 +CTX = 90624 +FRAC = LlamaCppBackend._GPU_PIN_VRAM_FRACTION # 0.97; usable = free - 0.03*total + + +def _backend( + vocab = 248320, + embd = 5120, + kv_fixed_mib = 0, +): + """Backend with the dims the compute buffer reads; KV mocked to a fixed size so the + only slot-dependent term is the compute buffer (485 MiB/slot f32 output x 1.15).""" + b = LlamaCppBackend.__new__(LlamaCppBackend) + b._vocab_size = vocab + b._embedding_length = embd + b._key_length_mla = None + b._estimate_kv_cache_bytes = lambda ctx, t = None, **k: kv_fixed_mib * MIB + b._can_estimate_kv = lambda: True + return b + + +def _run( + b, + n_parallel, + base_mib, + gpus, + total_by_idx, + overhead_mib = 0, +): + return b._slots_that_fit_on_gpu( + n_parallel, + CTX, + gpus, + total_by_idx, + int(base_mib * MIB), + "q8_0", + FRAC, + int(overhead_mib * MIB), + 1, + 512, + ) + + +class TestSlotsThatFitOnGpu: + """Compute-buffer per slot (vocab 248320, embd 5120): cb(1)=46, cb(2)=604, cb(3)=1162, + cb(4)=1719 MiB. Single 24 GB card usable = 24576 - 0.03*24576 = 23839 MiB.""" + + def test_reduces_to_largest_fitting_slot(self): + # base+KV = 22500: par4 (24219) over 23839, par3 (23662) fits -> 3 slots on GPU. + gi, use_fit, slots = _run(_backend(), 4, 22500, [(0, 24576)], {0: 24576}) + assert use_fit is False and gi == [0] and slots == 3 + + def test_floor_when_only_one_slot_fits(self): + # base 23400: par2 (24004) over, par1 (23446) fits -> drop all the way to 1. + gi, use_fit, slots = _run(_backend(), 4, 23400, [(0, 24576)], {0: 24576}) + assert use_fit is False and gi == [0] and slots == 1 + + def test_none_fit_stays_offload(self): + # Even a single slot (24046) exceeds usable -> genuine offload, unchanged. + gi, use_fit, slots = _run(_backend(), 4, 24000, [(0, 24576)], {0: 24576}) + assert use_fit is True and gi is None and slots == 4 + + def test_roomy_would_keep_all_but_helper_only_reduces(self): + # On a roomy card par4 fits, so load_model never calls this helper; if called it + # still only searches < n_parallel and never raises the count above the request. + gi, use_fit, slots = _run(_backend(), 4, 5000, [(0, 183000)], {0: 183000}) + assert use_fit is False and slots == 3 and slots < 4 + + def test_single_slot_request_is_noop(self): + # n_parallel == 1: nothing to reduce (range empty) -> report offload unchanged. + gi, use_fit, slots = _run(_backend(), 1, 22500, [(0, 24576)], {0: 24576}) + assert use_fit is True and gi is None and slots == 1 + + def test_multi_gpu_reduces_across_devices(self): + # Needs 2 GPUs: usable/GPU = 23839, cumulative 47677. base+KV 46200: par4 (47919) + # over, par3 (47362) fits across both -> 3 slots spanning [0, 1]. + gi, use_fit, slots = _run( + _backend(), 4, 46200, [(0, 24576), (1, 24576)], {0: 24576, 1: 24576} + ) + assert use_fit is False and gi == [0, 1] and slots == 3 + + def test_kv_counted_per_candidate(self): + # A non-zero (slot-independent) KV shifts the threshold: with 3000 MiB KV and + # base 19500 (= 22500 total at par-independent terms) the same par3 fit holds. + gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576}) + assert use_fit is False and slots == 3 diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 1f09f9a091..0d71b89d87 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -746,10 +746,14 @@ def test_tp_plan_weighted_split_on_asymmetric_big_model(): b, (ec, mac, gi, ts) = _plan(50) reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB assert gi == [0, 1] - # split weighted by (usable - buffer); with no totals usable is free*frac + # split weighted by (usable - flat buffer - per-device context compute); with + # no totals usable is free*frac. The per-device cc is subtracted so the smaller + # card isn't weighted above its real usable budget (see below). + cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024) + assert cc_per_dev > 0 assert ts == [ - int(48000 * _CTX_FIT_VRAM_FRACTION - reserve), - int(24000 * _CTX_FIT_VRAM_FRACTION - reserve), + int(48000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev), + int(24000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev), ] assert ec < 131072 # capped below native @@ -819,6 +823,75 @@ def test_tp_plan_mtp_reserves_extra_and_shrinks_context(): assert ec_mtp < ec_no +def test_tp_plan_reserves_context_linear_compute_buffer(): + # Tensor mode replicates the compute graph on every device; measured on + # Qwen3.5-9B at f16 the per-device buffer grows ~n_ubatch*2 B/token (~1024 + # B/tok), so the fit must reserve n_dev x that on top of the flat reserve or + # it over-pins and OOMs at high context. The chosen KV must leave room for it. + b, (ec, mac, gi, ts) = _plan(50) + cc = len(gi) * b._compute_buffer_ctx_bytes(ec, None, "f16") + assert cc > 0 + assert b._estimate_kv_cache_bytes(ec) + cc <= _kv_budget_b(50) + + +def test_tp_plan_context_shrinks_vs_compute_unaware(): + # With the context-linear term the pinned context is strictly below what a + # KV-only (compute-unaware) fit at the same budget would allow. + b, (ec, *_r) = _plan(50) + b2 = _kv_seeded_backend() + b2._embedding_length = 0 # kills the context-linear compute term (returns 0) + ec_naive, *_r2 = b2._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec < ec_naive + + +def test_tp_plan_soft_overhead_shrinks_context(): + # The CUDA-ctx / mmproj / MTP-draft reserve the layer path folds into the fit + # budget (model_size_fit) must also shrink the tensor context. Tensor mode has + # no --fit valve, so an unreserved overshoot OOMs at startup instead of + # offloading. A non-zero soft_overhead must pin a strictly smaller context. + b = _kv_seeded_backend() + ec_no, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + ec_soft, *_r2 = b._plan_tensor_parallel( + _ASYM, int(50 * _GB), 131072, soft_overhead_bytes = 2 * _GB + ) + assert 2048 < ec_soft < ec_no + + +def test_tp_plan_soft_overhead_reserved_against_budget(): + # The pinned context must leave the whole soft reserve free on top of KV and + # the replicated context compute, so the real footprint stays within the pool. + b = _kv_seeded_backend() + soft = 2 * _GB + ec, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072, soft_overhead_bytes = soft) + cc = len(_ASYM) * b._compute_buffer_ctx_bytes(ec, None, None) + assert b._estimate_kv_cache_bytes(ec) + cc + soft <= _kv_budget_b(50) + + +def test_tp_plan_weighted_split_keeps_small_gpu_within_budget(): + # Regression: the weighted split must subtract each device's replicated context + # compute (cc_bytes/n_dev), not just the flat reserve. Otherwise the smaller + # card is weighted above its usable budget and OOMs at launch. Model the split: + # llama.cpp distributes weights+KV by the tensor-split weights; every device + # also holds the flat reserve plus its per-device context compute. + b, (ec, mac, gi, ts) = _plan(50) + assert ts is not None and len(ts) == len(gi) == 2 + reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024) + free_by_idx = {0: 48000, 1: 24000} + split_content_mib = (int(50 * _GB) + b._estimate_kv_cache_bytes(ec)) / (1024 * 1024) + total_weight = sum(ts) + for w, idx in zip(ts, gi): + placed = split_content_mib * w / total_weight + usable = free_by_idx[idx] * _CTX_FIT_VRAM_FRACTION + assert placed + reserve + cc_per_dev <= usable + 1 # +1 MiB for int rounding + + # Lock the regression: under the old formula (flat reserve only) the smaller + # card was placed over its budget; the cc term is what pulls it back. + old_adj = [int(free_by_idx[i] * _CTX_FIT_VRAM_FRACTION - reserve) for i in gi] + old_small_placed = split_content_mib * old_adj[1] / sum(old_adj) + assert old_small_placed + reserve + cc_per_dev > free_by_idx[1] * _CTX_FIT_VRAM_FRACTION + + def test_tp_plan_no_kv_metadata_floors_context(): b = LlamaCppBackend() # no KV metadata -> can't size safely ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) From 01f7e14988d081643c48286b23d4f1f5d4ed9112 Mon Sep 17 00:00:00 2001 From: ramisworld Date: Sat, 4 Jul 2026 06:10:04 +1200 Subject: [PATCH 20/23] Fix Studio custom folders on Linux external drives (#6799) * Fix external drive custom folder selection * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/backend/tests/test_linux_external_media_paths.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep legacy media scan validation strict * Apply sensitive-dir denylist to legacy folder browser for PR #6799 The legacy /api/models browse endpoint gained the new /run/media mount roots in its allowlist but not the credential/config guard that scan-folder registration and the Hub browser already enforce. Filter sensitive names during enumeration and reject them in _resolve_browse_target so .ssh, .aws, .config, etc. under allowlisted roots stay unbrowseable, matching the Hub browser. Add a public contains_sensitive_path_component helper and cover the legacy resolver with a regression test. * Trim redundant comments in PR #6799 changes * Skip sensitive Linux media roots * Reject sensitive dirs at exact browse roots for PR #6799 Both _resolve_browse_target functions only checked contains_sensitive_path_component while walking descendant parts, so requesting an allowlisted root itself (empty relative path) returned it unchecked. A pre-existing scan-folder row under ~/.ssh, ~/.aws, ~/.config, etc. (registerable before the denylist was added) is re-added to the allowlist on upgrade and could then be browsed. Check the resolved target once before returning in both the legacy and Hub browsers, and cover the root case in both test suites. * fix: avoid unused path helper reexports * fix: import sensitive path helpers directly * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: imagineer99 Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- .../hub/services/models/folder_browser.py | 10 + studio/backend/hub/storage/scan_folders.py | 36 +-- .../backend/hub/tests/test_model_services.py | 28 ++ studio/backend/routes/models.py | 22 +- studio/backend/storage/studio_db.py | 22 +- .../tests/test_linux_external_media_paths.py | 287 ++++++++++++++++++ studio/backend/utils/paths/external_media.py | 100 ++++++ studio/backend/utils/paths/sensitive.py | 46 +++ 8 files changed, 520 insertions(+), 31 deletions(-) create mode 100644 studio/backend/tests/test_linux_external_media_paths.py create mode 100644 studio/backend/utils/paths/external_media.py create mode 100644 studio/backend/utils/paths/sensitive.py diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py index 9b0b46509b..eb137127fb 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -27,6 +27,7 @@ from hub.utils.paths import ( studio_root, well_known_model_dirs, ) +from utils.paths.external_media import linux_run_media_mount_roots from hub.services.models.common import _safe_is_dir from hub.services.models.local_inventory import _resolve_hf_cache_dir @@ -175,6 +176,8 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) + for p in linux_run_media_mount_roots(): + _add(p) _add(_resolve_hf_cache_dir()) try: _add(hf_default_cache_dir()) @@ -346,6 +349,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa ) current = resolved_child + if contains_sensitive_path_component(str(current)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -485,6 +493,8 @@ def browse_folders_response( # Home first as the safe fallback. _add_sug(Path.home()) + for p in linux_run_media_mount_roots(): + _add_sug(p) # The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default. try: _add_sug(_resolve_hf_cache_dir()) diff --git a/studio/backend/hub/storage/scan_folders.py b/studio/backend/hub/storage/scan_folders.py index 85f515da00..fdb15c7c3c 100644 --- a/studio/backend/hub/storage/scan_folders.py +++ b/studio/backend/hub/storage/scan_folders.py @@ -16,37 +16,14 @@ from datetime import datetime, timezone from storage.studio_db import get_connection from hub.utils.paths import normalize_path +from utils.paths.external_media import is_linux_run_media_path +from utils.paths.sensitive import ( + contains_sensitive_path_component as _shared_contains_sensitive_path_component, +) _schema_lock = threading.Lock() _schema_ready = False -_SENSITIVE_PATH_COMPONENTS = { - ".aws", - ".azure", - ".config", - ".docker", - ".gcloud", - ".gnupg", - ".huggingface", - ".kaggle", - ".kube", - ".modelscope", - ".ngc", - ".local", - ".mozilla", - ".pki", - ".thunderbird", - ".ssh", - ".1password", - ".bitwarden", - ".password-store", - "1password", - "bitwarden", - "keychains", - "keyrings", - "mozilla", - "thunderbird", -} def _denied_path_prefixes() -> list[str]: @@ -76,8 +53,7 @@ def _denied_path_prefixes() -> list[str]: def _contains_sensitive_path_component(path: str) -> bool: - parts = os.path.normpath(path).split(os.sep) - return any(part.lower() in _SENSITIVE_PATH_COMPONENTS for part in parts) + return _shared_contains_sensitive_path_component(path) def contains_sensitive_path_component(path: str) -> bool: @@ -142,6 +118,8 @@ def add_scan_folder(path: str) -> dict: check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue raise ValueError(f"Path under {prefix} is not allowed") conn = get_connection() diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index e22aaba282..f05d8359ec 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -168,6 +168,16 @@ def test_resolve_browse_target_rejects_sensitive_dir(tmp_path): assert exc_info.value.status_code == 403 +def test_resolve_browse_target_rejects_sensitive_root(tmp_path): + ssh = tmp_path / "home" / ".ssh" + ssh.mkdir(parents = True) + + with pytest.raises(HTTPException) as exc_info: + folder_browser._resolve_browse_target(str(ssh), [ssh]) + + assert exc_info.value.status_code == 403 + + def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): home = tmp_path / "home" (home / ".ssh").mkdir(parents = True) @@ -181,6 +191,24 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): assert ".ssh" not in names +def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): + home = tmp_path / "home" + media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB" + model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6" + home.mkdir() + model_dir.mkdir(parents = True) + monkeypatch.setattr(folder_browser.Path, "home", lambda: home) + monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root]) + monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf") + monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: []) + monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: []) + + allowlist = folder_browser._build_browse_allowlist() + + assert media_root.resolve() in allowlist + assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve() + + def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path): # The endpoint creates the cache dir on demand so the desktop "Open folder" # action works even before the first download. diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1501868860..c23ab1d428 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1202,6 +1202,7 @@ def _build_browse_allowlist() -> list[Path]: legacy_hf_cache_dir, well_known_model_dirs, ) + from utils.paths.external_media import linux_run_media_mount_roots from storage.studio_db import list_scan_folders candidates: list[Path] = [] @@ -1217,6 +1218,8 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) + for p in linux_run_media_mount_roots(): + _add(p) _add(_resolve_hf_cache_dir()) try: _add(hf_default_cache_dir()) @@ -1336,6 +1339,8 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]: def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path: """Resolve a requested browse path by walking from trusted allowlist roots.""" + from storage.studio_db import contains_sensitive_path_component + requested_path = _normalize_browse_request_path(path) resolved_roots: list[Path] = [] seen_roots: set[str] = set() @@ -1386,8 +1391,18 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa "under your home folder." ), ) + if contains_sensitive_path_component(str(resolved_child)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) current = resolved_child + if contains_sensitive_path_component(str(current)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -1435,7 +1450,8 @@ async def browse_folders( then hidden (if ``show_hidden=true``). """ from utils.paths import hf_default_cache_dir, well_known_model_dirs - from storage.studio_db import list_scan_folders + from utils.paths.external_media import linux_run_media_mount_roots + from storage.studio_db import contains_sensitive_path_component, list_scan_folders # Build once; the sandbox check and suggestion chips share it. allowed_roots = _build_browse_allowlist() @@ -1488,6 +1504,8 @@ async def browse_folders( is_hidden = name.startswith(".") if is_hidden and not show_hidden: continue + if contains_sensitive_path_component(name): + continue entries.append( BrowseEntry( name = name, @@ -1541,6 +1559,8 @@ async def browse_folders( # Home first -- the safe fallback when everything else is cold. _add_sug(Path.home()) + for p in linux_run_media_mount_roots(): + _add_sug(p) # The HF cache root the process is actually using. try: _add_sug(hf_default_cache_dir()) diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index ba9f5b9cbc..41a9adcc29 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -22,7 +22,15 @@ logger = logging.getLogger(__name__) from typing import Any, Iterable, Optional -from utils.paths import project_workspaces_root, studio_db_path, ensure_dir +from utils.paths import ( + ensure_dir, + project_workspaces_root, + studio_db_path, +) +from utils.paths.external_media import is_linux_run_media_path +from utils.paths.sensitive import ( + contains_sensitive_path_component as _shared_contains_sensitive_path_component, +) from utils.training_runs import extract_project_name @@ -61,6 +69,14 @@ def _denied_path_prefixes() -> list[str]: return [] +def _contains_sensitive_path_component(path: str) -> bool: + return _shared_contains_sensitive_path_component(path) + + +def contains_sensitive_path_component(path: str) -> bool: + return _contains_sensitive_path_component(path) + + _schema_lock = threading.Lock() _schema_ready = False _SQLITE_IN_CHUNK_SIZE = 900 @@ -896,6 +912,8 @@ def add_scan_folder(path: str) -> dict: raise ValueError("Path must be a directory, not a file") if not os.access(normalized, os.R_OK | os.X_OK): raise ValueError("Path is not readable") + if _contains_sensitive_path_component(normalized): + raise ValueError("Credential or configuration directories are not allowed") # Windows: normcase for the denylist check but store original casing # so consumers see the native drive-letter casing (e.g. C:\Models). @@ -903,6 +921,8 @@ def add_scan_folder(path: str) -> dict: check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue raise ValueError(f"Path under {prefix} is not allowed") conn = get_connection() diff --git a/studio/backend/tests/test_linux_external_media_paths.py b/studio/backend/tests/test_linux_external_media_paths.py new file mode 100644 index 0000000000..c763248f6a --- /dev/null +++ b/studio/backend/tests/test_linux_external_media_paths.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import ast +import os +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Optional + +import pytest + +from hub.storage import scan_folders +from storage import studio_db +from utils.paths import external_media + + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +class _ExistingScanFolderConn: + def __init__(self): + self.params = () + + def execute( + self, + _sql, + params = (), + ): + self.params = params + return self + + def fetchone(self): + return {"id": 1, "path": self.params[0], "created_at": "fake"} + + def commit(self): + pass + + def close(self): + pass + + +class _HTTPException(Exception): + def __init__(self, status_code: int, detail: str): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +def _stub_linux_path_checks(monkeypatch, module): + monkeypatch.setattr(module.platform, "system", lambda: "Linux") + monkeypatch.setattr(module.os.path, "realpath", os.path.normpath) + monkeypatch.setattr(module.os.path, "expanduser", lambda p: p) + monkeypatch.setattr(module.os.path, "exists", lambda _p: True) + monkeypatch.setattr(module.os.path, "isdir", lambda _p: True) + monkeypatch.setattr(module.os, "access", lambda _p, _mode: True) + + +def _stub_hub_scan_folder_db(monkeypatch): + monkeypatch.setattr(scan_folders, "_ensure_schema", lambda _conn: None) + monkeypatch.setattr(scan_folders, "get_connection", _ExistingScanFolderConn) + + +def _stub_legacy_scan_folder_db(monkeypatch): + monkeypatch.setattr(studio_db, "get_connection", _ExistingScanFolderConn) + + +def test_linux_run_media_policy_accepts_mounted_volume_descendants(monkeypatch): + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB") + assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6") + + +@pytest.mark.parametrize( + "path", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_linux_run_media_policy_rejects_unrelated_run_paths(monkeypatch, path): + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + assert not external_media.is_linux_run_media_path(path) + + +def test_linux_run_media_mount_roots_lists_readable_volume_roots(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + mount = base / "dspofu" / "nvmeB" + sensitive_mount = base / "dspofu" / ".ssh" + sensitive_aws_mount = base / "dspofu" / ".aws" + other_user_mount = base / "other" / "backup" + incomplete = base / "dspofu-only" + mount.mkdir(parents = True) + sensitive_mount.mkdir() + sensitive_aws_mount.mkdir() + other_user_mount.mkdir(parents = True) + incomplete.mkdir() + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [mount.resolve()] + + +def test_linux_run_media_mount_roots_skips_sensitive_resolved_volume_name(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + normal_mount = base / "dspofu" / "nvmeB" + sensitive_target = base / "dspofu" / ".config" + normal_mount.mkdir(parents = True) + sensitive_target.mkdir() + alias = base / "dspofu" / "config-alias" + alias.symlink_to(sensitive_target, target_is_directory = True) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [normal_mount.resolve()] + + +def test_linux_run_media_mount_roots_skips_sensitive_resolved_descendant(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + normal_mount = base / "dspofu" / "nvmeB" + sensitive_descendant = normal_mount / ".ssh" / "models" + sensitive_descendant.mkdir(parents = True) + alias = base / "dspofu" / "models-alias" + alias.symlink_to(sensitive_descendant, target_is_directory = True) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [normal_mount.resolve()] + + +def test_hub_scan_folder_accepts_linux_run_media_mount(monkeypatch): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + + row = scan_folders.add_scan_folder(target) + + assert row["path"] == target + + +@pytest.mark.parametrize( + "target", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_hub_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Path under /run is not allowed"): + scan_folders.add_scan_folder(target) + + +def test_hub_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Credential or configuration"): + scan_folders.add_scan_folder("/run/media/dspofu/nvmeB/.ssh/models") + + +def test_legacy_scan_folder_accepts_linux_run_media_mount(monkeypatch): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + + row = studio_db.add_scan_folder(target) + + assert row["path"] == target + + +@pytest.mark.parametrize( + "target", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_legacy_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Path under /run is not allowed"): + studio_db.add_scan_folder(target) + + +def test_legacy_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Credential or configuration"): + studio_db.add_scan_folder("/run/media/dspofu/nvmeB/.aws/models") + + +def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): + tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")) + function_names = { + "_build_browse_allowlist", + "_browse_relative_parts", + "_is_path_inside_allowlist", + "_match_browse_child", + "_normalize_browse_request_path", + "_resolve_browse_target", + } + functions = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in function_names + ] + module = ast.Module(body = functions, type_ignores = []) + ast.fix_missing_locations(module) + + home = tmp_path / "home" + media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB" + model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6" + home.mkdir() + model_dir.mkdir(parents = True) + (media_root / ".ssh").mkdir() + + fake_paths = SimpleNamespace( + hf_default_cache_dir = lambda: tmp_path / "missing-default-hf", + legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf", + well_known_model_dirs = lambda: [], + studio_root = lambda: tmp_path / "missing-studio", + outputs_root = lambda: tmp_path / "missing-outputs", + exports_root = lambda: tmp_path / "missing-exports", + ) + fake_external_media = SimpleNamespace(linux_run_media_mount_roots = lambda: [media_root]) + fake_studio_db = SimpleNamespace( + list_scan_folders = lambda: [], + contains_sensitive_path_component = studio_db.contains_sensitive_path_component, + ) + monkeypatch.setitem(sys.modules, "utils.paths", fake_paths) + monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db) + + ns = { + "HTTPException": _HTTPException, + "os": os, + "Path": Path, + "Optional": Optional, + "_safe_is_dir": lambda p: Path(p).is_dir(), + "_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf", + "logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None), + } + exec(compile(module, "", "exec"), ns) + + allowlist = ns["_build_browse_allowlist"]() + + assert media_root.resolve() in allowlist + assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve() + + with pytest.raises(_HTTPException) as exc: + ns["_resolve_browse_target"](str(media_root / ".ssh"), allowlist) + assert exc.value.status_code == 403 + + ssh_root = media_root / ".ssh" + with pytest.raises(_HTTPException) as exc_root: + ns["_resolve_browse_target"](str(ssh_root), [ssh_root]) + assert exc_root.value.status_code == 403 diff --git a/studio/backend/utils/paths/external_media.py b/studio/backend/utils/paths/external_media.py new file mode 100644 index 0000000000..1f1754664f --- /dev/null +++ b/studio/backend/utils/paths/external_media.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""External media path helpers.""" + +from __future__ import annotations + +import getpass +import os +import platform +from pathlib import Path + +from utils.paths.sensitive import ( + contains_sensitive_path_component, + is_sensitive_path_component, +) + + +def _is_linux_media_mount_path(path: str, media_root: Path | str) -> bool: + normalized = os.path.normpath(os.path.realpath(os.path.expanduser(path))) + root = os.path.normpath(os.path.realpath(os.path.expanduser(str(media_root)))) + try: + rel = os.path.relpath(normalized, root) + except ValueError: + return False + if rel == "." or rel == ".." or rel.startswith(f"..{os.sep}"): + return False + parts = [part for part in rel.split(os.sep) if part] + return len(parts) >= 2 and all(part not in (".", "..") for part in parts[:2]) + + +def is_linux_run_media_path(path: str) -> bool: + """True for Linux removable-media paths under /run/media//.""" + if platform.system() != "Linux": + return False + return _is_linux_media_mount_path(path, "/run/media") + + +def _current_username() -> str | None: + try: + user = getpass.getuser().strip() + except Exception: + return None + return user or None + + +def _contains_sensitive_media_component(path: Path, media_root: Path) -> bool: + try: + rel = path.relative_to(media_root) + except ValueError: + rel = path + return contains_sensitive_path_component(str(rel)) + + +def linux_run_media_mount_roots( + base: Path | str = "/run/media", *, user: str | None = None +) -> list[Path]: + """Readable /run/media// roots for the folder browser.""" + if platform.system() != "Linux": + return [] + user = user or _current_username() + if not user or user in (".", "..") or os.sep in user: + return [] + base_path = Path(base) + try: + resolved_base = base_path.resolve() + except (OSError, RuntimeError, ValueError): + return [] + + roots: list[Path] = [] + seen: set[str] = set() + user_dir = base_path / user + try: + if not user_dir.is_dir(): + return [] + volume_dirs = list(user_dir.iterdir()) + except (OSError, RuntimeError, ValueError): + return [] + for volume_dir in volume_dirs: + if is_sensitive_path_component(volume_dir.name): + continue + try: + resolved = volume_dir.resolve() + except (OSError, RuntimeError, ValueError): + continue + if not _is_linux_media_mount_path(str(resolved), resolved_base): + continue + if _contains_sensitive_media_component(resolved, resolved_base): + continue + key = os.path.normcase(os.path.realpath(str(resolved))) + if key in seen: + continue + try: + is_dir = resolved.is_dir() + except OSError: + continue + if is_dir and os.access(resolved, os.R_OK | os.X_OK): + seen.add(key) + roots.append(resolved) + return roots diff --git a/studio/backend/utils/paths/sensitive.py b/studio/backend/utils/paths/sensitive.py new file mode 100644 index 0000000000..7d32a5f4cf --- /dev/null +++ b/studio/backend/utils/paths/sensitive.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared sensitive path-component policy.""" + +from __future__ import annotations + +import os + + +SENSITIVE_PATH_COMPONENTS = { + ".aws", + ".azure", + ".config", + ".docker", + ".gcloud", + ".gnupg", + ".huggingface", + ".kaggle", + ".kube", + ".modelscope", + ".ngc", + ".local", + ".mozilla", + ".pki", + ".thunderbird", + ".ssh", + ".1password", + ".bitwarden", + ".password-store", + "1password", + "bitwarden", + "keychains", + "keyrings", + "mozilla", + "thunderbird", +} + + +def is_sensitive_path_component(name: str) -> bool: + return name.lower() in SENSITIVE_PATH_COMPONENTS + + +def contains_sensitive_path_component(path: str) -> bool: + parts = os.path.normpath(path).split(os.sep) + return any(is_sensitive_path_component(part) for part in parts) From c356427f30e8abadd38fbc722e494321f0b5b64b Mon Sep 17 00:00:00 2001 From: Ayushman <139611211+InfoSage05@users.noreply.github.com> Date: Fri, 3 Jul 2026 23:54:29 +0530 Subject: [PATCH 21/23] Guard Windows ROCm torchao override skip (#6837) * Fix: skip fp16/bf16 validation for full finetuning in RL trainers When doing full finetuning (FFT) of a bfloat16 model, the fp16/bf16 mismatch validation fires before the corrective logic runs, causing a misleading error even though the code would properly handle it downstream. Skip the validation when full_finetuning is active. Fixes #6731 * Fix: auto-correct fp16/bf16 mismatches for full finetuning before validation Instead of entirely skipping validation (which could let mismatches through when mixed_precision_dtype is float32), auto-correct explicit fp16/bf16 settings that conflict with the model's dtype for FFT. This way the existing validation still catches real mismatches for non-FFT cases, and the corrective logic below handles the normalized settings. Fixes the issue raised in Codex review of PR #6813. * Guard Windows ROCm torchao override skip Detect installed ROCm torch directly before applying the torchao override so Windows ROCm environments never install the crashing torchao package even if the earlier ROCm-installed flag is missing. * Update unsloth/models/rl.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/install_python_stack.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Harden ROCm probe and sync RL precision flags Tolerate stray stdout noise when probing Windows ROCm torch installs by checking the last non-empty output line, matching the existing torch version probe behavior. Also keep args.fp16 and args.bf16 synchronized with the full-finetuning precision auto-corrections in the RL trainer patch so downstream eval settings see a consistent TrainingArguments state. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add MLX trainer compatibility shims Patch imported MLXTrainer and MLXTrainingConfig objects to preserve the expected dataclass field ordering and to provide a _train_dataset_for_batches fallback when older trainers or test doubles only expose train_dataset. Also add focused worker tests covering both compatibility paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope PR to Windows ROCm torchao guard * Restore PR scope to Windows ROCm guard * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: cover Windows ROCm torchao skip behavior * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Ayushman Paul Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 --- studio/backend/tests/test_torchao_select.py | 59 +++++++++++++++++-- studio/install_python_stack.py | 36 +++++++++++- tests/studio/install/test_rocm_support.py | 64 +++++++++++++++++++++ 3 files changed, 152 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index a99eb4c45c..2d3dc5fbff 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -12,6 +12,7 @@ from __future__ import annotations import sys from pathlib import Path +from unittest.mock import MagicMock import pytest @@ -71,12 +72,58 @@ def test_default_spec_matches_table(monkeypatch): assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC -def test_skips_torchao_on_windows_rocm(): +@pytest.mark.parametrize( + ("rocm_windows_torch_installed", "installed_torch_is_windows_rocm"), + [ + (True, False), + (False, True), + ], +) +def test_skips_torchao_on_windows_rocm( + monkeypatch, tmp_path, rocm_windows_torch_installed, installed_torch_is_windows_rocm +): """The overrides step must skip torchao on Windows ROCm: no working build exists there (it imports an absent c10d backend and crashes transformers.quantizers), so the installer skips it and relies on the runtime stub instead.""" - source = _INSTALL_SCRIPT.read_text(encoding = "utf-8") - # Branches on the Windows-ROCm marker set by _ensure_rocm_torch ... - assert "elif _rocm_windows_torch_installed:" in source - # ... and reports the skip in the progress label. - assert "dependency overrides (skipped, Windows ROCm)" in source + mod = _load_module(monkeypatch) + installed_specs: list[str] = [] + progress_labels: list[str] = [] + + def _record_pip_install(*args, **kwargs): + installed_specs.extend(str(arg) for arg in args) + return 0 + + unstructured_plugin = tmp_path / "unstructured" + github_plugin = tmp_path / "github" + unstructured_plugin.mkdir() + github_plugin.mkdir() + + subprocess_result = MagicMock() + subprocess_result.returncode = 0 + subprocess_result.stdout = "" + + monkeypatch.setenv("SKIP_STUDIO_BASE", "1") + monkeypatch.setattr(mod, "IS_WINDOWS", True) + monkeypatch.setattr(mod, "IS_MACOS", False) + monkeypatch.setattr(mod, "IS_MAC_ARM", False) + monkeypatch.setattr(mod, "NO_TORCH", False) + monkeypatch.setattr(mod, "_rocm_windows_torch_installed", rocm_windows_torch_installed) + monkeypatch.setattr( + mod, "_installed_torch_is_windows_rocm", lambda: installed_torch_is_windows_rocm + ) + monkeypatch.setattr(mod, "_bootstrap_uv", lambda: False) + monkeypatch.setattr(mod, "_repair_bad_anyio", lambda: None) + monkeypatch.setattr(mod, "_ensure_rocm_torch", lambda: None) + monkeypatch.setattr(mod, "_ensure_cuda_torch", lambda: None) + monkeypatch.setattr(mod, "_has_usable_nvidia_gpu", lambda: True) + monkeypatch.setattr(mod, "run", lambda *args, **kwargs: None) + monkeypatch.setattr(mod, "pip_install", _record_pip_install) + monkeypatch.setattr(mod, "_progress", lambda label: progress_labels.append(label)) + monkeypatch.setattr(mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin) + monkeypatch.setattr(mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin) + monkeypatch.setattr(mod.subprocess, "run", lambda *args, **kwargs: subprocess_result) + + assert mod.install_python_stack() == 0 + + assert not any(spec.startswith("torchao") for spec in installed_specs) + assert "dependency overrides (skipped, Windows ROCm)" in progress_labels diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 37805e2a57..439d3ffe7b 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -181,6 +181,40 @@ def _probe_installed_torch_version() -> str | None: return lines[-1] if lines else None +def _installed_torch_is_windows_rocm() -> bool: + """Return True when the target venv currently has a Windows ROCm torch build. + + This is a belt-and-suspenders guard for the torchao override step: if the + earlier ROCm install path failed to set _rocm_windows_torch_installed but the + venv already contains a ROCm torch wheel, still skip torchao because it + crashes on import on Windows ROCm. + """ + if not IS_WINDOWS: + return False + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys, torch; " + "hip = getattr(getattr(torch, 'version', None), 'hip', None) or ''; " + "ver = getattr(torch, '__version__', '').lower(); " + "sys.stdout.write('yes' if (hip or 'rocm' in ver or 'rocmsdk' in ver) else '')" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 90, + **_windows_hidden_subprocess_kwargs(), + ) + except (OSError, subprocess.TimeoutExpired): + return False + lines = [line.strip() for line in (probe.stdout or "").splitlines() if line.strip()] + return probe.returncode == 0 and bool(lines and lines[-1] == "yes") + + # constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install # from before the cap existed can already be stuck at 4.14+, which later # constrained installs won't touch since it already satisfies mcp/fastmcp. @@ -2256,7 +2290,7 @@ def install_python_stack() -> int: # (no working build; see below). if NO_TORCH: _progress("dependency overrides (skipped, no torch)") - elif _rocm_windows_torch_installed: + elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm(): # No working Windows ROCm torchao build: it imports an absent c10d backend # and crashes transformers.quantizers. Studio stubs it at runtime, so # installing it only ships a package that crashes on import -- skip it. diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 092e1803f8..c8f2053946 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2296,6 +2296,70 @@ class TestRocmTorchInstalledEnvVar: mock_bnb.assert_not_called() +class TestWindowsRocmTorchaoGuard: + """Verify the torchao skip can detect an installed Windows ROCm torch build.""" + + def test_installed_torch_is_windows_rocm_accepts_rocm_probe(self): + rv = MagicMock() + rv.returncode = 0 + rv.stdout = "yes" + with ( + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod.subprocess, "run", return_value = rv), + ): + assert stack_mod._installed_torch_is_windows_rocm() is True + + def test_installed_torch_is_windows_rocm_rejects_non_rocm_probe(self): + rv = MagicMock() + rv.returncode = 0 + rv.stdout = "" + with ( + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod.subprocess, "run", return_value = rv), + ): + assert stack_mod._installed_torch_is_windows_rocm() is False + + def test_installed_torch_is_windows_rocm_is_non_windows_noop(self): + with patch.object(stack_mod, "IS_WINDOWS", False): + assert stack_mod._installed_torch_is_windows_rocm() is False + + @patch.object(stack_mod, "_repair_bad_anyio") + @patch.object(stack_mod, "_ensure_rocm_torch") + @patch.object(stack_mod, "_ensure_cuda_torch") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) + @patch.object(stack_mod, "run") + @patch.object(stack_mod, "pip_install") + def test_install_python_stack_skips_torchao_when_windows_rocm_torch_is_installed( + self, mock_pip, mock_run, mock_has_nvidia, mock_cuda, mock_rocm, mock_anyio, tmp_path + ): + unstructured_plugin = tmp_path / "unstructured" + github_plugin = tmp_path / "github" + unstructured_plugin.mkdir() + github_plugin.mkdir() + + subprocess_result = MagicMock() + subprocess_result.returncode = 0 + subprocess_result.stdout = "" + + with ( + patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}), + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod, "IS_MACOS", False), + patch.object(stack_mod, "IS_MAC_ARM", False), + patch.object(stack_mod, "NO_TORCH", False), + patch.object(stack_mod, "_rocm_windows_torch_installed", False), + patch.object(stack_mod, "_bootstrap_uv", return_value = False), + patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = True), + patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin), + patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin), + patch.object(stack_mod.subprocess, "run", return_value = subprocess_result), + ): + assert stack_mod.install_python_stack() == 0 + + installed_specs = [str(arg) for call in mock_pip.call_args_list for arg in call.args] + assert not any("torchao" in arg for arg in installed_specs) + + # TEST: worker.py -- Windows ROCm patches (source-level checks) From fefa2187f077185b8c13ee8eae28333ab9b5ad8f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 01:30:06 +0000 Subject: [PATCH 22/23] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/diffusion_dit_trainer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 1cbdd93c53..d6288ce48c 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -279,7 +279,13 @@ def _apply_fp8_training(transformer, on_event) -> bool: def _pick_auto_precision( - prequant, device, free_gb, dense_gb, capability, has_fp8, has_torchao = True + prequant, + device, + free_gb, + dense_gb, + capability, + has_fp8, + has_torchao = True, ) -> str: """Pure policy for base_precision="auto": nf4 for a prequant base or no CUDA; else the fastest dense mode whose weights + headroom (activations, optimizer, cache) fit the From bbca561d38875b894564036511f7d7fc5bdf8c75 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 02:20:15 +0000 Subject: [PATCH 23/23] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/sd_cpp_backend.py | 5 ++++- studio/backend/tests/test_sd_cpp_backend.py | 14 +++++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 7632551661..d8cb21d409 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -893,7 +893,10 @@ class SdCppDiffusionBackend: lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}" materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage) lora_payload = [ - {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)} + { + "path": f"{lora_stage.name}/{Path(m.path).name}", + "multiplier": float(m.weight), + } for m in materialized ] try: diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 0b52ea6e14..be2fc22a4b 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -691,7 +691,11 @@ def _fake_materialize(resolved, dest): return out -def _patch_lora(monkeypatch, resolved, supported = True): +def _patch_lora( + monkeypatch, + resolved, + supported = True, +): from core.inference import diffusion_lora as dl monkeypatch.setattr(dl, "supports_lora", lambda **k: supported) @@ -706,7 +710,9 @@ def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch): eng = _FakeEngine() b = _loaded_backend(engine = eng) # mode = "oneshot" - _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]) + _patch_lora( + monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)] + ) b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)]) _, params, _, _ = eng.calls[0] assert params.lora_dir is not None and params.lora_apply_mode == "auto" @@ -725,7 +731,9 @@ def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tm servers: list = [] _run_server_load(monkeypatch, b, servers) servers[0].lora_dir = str(tmp_path) - _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]) + _patch_lora( + monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)] + ) b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)]) payload = servers[0].payloads[0] assert "lora" in payload and len(payload["lora"]) == 1