From f899834e58a735a7fa7d5341bf944476b32a835e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 19:33:05 -0700 Subject: [PATCH 001/210] DeepSeek-V4: eager attention and trainable FP8 grouped experts (#7042) * DeepSeek-V4: eager attention and trainable FP8 grouped experts deepseek_v4 ships a custom attention that is not compatible with the sdpa and flash paths, so add it to _EAGER_ONLY_PREFIXES to load with eager. Its fused experts load as FP8GroupedLinear, whose forward calls a grouped matmul kernel with no autograd formula, so loss.backward() fails during finetuning. Patch the forward to dequantize the frozen fp8 weight and run a differentiable grouped matmul while training, keeping the fused fp8 kernel for inference. * DeepSeek-V4: exclude sdpa/flash and stream fp8 grouped backward Add deepseek_v4 to _SDPA_EXCLUDED_MODELS and _FLASH_EXCLUDED_MODELS so an explicit attn_implementation=sdpa/flash request downgrades to eager instead of raising (the model has no sdpa/flash kernel), matching the eager-only default. Replace the FP8GroupedLinear training bmm with a custom autograd Function that saves only the fp8 weight + scale rather than a full bf16 dequantized copy, so no dequantized grouped weight is retained per layer, and unwrap tensor-parallel shards before dequant. Bit-exact forward and grad with the previous path. * FP8 grouped: consistent checkpointing math and block-size-aware dequant Gate the differentiable training path on self.training rather than torch.is_grad_enabled(), so a gradient-checkpointed segment runs the same bmm math in its no-grad forward and its grad recompute instead of mixing the fused fp8 kernel with bmm. Dequantize with the layer's own block_size via _blockwise_weight_dequant_any_shape so non-128 or rectangular fp8 blocks are scaled correctly instead of assuming 128x128. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/kernels/fp8.py | 68 ++++++++++++++++++++++++++++++++++++++++ unsloth/models/_utils.py | 8 +++-- 2 files changed, 73 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 4efc4bd5d3..7a57f91ce1 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -42,6 +42,11 @@ except: "Unsloth: FP8 models need importing FP8Linear from `transformers.integrations.finegrained_fp8` but we don't see it." ) +try: + from transformers.integrations.finegrained_fp8 import FP8GroupedLinear +except: + FP8GroupedLinear = None + try: from transformers.integrations.fbgemm_fp8 import FbgemmFp8Linear except: @@ -688,3 +693,66 @@ if FbgemmFp8Linear is not None: FbgemmFp8Linear.forward = module_forward_patch(fbgemm_fp8_linear, "weight_scale") if FP8Linear is not None: FP8Linear.forward = module_forward_patch(fp8_block_quant_linear, "weight_scale_inv") + +# FP8GroupedLinear's fused grouped matmul has no autograd formula, so training +# backward fails. In training, use a custom autograd Function: dequant the frozen +# fp8 weight for a differentiable bmm, saving only the fp8 weight + scale and +# unwrapping TP shards; eval keeps the fused kernel. Gate on self.training (not +# is_grad_enabled) so the grad-checkpoint no-grad forward and its recompute match. +if FP8GroupedLinear is not None: + _fp8_grouped_forward_orig = FP8GroupedLinear.forward + + def _fp8_to_local(t): + dt = getattr(getattr(torch, "distributed", None), "tensor", None) + DTensor = getattr(dt, "DTensor", None) if dt is not None else None + return t.to_local() if DTensor is not None and isinstance(t, DTensor) else t + + def _fp8_grouped_dequant(weight, scale_inv, block_size, dtype): + # Honor the layer's block size; weight_dequant would assume 128 and mis-scale. + if block_size is not None and len(block_size) == 2: + return _blockwise_weight_dequant_any_shape(weight, scale_inv.float(), block_size, dtype) + return weight_dequant(weight, scale_inv.float()).to(dtype) + + class _FP8GroupedMM(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, scale_inv, n_groups, block_size, bias): + weight, scale_inv = _fp8_to_local(weight), _fp8_to_local(scale_inv) + hidden = x.shape[-1] + W = _fp8_grouped_dequant(weight, scale_inv, block_size, x.dtype) + out_per = W.shape[0] // n_groups + xg = x.reshape(-1, n_groups, hidden).transpose(0, 1) + y = torch.bmm(xg, W.view(n_groups, out_per, hidden).transpose(1, 2)) + y = y.transpose(0, 1).reshape(*x.shape[:-2], n_groups, out_per) + if bias is not None: + y = y + bias.view(n_groups, out_per) + ctx.save_for_backward(weight, scale_inv) + ctx.n_groups, ctx.out_per, ctx.x_shape = n_groups, out_per, x.shape + ctx.dtype, ctx.has_bias, ctx.block_size = x.dtype, bias is not None, block_size + return y + + @staticmethod + def backward(ctx, grad_y): + weight, scale_inv = ctx.saved_tensors + ng, out_per, hidden = ctx.n_groups, ctx.out_per, ctx.x_shape[-1] + W = _fp8_grouped_dequant(weight, scale_inv, ctx.block_size, ctx.dtype).view( + ng, out_per, hidden + ) + gy = grad_y.reshape(-1, ng, out_per).transpose(0, 1) + grad_x = torch.bmm(gy, W).transpose(0, 1).reshape(ctx.x_shape) + grad_bias = gy.sum(1).reshape(-1) if ctx.has_bias else None + return grad_x, None, None, None, None, grad_bias + + def _fp8_grouped_forward(self, x): + if self.weight.element_size() > 1 or not self.training: + return _fp8_grouped_forward_orig(self, x) + bias = self.bias if self.has_bias else None + return _FP8GroupedMM.apply( + x, + self.weight, + self.weight_scale_inv, + self.n_groups, + getattr(self, "block_size", None), + bias, + ) + + FP8GroupedLinear.forward = _fp8_grouped_forward diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index fa0e0b1c49..ae3b94ecc8 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -424,7 +424,7 @@ def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_len # access on some GPU architectures (B200). Falls back to eager safely. _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2") -_SDPA_EXCLUDED_MODELS = ("gpt_oss",) +_SDPA_EXCLUDED_MODELS = ("gpt_oss", "deepseek_v4") # The loader (loader.py) forces supports_sdpa=False for these because their bundled # SDPA modules are wrong. Kept here, not in loader.py, so _is_sdpa_excluded can honor # them without a loader -> _utils import cycle (loader.py already imports from _utils @@ -437,8 +437,10 @@ DISABLE_SDPA_MODEL_NAMES = [ "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore "gpt_oss", ] -_FLASH_EXCLUDED_MODELS = ("gpt_oss",) -_EAGER_ONLY_PREFIXES = ("gemma3n",) +_FLASH_EXCLUDED_MODELS = ("gpt_oss", "deepseek_v4") +# deepseek_v4's custom attention is sdpa/flash-incompatible; force eager, and +# excluded above so an explicit sdpa/flash request cannot re-enable the crash. +_EAGER_ONLY_PREFIXES = ("gemma3n", "deepseek_v4") _FLASH_ATTENTION_MAX_HEAD_DIM = 256 _FLASH_ATTENTION_DISABLED_WARNED = set() From 275bad1f644c730b9d3ed354bfc3a6cc6fdcb0d9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 11 Jul 2026 21:29:19 -0700 Subject: [PATCH 002/210] Studio: fix the manual response-template markers that never match their rendered templates (#7062) * Fix broken manual response-template markers in Studio's fallback table Six template families in TEMPLATE_TO_RESPONSES_MAPPER shipped markers that never match what their chat templates actually render, so the manual train_on_completions path masked every assistant token and the run died on the all-labels-masked safety net: - mistral, llama: '[INST] ' / ' [/INST]' - the surrounding spaces fold into the neighbouring tokens ('[INST]'/'[/INST]' are single special tokens in Mistral v0.3, SentencePiece pieces in Llama-2), so the padded strings never match. Now '[INST]' / '[/INST]'. - starling: trailing space after 'GPT4 Correct Assistant:' folds into the next content token. Now no trailing space. - glm: '[gMASK]' renders once at text start, never before later user turns, and '' is generation scaffolding rendered as a lone '' on non-final turns. Now '<|user|>' / '<|assistant|>'. - qwen3-thinking: '' is stripped from non-final assistant turns (Qwen3-Thinking-2507) and never rendered by QwQ. Now the bare assistant header, matching the other qwen entries. - zephyr: role tags are plain text and SentencePiece tokenizes them differently at text start than after '' + newline mid-conversation; the markers need the leading newline anchor. Now '\n<|user|>\n' / '\n<|assistant|>\n'. Validated token-level on each family's representative tokenizer with a two-turn fixture plus system message: user and system content fully masked, every assistant turn trained, and the final EOS label never -100. The fixed mistral, llama, starling and glm markers produce labels identical to zoo auto-detection; qwen3-thinking differs only in one turn-separator newline token. All 22 unchanged entries produce byte-identical labels to before this change. Adds tests/test_response_template_markers.py pinning the fixed and key unchanged marker literals (dependency-free) plus token-level masking checks that skip when tokenizers or unsloth_zoo are unavailable offline. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close tokenizer config handle and read it as UTF-8 Chat templates in tokenizer_config.json are rarely ASCII-only, so the default locale codec could fail the GLM fallback loader on Windows. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments * Anchor the llama marker on and harden the marker test On transformers 5.x llama-2 tokenizes [INST] after as a bare left bracket while the standalone encoding gives the space-prefixed piece, so the unanchored marker missed every turn boundary and later user turns leaked into training; 4.57 masked this. Anchoring on [INST] matches both tokenizations, verified token-level under 4.57.6 and 5.5.0. The test now unwraps the BatchEncoding that apply_chat_template returns on 5.x before indexing, and the latent trailing spaces in the unreachable unsloth and vicuna entries are dropped for table consistency. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../tests/test_response_template_markers.py | 216 ++++++++++++++++++ .../backend/utils/datasets/model_mappings.py | 47 ++-- 2 files changed, 248 insertions(+), 15 deletions(-) create mode 100644 studio/backend/tests/test_response_template_markers.py diff --git a/studio/backend/tests/test_response_template_markers.py b/studio/backend/tests/test_response_template_markers.py new file mode 100644 index 0000000000..8c813e62f2 --- /dev/null +++ b/studio/backend/tests/test_response_template_markers.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""TEMPLATE_TO_RESPONSES_MAPPER markers must match what the templates render. + +The manual instruction/response markers are the fallback for +train_on_completions when auto-detection is unavailable, so a marker that +never matches the rendered chat template masks every assistant token and the +run dies on the all-labels-masked safety net. Six template families shipped +such markers: + + mistral - "[INST] " / " [/INST]": the surrounding spaces fold into + the neighbouring tokens ("[INST]" is a single special + token in Mistral v0.3), so the padded strings never match. + llama - same space folding, plus llama-2 tokenizes [INST] after + as bare "[" on transformers 5.x while the standalone + encoding gives "▁[", so the marker must anchor on . + starling - trailing space after "GPT4 Correct Assistant:" folds + into the next content token ("▁Hello"). + glm - "[gMASK]" renders once at text start, never before + later user turns; "" is generation scaffolding + that non-final turns render as a lone "". + qwen3-thinking - "" is stripped from non-final assistant turns + (Qwen3-Thinking-2507) or never rendered (QwQ). + zephyr - role tags are plain text, and SentencePiece tokenizes + "<|assistant|>" differently at text start than after + "\\n" mid-conversation; the markers need the leading + newline anchor to tokenize like a real turn boundary. + +Literal assertions run everywhere; the token-level masking checks need the +representative tokenizers plus unsloth_zoo and skip when either is +unavailable (offline CI). +""" + +from __future__ import annotations + +import importlib.util +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) + +# model_mappings is dependency-free: load it directly so these tests run +# without the studio venv / package import side effects. +_MM_PATH = Path(_BACKEND_DIR) / "utils" / "datasets" / "model_mappings.py" +_mm_spec = importlib.util.spec_from_file_location("_marker_test_mm", _MM_PATH) +model_mappings = importlib.util.module_from_spec(_mm_spec) +_mm_spec.loader.exec_module(model_mappings) + +T2R = model_mappings.TEMPLATE_TO_RESPONSES_MAPPER + + +# ── Fixed entries: markers derived from what each representative tokenizer +# actually renders (see PR for the token-level derivation). ── +EXPECTED_FIXED = { + "mistral": {"instruction": "[INST]", "response": "[/INST]"}, + "llama": {"instruction": "[INST]", "response": "[/INST]"}, + "starling": {"instruction": "GPT4 Correct User:", "response": "GPT4 Correct Assistant:"}, + "glm": {"instruction": "<|user|>", "response": "<|assistant|>"}, + "qwen3-thinking": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "zephyr": {"instruction": "\n<|user|>\n", "response": "\n<|assistant|>\n"}, +} + +# Spot-pin some known-good entries so a refactor cannot silently change them. +EXPECTED_UNCHANGED = { + "qwen3": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "llama-3.1": { + "instruction": "<|start_header_id|>user<|end_header_id|>\n\n", + "response": "<|start_header_id|>assistant<|end_header_id|>\n\n", + }, + "phi-4": { + "instruction": "<|im_start|>user<|im_sep|>", + "response": "<|im_start|>assistant<|im_sep|>", + }, + "gemma-3": {"instruction": "user\n", "response": "model\n"}, + "gpt-oss": { + "instruction": "<|start|>user<|message|>", + "response": "<|start|>assistant<|channel|>final<|message|>", + }, +} + + +@pytest.mark.parametrize("template", sorted(EXPECTED_FIXED)) +def test_fixed_marker_literals(template): + assert T2R[template] == EXPECTED_FIXED[template] + + +@pytest.mark.parametrize("template", sorted(EXPECTED_UNCHANGED)) +def test_unchanged_marker_literals(template): + assert T2R[template] == EXPECTED_UNCHANGED[template] + + +def test_no_marker_is_empty_or_whitespace(): + for template, parts in T2R.items(): + assert parts["instruction"].strip(), template + assert parts["response"].strip(), template + + +# ── Token-level checks: markers must select exactly the assistant turns on a +# rendered two-turn fixture, and the final EOS label must never be -100. ── + +REPRESENTATIVES = { + "mistral": ["unsloth/mistral-7b-instruct-v0.3"], + "llama": ["unsloth/llama-2-7b-chat"], + "starling": ["unsloth/Starling-LM-7B-beta"], + "glm": ["unsloth/GLM-4.7-Flash"], + "qwen3-thinking": ["unsloth/Qwen3-4B-Thinking-2507", "Qwen/QwQ-32B"], + "zephyr": ["unsloth/zephyr-sft"], +} + +FIXTURE = [ + {"role": "user", "content": "zebra alpha question one?"}, + {"role": "assistant", "content": "grape reply number one."}, + {"role": "user", "content": "zebra beta question two?"}, + {"role": "assistant", "content": "grape reply number two."}, +] + + +def _load_tokenizer(repo): + try: + from transformers import AutoTokenizer + except Exception as e: # pragma: no cover + pytest.skip(f"transformers unavailable: {e}") + try: + return AutoTokenizer.from_pretrained(repo) + except OSError as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + except Exception: + # Tokenizer class newer than this transformers (e.g. GLM-4.7's + # TokenizersBackend): build directly from tokenizer.json. + try: + import json as _json + from huggingface_hub import hf_hub_download + from transformers import PreTrainedTokenizerFast + + with open(hf_hub_download(repo, "tokenizer_config.json"), encoding = "utf-8") as f: + cfg = _json.load(f) + tok_file = hf_hub_download(repo, "tokenizer.json") + + def _tokval(v): + return v["content"] if isinstance(v, dict) else v + + return PreTrainedTokenizerFast( + tokenizer_file = tok_file, + chat_template = cfg.get("chat_template"), + **{ + k: _tokval(cfg[k]) + for k in ("bos_token", "eos_token", "pad_token", "unk_token") + if cfg.get(k) is not None + }, + ) + except Exception as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + + +def _train_on_responses_only(): + try: + from unsloth_zoo.dataset_utils import train_on_responses_only + except Exception as e: + pytest.skip(f"unsloth_zoo unavailable: {e}") + return train_on_responses_only + + +@pytest.mark.parametrize( + "template,repo", + [(t, r) for t, repos in sorted(REPRESENTATIVES.items()) for r in repos], +) +def test_fixed_markers_token_level(template, repo): + tor = _train_on_responses_only() + tok = _load_tokenizer(repo) + parts = T2R[template] + + msgs = [{"role": "system", "content": "You are a terse assistant."}] + FIXTURE + try: + ids = tok.apply_chat_template(msgs, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] # transformers 5.x returns a BatchEncoding + except Exception: + ids = tok.apply_chat_template(FIXTURE, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] + + fn = tor( + None, + instruction_part = parts["instruction"], + response_part = parts["response"], + tokenizer = tok, + return_function = True, + ) + labels = fn({"input_ids": [list(ids)]})["labels"][0] + + n = len(ids) + trained = tok.decode([ids[i] for i in range(n) if labels[i] != -100]) + masked = tok.decode([ids[i] for i in range(n) if labels[i] == -100]) + + # User and system content fully masked + assert "question one" not in trained and "question one" in masked + assert "question two" not in trained and "question two" in masked + assert "terse assistant" not in trained + # EVERY assistant turn trained, not just the last + assert "reply number one" in trained + assert "reply number two" in trained + # The final EOS (last non-whitespace token) must never be -100, or the + # fine-tuned model never learns to stop generating. + i = n - 1 + while i > 0 and tok.decode([ids[i]]).strip() == "": + i -= 1 + assert labels[i] != -100, f"final token {tok.convert_ids_to_tokens(int(ids[i]))!r} is masked" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 9d2c983aed..65ba4b4688 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -485,9 +485,11 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No "" suffix: Qwen3-Thinking-2507 strips it from non-final turns + # and QwQ renders none, so a marker holding it masks those responses. "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", @@ -525,29 +527,39 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user<|im_sep|>", "response": "<|im_start|>assistant<|im_sep|>", }, + # No surrounding spaces: in Mistral v0.3 they fold into neighbouring text + # tokens ("[INST]"/"[/INST]" are single special tokens), so padded strings + # never match and everything masks. Same for Llama-2's SentencePiece. "mistral": { - "instruction": "[INST] ", - "response": " [/INST]", + "instruction": "[INST]", + "response": "[/INST]", }, "llama": { - "instruction": "[INST] ", - "response": " [/INST]", + # -anchored: llama-2 tokenizes [INST] after as bare "[" on + # transformers 5.x (standalone gives space-prefixed "▁["), so an + # unanchored marker misses every turn boundary there. + "instruction": "[INST]", + "response": "[/INST]", }, "chatml": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # Leading newline required: Zephyr's role tags are plain text, and + # SentencePiece tokenizes "<|assistant|>" differently at text start than + # after "\n". Without the "\n" anchor the markers never match real + # turns, so every assistant token masks. "zephyr": { - "instruction": "<|user|>\n", - "response": "<|assistant|>\n", + "instruction": "\n<|user|>\n", + "response": "\n<|assistant|>\n", }, "unsloth": { - "instruction": ">>> User: ", - "response": ">>> Assistant: ", + "instruction": ">>> User:", + "response": ">>> Assistant:", }, "vicuna": { - "instruction": "USER: ", - "response": "ASSISTANT: ", + "instruction": "USER:", + "response": "ASSISTANT:", }, "alpaca": { "instruction": "### Instruction:\n", @@ -573,16 +585,21 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No trailing space: SentencePiece folds it into the next content token + # ("▁Hello"), so the padded marker never matches and masks everything. "starling": { - "instruction": "GPT4 Correct User: ", - "response": "GPT4 Correct Assistant: ", + "instruction": "GPT4 Correct User:", + "response": "GPT4 Correct Assistant:", }, "yi-chat": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # "[gMASK]" appears once at text start, so a marker holding it matches + # no later user turn; "" is scaffolding GLM-4.x renders as a lone + # "" on non-final turns, so "<|assistant|>" never matches. "glm": { - "instruction": "[gMASK]<|user|>", - "response": "<|assistant|>", + "instruction": "<|user|>", + "response": "<|assistant|>", }, } From 935474c20aabc2aadb1da17338959c7c6f9bdafe Mon Sep 17 00:00:00 2001 From: WinkleMad Date: Sun, 12 Jul 2026 17:36:11 +0530 Subject: [PATCH 003/210] Fix SyntheticDataKit.chunk_data emitting chunks over max_tokens (#7073) * Fix SyntheticDataKit.chunk_data emitting chunks over max_tokens The multi-chunk path built boundaries from np.linspace(..., n_chunks), but pairing boundaries[:-1] with boundaries[1:] turns N points into N-1 ranges, so it produced one fewer, oversized chunk: every chunk exceeded max_tokens and a document just over the threshold came back as a single unsplit chunk. Use n_chunks + 1 points so exactly n_chunks ranges are emitted, each within max_tokens. Also base n_chunks on the non-overlapped span: consecutive chunks overlap by overlap, so covering length needs ceil((length - overlap) / stride) chunks, not ceil(length / stride). The looser count over-counted by one just past a stride multiple (a 673-token doc became 3 chunks of ~267 instead of 2 of ~369), emitting an extra redundant chunk. Coverage and overlap are unchanged and every chunk still stays within max_tokens. * Condense chunk_data comments and clarify over-split test for PR #7073 --------- Co-authored-by: danielhanchen --- tests/test_synthetic_chunk_data.py | 35 ++++++++++++++++++++++++++++++ unsloth/dataprep/synthetic.py | 9 ++++++-- 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/tests/test_synthetic_chunk_data.py b/tests/test_synthetic_chunk_data.py index abc2c01443..17e5228e32 100644 --- a/tests/test_synthetic_chunk_data.py +++ b/tests/test_synthetic_chunk_data.py @@ -129,6 +129,39 @@ def test_chunk_data_uninitialized_error_names_real_class(): os.unlink(path) +def test_chunk_data_chunks_do_not_exceed_max_tokens(): + # Every chunk must fit within max_tokens. The old multi-chunk path emitted one + # fewer, oversized chunk, and a doc just over the threshold came back unsplit. + kit = _make_kit(max_seq_length = 2048, max_generation_tokens = 760, overlap = 64) + max_tokens = 2048 - 760 * 2 - 128 # 400 + + for n_words in (500, 2000): + out, contents = _chunk("word " * n_words, kit = kit) + assert ( + len(out) >= 2 + ), f"a {n_words}-token doc (> max_tokens={max_tokens}) must be split, got {len(out)}" + for content in contents: + n_tokens = len(content.split()) + assert ( + n_tokens <= max_tokens + ), f"chunk has {n_tokens} tokens, exceeding max_tokens={max_tokens}" + + +def test_chunk_data_does_not_over_split(): + # n_chunks must be the minimum count: ceil((length - overlap) / stride), not + # ceil(length / stride) which over-splits just past a stride multiple. At 673 + # tokens (max_tokens=400, overlap=64) the tight count gives 2 chunks (~369+368). + kit = _make_kit(max_seq_length = 2048, max_generation_tokens = 760, overlap = 64) + max_tokens = 2048 - 760 * 2 - 128 # 400 + out, contents = _chunk("word " * 673, kit = kit) + assert len(out) == 2, f"673-token doc should yield the minimal 2 chunks, got {len(out)}" + for content in contents: + n_tokens = len(content.split()) + assert ( + n_tokens <= max_tokens + ), f"chunk has {n_tokens} tokens, exceeding max_tokens={max_tokens}" + + if __name__ == "__main__": test_chunk_data_keeps_single_chunk_document() test_chunk_data_still_splits_long_document() @@ -136,4 +169,6 @@ if __name__ == "__main__": test_chunk_data_short_document_is_not_split_into_fragments() test_chunk_data_rejects_overlap_not_smaller_than_chunk() test_chunk_data_uninitialized_error_names_real_class() + test_chunk_data_chunks_do_not_exceed_max_tokens() + test_chunk_data_does_not_over_split() print("OK") diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py index 6f025343f5..9c529d5a03 100644 --- a/unsloth/dataprep/synthetic.py +++ b/unsloth/dataprep/synthetic.py @@ -425,8 +425,13 @@ class SyntheticDataKit: else: # length > max_tokens > overlap here, so length - overlap > 0 and the # linspace boundaries below are always non-negative. - n_chunks = int(np.ceil(length / (max_tokens - self.overlap))) - boundaries = np.ceil(np.linspace(0, length - self.overlap, n_chunks)).astype(int) + # Minimal count: overlapping chunks cover `length` in + # ceil((length - overlap) / stride) chunks, not ceil(length / stride) + # which over-splits just past a stride multiple. + n_chunks = int(np.ceil((length - self.overlap) / (max_tokens - self.overlap))) + # n_chunks + 1 points: [:-1]/[1:] pairing yields n_chunks ranges; using + # n_chunks points gave one fewer, oversized chunk (over max_tokens). + boundaries = np.ceil(np.linspace(0, length - self.overlap, n_chunks + 1)).astype(int) boundaries = np.stack((boundaries[:-1], (boundaries + self.overlap)[1:])).T boundaries = np.minimum(boundaries, length).tolist() From 2a22da9fd7cc8bc843ae6aa3a480b8861221729a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 18:19:33 -0700 Subject: [PATCH 004/210] Studio: startup loading banner and mute the benign bitsandbytes ROCm warning (#7085) * Studio: startup loading banner and mute the benign bitsandbytes ROCm warning * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: shorten startup banner wording * [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> --- studio/backend/core/training/worker.py | 9 +++++++++ studio/backend/main.py | 8 ++++++++ studio/backend/run.py | 11 +++++++++++ 3 files changed, 28 insertions(+) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 5fb8fc2bb6..130e6ece64 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2579,6 +2579,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> _bnb_rocm_ver, ) + # Setting BNB_ROCM_VERSION makes bitsandbytes log a benign override + # notice on import; drop only that record so real errors and mismatch + # warnings still show. + if os.environ.get("BNB_ROCM_VERSION"): + import logging as _logging + _logging.getLogger("bitsandbytes.cextension").addFilter( + lambda _r: "environment variable detected" not in _r.getMessage() + ) + # Parse HIP version for the kernel-fix gate below, falling back to # the rocm version embedded in torch.__version__ when version.hip is # unset (AMD SDK / Radeon wheels). diff --git a/studio/backend/main.py b/studio/backend/main.py index 8762c43195..332f4e7d1f 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -159,6 +159,14 @@ if sys.platform == "win32": _bnb_rocm_ver_final, ) + # Setting BNB_ROCM_VERSION makes bitsandbytes log a benign override notice on + # import; drop only that record so real errors and mismatch warnings show. + if os.environ.get("BNB_ROCM_VERSION"): + import logging as _logging + _logging.getLogger("bitsandbytes.cextension").addFilter( + lambda _r: "environment variable detected" not in _r.getMessage() + ) + # ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ────── # In WSL the AMD GPU is reached via the ROCDXG bridge (librocdxg.so over # /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_DETECTION=1 is set BEFORE diff --git a/studio/backend/run.py b/studio/backend/run.py index 2cc6c4a93e..3efeac960e 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1156,6 +1156,15 @@ def run_server( from threading import Thread, Event import uvicorn + # `from main import app` below loads torch/unsloth/transformers (~2 min cold, + # silent), so print a flushed heads-up (piped stdout is block-buffered). + if not silent: + print( + "Loading Unsloth Studio, please wait... (this can take a few minutes)", + flush = True, + ) + print(" - loading PyTorch, Unsloth and Transformers...", flush = True) + import_started = time.perf_counter() from main import app, setup_frontend, _IS_COLAB @@ -1164,6 +1173,8 @@ def run_server( "Imported FastAPI app in %.1fms", (time.perf_counter() - import_started) * 1000, ) + if not silent: + print(" - Starting server...", flush = True) from utils.paths import ensure_studio_directories # Allow local stdio MCP servers on a loopback bind (the user's own machine), From ca979e9643d9641c1f709946514105109ed7c19e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 21:23:14 -0700 Subject: [PATCH 005/210] Studio: add UNSLOTH_SKIP_AUTOSTART installer flag (#7093) * Studio: add installer autostart opt-out * CI: run installer autostart tests cross-platform * Tests: combine Studio installer skip flags --- .../workflows/cross-platform-parity-ci.yml | 14 +- README.md | 8 ++ install.ps1 | 8 +- install.sh | 13 +- tests/sh/test_strixhalo_wsl_reroute.sh | 10 +- tests/test_installer_skip_autostart.py | 126 ++++++++++++++++++ 6 files changed, 169 insertions(+), 10 deletions(-) create mode 100644 tests/test_installer_skip_autostart.py diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml index 4632794587..bb7dcbf8e4 100644 --- a/.github/workflows/cross-platform-parity-ci.yml +++ b/.github/workflows/cross-platform-parity-ci.yml @@ -1,7 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Runs tests/python/test_cross_platform_parity.py on Windows and macOS. +# Runs installer parity and autostart opt-out tests on Windows and macOS. # # Why: that test is the guard that install.sh and install.ps1 stay in # sync, but today it only runs on ubuntu-latest (auto-discovered by @@ -21,6 +21,7 @@ on: paths: - 'install.sh' - 'install.ps1' + - 'tests/test_installer_skip_autostart.py' - 'tests/python/test_cross_platform_parity.py' - '.github/workflows/cross-platform-parity-ci.yml' push: @@ -28,6 +29,7 @@ on: paths: - 'install.sh' - 'install.ps1' + - 'tests/test_installer_skip_autostart.py' - 'tests/python/test_cross_platform_parity.py' - '.github/workflows/cross-platform-parity-ci.yml' workflow_dispatch: @@ -57,5 +59,11 @@ jobs: python-version: '3.12' cache: 'pip' - run: python -m pip install -U pip pytest - - name: Cross-platform parity test - run: python -m pytest tests/python/test_cross_platform_parity.py -q + - name: Cross-platform parity tests + env: + UNSLOTH_NO_TORCH: '1' + run: >- + python -m pytest + tests/python/test_cross_platform_parity.py + tests/test_installer_skip_autostart.py + -q diff --git a/README.md b/README.md index 849ee2e87b..5f1630e2ba 100644 --- a/README.md +++ b/README.md @@ -230,6 +230,14 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh $env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex ``` +Skip the post-install prompt that starts Studio (useful for automated installs): +```bash +curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_SKIP_AUTOSTART=1 sh +``` +```powershell +$env:UNSLOTH_SKIP_AUTOSTART=1; irm https://unsloth.ai/install.ps1 | iex +``` + Pin the Python version: ```bash curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh diff --git a/install.ps1 b/install.ps1 index 0797cd3868..100a3177ba 100644 --- a/install.ps1 +++ b/install.ps1 @@ -6,6 +6,7 @@ # irm | iex cannot forward arguments, so web installs take options as env vars set # before the pipe (flags still work via .\install.ps1): # $env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex # skip PyTorch (GGUF-only) +# $env:UNSLOTH_SKIP_AUTOSTART=1; irm https://unsloth.ai/install.ps1 | iex # do not prompt to launch # $env:UNSLOTH_PYTHON='3.12'; irm https://unsloth.ai/install.ps1 | iex # pin Python version # $env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex # .\install.ps1 --no-torch # equivalent flag @@ -98,6 +99,7 @@ function Install-UnslothStudio { $RepoRoot = "" $TauriMode = $false $SkipTorch = $false + $SkipAutostart = $false $ShortcutsOnly = $false $WithLlamaCppDir = "" $argList = $args @@ -130,6 +132,7 @@ function Install-UnslothStudio { # Env-var equivalent for web installs; an explicit flag still wins. if ($env:UNSLOTH_NO_TORCH -in @('1', 'true', 'yes', 'on')) { $SkipTorch = $true } + if ($env:UNSLOTH_SKIP_AUTOSTART -in @('1', 'true', 'yes', 'on')) { $SkipAutostart = $true } # Propagate to child processes so they also respect verbose mode. # Process-scoped -- does not persist. @@ -2612,9 +2615,10 @@ exit 0 # Diagnostic only; never block install on a probe failure. } - # In interactive terminals, ask the user before starting Studio. + # In interactive terminals, ask the user before starting Studio unless the + # caller explicitly disabled the post-install prompt. # In non-interactive environments (CI, Docker) just print instructions. - $IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected) + $IsInteractive = (-not $SkipAutostart) -and [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected) if ($IsInteractive) { Write-Host "" $reply = Read-Host " Start Unsloth Studio now? [Y/n]" diff --git a/install.sh b/install.sh index 3f4ea92387..3bf6fd1855 100755 --- a/install.sh +++ b/install.sh @@ -8,8 +8,9 @@ # # Piped installs take options as env vars after the pipe (a bare `| sh --no-torch` # makes sh reject --no-torch as its own option). Flags still work via ./install.sh: -# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh # skip PyTorch (GGUF-only) -# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh # pin Python version +# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh # skip PyTorch (GGUF-only) +# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_SKIP_AUTOSTART=1 sh # do not prompt to launch +# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh # pin Python version # curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh # Equivalent flags: ./install.sh --no-torch --python 3.12 (or pipe them: sh -s -- --no-torch) # @@ -49,6 +50,7 @@ PACKAGE_NAME="unsloth" TAURI_MODE=false _USER_PYTHON="" _NO_TORCH_FLAG=false +_SKIP_AUTOSTART=false _VERBOSE=false _SHORTCUTS_ONLY=false _next_is_package=false @@ -88,6 +90,7 @@ done # Env-var equivalents for piped installs; an explicit flag still wins. case "${UNSLOTH_NO_TORCH:-}" in 1|true|TRUE|yes|YES|on|ON) _NO_TORCH_FLAG=true ;; esac +case "${UNSLOTH_SKIP_AUTOSTART:-}" in 1|true|TRUE|yes|YES|on|ON) _SKIP_AUTOSTART=true ;; esac [ -z "$_USER_PYTHON" ] && [ -n "${UNSLOTH_PYTHON:-}" ] && _USER_PYTHON="$UNSLOTH_PYTHON" if [ "$_VERBOSE" = true ]; then @@ -1631,6 +1634,7 @@ _maybe_reroute_strixhalo_to_2404() { # Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the # GPU instead of falling back to the desktop-app prompt path. [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1" + [ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1" _rr_args="" [ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")" [ -n "$_USER_PYTHON" ] && _rr_args="$_rr_args --python $(_rr_q "$_USER_PYTHON")" @@ -3223,9 +3227,10 @@ printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!" printf " ${C_DIM}%s${C_RST}\n" "$RULE" echo "" -# In interactive terminals, ask the user before starting Studio. +# In interactive terminals, ask the user before starting Studio unless the +# caller explicitly disabled the post-install prompt. # In non-interactive environments (Docker, CI, cloud-init) just print instructions. -if [ -t 1 ]; then +if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then echo "" printf " Start Unsloth Studio now? [Y/n] " # No readable answer (closed/EOF tty) defaults to no; Enter is still yes. diff --git a/tests/sh/test_strixhalo_wsl_reroute.sh b/tests/sh/test_strixhalo_wsl_reroute.sh index 7edf475ccb..9c9b04cf21 100644 --- a/tests/sh/test_strixhalo_wsl_reroute.sh +++ b/tests/sh/test_strixhalo_wsl_reroute.sh @@ -309,7 +309,15 @@ if [ "$_rc" = "2" ]; then echo " PASS: tauri child exit 2 -> reroute propagates assert_absent "tauri exit 2 -> not a CPU fallback" "$_out" "__NOROUTE__" rm -rf "$_d" -# 27) Non-tauri mode: a child exit 2 is just a failure -> CPU fallback, not propagated. +# 27) The post-install autostart opt-out must reach the target distro, where the +# final launch prompt is evaluated. +_d=$(make_fixture 1 strix 0 26.04 1) +_out=$(run_func "$_d" _SKIP_AUTOSTART=true UNSLOTH_SKIP_AUTOSTART= \ + UNSLOTH_WSL_REROUTE_CMD='echo skip=[$UNSLOTH_SKIP_AUTOSTART]') +assert_contains "UNSLOTH_SKIP_AUTOSTART forwarded to reroute" "$_out" "skip=[1]" +rm -rf "$_d" + +# 28) Non-tauri mode: a child exit 2 is just a failure -> CPU fallback, not propagated. _d=$(make_fixture 1 strix 0 26.04 1) _rc=0 _out=$(run_func "$_d" UNSLOTH_WSL_REROUTE_CMD='exit 2') || _rc=$? diff --git a/tests/test_installer_skip_autostart.py b/tests/test_installer_skip_autostart.py new file mode 100644 index 0000000000..8c283f4458 --- /dev/null +++ b/tests/test_installer_skip_autostart.py @@ -0,0 +1,126 @@ +"""Regression tests for the installers' post-install autostart opt-out.""" + +from __future__ import annotations + +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +INSTALL_SH = REPO_ROOT / "install.sh" +INSTALL_PS1 = REPO_ROOT / "install.ps1" +README = REPO_ROOT / "README.md" + +TRUTHY_VALUES = ("1", "true", "TRUE", "yes", "YES", "on", "ON") +FALSEY_VALUES = ("", "0", "false", "no", "off", "anything-else") + + +def _extract(pattern: str, source: str) -> str: + match = re.search(pattern, source, flags = re.DOTALL) + assert match is not None, f"installer block not found: {pattern}" + return match.group(0) + + +@pytest.mark.parametrize( + ("value", "expected"), + [(value, "true") for value in TRUTHY_VALUES] + [(value, "false") for value in FALSEY_VALUES], +) +@pytest.mark.skipif(shutil.which("sh") is None, reason = "POSIX shell is unavailable") +def test_posix_skip_autostart_value_parsing_with_no_torch(value: str, expected: str): + source = INSTALL_SH.read_text(encoding = "utf-8") + no_torch_parser = _extract( + r'case "\$\{UNSLOTH_NO_TORCH:-\}" in.*?esac', + source, + ) + autostart_parser = _extract( + r'case "\$\{UNSLOTH_SKIP_AUTOSTART:-\}" in.*?esac', + source, + ) + env = os.environ.copy() + env["UNSLOTH_NO_TORCH"] = "1" + env["UNSLOTH_SKIP_AUTOSTART"] = value + result = subprocess.run( + [ + "sh", + "-c", + ( + f"_NO_TORCH_FLAG=false\n_SKIP_AUTOSTART=false\n{no_torch_parser}\n" + f'{autostart_parser}\nprintf "%s %s" "$_NO_TORCH_FLAG" "$_SKIP_AUTOSTART"' + ), + ], + check = True, + capture_output = True, + text = True, + env = env, + ) + assert result.stdout == f"true {expected}" + + +def test_posix_skip_autostart_bypasses_only_the_interactive_prompt(): + source = INSTALL_SH.read_text(encoding = "utf-8") + gate = 'if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then' + assert gate in source + assert source.index(gate) < source.index("Start Unsloth Studio now? [Y/n]") + assert source.count("Start Unsloth Studio now? [Y/n]") == 1 + assert source.index("Start Unsloth Studio now? [Y/n]") < source.index( + 'step "launch" "manual commands:"' + ) + assert "export UNSLOTH_SKIP_AUTOSTART=1" in source + + +@pytest.mark.skipif(shutil.which("pwsh") is None, reason = "PowerShell is unavailable") +@pytest.mark.parametrize( + ("value", "expected"), + [(value, "True") for value in TRUTHY_VALUES] + [(value, "False") for value in FALSEY_VALUES], +) +def test_windows_skip_autostart_value_parsing_with_no_torch(value: str, expected: str): + source = INSTALL_PS1.read_text(encoding = "utf-8") + parser = _extract( + r"\$SkipTorch = \$false\s+\$SkipAutostart = \$false\s+.*?" + r"if \(\$env:UNSLOTH_NO_TORCH -in @\('1', 'true', 'yes', 'on'\)\) " + r"\{ \$SkipTorch = \$true \}\s+" + r"if \(\$env:UNSLOTH_SKIP_AUTOSTART -in @\('1', 'true', 'yes', 'on'\)\) " + r"\{ \$SkipAutostart = \$true \}", + source, + ) + env = os.environ.copy() + env["UNSLOTH_NO_TORCH"] = "1" + env["UNSLOTH_SKIP_AUTOSTART"] = value + result = subprocess.run( + [ + "pwsh", + "-NoProfile", + "-NonInteractive", + "-Command", + f'{parser}; "$SkipTorch $SkipAutostart"', + ], + check = True, + capture_output = True, + text = True, + env = env, + ) + assert result.stdout.strip() == f"True {expected}" + + +def test_windows_skip_autostart_bypasses_only_the_interactive_prompt(): + source = INSTALL_PS1.read_text(encoding = "utf-8") + gate = ( + "$IsInteractive = (-not $SkipAutostart) -and " + "[Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)" + ) + assert gate in source + assert source.index(gate) < source.index("Start Unsloth Studio now? [Y/n]") + assert source.count("Start Unsloth Studio now? [Y/n]") == 1 + assert source.index("Start Unsloth Studio now? [Y/n]") < source.index( + 'step "launch" "manual commands:"' + ) + + +def test_skip_autostart_is_documented_for_all_installers(): + readme = README.read_text(encoding = "utf-8") + assert "UNSLOTH_SKIP_AUTOSTART=1 sh" in readme + assert "$env:UNSLOTH_SKIP_AUTOSTART=1" in readme From 9e77c1e663ceba036c4007bbfb21ce707e61b933 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 00:15:04 -0700 Subject: [PATCH 006/210] Studio: remove AGENTS.md and CLAUDE.md from install artifacts (#7096) * Studio: remove AGENTS.md from install artifacts * Studio: prune CLAUDE.md from install artifacts * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Studio instruction cleanup edge cases * Trim Studio cleanup comments * Make Studio cleanup safe on PowerShell 5.1 * Fix Studio cleanup ownership boundaries * Simplify Windows link detection --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../src/i18n/{AGENTS.md => README.md} | 4 +- studio/install_llama_prebuilt.py | 58 +++- studio/setup.ps1 | 38 +++ studio/setup.sh | 20 ++ .../test_install_llama_prebuilt_logic.py | 256 ++++++++++++++++++ 5 files changed, 373 insertions(+), 3 deletions(-) rename studio/frontend/src/i18n/{AGENTS.md => README.md} (94%) diff --git a/studio/frontend/src/i18n/AGENTS.md b/studio/frontend/src/i18n/README.md similarity index 94% rename from studio/frontend/src/i18n/AGENTS.md rename to studio/frontend/src/i18n/README.md index 35c025cd0b..ba8cce5d8d 100644 --- a/studio/frontend/src/i18n/AGENTS.md +++ b/studio/frontend/src/i18n/README.md @@ -1,4 +1,4 @@ -# i18n Contribution Instructions +# i18n Contribution Guide - `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. @@ -9,4 +9,4 @@ - 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 +- 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. diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index ca1fe79efa..cf9b2cf45f 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -20,6 +20,7 @@ import re import shutil import site import socket +import stat import struct import subprocess import sys @@ -4385,6 +4386,48 @@ def copy_directory_contents(source_dir: Path, destination: Path) -> None: shutil.copy2(item, target) +def _is_link_or_junction(path: Path) -> bool: + """Return whether ``path`` redirects to another filesystem location.""" + if os.name == "nt": + try: + attributes = getattr(path.lstat(), "st_file_attributes", 0) + except OSError: + return True + return bool(attributes & stat.FILE_ATTRIBUTE_REPARSE_POINT) + try: + return path.is_symlink() + except OSError: + return True + + +def remove_agent_instruction_files(root: Path) -> int: + """Best-effort removal inside a managed tree without following links.""" + if _is_link_or_junction(root) or not root.is_dir(): + return 0 + + removed = 0 + for current_dir, dirnames, filenames in os.walk(root, topdown = True, followlinks = False): + current_path = Path(current_dir) + # followlinks=False still follows Windows junctions. + if current_path != root and _is_link_or_junction(current_path): + dirnames.clear() + continue + dirnames[:] = [ + dirname for dirname in dirnames if not _is_link_or_junction(current_path / dirname) + ] + for filename in sorted({"AGENTS.md", "CLAUDE.md"}.intersection(filenames)): + candidate = current_path / filename + try: + candidate.unlink() + except FileNotFoundError: + continue + except OSError as exc: + log(f"could not remove contributor-only instruction {candidate}: {exc}") + else: + removed += 1 + return removed + + def hydrate_source_tree( source_ref: str, install_dir: Path, @@ -4447,6 +4490,9 @@ def hydrate_source_tree( "upstream source archive was missing required repo files: " + ", ".join(missing) ) copy_directory_contents(source_root, install_dir) + removed = remove_agent_instruction_files(install_dir) + if removed: + log(f"removed {removed} contributor-only agent instruction file(s) from staged source") except PrebuiltFallback: raise except Exception as exc: @@ -6824,6 +6870,7 @@ def install_prebuilt( override_has_rocm: bool = False, override_rocm_gfx: str | None = None, force_cpu: bool = False, + instruction_cleanup_root: Path | None = None, ) -> None: host = detect_host() host = _apply_host_overrides( @@ -6836,8 +6883,15 @@ def install_prebuilt( host, published_repo, published_release_tag, force_cpu = force_cpu ) choice: AssetChoice | None = None + cleanup_root = install_dir if instruction_cleanup_root is None else instruction_cleanup_root try: with install_lock(install_lock_path(install_dir)): + if (install_dir / "UNSLOTH_PREBUILT_INFO.json").is_file(): + removed = remove_agent_instruction_files(cleanup_root) + if removed: + log( + f"removed {removed} contributor-only agent instruction file(s) from install" + ) if install_dir.exists(): log( f"existing llama.cpp install detected at {install_dir}; validating staged prebuilt update before replacement" @@ -7168,14 +7222,16 @@ def main() -> int: # Install path only: route status logs to stdout (see _LOG_TO_STDOUT note). global _LOG_TO_STDOUT _LOG_TO_STDOUT = True + install_arg = Path(args.install_dir).expanduser() install_prebuilt( - install_dir = Path(args.install_dir).expanduser().resolve(), + install_dir = install_arg.resolve(), llama_tag = args.llama_tag, published_repo = args.published_repo, published_release_tag = args.published_release_tag or "", override_has_rocm = args.has_rocm, override_rocm_gfx = args.rocm_gfx, force_cpu = args.cpu_fallback, + instruction_cleanup_root = install_arg.absolute(), ) return EXIT_SUCCESS diff --git a/studio/setup.ps1 b/studio/setup.ps1 index db01a1ecad..2c828281d2 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -203,6 +203,31 @@ function New-UnslothTemporaryFile { return Get-Item -LiteralPath $tempPath } +function Remove-AgentInstructionFiles { + param([string[]]$Roots) + + foreach ($root in $Roots) { + if (-not $root) { continue } + $item = Get-Item -LiteralPath $root -Force -ErrorAction SilentlyContinue + if (-not $item -or -not $item.PSIsContainer) { continue } + if ($item.Attributes -band [System.IO.FileAttributes]::ReparsePoint) { continue } + $pending = New-Object System.Collections.Stack + $pending.Push($item) + while ($pending.Count -gt 0) { + $current = $pending.Pop() + foreach ($child in @(Get-ChildItem -LiteralPath $current.FullName -Force -ErrorAction SilentlyContinue)) { + if ($child.PSIsContainer) { + if (-not ($child.Attributes -band [System.IO.FileAttributes]::ReparsePoint)) { + $pending.Push($child) + } + } elseif ($child.Name -in @("AGENTS.md", "CLAUDE.md")) { + Remove-Item -LiteralPath $child.FullName -Force -ErrorAction SilentlyContinue + } + } + } + } +} + function Get-InstalledLlamaPrebuiltRelease { param([string]$InstallDir) @@ -2314,6 +2339,8 @@ if ((Test-Path $OxcValidatorDir) -and $NodeSource -ne "skip" -and (Get-Command n substep "OXC validator runtime skipped (no npm found); code validation degrades until Node is available" "Yellow" } +Remove-AgentInstructionFiles -Roots @($FrontendDir, $OxcValidatorDir) + # ========================================================================== # PHASE 3: Python environment + dependencies # ========================================================================== @@ -3259,6 +3286,7 @@ if ($LocalLlamaCppSrc) { if ($LASTEXITCODE -ne 0) { substep "Could not create directory junction; copying instead..." "Yellow" Copy-Item -Recurse -LiteralPath $ResolvedLocal -Destination $LlamaCppDir + Remove-AgentInstructionFiles -Roots @($LlamaCppDir) } Write-Host "" step "llama.cpp" "linked local directory: $ResolvedLocal" @@ -4024,6 +4052,16 @@ if ($LocalLlamaCppLinked) { } } +$llamaCppItem = Get-Item -LiteralPath $LlamaCppDir -Force -ErrorAction SilentlyContinue +$llamaCppIsLink = $llamaCppItem -and ($llamaCppItem.Attributes -band [System.IO.FileAttributes]::ReparsePoint) +if (-not $llamaCppIsLink -and ( + -not $StudioHomeIsCustom -or + (Test-Path -LiteralPath (Join-Path $LlamaCppDir $StudioOwnedMarker) -PathType Leaf) -or + (Test-StudioOwnedAdoptable $LlamaCppDir) + )) { + Remove-AgentInstructionFiles -Roots @($LlamaCppDir) +} + # ───────────────────────────────────────────── # Footer # ───────────────────────────────────────────── diff --git a/studio/setup.sh b/studio/setup.sh index d244e3cdcf..3e81b332ad 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -74,6 +74,16 @@ verbose_substep() { return 0 } +_remove_agent_instruction_files() { + local _root + for _root in "$@"; do + [ -d "$_root" ] || continue + [ -L "$_root" ] && continue + find "$_root" -type f \( -name 'AGENTS.md' -o -name 'CLAUDE.md' \) \ + -exec rm -f {} + 2>/dev/null || true + done +} + # ── Corporate-mirror / proxy escape hatch for the frontend npm/bun install (#6491) ── # studio/frontend/.npmrc pins registry=https://registry.npmjs.org/ as a supply-chain # lock. A project-level pin overrides a corporate user's ~/.npmrc proxy, so the install @@ -847,6 +857,8 @@ elif [ -d "$_OXC_DIR" ] && [ "${NODE_SOURCE:-}" != skip ]; then substep "OXC validator runtime skipped (no npm found); code validation degrades until Node is available" "$C_WARN" fi +_remove_agent_instruction_files "$SCRIPT_DIR/frontend" "$_OXC_DIR" + # ── Python venv + deps ── [ -d "$REPO_ROOT/.venv" ] && rm -rf "$REPO_ROOT/.venv" @@ -1919,6 +1931,14 @@ if [ "$_LLAMA_CPP_DEGRADED" = true ] \ fi fi +if [ ! -L "$LLAMA_CPP_DIR" ] && { + [ "$_STUDIO_HOME_IS_CUSTOM" != true ] || + [ -f "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" ] || + _studio_owned_adoptable "$LLAMA_CPP_DIR" +}; then + _remove_agent_instruction_files "$LLAMA_CPP_DIR" +fi + # ── Footer ── if [ "$_LLAMA_ONLY" = "1" ]; then echo "" diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index c852d7c495..e4b9c7df85 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -3,6 +3,7 @@ import importlib.util import io import json import os +import subprocess import sys import tarfile import zipfile @@ -30,6 +31,7 @@ AssetChoice = INSTALL_LLAMA_PREBUILT.AssetChoice ApprovedArtifactHash = INSTALL_LLAMA_PREBUILT.ApprovedArtifactHash ApprovedReleaseChecksums = INSTALL_LLAMA_PREBUILT.ApprovedReleaseChecksums hydrate_source_tree = INSTALL_LLAMA_PREBUILT.hydrate_source_tree +remove_agent_instruction_files = INSTALL_LLAMA_PREBUILT.remove_agent_instruction_files validate_prebuilt_choice = INSTALL_LLAMA_PREBUILT.validate_prebuilt_choice activate_install_tree = INSTALL_LLAMA_PREBUILT.activate_install_tree activate_staged_dir = INSTALL_LLAMA_PREBUILT.activate_staged_dir @@ -203,6 +205,178 @@ def test_extract_archive_rejects_zip_symlink_entry(tmp_path: Path): extract_archive(archive_path, tmp_path / "extract") +def test_remove_agent_instruction_files_does_not_follow_links(tmp_path: Path): + managed = tmp_path / "managed" + nested = managed / "nested" + external = tmp_path / "external" + nested.mkdir(parents = True) + external.mkdir() + (managed / "AGENTS.md").write_text("managed root", encoding = "utf-8") + (nested / "AGENTS.md").write_text("managed nested", encoding = "utf-8") + (managed / "CLAUDE.md").write_text("managed Claude root", encoding = "utf-8") + (nested / "CLAUDE.md").write_text("managed Claude nested", encoding = "utf-8") + (external / "AGENTS.md").write_text("user owned", encoding = "utf-8") + (external / "CLAUDE.md").write_text("user-owned Claude", encoding = "utf-8") + try: + (managed / "external-link").symlink_to(external, target_is_directory = True) + linked_root = tmp_path / "linked-root" + linked_root.symlink_to(external, target_is_directory = True) + except OSError as exc: + pytest.skip(f"directory symlinks unavailable: {exc}") + + assert remove_agent_instruction_files(managed) == 4 + assert not list(managed.rglob("AGENTS.md")) + assert not list(managed.rglob("CLAUDE.md")) + assert (external / "AGENTS.md").read_text(encoding = "utf-8") == "user owned" + assert (external / "CLAUDE.md").read_text(encoding = "utf-8") == "user-owned Claude" + + assert remove_agent_instruction_files(linked_root) == 0 + assert (external / "AGENTS.md").exists() + assert (external / "CLAUDE.md").exists() + + +@pytest.mark.skipif(os.name != "nt", reason = "Windows junction behavior") +def test_remove_agent_instruction_files_does_not_follow_windows_junctions(tmp_path: Path): + managed = tmp_path / "managed" + external = tmp_path / "external" + managed.mkdir() + external.mkdir() + (external / "AGENTS.md").write_text("user owned", encoding = "utf-8") + (external / "CLAUDE.md").write_text("user-owned Claude", encoding = "utf-8") + + nested_junction = managed / "external-junction" + root_junction = tmp_path / "linked-root" + for junction in (nested_junction, root_junction): + result = subprocess.run( + ["cmd", "/d", "/c", "mklink", "/J", str(junction), str(external)], + capture_output = True, + text = True, + check = False, + ) + if result.returncode != 0: + pytest.skip(f"directory junctions unavailable: {result.stderr or result.stdout}") + + assert remove_agent_instruction_files(managed) == 0 + assert remove_agent_instruction_files(root_junction) == 0 + assert (external / "AGENTS.md").read_text(encoding = "utf-8") == "user owned" + assert (external / "CLAUDE.md").read_text(encoding = "utf-8") == "user-owned Claude" + + +def test_remove_agent_instruction_files_prunes_linklike_directories( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + managed = tmp_path / "managed" + simulated_junction = managed / "simulated-junction" + simulated_junction.mkdir(parents = True) + agents = simulated_junction / "AGENTS.md" + claude = simulated_junction / "CLAUDE.md" + agents.write_text("external instructions", encoding = "utf-8") + claude.write_text("external Claude instructions", encoding = "utf-8") + real_is_link_or_junction = INSTALL_LLAMA_PREBUILT._is_link_or_junction + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "_is_link_or_junction", + lambda path: path == simulated_junction or real_is_link_or_junction(path), + ) + + assert remove_agent_instruction_files(managed) == 0 + assert agents.exists() + assert claude.exists() + + +def test_remove_agent_instruction_files_continues_after_unlink_error( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +): + managed = tmp_path / "managed" + managed.mkdir() + blocked = managed / "AGENTS.md" + removable = managed / "CLAUDE.md" + blocked.write_text("blocked", encoding = "utf-8") + removable.write_text("remove me", encoding = "utf-8") + real_unlink = Path.unlink + + def selective_unlink(path: Path, *args, **kwargs): + if path == blocked: + raise PermissionError(errno.EACCES, "Access is denied", str(path)) + return real_unlink(path, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", selective_unlink) + + assert remove_agent_instruction_files(managed) == 1 + assert blocked.exists() + assert not removable.exists() + captured = capsys.readouterr() + assert "could not remove contributor-only instruction" in captured.out + captured.err + + +def test_main_resolves_linked_install_path_and_preserves_cleanup_root( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + target = tmp_path / "target" + linked_root = tmp_path / "linked-root" + target.mkdir() + try: + linked_root.symlink_to(target, target_is_directory = True) + except OSError as exc: + pytest.skip(f"directory symlinks unavailable: {exc}") + + received = {} + monkeypatch.setattr( + sys, + "argv", + ["install_llama_prebuilt.py", "--install-dir", str(linked_root)], + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "install_prebuilt", + lambda **kwargs: received.update(kwargs), + ) + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_LOG_TO_STDOUT", False) + + assert INSTALL_LLAMA_PREBUILT.main() == 0 + assert received["install_dir"] == target.resolve() + assert received["instruction_cleanup_root"] == linked_root.absolute() + assert received["instruction_cleanup_root"].is_symlink() + + +def test_install_prebuilt_uses_explicit_instruction_cleanup_root( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "target" + linked_root = tmp_path / "linked-root" + install_dir.mkdir() + (install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text("{}", encoding = "utf-8") + try: + linked_root.symlink_to(install_dir, target_is_directory = True) + except OSError as exc: + pytest.skip(f"directory symlinks unavailable: {exc}") + + cleanup_roots = [] + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", linux_host) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "remove_agent_instruction_files", + lambda root: cleanup_roots.append(root) or 0, + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "resolve_simple_install_release_plans", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("stop after cleanup")), + ) + + with pytest.raises(RuntimeError, match = "stop after cleanup"): + install_prebuilt( + install_dir.resolve(), + "latest", + "unslothai/llama.cpp", + "", + instruction_cleanup_root = linked_root.absolute(), + ) + + assert cleanup_roots == [linked_root.absolute()] + + def test_hydrate_source_tree_extracts_upstream_archive_contents( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): @@ -224,6 +398,26 @@ def test_hydrate_source_tree_extracts_upstream_archive_contents( f"llama.cpp-{upstream_tag}/gguf-py/gguf/__init__.py", b"__all__ = []\n", ) + add_bytes_to_tar( + archive, + f"llama.cpp-{upstream_tag}/AGENTS.md", + b"upstream contributor instructions\n", + ) + add_bytes_to_tar( + archive, + f"llama.cpp-{upstream_tag}/examples/AGENTS.md", + b"nested contributor instructions\n", + ) + add_bytes_to_tar( + archive, + f"llama.cpp-{upstream_tag}/CLAUDE.md", + b"Claude contributor instructions\n", + ) + add_bytes_to_tar( + archive, + f"llama.cpp-{upstream_tag}/examples/CLAUDE.md", + b"nested Claude contributor instructions\n", + ) source_urls = set(INSTALL_LLAMA_PREBUILT.upstream_source_archive_urls(upstream_tag)) @@ -244,6 +438,8 @@ def test_hydrate_source_tree_extracts_upstream_archive_contents( assert (install_dir / "convert_hf_to_gguf.py").exists() assert (install_dir / "gguf-py" / "gguf" / "__init__.py").exists() assert not (install_dir / f"llama.cpp-{upstream_tag}").exists() + assert not list(install_dir.rglob("AGENTS.md")) + assert not list(install_dir.rglob("CLAUDE.md")) def test_release_asset_download_url(): @@ -644,6 +840,29 @@ def test_activate_install_tree_restores_existing_install_after_activation_failur assert "restored previous install from rollback path" in output +def test_activate_install_tree_preserves_symlink_to_resolved_target( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "target" + linked_root = tmp_path / "linked-root" + staging_dir = tmp_path / "staging" + install_dir.mkdir() + staging_dir.mkdir() + (install_dir / "old.txt").write_text("old", encoding = "utf-8") + (staging_dir / "new.txt").write_text("new", encoding = "utf-8") + try: + linked_root.symlink_to(install_dir, target_is_directory = True) + except OSError as exc: + pytest.skip(f"directory symlinks unavailable: {exc}") + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "confirm_install_tree", lambda *_args: None) + + activate_install_tree(staging_dir, linked_root.resolve(), linux_host()) + + assert linked_root.is_symlink() + assert (linked_root / "new.txt").read_text(encoding = "utf-8") == "new" + assert not (linked_root / "old.txt").exists() + + def test_activate_install_tree_cleans_all_paths_when_rollback_restore_fails( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ): @@ -2007,6 +2226,14 @@ def test_install_prebuilt_skips_download_when_existing_install_matches( approved_checksums = checksums, prebuilt_fallback_used = False, ) + (install_dir / "AGENTS.md").write_text("old root instructions", encoding = "utf-8") + nested_agents = install_dir / "examples" / "AGENTS.md" + nested_agents.parent.mkdir() + nested_agents.write_text("old nested instructions", encoding = "utf-8") + (install_dir / "CLAUDE.md").write_text("old Claude instructions", encoding = "utf-8") + (nested_agents.parent / "CLAUDE.md").write_text( + "old nested Claude instructions", encoding = "utf-8" + ) monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: host) monkeypatch.setattr( @@ -2026,6 +2253,35 @@ def test_install_prebuilt_skips_download_when_existing_install_matches( ) install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "") + assert not list(install_dir.rglob("AGENTS.md")) + assert not list(install_dir.rglob("CLAUDE.md")) + + +def test_setup_scripts_prune_agent_files_without_shipping_a_repo_copy(): + setup_sh = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8") + setup_ps1 = (PACKAGE_ROOT / "studio" / "setup.ps1").read_text(encoding = "utf-8") + + assert '_remove_agent_instruction_files "$SCRIPT_DIR/frontend" "$_OXC_DIR"' in setup_sh + assert '_remove_agent_instruction_files "$LLAMA_CPP_DIR"' in setup_sh + assert "-name 'CLAUDE.md'" in setup_sh + assert 'if [ ! -L "$LLAMA_CPP_DIR" ] && {' in setup_sh + assert '${_LOCAL_LLAMA_CPP_LINKED:-false}" != true' not in setup_sh + assert "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" in setup_sh + assert '_studio_owned_adoptable "$LLAMA_CPP_DIR"' in setup_sh + assert "Remove-AgentInstructionFiles -Roots @($FrontendDir, $OxcValidatorDir)" in setup_ps1 + assert '"CLAUDE.md"' in setup_ps1 + assert '-Include "AGENTS.md", "CLAUDE.md"' not in setup_ps1 + assert '$child.Name -in @("AGENTS.md", "CLAUDE.md")' in setup_ps1 + assert "$llamaCppIsLink" in setup_ps1 + assert "if (-not $LocalLlamaCppLinked)" not in setup_ps1 + assert "Join-Path $LlamaCppDir $StudioOwnedMarker" in setup_ps1 + assert "Test-StudioOwnedAdoptable $LlamaCppDir" in setup_ps1 + assert ( + "Copy-Item -Recurse -LiteralPath $ResolvedLocal -Destination $LlamaCppDir\n" + " Remove-AgentInstructionFiles -Roots @($LlamaCppDir)" + ) in setup_ps1 + assert not (PACKAGE_ROOT / "studio" / "frontend" / "src" / "i18n" / "AGENTS.md").exists() + assert (PACKAGE_ROOT / "studio" / "frontend" / "src" / "i18n" / "README.md").is_file() def test_install_prebuilt_does_not_skip_unhealthy_existing_install( From c570180a32bd5aa97be1e8af6328aef45c8da75d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 01:46:23 -0700 Subject: [PATCH 007/210] Tighten Studio instruction-file cleanup boundaries (#7097) * Handle linked instruction files in Bash cleanup * Limit instruction cleanup to managed dependencies * Make Bash cleanup test portable * Run junction cleanup regression on Windows * Keep instruction cleanup CI focused --- studio/setup.ps1 | 5 ++- studio/setup.sh | 6 ++- .../test_install_llama_prebuilt_logic.py | 43 ++++++++++++++++++- 3 files changed, 49 insertions(+), 5 deletions(-) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 2c828281d2..b19167c478 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2339,7 +2339,10 @@ if ((Test-Path $OxcValidatorDir) -and $NodeSource -ne "skip" -and (Get-Command n substep "OXC validator runtime skipped (no npm found); code validation degrades until Node is available" "Yellow" } -Remove-AgentInstructionFiles -Roots @($FrontendDir, $OxcValidatorDir) +Remove-AgentInstructionFiles -Roots @( + (Join-Path $FrontendDir "node_modules"), + (Join-Path $OxcValidatorDir "node_modules") +) # ========================================================================== # PHASE 3: Python environment + dependencies diff --git a/studio/setup.sh b/studio/setup.sh index 3e81b332ad..64d4f852b5 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -79,7 +79,7 @@ _remove_agent_instruction_files() { for _root in "$@"; do [ -d "$_root" ] || continue [ -L "$_root" ] && continue - find "$_root" -type f \( -name 'AGENTS.md' -o -name 'CLAUDE.md' \) \ + find "$_root" \( -type f -o -type l \) \( -name 'AGENTS.md' -o -name 'CLAUDE.md' \) \ -exec rm -f {} + 2>/dev/null || true done } @@ -857,7 +857,9 @@ elif [ -d "$_OXC_DIR" ] && [ "${NODE_SOURCE:-}" != skip ]; then substep "OXC validator runtime skipped (no npm found); code validation degrades until Node is available" "$C_WARN" fi -_remove_agent_instruction_files "$SCRIPT_DIR/frontend" "$_OXC_DIR" +_remove_agent_instruction_files \ + "$SCRIPT_DIR/frontend/node_modules" \ + "$_OXC_DIR/node_modules" # ── Python venv + deps ── diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index e4b9c7df85..e995e5033e 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -3,6 +3,7 @@ import importlib.util import io import json import os +import shutil import subprocess import sys import tarfile @@ -2261,14 +2262,25 @@ def test_setup_scripts_prune_agent_files_without_shipping_a_repo_copy(): setup_sh = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8") setup_ps1 = (PACKAGE_ROOT / "studio" / "setup.ps1").read_text(encoding = "utf-8") - assert '_remove_agent_instruction_files "$SCRIPT_DIR/frontend" "$_OXC_DIR"' in setup_sh + assert ( + "_remove_agent_instruction_files \\\n" + ' "$SCRIPT_DIR/frontend/node_modules" \\\n' + ' "$_OXC_DIR/node_modules"' + ) in setup_sh + assert '_remove_agent_instruction_files "$SCRIPT_DIR/frontend" "$_OXC_DIR"' not in setup_sh assert '_remove_agent_instruction_files "$LLAMA_CPP_DIR"' in setup_sh assert "-name 'CLAUDE.md'" in setup_sh assert 'if [ ! -L "$LLAMA_CPP_DIR" ] && {' in setup_sh assert '${_LOCAL_LLAMA_CPP_LINKED:-false}" != true' not in setup_sh assert "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" in setup_sh assert '_studio_owned_adoptable "$LLAMA_CPP_DIR"' in setup_sh - assert "Remove-AgentInstructionFiles -Roots @($FrontendDir, $OxcValidatorDir)" in setup_ps1 + assert ( + "Remove-AgentInstructionFiles -Roots @(\n" + ' (Join-Path $FrontendDir "node_modules"),\n' + ' (Join-Path $OxcValidatorDir "node_modules")\n' + ")" + ) in setup_ps1 + assert "Remove-AgentInstructionFiles -Roots @($FrontendDir, $OxcValidatorDir)" not in setup_ps1 assert '"CLAUDE.md"' in setup_ps1 assert '-Include "AGENTS.md", "CLAUDE.md"' not in setup_ps1 assert '$child.Name -in @("AGENTS.md", "CLAUDE.md")' in setup_ps1 @@ -2284,6 +2296,33 @@ def test_setup_scripts_prune_agent_files_without_shipping_a_repo_copy(): assert (PACKAGE_ROOT / "studio" / "frontend" / "src" / "i18n" / "README.md").is_file() +def test_setup_sh_cleanup_unlinks_instruction_symlink_only(tmp_path: Path): + if shutil.which("bash") is None: + pytest.skip("bash is not available") + + setup_sh = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8") + start = setup_sh.index("_remove_agent_instruction_files() {") + end = setup_sh.index("\n}\n", start) + 2 + function = setup_sh[start:end] + managed = tmp_path / "managed" + external = tmp_path / "external.md" + managed.mkdir() + external.write_text("external", encoding = "utf-8") + instruction = managed / "AGENTS.md" + try: + instruction.symlink_to(external) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + subprocess.run( + ["bash", "-c", function + '\n_remove_agent_instruction_files "$1"', "bash", str(managed)], + check = True, + ) + + assert not os.path.lexists(instruction) + assert external.read_text(encoding = "utf-8") == "external" + + def test_install_prebuilt_does_not_skip_unhealthy_existing_install( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): From cc8599207cce0d22d52a46bf4a26bab10b8cfbf9 Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Tue, 14 Jul 2026 01:42:40 +0800 Subject: [PATCH 008/210] Fix Studio user-message overflow for long unbroken text (#7100) --- studio/frontend/src/components/assistant-ui/thread.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index d987092c48..230a1fb40e 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -4006,7 +4006,7 @@ const UserMessage: FC = () => {
-
+
From 2573dbdd6bfc74ce12bb6ca64dbbc85117bc86ae Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Tue, 14 Jul 2026 02:08:45 +0800 Subject: [PATCH 009/210] fix(studio): use writable recipe artifact path (#7044) --- studio/backend/core/data_recipe/service.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 9d8ca5cfcc..4647dc098d 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -9,6 +9,8 @@ import os from pathlib import Path from typing import Any +from utils.paths import recipe_datasets_root + from .jsonable import to_jsonable from .local_callable_validators import ( register_oxc_local_callable_validators, @@ -277,6 +279,11 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None = _apply_data_designer_image_context_patch() from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports] + if artifact_path is None: + # DataDesigner defaults to cwd/artifacts; packaged Studio can run with + # cwd=/, so keep default callers on Studio's writable recipe artifact root. + artifact_path = str(recipe_datasets_root()) + recipe = _strip_frontend_model_config_metadata(recipe) model_providers = build_model_providers(recipe) _validate_recipe_runtime_support(recipe, model_providers) From a337c72753b2aba50a19613c751d15e315ad1cac Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Tue, 14 Jul 2026 02:13:50 +0800 Subject: [PATCH 010/210] Fix Studio auto-titles for reasoning models (#7098) --- .../src/features/chat/runtime-provider.tsx | 9 +++-- tests/studio/test_chat_title_generation.py | 36 ++++++++++++++++--- 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index b545695f9e..634e5ec8b0 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -77,6 +77,7 @@ const pendingRunStartReadyByMessageId = new Map>(); type TitleResponse = { choices?: Array<{ + finish_reason?: string | null; message?: { content?: string; }; @@ -474,6 +475,8 @@ async function generateTitleWithModel(payload: { max_tokens: 24, top_k: 20, repetition_penalty: 1.0, + enable_thinking: false, + reasoning_effort: "none", messages: [ { role: "system", @@ -489,8 +492,10 @@ async function generateTitleWithModel(payload: { .json() .catch(() => null)) as TitleResponse | null; if (!response.ok) return null; - const raw: string | undefined = body?.choices?.[0]?.message?.content; - if (!raw) return null; + const choice = body?.choices?.[0]; + if (choice?.finish_reason === "length") return null; + const raw: string | undefined = choice?.message?.content; + if (!raw || /<\/?think>/i.test(raw)) return null; return normalizeTitle(raw); } diff --git a/tests/studio/test_chat_title_generation.py b/tests/studio/test_chat_title_generation.py index 4f82b4ec03..b568a51400 100644 --- a/tests/studio/test_chat_title_generation.py +++ b/tests/studio/test_chat_title_generation.py @@ -65,6 +65,8 @@ def test_title_model_payload_includes_optional_assistant_reply(): assert "if (assistant)" in block assert "parts.push(`Assistant: ${assistant}`);" in block assert 'parts.join("\\n")' in block + assert "enable_thinking: false" in block + assert 'reasoning_effort: "none"' in block def test_generate_title_passes_first_assistant_reply_after_first_user(): @@ -81,6 +83,25 @@ def test_generate_title_passes_first_assistant_reply_after_first_user(): assert "assistantText," in block +def test_tool_call_only_first_assistant_still_uses_first_user_message(): + source = RUNTIME_TSX.read_text() + extract_block = " ".join(_balanced_block(source, "function extractTextParts").split()) + generate_block = " ".join(_balanced_block(source, "async generateTitle(remoteId").split()) + + assert ( + '.filter((p): p is Extract => p.type === "text")' + in extract_block + ) + assert ( + "const userText = extractTextParts(firstUser) || defaultTitle; const assistantText = extractTextParts(firstAssistant);" + in generate_block + ) + assert ( + "(await generateTitleWithModel({ userText, assistantText, })) || fallbackTitleFromUserText(userText);" + in generate_block + ) + + def test_auto_title_disabled_uses_deterministic_user_text_fallback(): block = _balanced_block( RUNTIME_TSX.read_text(), @@ -93,12 +114,18 @@ def test_auto_title_disabled_uses_deterministic_user_text_fallback(): def test_model_failure_still_falls_back_to_user_text(): - block = _balanced_block( - RUNTIME_TSX.read_text(), - "async generateTitle(remoteId", + source = RUNTIME_TSX.read_text() + model_block = _source_until( + source, + "async function generateTitleWithModel", + "\nconst inflightTitleByKey", ) + generate_block = _balanced_block(source, "async generateTitle(remoteId") - assert "})) || fallbackTitleFromUserText(userText);" in block + assert "finish_reason?: string | null;" in source + assert 'if (choice?.finish_reason === "length") return null;' in model_block + assert r"if (!raw || /<\/?think>/i.test(raw)) return null;" in model_block + assert "})) || fallbackTitleFromUserText(userText);" in generate_block def test_title_normalizer_still_enforces_output_constraints(): @@ -113,3 +140,4 @@ def test_title_normalizer_still_enforces_output_constraints(): assert 'replace(/[.!?:;,]+/g, " ")' in block assert 'title.split(" ").filter(Boolean).slice(0, 6)' in block assert "joined.length > 60" in block + assert "return normalizeTitle(raw);" in block From 85f5292097638d7f005945005db8b083c6794d3c Mon Sep 17 00:00:00 2001 From: oobabooga Date: Mon, 13 Jul 2026 15:43:13 -0300 Subject: [PATCH 011/210] Studio: resync model state after a llama.cpp update unloads it (#6998) --- .../src/components/llama-update-banner.tsx | 13 +- .../src/features/chat/chat-settings-sheet.tsx | 6 +- .../chat/hooks/use-chat-model-runtime.ts | 144 +++++++++++------- studio/frontend/src/features/chat/index.ts | 5 +- .../src/hooks/use-llama-update-check.ts | 137 +++++++++++++++-- 5 files changed, 231 insertions(+), 74 deletions(-) diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 840383de90..3db15ffe30 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { Button } from "@/components/ui/button"; +import { resyncInferenceStatusAfterServerModelChange } from "@/features/chat"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref"; import { toast } from "@/lib/toast"; @@ -81,9 +82,14 @@ export function LlamaUpdateBanner({ positioned = true, }: LlamaUpdateBannerProps): ReactElement | null { const showBannerPref = useShowLlamaUpdateBanner(); + // Not gated on showBannerPref: this hook instance is the app-wide listener + // for a cross-tab reload_required resync (the settings-sheet's own instance + // only runs during an MTP-fallback rebuild), so muting the banner must not + // also silence that resync -- it only suppresses the UI below. const { status, visible, applying, apply, dismiss, snooze } = useLlamaUpdateCheck({ - enabled: enabled && showBannerPref, + enabled, + onReloadRequired: resyncInferenceStatusAfterServerModelChange, }); async function handleUpdate() { @@ -102,7 +108,10 @@ export function LlamaUpdateBanner({ } const show = - visible && status != null && (status.update_available || applying); + showBannerPref && + visible && + status != null && + (status.update_available || applying); const sizeBytes = status?.update_size_bytes ?? null; const sizeLabel = sizeBytes && sizeBytes > 0 diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index ea6c409b40..a547b88e57 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -80,6 +80,7 @@ import { Fragment, type ReactNode } from "react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "@/lib/toast"; import { OpenAICodeExecSection } from "./components/openai-code-exec-section"; +import { resyncInferenceStatusAfterServerModelChange } from "./hooks/use-chat-model-runtime"; import { type ExternalProviderConfig, getExternalProviderApiKey, @@ -573,7 +574,10 @@ export function ChatSettingsPanel({ status: llamaUpdateStatus, applying: llamaUpdating, apply: applyLlamaUpdate, - } = useLlamaUpdateCheck({ enabled: mtpUpdatable }); + } = useLlamaUpdateCheck({ + enabled: mtpUpdatable, + onReloadRequired: resyncInferenceStatusAfterServerModelChange, + }); const handleMtpUpdate = useCallback(async () => { const result = await applyLlamaUpdate(); if (result.ok) { diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 0a4342f208..a659b7f83e 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -245,18 +245,98 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string { return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`; } +/** + * Reconcile the chat runtime store against `/api/inference/status`: refresh the + * models/loras catalogs and either re-pin the active checkpoint or clear the + * loaded-model flags when nothing is loaded. Module-level so it can run outside + * a React render (e.g. the imperative resync below); `useChatModelRuntime.refresh` + * is a thin wrapper over it. External selections are left untouched since they + * have no backend mirror. + */ +async function syncInferenceStatusToStore(options?: { + signal?: AbortSignal; + includeLoras?: boolean; +}): Promise { + const signal = options?.signal; + const includeLoras = options?.includeLoras ?? true; + const { setModels, setLoras, setCheckpoint, setModelsError } = + useChatRuntimeStore.getState(); + setModelsError(null); + try { + const [listRes, statusRes, lorasRes] = await Promise.all([ + listModels(), + getInferenceStatus(), + includeLoras ? listLoras() : Promise.resolve(null), + ]); + + // Cancellation can land while the requests above are in flight. Bail + // before writing backend state back -- cancelLoading already cleared it. + if (signal?.aborted) return; + + setModels(listRes.models.map(toChatModelSummary)); + if (lorasRes) { + setLoras(lorasRes.loras.map(toLoraSummary)); + } + + const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint; + const isExternalSelectionActive = isExternalModelId(selectedCheckpoint); + if (statusRes.active_model && !isExternalSelectionActive) { + const checkpointId = resolveInferenceCheckpointId(statusRes); + if (checkpointId) { + setCheckpoint(checkpointId, statusRes.gguf_variant); + applyActiveModelStatusToStore(statusRes, { + previousCheckpoint: selectedCheckpoint, + }); + // setModels(listRes...) above used catalog data, which omits audio + // capability. Re-apply live status so attach gates survive a refresh. + syncModelCapabilities(checkpointId, statusRes); + } + } else if (!statusRes.active_model && !isExternalSelectionActive) { + useChatRuntimeStore.setState({ + modelRequiresTrustRemoteCode: false, + loadedIsMultimodal: false, + loadedIsDiffusion: false, + }); + } + } catch (error) { + if (signal?.aborted) return; + const message = + error instanceof Error ? error.message : "Failed to load models"; + setModelsError(message); + toast.error("Failed to refresh models", { + description: message, + }); + } +} + +/** + * Reconcile the UI after the SERVER unloaded the active model out from under it + * (e.g. a llama.cpp update unloads the running model to swap the binary): the + * model selector drops to "select model" instead of pointing at a model that now + * 400s on send. Imperative so the global llama-update banner (which has no + * chat-runtime handle) can call it. + * + * Only a LOCAL selection points at the unloaded model. An external-provider + * selection has no llama.cpp mirror and still works, so clearing it (which also + * wipes its persisted id) would drop a valid, unrelated model; skip the clear so + * the refresh below leaves it intact. + */ +export async function resyncInferenceStatusAfterServerModelChange(): Promise { + if (!isExternalModelId(useChatRuntimeStore.getState().params.checkpoint)) { + useChatRuntimeStore.getState().clearCheckpoint(); + } + await syncInferenceStatusToStore(); +} + export function useChatModelRuntime() { const params = useChatRuntimeStore((state) => state.params); const models = useChatRuntimeStore((state) => state.models); const loras = useChatRuntimeStore((state) => state.loras); - const setModels = useChatRuntimeStore((state) => state.setModels); - const setLoras = useChatRuntimeStore((state) => state.setLoras); const setParams = useChatRuntimeStore((state) => state.setParams); const setModelsError = useChatRuntimeStore((state) => state.setModelsError); const setLastModelLoadError = useChatRuntimeStore( (state) => state.setLastModelLoadError, ); - const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint); const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint); const [loadingModel, setLoadingModel] = useState<{ @@ -313,59 +393,11 @@ export function useChatModelRuntime() { [], ); - const refresh = useCallback(async (options?: { - signal?: AbortSignal; - includeLoras?: boolean; - }) => { - const signal = options?.signal; - const includeLoras = options?.includeLoras ?? true; - setModelsError(null); - try { - const [listRes, statusRes, lorasRes] = await Promise.all([ - listModels(), - getInferenceStatus(), - includeLoras ? listLoras() : Promise.resolve(null), - ]); - - // Cancellation can land while the requests above are in flight. Bail - // before writing backend state back -- cancelLoading already cleared it. - if (signal?.aborted) return; - - setModels(listRes.models.map(toChatModelSummary)); - if (lorasRes) { - setLoras(lorasRes.loras.map(toLoraSummary)); - } - - const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint; - const isExternalSelectionActive = isExternalModelId(selectedCheckpoint); - if (statusRes.active_model && !isExternalSelectionActive) { - const checkpointId = resolveInferenceCheckpointId(statusRes); - if (checkpointId) { - setCheckpoint(checkpointId, statusRes.gguf_variant); - applyActiveModelStatusToStore(statusRes, { - previousCheckpoint: selectedCheckpoint, - }); - // setModels(listRes...) above used catalog data, which omits audio - // capability. Re-apply live status so attach gates survive a refresh. - syncModelCapabilities(checkpointId, statusRes); - } - } else if (!statusRes.active_model && !isExternalSelectionActive) { - useChatRuntimeStore.setState({ - modelRequiresTrustRemoteCode: false, - loadedIsMultimodal: false, - loadedIsDiffusion: false, - }); - } - } catch (error) { - if (signal?.aborted) return; - const message = - error instanceof Error ? error.message : "Failed to load models"; - setModelsError(message); - toast.error("Failed to refresh models", { - description: message, - }); - } - }, [setCheckpoint, setLoras, setModels, setModelsError, setParams]); + const refresh = useCallback( + (options?: { signal?: AbortSignal; includeLoras?: boolean }) => + syncInferenceStatusToStore(options), + [], + ); const cancelLoading = useCallback(() => { const model = loadingModelRef.current; diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index d070ed15de..3099884645 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -25,7 +25,10 @@ export { usePlusMenuPrefsStore, type PlusMenuItemId, } from "./stores/plus-menu-prefs-store"; -export { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; +export { + useChatModelRuntime, + resyncInferenceStatusAfterServerModelChange, +} from "./hooks/use-chat-model-runtime"; export { customProviderDisplayName, isExternalModelId, diff --git a/studio/frontend/src/hooks/use-llama-update-check.ts b/studio/frontend/src/hooks/use-llama-update-check.ts index d8f894b1b1..e3a735ffde 100644 --- a/studio/frontend/src/hooks/use-llama-update-check.ts +++ b/studio/frontend/src/hooks/use-llama-update-check.ts @@ -22,6 +22,9 @@ export interface LlamaUpdateJob { error: string | null; // Download fraction while running, 1 on success. progress: number | null; + // Set once the job leaves "running"; identifies a completed job so a + // repeated fetch of the same success can be told apart from the next one. + finished_at: string | null; } export interface LlamaUpdateStatus { @@ -34,10 +37,24 @@ export interface LlamaUpdateStatus { job: LlamaUpdateJob; } +function parseJob(value: unknown): LlamaUpdateJob { + const job = (value ?? {}) as Record; + return { + state: (job.state as LlamaUpdateJob["state"]) ?? "idle", + message: typeof job.message === "string" ? job.message : "", + from_tag: typeof job.from_tag === "string" ? job.from_tag : null, + to_tag: typeof job.to_tag === "string" ? job.to_tag : null, + reload_required: + typeof job.reload_required === "boolean" ? job.reload_required : null, + error: typeof job.error === "string" ? job.error : null, + progress: typeof job.progress === "number" ? job.progress : null, + finished_at: typeof job.finished_at === "string" ? job.finished_at : null, + }; +} + function parseStatus(value: unknown): LlamaUpdateStatus | null { if (!value || typeof value !== "object") return null; const s = value as Record; - const job = (s.job ?? {}) as Record; return { supported: s.supported === true, update_available: s.update_available === true, @@ -45,19 +62,34 @@ function parseStatus(value: unknown): LlamaUpdateStatus | null { latest_tag: typeof s.latest_tag === "string" ? s.latest_tag : null, update_size_bytes: typeof s.update_size_bytes === "number" ? s.update_size_bytes : null, - job: { - state: (job.state as LlamaUpdateJob["state"]) ?? "idle", - message: typeof job.message === "string" ? job.message : "", - from_tag: typeof job.from_tag === "string" ? job.from_tag : null, - to_tag: typeof job.to_tag === "string" ? job.to_tag : null, - reload_required: - typeof job.reload_required === "boolean" ? job.reload_required : null, - error: typeof job.error === "string" ? job.error : null, - progress: typeof job.progress === "number" ? job.progress : null, - }, + job: parseJob(s.job), }; } +// The backend job persists as "success" until the next update starts (it's a +// single in-memory record, not per-tab), so a fresh mount -- a new tab, or a +// page reload of a tab that already resynced -- would otherwise replay the +// same completed job forever. Persist the handled marker outside React state +// so it survives both, and is shared across tabs in this browser. +const HANDLED_RELOAD_STORAGE_KEY = "unsloth_llama_update_reload_handled_at"; + +function getHandledReloadAt(): string | null { + try { + return localStorage.getItem(HANDLED_RELOAD_STORAGE_KEY); + } catch { + return null; + } +} + +function setHandledReloadAt(finishedAt: string | null): void { + if (!finishedAt) return; + try { + localStorage.setItem(HANDLED_RELOAD_STORAGE_KEY, finishedAt); + } catch { + // storage unavailable + } +} + async function fetchStatus( forceRefresh = false, ): Promise { @@ -78,6 +110,14 @@ const recheckStatus = () => fetchStatus(true); interface UseLlamaUpdateCheckOptions { enabled?: boolean; + /** + * Called when a completed update reports `reload_required` (i.e. it unloaded + * the active model server-side). Consumers use it to resync the chat runtime + * so the model selector drops to "select model" instead of pointing at a + * model that now 400s on send. Fires for both this tab's own apply() and a + * cross-tab update mirrored through the background poll. + */ + onReloadRequired?: () => void; } export interface LlamaApplyResult { @@ -90,12 +130,25 @@ export interface LlamaApplyResult { /** Tracks llama.cpp update visibility and apply progress. */ export function useLlamaUpdateCheck({ enabled = true, + onReloadRequired, }: UseLlamaUpdateCheckOptions = {}) { const [status, setStatus] = useState(null); const [visible, setVisible] = useState(false); const [applying, setApplying] = useState(false); const pollTimer = useRef | null>(null); const snoozeTimer = useRef | null>(null); + // Read through a ref so startJobPoll stays stable (apply/surfaceIfAvailable + // depend on it) while still calling the latest callback. + const onReloadRequiredRef = useRef(onReloadRequired); + useEffect(() => { + onReloadRequiredRef.current = onReloadRequired; + }, [onReloadRequired]); + // Fires the callback once per completed job, whether this tab watched it run + // or only saw the persisted "success" after the fact (e.g. another tab + // applied it). Keyed by finished_at and seeded from localStorage so a fresh + // mount (new tab, or a page reload of a tab that already resynced) doesn't + // replay a job some tab already handled. + const reloadNotifiedForRef = useRef(getHandledReloadAt()); const clearPollTimer = useCallback(() => { if (pollTimer.current) { @@ -104,6 +157,25 @@ export function useLlamaUpdateCheck({ } }, []); + // Shared by the poll path (this tab watched the job run), the surface path + // (this tab only saw the persisted success), and apply()'s stale-click path + // (the job came back embedded in a "not started" response) so none of them + // can drop or double-fire the notification. + const notifyReloadIfNeeded = useCallback( + (job: Pick) => { + if ( + job.state === "success" && + job.reload_required && + job.finished_at !== reloadNotifiedForRef.current + ) { + reloadNotifiedForRef.current = job.finished_at; + setHandledReloadAt(job.finished_at); + onReloadRequiredRef.current?.(); + } + }, + [], + ); + // Used by apply() and another-tab job tracking. const startJobPoll = useCallback( (onDone?: (result: LlamaApplyResult) => void) => { @@ -118,6 +190,12 @@ export function useLlamaUpdateCheck({ if (s.job.state === "success") { setVisible(false); void refreshHardwareInfo(); + // The update unloads the running model server-side, so the chat + // runtime still points at a model that now 400s on send. Let the + // consumer drop the selector to "select model" instead of waiting for + // a page reload. Fires here (not just from apply's onDone) so a + // cross-tab update mirrored through this poll is covered too. + notifyReloadIfNeeded(s.job); onDone?.({ ok: true, tag: s.job.to_tag, @@ -131,7 +209,7 @@ export function useLlamaUpdateCheck({ } }, JOB_POLL_INTERVAL_MS); }, - [clearPollTimer], + [clearPollTimer, notifyReloadIfNeeded], ); const surfaceIfAvailable = useCallback( @@ -145,11 +223,16 @@ export function useLlamaUpdateCheck({ if (!pollTimer.current) startJobPoll(); return; } + // A completed job persists as "success" until the next update starts, so + // a tab that missed the running window entirely (mounted, or only checks + // hourly and misses both the running and just-finished moments) still + // needs to resync here, not just from the poll path above. + notifyReloadIfNeeded(next.job); if (next.update_available) { setVisible(true); } }, - [startJobPoll], + [startJobPoll, notifyReloadIfNeeded], ); useEffect(() => { @@ -185,6 +268,26 @@ export function useLlamaUpdateCheck({ }; }, [enabled, surfaceIfAvailable, clearPollTimer]); + // Cross-tab nudge: a tab that only checks hourly would otherwise stay + // pointed at a server-unloaded model for up to an hour after a DIFFERENT + // open tab applies an update. The storage event only fires in other tabs + // (never the one that wrote it), so this recheck fires promptly there + // without this tab redundantly re-triggering itself. + useEffect(() => { + if (!enabled) return; + const onStorage = (event: StorageEvent) => { + if ( + event.key === HANDLED_RELOAD_STORAGE_KEY && + event.newValue && + event.newValue !== reloadNotifiedForRef.current + ) { + recheckStatus().then(surfaceIfAvailable); + } + }; + window.addEventListener("storage", onStorage); + return () => window.removeEventListener("storage", onStorage); + }, [enabled, surfaceIfAvailable]); + const dismiss = useCallback(() => { setVisible(false); }, []); @@ -206,6 +309,7 @@ export function useLlamaUpdateCheck({ started?: boolean; reason?: string | null; message?: string | null; + job?: unknown; } | null = null; try { const res = await authFetch("/api/llama/update", { method: "POST" }); @@ -229,6 +333,11 @@ export function useLlamaUpdateCheck({ action.started === false && action.reason !== "already_running" ) { + // A stale banner's click can land after another tab already applied the + // update (e.g. "up_to_date"): the response still carries that tab's + // completed job, so process reload_required here too, not just from the + // poll path -- otherwise this rejection silently drops it. + notifyReloadIfNeeded(parseJob(action.job)); setApplying(false); return { ok: false, @@ -239,7 +348,7 @@ export function useLlamaUpdateCheck({ return await new Promise((resolve) => startJobPoll(resolve), ); - }, [applying, startJobPoll]); + }, [applying, startJobPoll, notifyReloadIfNeeded]); return { status: enabled ? status : null, From f60b982a09ed258b9469ef1cbca54c466d9f0cb1 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Mon, 13 Jul 2026 17:34:25 -0300 Subject: [PATCH 012/210] Studio: Fix torch_dtype deprecation warning on startup and ASR load (#6999) --- studio/backend/core/inference/inference.py | 3 +- studio/backend/core/rag/embeddings.py | 5 +- .../backend/tests/test_transformers_dtype.py | 77 +++++++++++++++++++ studio/backend/utils/transformers_dtype.py | 51 ++++++++++++ 4 files changed, 132 insertions(+), 4 deletions(-) create mode 100644 studio/backend/tests/test_transformers_dtype.py create mode 100644 studio/backend/utils/transformers_dtype.py diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 7e69e05124..bb4be39ef6 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -15,6 +15,7 @@ from pathlib import Path from typing import Optional, Union, Generator, Tuple from utils.models import ModelConfig, get_base_model_from_lora from utils.paths import is_model_cached +from utils.transformers_dtype import dtype_kwargs from utils.utils import format_error_message from utils.hardware import ( get_device, @@ -440,7 +441,7 @@ class InferenceBackend: feature_extractor = tokenizer.feature_extractor, processor = tokenizer, return_language = True, - torch_dtype = torch.float16, + **dtype_kwargs(torch.float16), ) self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = tokenizer diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 47d26209b4..b0ecedd593 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -21,6 +21,7 @@ from functools import lru_cache from typing import Callable from utils.hardware.hardware import DeviceType, get_device +from utils.transformers_dtype import dtype_kwargs from . import config @@ -157,9 +158,7 @@ def _get(model_name: str | None = None): device = _device() logger.info("loading embedding model %s on %s", name, device) _guard_model_security(name) - _model = SentenceTransformer( - name, device = device, model_kwargs = {"torch_dtype": "float16"} - ) + _model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16")) _name = name return _model diff --git a/studio/backend/tests/test_transformers_dtype.py b/studio/backend/tests/test_transformers_dtype.py new file mode 100644 index 0000000000..28629e5620 --- /dev/null +++ b/studio/backend/tests/test_transformers_dtype.py @@ -0,0 +1,77 @@ +# 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 version-safe torch_dtype/dtype kwarg helper.""" + +import sys +import types + +import pytest + +from utils.transformers_dtype import _has_torch_dtype_kwarg, dtype_kwargs + + +@pytest.fixture(autouse = True) +def _clear_cache(): + _has_torch_dtype_kwarg.cache_clear() + yield + _has_torch_dtype_kwarg.cache_clear() + + +def _stub_transformers(monkeypatch, version): + stub = types.ModuleType("transformers") + stub.__version__ = version + monkeypatch.setitem(sys.modules, "transformers", stub) + + +def test_old_transformers_uses_torch_dtype(monkeypatch): + _stub_transformers(monkeypatch, "4.51.3") + assert _has_torch_dtype_kwarg() is True + assert dtype_kwargs("float16") == {"torch_dtype": "float16"} + + +def test_new_transformers_uses_dtype(monkeypatch): + _stub_transformers(monkeypatch, "4.57.6") + assert _has_torch_dtype_kwarg() is False + assert dtype_kwargs("float16") == {"dtype": "float16"} + + +def test_rename_boundary_uses_dtype(monkeypatch): + _stub_transformers(monkeypatch, "4.56.0") + assert _has_torch_dtype_kwarg() is False + + +def test_just_below_boundary_uses_torch_dtype(monkeypatch): + _stub_transformers(monkeypatch, "4.55.4") + assert _has_torch_dtype_kwarg() is True + + +@pytest.mark.parametrize("version", ["4.56.0.dev0", "4.56.0rc1"]) +def test_rename_prerelease_uses_dtype(monkeypatch, version): + """A pre-release of the rename version sorts *below* ``4.56.0`` but already + accepts (and prefers) ``dtype``; the release-tuple check must not fall back to + the legacy name there, or it re-emits the deprecation warning it suppresses.""" + _stub_transformers(monkeypatch, version) + assert _has_torch_dtype_kwarg() is False + + +def test_malformed_version_prefers_modern_name(monkeypatch): + """A non-PEP440 __version__ raises InvalidVersion; the except branch must + swallow it and default to the modern name rather than crash the embedder warm-up.""" + _stub_transformers(monkeypatch, "not-a-version") + assert _has_torch_dtype_kwarg() is False + assert dtype_kwargs("float16") == {"dtype": "float16"} + + +def test_missing_transformers_prefers_modern_name(monkeypatch): + monkeypatch.delitem(sys.modules, "transformers", raising = False) + real_import = __import__ + + def _raise(name, *args, **kwargs): + if name == "transformers": + raise ImportError("no transformers") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr("builtins.__import__", _raise) + assert _has_torch_dtype_kwarg() is False + assert dtype_kwargs("float16") == {"dtype": "float16"} diff --git a/studio/backend/utils/transformers_dtype.py b/studio/backend/utils/transformers_dtype.py new file mode 100644 index 0000000000..daeb6e2452 --- /dev/null +++ b/studio/backend/utils/transformers_dtype.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Version-safe fp-dtype kwarg for transformers/sentence-transformers loads. + +transformers renamed the ``torch_dtype`` kwarg to ``dtype`` in 4.56.0, and emits +``torch_dtype is deprecated! Use dtype instead!`` when the old name is passed. But our floor (``transformers>=4.51.3``) predates ``dtype`` and only +accepts ``torch_dtype``, so a bare rename would ``TypeError`` on the floor. Pick +the name the installed version accepts instead. + +Answers the same question as ``unsloth_zoo.hf_utils.HAS_TORCH_DTYPE`` but derives +it independently, for two reasons. It uses a ``packaging.version`` check rather +than that constant's ``"torch_dtype" in PretrainedConfig.__doc__`` sniffing, which +raises ``TypeError`` under ``python -OO`` / ``PYTHONOPTIMIZE=2`` (docstrings are +stripped to ``None``, and ``"torch_dtype" in None`` is a type error). And it avoids +importing the constant at all: the RAG +embedder warms here at startup in the lean main process, and reading it would run +``unsloth_zoo``'s package ``__init__`` (torch import, GPU/Pytorch checks, the +patching banner) as a side effect. The embedder is deliberately torch-optional (it +degrades to the ``llama-server`` GGUF backend), so it must not drag in that +heavyweight import just to read one bool. +""" + +from functools import lru_cache + + +@lru_cache(maxsize = 1) +def _has_torch_dtype_kwarg() -> bool: + """True if the installed transformers still expects the legacy ``torch_dtype`` + name (i.e. predates the ``dtype`` rename). False when ``dtype`` is the accepted + name, or when transformers is missing/broken (prefer the modern name).""" + try: + import transformers + from packaging.version import Version + + # Compare on the release tuple so a pre-release of the rename version + # (``4.56.0.dev0``/``rc1``, which sort *below* ``4.56.0``) still counts as + # new and picks ``dtype`` -- those builds already accept it, and picking + # ``torch_dtype`` there would re-emit the very warning this suppresses. + return Version(transformers.__version__).release < (4, 56, 0) + except Exception: + return False + + +def dtype_kwargs(value) -> dict: + """``{"torch_dtype": value}`` on old transformers, ``{"dtype": value}`` on new. + + Splat into a load call (``pipeline(..., **dtype_kwargs(torch.float16))``) or use + directly as ``model_kwargs`` (``model_kwargs = dtype_kwargs("float16")``). + """ + return {"torch_dtype" if _has_torch_dtype_kwarg() else "dtype": value} From 76d7088e0ae032f537eca2a781e1b7578d2701f5 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Mon, 13 Jul 2026 19:18:20 -0300 Subject: [PATCH 013/210] Studio: Show Run button for downloaded non-GGUF models in the Model Hub (#7001) --- .../hub/catalog/local-on-device-card.tsx | 3 +- .../features/hub/catalog/model-inspector.tsx | 39 +++++++++++++++++-- .../hub/catalog/safetensors-download-card.tsx | 16 +++++--- .../src/features/hub/lib/hub-feature-flags.ts | 9 ++++- .../src/features/hub/lib/unsloth-support.ts | 3 ++ 5 files changed, 57 insertions(+), 13 deletions(-) diff --git a/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx b/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx index e40f5fd5ac..3c020a7199 100644 --- a/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx +++ b/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx @@ -4,6 +4,7 @@ import { TrainIcon } from "../components/train-icon"; import { HUB_GGUF_RUN_ACTIONS_VISIBLE, + HUB_NON_GGUF_RUN_ACTIONS_VISIBLE, HUB_POST_DOWNLOAD_ACTIONS_VISIBLE, } from "../lib/hub-feature-flags"; import { @@ -410,7 +411,7 @@ export function LocalOnDeviceCard({ const showOldCacheHint = source === "hf_cache" && !!unsupportedReason; const runActionsVisible = isGguf ? HUB_GGUF_RUN_ACTIONS_VISIBLE - : HUB_POST_DOWNLOAD_ACTIONS_VISIBLE; + : HUB_NON_GGUF_RUN_ACTIONS_VISIBLE; return (
diff --git a/studio/frontend/src/features/hub/catalog/model-inspector.tsx b/studio/frontend/src/features/hub/catalog/model-inspector.tsx index 0f244bce53..a67b2a57cc 100644 --- a/studio/frontend/src/features/hub/catalog/model-inspector.tsx +++ b/studio/frontend/src/features/hub/catalog/model-inspector.tsx @@ -263,17 +263,22 @@ type VramInfo = { est: number; status: "fits" | "tight" | "exceeds" } | null; function ModelStatusChips({ isDataset, isGguf, + chatOnly, unslothSupport, vramInfo, }: { isDataset: boolean; isGguf: boolean; + chatOnly: boolean; unslothSupport: UnslothSupport; vramInfo: VramInfo; }) { const showUnsupported = !isDataset && unslothSupport.status === "unsupported"; + // The format-unsupported chip already explains itself; this one covers the + // supported-format model a chat-only host still can't run. + const showChatOnly = !isDataset && !isGguf && chatOnly && !showUnsupported; const showVram = !isDataset && vramInfo && !isGguf; - if (!showUnsupported && !showVram) return null; + if (!showUnsupported && !showChatOnly && !showVram) return null; const vramTone = vramInfo ? vramInfo.status === "exceeds" @@ -323,6 +328,26 @@ function ModelStatusChips({ )} + {showChatOnly && ( + + + + + + + + This device has no supported GPU or usable MLX, so only GGUF models + can run here. + + Still downloadable to your Hugging Face cache. + + + + )} {showVram && vramInfo && ( @@ -406,6 +431,7 @@ export const ModelInspector = memo(function ModelInspector({ onSearchHub, } = actions; const deviceType = usePlatformStore((s) => s.deviceType); + const chatOnly = usePlatformStore((s) => s.isChatOnly()); const hfToken = useHfTokenStore((s) => s.token); const datasetRepoId = isDataset && model?.hubRepoId ? model.hubRepoId : null; const datasetSize = useDatasetSize(datasetRepoId, { @@ -504,15 +530,19 @@ export const ModelInspector = memo(function ModelInspector({ const paramsLabel = model.totalParams ? formatCompact(model.totalParams) : "N/A"; - const trainingSupported = unslothSupport.status !== "unsupported"; + const unslothSupported = unslothSupport.status !== "unsupported"; + // Chat-only hosts (no supported GPU / usable MLX) run inference only through + // llama.cpp, so only GGUF is loadable. const canRunModel = - !isDataset && (model.runtimeCapabilities?.canChat ?? true); + !isDataset && + (model.runtimeCapabilities?.canChat ?? true) && + (model.isGguf || (!chatOnly && unslothSupported)); const canTrainModel = !isDataset && (model.runtimeCapabilities?.canTrain ?? false) && model.modelFormat !== "gguf" && model.modelFormat !== "adapter" && - trainingSupported; + unslothSupported; const languages = parseLanguageTags(model.tags); const datasetSizeBytes = @@ -765,6 +795,7 @@ export const ModelInspector = memo(function ModelInspector({ diff --git a/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx b/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx index ccd61c5f04..424afa2e05 100644 --- a/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx +++ b/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx @@ -16,7 +16,10 @@ import { PlayIcon, } from "@hugeicons/core-free-icons"; import { TrainIcon } from "../components/train-icon"; -import { HUB_POST_DOWNLOAD_ACTIONS_VISIBLE } from "../lib/hub-feature-flags"; +import { + HUB_NON_GGUF_RUN_ACTIONS_VISIBLE, + HUB_POST_DOWNLOAD_ACTIONS_VISIBLE, +} from "../lib/hub-feature-flags"; import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useState } from "react"; import { useHfTokenStore } from "../stores/hf-token-store"; @@ -158,6 +161,7 @@ export function SafetensorsDownloadCard({ const showActionPair = isDownloaded && !downloading && (canRun || !!onTrain); const showUnavailableAction = isDownloaded && !downloading && !canRun && !onTrain; + const trainActionVisible = !!onTrain && HUB_POST_DOWNLOAD_ACTIONS_VISIBLE; const canDelete = (isDownloaded || isPartial) && !downloading && @@ -238,17 +242,17 @@ export function SafetensorsDownloadCard({ )}
- {/* Divider sits above the Download CTA; in the action-pair state it hides with the pair. */} - {(!showActionPair || HUB_POST_DOWNLOAD_ACTIONS_VISIBLE) && } + {/* Info/actions hairline; dropped for the run action row (no divider before + Run, as in the GGUF card's Run CTA), restored when the Train pair ships. */} + {(!showActionPair || trainActionVisible) && } {showActionPair ? ( From ed427027305ca2eaee3ad4488621a98a19c2259b Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 14 Jul 2026 00:01:25 -0300 Subject: [PATCH 016/210] Probe xformers support on sm_120 instead of disabling it by version (#6828) --- .../test_attention_dispatch_dora_dtype.py | 56 ++++++++++ tests/utils/test_xformers_capability_gate.py | 102 ++++++++++++++++++ unsloth/utils/attention_dispatch.py | 66 +++++++++--- 3 files changed, 210 insertions(+), 14 deletions(-) create mode 100644 tests/utils/test_xformers_capability_gate.py diff --git a/tests/utils/test_attention_dispatch_dora_dtype.py b/tests/utils/test_attention_dispatch_dora_dtype.py index a6586dea8d..f0496ee6a8 100644 --- a/tests/utils/test_attention_dispatch_dora_dtype.py +++ b/tests/utils/test_attention_dispatch_dora_dtype.py @@ -71,3 +71,59 @@ def test_varlen_flash_downcasts_fp32_qkv(monkeypatch): def test_bf16_qkv_left_untouched(monkeypatch): # Standard LoRA path (already bf16) must not be altered. assert _run(monkeypatch, torch.bfloat16, ad.FLASH_DENSE) == torch.bfloat16 + + +def _run_xformers(monkeypatch, qkv_dtype, fp32_unsupported): + # Same #1013 fp32 downcast, but for the xformers backend. On sm_100+ (B200, sm_120) + # xformers' fp32-capable cutlass op is capability-rejected and only its flash-2 op + # runs (fp16/bf16 only), so fp32 must be downcast there too or the op raises. + captured = {} + + def fake_xformers_attention( + Q, + K, + V, + attn_bias = None, + **kwargs, + ): + captured["dtype"] = Q.dtype + # Mirror the flash-2 op's real dtype constraint so an unfixed dispatch fails loudly. + if fp32_unsupported and Q.dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError("no operator found for memory_efficient_attention with fp32") + return torch.zeros_like(Q) + + monkeypatch.setattr(ad, "_XFORMERS_FP32_UNSUPPORTED", fp32_unsupported, raising = False) + monkeypatch.setattr(ad, "xformers_attention", fake_xformers_attention, raising = False) + monkeypatch.setattr( + ad, "build_xformers_block_causal_mask", lambda *a, **k: object(), raising = False + ) + + bsz, n_heads, q_len, head_dim = 1, 2, 4, 8 + Q = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype) + K = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype) + V = torch.randn(bsz, n_heads, q_len, head_dim, dtype = qkv_dtype) + + config = ad.AttentionConfig(backend = ad.XFORMERS, n_kv_heads = n_heads, n_groups = 1) + context = ad.AttentionContext( + bsz = bsz, + q_len = q_len, + kv_seq_len = q_len, + n_heads = n_heads, + head_dim = head_dim, + requires_grad = False, + seq_info = None, + attention_mask = None, + causal_mask = None, + ) + ad.run_attention(config = config, context = context, Q = Q, K = K, V = V) + return captured["dtype"] + + +def test_xformers_downcasts_fp32_qkv_on_sm100_plus(monkeypatch): + # sm_100+ (fp32 op gone): fp32 DoRA output must be downcast, else the flash-2 op raises. + assert _run_xformers(monkeypatch, torch.float32, True) in (torch.bfloat16, torch.float16) + + +def test_xformers_leaves_fp32_qkv_below_sm100(monkeypatch): + # Below sm_100 the cutlass op handles fp32 natively, so it must be passed through as-is. + assert _run_xformers(monkeypatch, torch.float32, False) == torch.float32 diff --git a/tests/utils/test_xformers_capability_gate.py b/tests/utils/test_xformers_capability_gate.py new file mode 100644 index 0000000000..7514c623e3 --- /dev/null +++ b/tests/utils/test_xformers_capability_gate.py @@ -0,0 +1,102 @@ +"""Regression test for unslothai/unsloth#4631: xformers must not be blanket-disabled +on sm_120 GPUs where its kernel actually runs (a ~57% attention-memory saving over the +SDPA packed-mask fallback). The gate now probes the real op instead of guessing by the +compute-capability major version.""" + +import pytest +import torch +import unsloth # noqa: F401 + +from unsloth.utils import attention_dispatch as ad + + +@pytest.mark.parametrize( + "capability, probe_result, expect_disabled", + [ + ((8, 9), None, False), # Ada: below sm_120, never probed, always kept + ((9, 0), None, False), # Hopper: below sm_120, kept + ((10, 0), None, False), # Blackwell B200 (sm_100): below sm_120, kept + ((12, 0), True, False), # sm_120 where the kernel runs: keep xformers + ((12, 0), False, True), # sm_120 where the kernel can't run: fall back to SDPA + ], +) +def test_capability_gate(capability, probe_result, expect_disabled): + calls = {"n": 0} + + def probe(): + calls["n"] += 1 + return probe_result + + assert ad._xformers_disabled_for_capability(capability, probe = probe) is expect_disabled + # Below sm_120 the probe must not run at all (no import-time kernel launch there). + assert calls["n"] == (0 if capability[0] < 12 else 1) + + +@pytest.mark.skipif( + not (torch.cuda.is_available() and ad.HAS_XFORMERS), + reason = "needs a CUDA GPU with a working xformers build", +) +@pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12, + reason = "on real sm_120+ the probe legitimately returns False when the build ships no " + "sm_120 kernel, so asserting True there would be a false failure", +) +def test_probe_shapes_are_valid_on_working_gpu(): + # Guards against a malformed probe that raises on every GPU and would silently + # disable xformers on Blackwell even where it works. On a pre-sm_120 GPU with a + # functional xformers the real probe must succeed; sm_120+ is skipped above because + # there a False is a correct answer, not a malformed probe. + assert ad._xformers_runs_on_device() is True + + +@pytest.mark.parametrize( + "supports_bf16, expected_dtype", + [(True, torch.bfloat16), (False, torch.float16)], +) +def test_probe_dtype_follows_bf16_support(monkeypatch, supports_bf16, expected_dtype): + # Pre-Ampere GPUs (sm < 80: Turing/Volta, e.g. T4/V100) run xformers fine in + # float16 but have no bfloat16 attention kernel, so a hardcoded bf16 probe would + # raise there, get swallowed to False, and misreport a working xformers as broken. + # The probe must pick its dtype from SUPPORTS_BFLOAT16 (no Turing GPU needed here). + captured = {} + + def fake_zeros( + *args, + dtype = None, + **kwargs, + ): + captured["dtype"] = dtype + raise RuntimeError("stop after capturing the probe dtype") + + monkeypatch.setattr(ad, "SUPPORTS_BFLOAT16", supports_bf16) + monkeypatch.setattr(ad.torch, "zeros", fake_zeros) + ad._xformers_runs_on_device() # RuntimeError is swallowed; only the dtype matters + assert captured["dtype"] is expected_dtype + + +def test_probe_syncs_and_fails_on_deferred_async_error(monkeypatch): + # A CUDA kernel launch is async: xformers_attention can return before the GPU + # reports a failure. The probe must synchronize so a deferred launch/runtime error + # is caught and disables xformers here, instead of surfacing later on an unrelated + # CUDA call (unslothai/unsloth#6828 review). No GPU needed: everything is stubbed. + _bias = type( + "B", + (), + { + "BlockDiagonalCausalMask": type( + "M", (), {"from_seqlens": staticmethod(lambda seqlens: None)} + ) + }, + ) + monkeypatch.setattr(ad, "SUPPORTS_BFLOAT16", True) + monkeypatch.setattr(ad.torch, "zeros", lambda *a, **k: object()) + monkeypatch.setattr(ad, "xformers", type("X", (), {"attn_bias": _bias})) + monkeypatch.setattr(ad, "xformers_attention", lambda *a, **k: None) # "succeeds" + + def deferred_cuda_error(): + raise RuntimeError("CUDA error: an illegal memory access was encountered") + + monkeypatch.setattr(ad.torch.cuda, "synchronize", deferred_cuda_error) + # Without the synchronize the stubbed op returns cleanly and the probe wrongly + # reports True; the sync surfaces the deferred error so the probe returns False. + assert ad._xformers_runs_on_device() is False diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 68fb33dad9..eda6103d5b 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -35,12 +35,44 @@ if HAS_FLASH_ATTENTION: from flash_attn import flash_attn_func, flash_attn_varlen_func HAS_XFORMERS = xformers is not None -# xformers kernels (FA3, FA2, cutlass) only support compute capability <= 9.0. -# Disable xformers on newer GPUs (e.g. RTX 5070 Ti / sm_120) and fall back to SDPA. -if HAS_XFORMERS and torch.cuda.is_available(): - _cc = torch.cuda.get_device_capability() - if _cc[0] >= 12: + +def _xformers_runs_on_device() -> bool: + """One tiny attention forward; True iff the xformers kernel actually runs here.""" + try: + # Pre-Ampere GPUs (sm < 80: Turing/Volta) have no bfloat16 attention kernel + # but run xformers fine in float16, so pick the dtype the device supports. + dtype = torch.bfloat16 if SUPPORTS_BFLOAT16 else torch.float16 + q = torch.zeros((1, 8, 1, 64), device = "cuda", dtype = dtype) + attn_bias = xformers.attn_bias.BlockDiagonalCausalMask.from_seqlens([8]) + xformers_attention(q, q, q, attn_bias = attn_bias) + # Launches are async; synchronize so a deferred kernel failure fails the probe here. + torch.cuda.synchronize() + return True + except Exception: + return False + + +def _xformers_disabled_for_capability(capability, probe = _xformers_runs_on_device) -> bool: + # At sm_120 (RTX 50-series) xformers' cutlass op is capability-rejected (caps at + # sm_90) and its flash-2 op runs only if the build ships an sm_120 kernel, so run + # one real forward to decide. Below sm_120 xformers always works; skip the probe. + if capability[0] < 12: + return False + return not probe() + + +# FlashAttention always wins in select_attention_backend and nothing downgrades +# flash -> xformers, so when it's installed xformers is never selected: skip the probe. +if HAS_XFORMERS and not HAS_FLASH_ATTENTION and torch.cuda.is_available(): + if _xformers_disabled_for_capability(torch.cuda.get_device_capability()): HAS_XFORMERS = False + +# On sm_100+ (B200, sm_120) xformers' fp32-capable cutlass op is capability-rejected and +# only its fp16/bf16 flash-2 op runs, so fp32 Q/K/V (DoRA, #1013) must be downcast there; +# below sm_100 cutlass handles fp32 natively. Read once from device 0, like the probe gate. +_XFORMERS_FP32_UNSUPPORTED = ( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10 +) SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "") # PrefixGrouper kernel, resolved once when the env gate is on so PG-off users never load @@ -201,9 +233,13 @@ def run_attention( requires_grad = context.requires_grad sliding_window = context.sliding_window - # DoRA promotes q/k/v_proj outputs to fp32, which FlashAttention rejects, so - # downcast any fp32 Q/K/V to a flash-supported dtype (#1013). - if backend in (FLASH_DENSE, FLASH_VARLEN) and torch.float32 in ( + # DoRA promotes q/k/v_proj outputs to fp32, which FlashAttention rejects (and so does + # the xformers flash-2 op on sm_100+, see _XFORMERS_FP32_UNSUPPORTED), so downcast any + # fp32 Q/K/V to a supported dtype (#1013). + if ( + backend in (FLASH_DENSE, FLASH_VARLEN) + or (backend == XFORMERS and _XFORMERS_FP32_UNSUPPORTED) + ) and torch.float32 in ( Q.dtype, K.dtype, V.dtype, @@ -211,14 +247,16 @@ def run_attention( # Prefer the autocast dtype, else a non-fp32 input's dtype, then clamp. if torch.is_autocast_enabled(): try: - flash_dtype = torch.get_autocast_dtype("cuda") + downcast_dtype = torch.get_autocast_dtype("cuda") except (AttributeError, TypeError): - flash_dtype = torch.get_autocast_gpu_dtype() + downcast_dtype = torch.get_autocast_gpu_dtype() else: - flash_dtype = next((d for d in (Q.dtype, K.dtype, V.dtype) if d != torch.float32), None) - if flash_dtype not in (torch.float16, torch.bfloat16): - flash_dtype = torch.bfloat16 if SUPPORTS_BFLOAT16 else torch.float16 - Q, K, V = Q.to(flash_dtype), K.to(flash_dtype), V.to(flash_dtype) + downcast_dtype = next( + (d for d in (Q.dtype, K.dtype, V.dtype) if d != torch.float32), None + ) + if downcast_dtype not in (torch.float16, torch.bfloat16): + downcast_dtype = torch.bfloat16 if SUPPORTS_BFLOAT16 else torch.float16 + Q, K, V = Q.to(downcast_dtype), K.to(downcast_dtype), V.to(downcast_dtype) if backend == FLASH_VARLEN: Q_f = Q.transpose(1, 2).reshape(bsz * q_len, n_heads, head_dim) From fea7d9ba345262d42d2ea125bea330d65350cbe0 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Tue, 14 Jul 2026 13:00:21 +0530 Subject: [PATCH 017/210] Studio: render image content returned by MCP tools (#7081) * MCP image handling * clean upg * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: return MCP error results so image content is not dropped FastMCP client.call_tool raises ToolError by default on an is_error result, so it never reaches _flatten_result and any returned image is dropped. Pass raise_on_error=False so error results flow through _flatten_result and keep their images. Transport failures still raise and hit the existing handler. Add a regression test for the real path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: accept raise_on_error kwarg in MCP test fake clients The call_tool_sync fix passes raise_on_error=False to client.call_tool. Update the fake MCP clients patched into mcp_client._client so their call_tool signatures accept the keyword, keeping the stdio/servers MCP test suites green. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten MCP raise_on_error rationale comments * Studio: only strip MCP image sentinel when suffix is a valid image envelope * Studio: validate MCP image envelope in chat adapter and keep base64 out of exports * Studio: sanitize MCP images in all export formats and fall through to sandbox parser on invalid marker --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- studio/backend/core/inference/mcp_client.py | 35 +++- .../core/inference/tool_loop_controller.py | 24 +++ .../backend/tests/test_mcp_flatten_result.py | 177 ++++++++++++++++++ studio/backend/tests/test_mcp_servers.py | 14 +- studio/backend/tests/test_mcp_stdio_pr5863.py | 7 +- .../components/assistant-ui/tool-fallback.tsx | 53 +++++- .../src/features/chat/api/chat-adapter.ts | 52 +++++ .../chat/hooks/use-chat-search-index.ts | 34 +++- .../prompt-storage/prompt-storage-dialog.tsx | 16 +- 9 files changed, 400 insertions(+), 12 deletions(-) create mode 100644 studio/backend/tests/test_mcp_flatten_result.py diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index 7bd4a7d6e9..c6b8acfdc4 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -298,20 +298,48 @@ def invalidate_tool_cache(server_id: Optional[str] = None) -> None: _probe_cooloff_until.pop(server_id, None) +MCP_IMAGES_SENTINEL = "__MCP_IMAGES__:" +MAX_IMAGE_PAYLOAD_CHARS = 12_000_000 + + def _flatten_result(result: Any) -> str: parts = [] + images = [] + omitted = 0 + budget = MAX_IMAGE_PAYLOAD_CHARS for block in getattr(result, "content", None) or []: text = getattr(block, "text", None) if text: parts.append(str(text)) + continue + data = getattr(block, "data", None) + mime = getattr(block, "mimeType", None) + if data and isinstance(mime, str) and mime.startswith("image/"): + data = str(data) + if len(data) > budget: + omitted += 1 + continue + budget -= len(data) + images.append({"data": data, "mimeType": mime}) body = "\n".join(parts) if not body: structured = getattr(result, "structured_content", None) body = str(structured) if structured is not None else "" + if images or omitted: + notes = [] + if images: + n = len(images) + notes.append(f"{n} image{'s' if n > 1 else ''} attached; displayed to the user") + if omitted: + notes.append(f"{omitted} image{'s' if omitted > 1 else ''} omitted (too large)") + note = f"[{'; '.join(notes)}]" + body = f"{body}\n{note}" if body else note if getattr(result, "is_error", False): # "Error: " prefix triggers tool_call_parser's TOOL_ERROR_PREFIXES nudge. - return f"Error: {body}" if body else "Error: tool returned no content" + body = f"Error: {body}" if body else "Error: tool returned no content" + if images: + body += "\n" + MCP_IMAGES_SENTINEL + json.dumps(images) return body @@ -333,7 +361,10 @@ def call_tool_sync( async def _call() -> Any: async with _client(url, headers, use_oauth) as client: - return await client.call_tool(name, args) + # raise_on_error=False lets an is_error result (which may still carry + # image content) reach _flatten_result instead of FastMCP raising ToolError + # and dropping the images. Transport failures still raise (handled below). + return await client.call_tool(name, args, raise_on_error = False) async def _watch_cancel() -> None: # 50 ms cadence keeps cancellation responsive without busy-looping; diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py index cb751ede3d..f595531b90 100644 --- a/studio/backend/core/inference/tool_loop_controller.py +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -233,9 +233,33 @@ def is_tool_error(result: str) -> bool: return isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES) +def _strip_mcp_image_suffix(result: str) -> str: + """Drop a trailing __MCP_IMAGES__ envelope only when it is the valid JSON + image array appended by _flatten_result, so legit tool text that merely + mentions the marker is not truncated.""" + head, sep, payload = result.rpartition("\n__MCP_IMAGES__:") + if not sep: + return result + try: + images = json.loads(payload) + except (ValueError, RecursionError): + return result + if not isinstance(images, list) or not images: + return result + if not all( + isinstance(img, dict) + and isinstance(img.get("data"), str) + and isinstance(img.get("mimeType"), str) + for img in images + ): + return result + return head.rstrip() + + def strip_result_for_model(result: str) -> str: """Remove frontend-only sentinels (image paths, RAG source map) before feeding the result back to the model.""" + result = _strip_mcp_image_suffix(result) for sentinel in ("__IMAGES__:", "__RAG_SOURCES__:"): if sentinel in result: result = result.split(sentinel, 1)[0].rstrip() diff --git a/studio/backend/tests/test_mcp_flatten_result.py b/studio/backend/tests/test_mcp_flatten_result.py new file mode 100644 index 0000000000..7daee799f9 --- /dev/null +++ b/studio/backend/tests/test_mcp_flatten_result.py @@ -0,0 +1,177 @@ +# 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 contextlib +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference import mcp_client +from core.inference.mcp_client import ( + MAX_IMAGE_PAYLOAD_CHARS, + MCP_IMAGES_SENTINEL, + _flatten_result, + call_tool_sync, +) +from core.inference.tool_loop_controller import is_tool_error, strip_result_for_model + +PNG_B64 = "iVBORw0KGgoAAAANSUhEUg==" + + +def _text(value: str) -> SimpleNamespace: + return SimpleNamespace(type = "text", text = value) + + +def _image(data: str = PNG_B64, mime: str = "image/png") -> SimpleNamespace: + return SimpleNamespace(type = "image", data = data, mimeType = mime) + + +def _result( + *blocks, + is_error = False, + structured = None, +) -> SimpleNamespace: + return SimpleNamespace( + content = list(blocks), + is_error = is_error, + structured_content = structured, + ) + + +def test_text_only_result_unchanged(): + assert _flatten_result(_result(_text("hello"))) == "hello" + + +def test_image_only_result_keeps_image_and_notes_model(): + flat = _flatten_result(_result(_image())) + body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) + assert body == "[1 image attached; displayed to the user]" + assert json.loads(payload) == [{"data": PNG_B64, "mimeType": "image/png"}] + + +def test_text_plus_image_keeps_both(): + flat = _flatten_result(_result(_text("Took a screenshot"), _image())) + body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) + assert body == "Took a screenshot\n[1 image attached; displayed to the user]" + assert json.loads(payload)[0]["mimeType"] == "image/png" + + +def test_multiple_images_pluralized(): + flat = _flatten_result(_result(_image(), _image(mime = "image/jpeg"))) + body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) + assert "[2 images attached; displayed to the user]" in body + assert [img["mimeType"] for img in json.loads(payload)] == ["image/png", "image/jpeg"] + + +def test_strip_result_for_model_drops_image_payload(): + flat = _flatten_result(_result(_text("Took a screenshot"), _image())) + stripped = strip_result_for_model(flat) + assert stripped == "Took a screenshot\n[1 image attached; displayed to the user]" + assert PNG_B64 not in stripped + + +def test_strip_preserves_literal_mcp_sentinel_in_text(): + # A tool that legitimately returns text containing the marker (e.g. reading + # source/docs that quote it) must not be truncated: the suffix is not a + # valid JSON image array. + text = "before\n__MCP_IMAGES__: literal from source\nafter" + assert strip_result_for_model(text) == text + + +def test_strip_preserves_non_image_json_after_marker(): + text = 'log line\n__MCP_IMAGES__:["not", "image", "dicts"]' + assert strip_result_for_model(text) == text + + +def test_strip_removes_only_valid_terminal_envelope(): + text = ( + "Earlier mention: __MCP_IMAGES__: is documented here" + "\n[1 image attached; displayed to the user]" + '\n__MCP_IMAGES__:[{"data": "AAAA", "mimeType": "image/png"}]' + ) + assert strip_result_for_model(text) == ( + "Earlier mention: __MCP_IMAGES__: is documented here" + "\n[1 image attached; displayed to the user]" + ) + + +def test_strip_still_handles_images_and_rag_sentinels(): + assert strip_result_for_model("output\n__IMAGES__:['a.png']") == "output" + assert strip_result_for_model("answer\n__RAG_SOURCES__:[{}]") == "answer" + + +def test_error_result_keeps_error_prefix_and_images(): + flat = _flatten_result(_result(_text("boom"), _image(), is_error = True)) + assert flat.startswith("Error: boom") + assert is_tool_error(flat) + assert MCP_IMAGES_SENTINEL in flat + + +def test_image_only_error_no_longer_reports_no_content(): + flat = _flatten_result(_result(_image(), is_error = True)) + assert flat.startswith("Error: [1 image attached") + assert "tool returned no content" not in flat + + +def test_oversized_image_omitted_with_note(): + huge = "A" * (MAX_IMAGE_PAYLOAD_CHARS + 1) + flat = _flatten_result(_result(_image(data = huge))) + assert flat == "[1 image omitted (too large)]" + assert MCP_IMAGES_SENTINEL not in flat + + +def test_oversized_budget_shared_across_images(): + big = "A" * (MAX_IMAGE_PAYLOAD_CHARS - 10) + flat = _flatten_result(_result(_image(data = big), _image())) + body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) + assert "1 image attached" in body + assert "1 image omitted (too large)" in body + images = json.loads(payload) + assert len(images) == 1 and images[0]["data"] == big + + +def test_non_image_binary_block_still_ignored(): + flat = _flatten_result( + _result(SimpleNamespace(type = "audio", data = PNG_B64, mimeType = "audio/wav")) + ) + assert flat == "" + + +def test_structured_content_fallback_still_used(): + flat = _flatten_result(_result(structured = {"ok": True})) + assert flat == "{'ok': True}" + + +def test_call_tool_sync_passes_raise_on_error_false_and_keeps_error_images(monkeypatch): + # Guards that call_tool_sync passes raise_on_error=False, so an is_error result + # with image content reaches _flatten_result instead of FastMCP raising ToolError. + seen = {} + + class _FakeClient: + async def call_tool( + self, + name, + args, + raise_on_error = True, + ): + seen["raise_on_error"] = raise_on_error + return _result(_text("boom"), _image(), is_error = True) + + @contextlib.asynccontextmanager + async def _fake_client(url, headers, use_oauth): + yield _FakeClient() + + monkeypatch.setattr(mcp_client, "_client", _fake_client) + out = call_tool_sync("http://x", None, "take_screenshot", {}) + + assert seen["raise_on_error"] is False + assert out.startswith("Error: boom") + assert MCP_IMAGES_SENTINEL in out + assert is_tool_error(out) diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 6d26d075cf..24784ca34c 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -198,7 +198,12 @@ def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch): async def __aexit__(self, *args): return False - async def call_tool(self, name, args): + async def call_tool( + self, + name, + args, + raise_on_error = True, + ): import asyncio as _asyncio await _asyncio.sleep(30) # never finishes during the test @@ -520,7 +525,12 @@ def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch): async def __aexit__(self, *args): return False - async def call_tool(self, name, args): + async def call_tool( + self, + name, + args, + raise_on_error = True, + ): return "ran" monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient()) diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py index 15fe553fb2..1cb1211cf2 100644 --- a/studio/backend/tests/test_mcp_stdio_pr5863.py +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -93,7 +93,12 @@ class _RecordingClient: async def list_tools(self): return [_FakeTool("list_directory"), _FakeTool("write_file")] - async def call_tool(self, name, args): + async def call_tool( + self, + name, + args, + raise_on_error = True, + ): return _FakeResult(f"called {name}") diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index 20a22c3a4a..d0bc12706e 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -260,6 +260,30 @@ function ToolFallbackArgs({ ); } +interface McpImageResult { + text: string; + images: { data: string; mimeType: string }[]; +} + +function isMcpImageResult(val: unknown): val is McpImageResult { + if (typeof val !== "object" || val === null) { + return false; + } + const v = val as { text?: unknown; images?: unknown }; + return ( + typeof v.text === "string" && + Array.isArray(v.images) && + v.images.length > 0 && + v.images.every( + (img: unknown) => + typeof img === "object" && + img !== null && + typeof (img as { data?: unknown }).data === "string" && + typeof (img as { mimeType?: unknown }).mimeType === "string", + ) + ); +} + function ToolFallbackResult({ result, className, @@ -271,6 +295,8 @@ function ToolFallbackResult({ return null; } + const imageResult = isMcpImageResult(result) ? result : null; + return (

Result:

- 
-        {typeof result === "string" ? result : JSON.stringify(result, null, 2)}
-
+ {imageResult ? ( + <> + {imageResult.text && ( + 
+              {imageResult.text}
+
+ )} +
+ {imageResult.images.map((img, i) => ( + {`Tool + ))} +
+ + ) : ( + 
+          {typeof result === "string" ? result : JSON.stringify(result, null, 2)}
+
+ )}
); } diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index ca062bf93c..4fb7d9bd7e 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -901,6 +901,33 @@ function serializeAssistantToolCallPart( return entry; } +export interface McpImageToolResult { + text: string; + images: { data: string; mimeType: string }[]; +} + +export function isMcpImageToolResult( + val: unknown, +): val is McpImageToolResult { + if (typeof val !== "object" || val === null) { + return false; + } + const v = val as { text?: unknown; images?: unknown; sessionId?: unknown }; + return ( + typeof v.text === "string" && + v.sessionId === undefined && + Array.isArray(v.images) && + v.images.length > 0 && + v.images.every( + (img: unknown) => + typeof img === "object" && + img !== null && + typeof (img as { data?: unknown }).data === "string" && + typeof (img as { mimeType?: unknown }).mimeType === "string", + ) + ); +} + function serializeToolResultPart( part: ToolCallMessagePart, ): SerializedToolResult | null { @@ -920,6 +947,8 @@ function serializeToolResultPart( // content; serialise a sentinel JSON so legitimately empty tool // outputs still round-trip the follow-up turn to the provider. content = result.length > 0 ? result : JSON.stringify({ result: "" }); + } else if (isMcpImageToolResult(result)) { + content = result.text.length > 0 ? result.text : JSON.stringify({ result: "" }); } else { try { content = JSON.stringify(result); @@ -3196,9 +3225,12 @@ export function createOpenAIStreamAdapter( const rawResult = (toolEvent.result as string) ?? ""; const imgMarker = "\n__IMAGES__:"; const imgIdx = rawResult.lastIndexOf(imgMarker); + const mcpImgMarker = "\n__MCP_IMAGES__:"; + const mcpImgIdx = rawResult.lastIndexOf(mcpImgMarker); let parsedResult: | string | { text: string; images: string[]; sessionId: string } + | McpImageToolResult | { image_b64: string; image_mime: string; @@ -3208,6 +3240,24 @@ export function createOpenAIStreamAdapter( prompt?: string; }; const imageB64 = toolEvent.image_b64 as string | undefined; + // A valid MCP image envelope wins; an invalid marker falls + // through so a sandbox __IMAGES__ suffix still renders and + // legit text round-trips unchanged. + let mcpImages: McpImageToolResult | null = null; + if (mcpImgIdx !== -1) { + try { + const images = JSON.parse( + rawResult.slice(mcpImgIdx + mcpImgMarker.length), + ); + const candidate = { + text: rawResult.slice(0, mcpImgIdx), + images, + }; + if (isMcpImageToolResult(candidate)) mcpImages = candidate; + } catch { + // Not a valid envelope; fall through below. + } + } if ( toolCallParts[idx].toolName === "image_generation" && typeof imageB64 === "string" && @@ -3225,6 +3275,8 @@ export function createOpenAIStreamAdapter( background: toolEvent.background as string | undefined, prompt: toolEvent.prompt as string | undefined, }; + } else if (mcpImages !== null) { + parsedResult = mcpImages; } else if (imgIdx !== -1) { const text = rawResult.slice(0, imgIdx); // Fall back to "_default" to match the backend sandbox diff --git a/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts b/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts index 841f2c2a2f..9badd43bc3 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts @@ -28,13 +28,43 @@ const SEARCH_REBUILD_DEBOUNCE_MS = 300; // Keys whose values are base64 image/audio payloads, not searchable text. const BINARY_KEY = /b64|base64|^(images?|audio|video)$/i; +// Drop a trailing __MCP_IMAGES__ envelope only when it is the valid JSON image +// array appended by the backend, so legit tool text that merely mentions the +// marker stays searchable. (base64 runs below are scrubbed regardless.) +function stripMcpImageSuffix(value: string): string { + const marker = "\n__MCP_IMAGES__:"; + const idx = value.lastIndexOf(marker); + if (idx === -1) return value; + try { + const images: unknown = JSON.parse(value.slice(idx + marker.length)); + if ( + Array.isArray(images) && + images.length > 0 && + images.every( + (img) => + typeof img === "object" && + img !== null && + typeof (img as Record).data === "string" && + typeof (img as Record).mimeType === "string", + ) + ) { + return value.slice(0, idx); + } + } catch { + // Not a valid envelope; leave the text intact. + } + return value; +} + // Readable text from tool args/results, dropping base64 image/audio blobs so // they never bloat the index (object fields by key, plus data URLs / long // base64 runs and the "__IMAGES__" suffix inside strings). function searchableText(value: unknown, depth = 0): string { if (typeof value === "string") { - const cut = value.indexOf("\n__IMAGES__:"); - return (cut === -1 ? value : value.slice(0, cut)) + let text = stripMcpImageSuffix(value); + const cut = text.indexOf("\n__IMAGES__:"); + if (cut !== -1) text = text.slice(0, cut); + return text .replace(/data:[^;,\s]+;base64,[A-Za-z0-9+/=]+/g, " ") .replace(/[A-Za-z0-9+/]{120,}={0,2}/g, " "); } diff --git a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx index ee3a49526f..f4546a01a8 100644 --- a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx +++ b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx @@ -54,6 +54,7 @@ import { syncStoredChatMessages, } from "../utils/chat-history-storage"; import { notifyChatHistoryUpdated } from "../api/chat-api"; +import { isMcpImageToolResult } from "../api/chat-adapter"; import { usePlusMenuPrefsStore } from "../stores/plus-menu-prefs-store"; import type { ThreadRecord, MessageRecord } from "../types"; @@ -170,11 +171,14 @@ function contentBlocksToText(content: unknown): string { parts.push("[thinking]\n" + thinkText + "\n[/thinking]"); } } else if (p.type === "tool-call") { + // Keep base64 image payloads out of every export format: use the + // model-visible text for MCP image results (matches chat replay). + const result = isMcpImageToolResult(p.result) ? p.result.text : p.result; parts.push( JSON.stringify({ tool_call: p.toolName, args: p.args, - result: p.result, + result, }), ); } else if (p.type === "image") { @@ -299,7 +303,15 @@ function messageToOpenAI(msg: { role: unknown; content: unknown; attachments?: u const argsStr = p.args != null ? JSON.stringify(p.args) : (typeof p.argsText === "string" ? p.argsText : "{}"); toolCalls.push({ id, type: "function", function: { name, arguments: argsStr } }); if (p.result !== undefined && p.result !== null) { - const resultStr = typeof p.result === "string" ? p.result : JSON.stringify(p.result); + // Keep base64 image payloads out of exports: MCP image results carry + // their model-visible text alongside the data, so serialize the text + // (matching chat replay) instead of the full object. + const resultStr = + typeof p.result === "string" + ? p.result + : isMcpImageToolResult(p.result) + ? p.result.text + : JSON.stringify(p.result); toolResults.push({ role: "tool", tool_call_id: id, name, content: resultStr }); } } From 6e375a5b177cd39dfcbe6f1e01b287c5c0b83635 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Tue, 14 Jul 2026 01:38:04 -0700 Subject: [PATCH 018/210] Studio: add French, German, Spanish, Hindi, Arabic, Russian and Korean display languages (#7076) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Studio: add 7 display languages, complete and fix existing locales Adds fully translated French, German, Spanish, Hindi, Arabic, Russian and Korean locales. Fills in all missing keys for zh-CN (113), ja (71) and pt-BR (47), fixes translation errors found in review, and reorders the language dropdown by popularity. All overlays pass check-parity with zero missing keys and zero placeholder mismatches. * Studio: default display language to auto detect The language preference now defaults to auto and resolves against the browser language list, with exact tag match first and language subtag match second (pt-PT resolves to pt-BR, zh-TW to zh-CN). Auto detect is the first dropdown option and is translated in every locale. Explicit choices still persist and sync; personalization sync now round trips the preference instead of the resolved locale so auto stays auto across devices. Auto mode also follows browser languagechange events. * Studio: guard import.meta.env in translate for non-Vite contexts translate() read import.meta.env.DEV directly, which throws when the module runs outside Vite (SSR or Node tooling). Optional-chain it so the dev-only warning is skipped and translation still works everywhere. * Studio: RTL for Arabic, translate recipes, keep Traditional Chinese off zh-CN - Sync document dir from a per-locale dir field so Arabic mirrors the layout instead of rendering RTL text in an LTR shell. - Translate the recipes nav label in fr, de, ko and hi to match the other locales (Recettes, Rezepte, and native forms). - Detection no longer maps Traditional Chinese (zh-Hant / zh-TW / zh-HK / zh-MO) to Simplified zh-CN; those tags fall through to the next preferred language. Simplified tags (zh, zh-CN, zh-SG, zh-Hans) still resolve to zh-CN. * Studio: don't treat legacy synced English as an explicit language pick The old sync serialized the resolved locale on every save, so existing profiles carry appearance.language 'en' even when the user never chose a language. Hydrating that as a pinned locale forced non-English browsers back to English under the new Auto detect default. Payloads now carry version 2 (the preference itself); on hydrate a version 1 'en' maps to auto, while explicit picks and all version 2 values are kept as-is. * Studio: persist only known language codes from the locale table normalizePreference now returns a value re-derived from the LOCALES keys instead of the raw input. It stays functionally identical (the stored value was already whitelisted) but makes it explicit that only known, non-sensitive language codes are written to localStorage, and clears a false-positive clear-text-storage scan on the persistence path. * Studio i18n: fix Train label transliteration and tidy locale consistency - ja and hi: the nav and route Train label used the railway transliteration (トレイン and ट्रेन); switch to the training term already used everywhere else in each file (トレーニング, ट्रेनिंग). - zh-CN: keep VRAM in English to match every other locale and the PR's own keep-English rule, and drop an extra clause added to the upload size hint so it matches the English source. - hi: translate Recents to हाल के in the export and import section to match the sidebar label, and point users to the Configure tab by its translated name (कॉन्फ़िगर). - ru: reword the preview sharing hint to avoid the "disable to disable" repetition. i18n parity and the type checked build stay green. * Studio i18n: keep Arabic layout LTR until physical-direction CSS is converted Setting ar to dir rtl only mirrors the flex based shell, sidebar and settings dialog. The shared select, dialog and dropdown primitives use physical-direction utilities (right-2, top-5 right-5, ml-auto) that do not flip under dir rtl, so chevrons, close buttons and check marks land on the wrong side. Keep Arabic on an LTR layout for now, matching the original plan in this PR. Arabic text still renders right to left per element via bidi and chat content keeps dir auto, so nothing regresses. Full layout mirroring can follow once the physical-direction classes are converted to logical ones. * Studio i18n: do not let a generic zh after a Traditional tag pick Simplified navigator.languages can be a list like ['zh-TW', 'zh', 'en-US']. The zh-TW pass already falls through, but the bare zh then reached the language-subtag match and selected zh-CN, so Traditional Chinese users still got Simplified and the guard was defeated. detectLocale now remembers when a Traditional tag was seen and skips a later bare zh, so detection keeps falling through to the next non-Chinese language. A lone bare zh, and explicit zh-CN or zh-Hans fallbacks, still resolve to Simplified as before. * Studio i18n: collapse two locale comments to a single line The Arabic dir note in messages.ts and the bare-zh note in locale-store.ts were two lines each; tighten each to one. Comment only, no behavior change. * Studio i18n: translate Hindi strings that were left in English Seventeen hi.ts labels stayed in English while all the other locales translated them: the training parameter labels (Grad Accum, Grad Norm, Grad Checkpoint, Eval Loss, Clip p95/p99, Seed, Continued Pretraining), the API example labels (curl/Python/JavaScript + tools/advanced), Hugging Face token, the VRAM estimate and the training terminal start line. Parity only checks key/placeholder presence so it did not catch these. Brand and technical tokens (curl, Python, VRAM, Loss, p95/p99, Hugging Face, unsloth) stay in English as elsewhere. --------- Co-authored-by: danielhanchen --- .../profile/hooks/use-personalization-sync.ts | 48 +- .../settings/components/language-select.tsx | 14 +- studio/frontend/src/i18n/check-parity.ts | 16 +- studio/frontend/src/i18n/index.ts | 6 + studio/frontend/src/i18n/locale-store.ts | 169 ++- studio/frontend/src/i18n/locales/ar.ts | 1001 ++++++++++++++++ studio/frontend/src/i18n/locales/de.ts | 1042 ++++++++++++++++ studio/frontend/src/i18n/locales/en.ts | 1 + studio/frontend/src/i18n/locales/es.ts | 1043 +++++++++++++++++ studio/frontend/src/i18n/locales/fr.ts | 1038 ++++++++++++++++ studio/frontend/src/i18n/locales/hi.ts | 999 ++++++++++++++++ studio/frontend/src/i18n/locales/ja.ts | 101 +- studio/frontend/src/i18n/locales/ko.ts | 1002 ++++++++++++++++ studio/frontend/src/i18n/locales/pt-br.ts | 76 +- studio/frontend/src/i18n/locales/ru.ts | 999 ++++++++++++++++ studio/frontend/src/i18n/locales/zh-CN.ts | 171 ++- studio/frontend/src/i18n/messages.ts | 39 +- 17 files changed, 7684 insertions(+), 81 deletions(-) create mode 100644 studio/frontend/src/i18n/locales/ar.ts create mode 100644 studio/frontend/src/i18n/locales/de.ts create mode 100644 studio/frontend/src/i18n/locales/es.ts create mode 100644 studio/frontend/src/i18n/locales/fr.ts create mode 100644 studio/frontend/src/i18n/locales/hi.ts create mode 100644 studio/frontend/src/i18n/locales/ko.ts create mode 100644 studio/frontend/src/i18n/locales/ru.ts diff --git a/studio/frontend/src/features/profile/hooks/use-personalization-sync.ts b/studio/frontend/src/features/profile/hooks/use-personalization-sync.ts index eac1a64d7a..5dbb7b54bf 100644 --- a/studio/frontend/src/features/profile/hooks/use-personalization-sync.ts +++ b/studio/frontend/src/features/profile/hooks/use-personalization-sync.ts @@ -9,12 +9,12 @@ import { type Theme, } from "@/features/settings"; import { - DEFAULT_LOCALE, - getLocale, - isSupportedLocale, + DEFAULT_LOCALE_PREFERENCE, + getLocalePreference, + isLocalePreference, setLocale, - useLocale, - type Locale, + useLocalePreference, + type LocalePreference, } from "@/i18n"; import { useCallback, useEffect, useRef, useState } from "react"; import { @@ -25,6 +25,11 @@ import type { AvatarShape } from "../stores/user-profile-store"; const PUSH_DEBOUNCE_MS = 800; +// Version 2 payloads store the language preference ("auto" or a pinned +// locale). Version 1 always serialized the resolved locale, so its "en" is +// usually the old default rather than an explicit pick. +const PERSONALIZATION_VERSION = 2; + type ProfileSnapshot = { displayName: string; nickname: string; @@ -110,10 +115,10 @@ function profileSnapshot(): ProfileSnapshot { function payload( profile: ProfileSnapshot, theme: Theme, - language: Locale | null, + language: LocalePreference | null, ): PersonalizationWrite { return { - version: 1, + version: PERSONALIZATION_VERSION, profile: normalizeProfile(profile), appearance: { theme, language }, }; @@ -123,10 +128,23 @@ function serialized(data: PersonalizationWrite): string { return JSON.stringify(data); } +// Version 1 clients wrote language on every save, so a legacy "en" usually +// means the user never picked a language. Map it to auto; explicit picks of +// other locales (the old default was English) are kept. Version 2 payloads +// are trusted verbatim, so a deliberate English pick stays pinned. +export function remoteLanguagePreference( + version: unknown, + language: unknown, +): unknown { + const isLegacy = typeof version !== "number" || version < 2; + if (isLegacy && language === "en") return DEFAULT_LOCALE_PREFERENCE; + return language; +} + function hasLocalSettings( profile: ProfileSnapshot, theme: Theme, - language: Locale, + language: LocalePreference, ): boolean { return Boolean( profile.displayName || @@ -134,7 +152,7 @@ function hasLocalSettings( profile.avatarDataUrl || profile.avatarShape !== "circle" || theme !== "system" || - language !== DEFAULT_LOCALE, + language !== DEFAULT_LOCALE_PREFERENCE, ); } @@ -144,7 +162,7 @@ export function usePersonalizationSync(enabled: boolean): void { const avatarDataUrl = useUserProfileStore((s) => s.avatarDataUrl); const avatarShape = useUserProfileStore((s) => s.avatarShape); const { theme } = useTheme(); - const language = useLocale(); + const language = useLocalePreference(); const [hydratedGeneration, setHydratedGeneration] = useState(0); const authGenerationRef = useRef(0); const latestThemeRef = useRef(theme); @@ -191,8 +209,12 @@ export function usePersonalizationSync(enabled: boolean): void { avatarShape: remote.profile.avatarShape === "rounded" ? "rounded" : "circle", }; const nextTheme = remote.appearance.theme; - const nextLanguage = isSupportedLocale(remote.appearance.language) - ? remote.appearance.language + const remoteLanguage = remoteLanguagePreference( + remote.version, + remote.appearance.language, + ); + const nextLanguage = isLocalePreference(remoteLanguage) + ? remoteLanguage : latestLanguageRef.current; useUserProfileStore.setState(nextProfile); if (nextTheme !== latestThemeRef.current) setTheme(nextTheme); @@ -207,7 +229,7 @@ export function usePersonalizationSync(enabled: boolean): void { useUserProfileStore.setState(nextProfile); } const nextTheme = latestThemeRef.current; - const nextLanguage = getLocale(); + const nextLanguage = getLocalePreference(); const nextPayload = payload(nextProfile, nextTheme, nextLanguage); const nextSerialized = serialized(nextPayload); if (hasLocalSettings(nextProfile, nextTheme, nextLanguage)) { diff --git a/studio/frontend/src/features/settings/components/language-select.tsx b/studio/frontend/src/features/settings/components/language-select.tsx index 9d30e06147..01fe049a56 100644 --- a/studio/frontend/src/features/settings/components/language-select.tsx +++ b/studio/frontend/src/features/settings/components/language-select.tsx @@ -9,22 +9,23 @@ import { SelectValue, } from "@/components/ui/select"; import { + AUTO_LOCALE, LOCALES, - isSupportedLocale, + isLocalePreference, setLocale, useT, - useLocale, + useLocalePreference, } from "@/i18n"; export function LanguageSelect() { const t = useT(); - const locale = useLocale(); + const preference = useLocalePreference(); return (