From 79a431cd53cbae34c4026a4434af269d339f228b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 17 May 2026 01:35:28 -0700 Subject: [PATCH 1/8] tests: pinned-symbol canary for unsloth-zoo save_pretrained_merged guards (#5410) (#5433) * tests: pinned-symbol canary for unsloth-zoo save_pretrained_merged guards (#5410) unsloth#5410 was a class of silent-write bug in the save_pretrained_merged path that the existing CI matrix could not detect because the merge-helper tests were not wired through the upstream-drift suite. The full fix lives in unslothai/unsloth-zoo#647 (layout-aware MoE merge helpers, authoritative num_experts resolver, loud-fail counter, generation_config.json save). This PR adds the unsloth-side canary that watches for the four guards staying in place in unsloth-zoo so a future refactor cannot silently regress them. tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py fetches unsloth_zoo/saving_utils.py + tests/test_unsloth_zoo_lora_merge.py from unslothai/unsloth-zoo:main and asserts: - _MOE_MERGE_STATE / _reset_moe_merge_state / _record_moe_merge_fallback are still defined and a `raise RuntimeError(...MoE...)` still fires when fallback > 0. - _detect_moe_lora_layout exists and both "swapped" / "standard" branch labels are reachable in the source. - _resolve_num_experts_from_lora_stats is present AND its base_layer walk is bounded by `for _ in range(N):` (a cyclic ParamWrapper chain must not hang the merge). - merge_and_overwrite_lora still calls model.generation_config.save_pretrained(...). - tests/test_unsloth_zoo_lora_merge.py keeps the six PEFT 0.19+ standard-layout regression tests added in #647. - Local unsloth/save.py still names save_pretrained_merged and routes through merge_and_overwrite_lora (i.e. the entry point still reaches the upstream fix). While #647 is still open, the four symbol tests SKIP cleanly with a message naming #647. When #647 merges into unsloth-zoo main, the same tests automatically become hard gates and catch any future regression. The sixth test (local entry-point grep) passes today. CPU-only static fetch, ~0.1s. Wired into the existing peft-pinned-symbols job in .github/workflows/version-compat-ci.yml so it runs on every PR that touches unsloth/** and on the daily schedule. Local run: 1 passed, 5 skipped (expected; #647 open). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests/version_compat: relax MoE/generation_config regex to fit zoo#647 zoo#647 landed two layout changes that broke the pinned-symbol canary's exact-string regex matches but kept the underlying guarantees intact: - The post-loop MoE LoRA fallback `raise RuntimeError(...)` wraps the "MoE" wording onto a second line; the old `[^\n]*` did not cross newlines. Switch to `.*?` + re.DOTALL. - The generation_config save now binds the attr to a local var `gen_cfg = getattr(model, "generation_config", ...)` and calls `gen_cfg.save_pretrained(save_directory)`, so a literal `generation_config.save_pretrained(` substring no longer matches. Anchor on the conceptual operation: a `generation_config` mention followed (within a small char window) by a `.save_pretrained(` call. That is what the canary actually cares about. Verified locally: pytest tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py -> 2 passed (4 deselected) --------- Co-authored-by: Daniel Han-Chen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/version-compat-ci.yml | 1 + ..._unsloth_zoo_save_merged_pinned_symbols.py | 128 ++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100644 tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index 2fbdd15747..599b53df1d 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -127,6 +127,7 @@ jobs: run: | PYTHONPATH=. python -m pytest \ tests/version_compat/test_peft_pinned_symbols.py \ + tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py \ -v --tb=short st-pinned-symbols: diff --git a/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py b/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py new file mode 100644 index 0000000000..19faa51119 --- /dev/null +++ b/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Pinned-symbol canary for unsloth-zoo save_pretrained_merged guards +(unslothai/unsloth-zoo#647 / unslothai/unsloth#5410). Skips until #647 +lands, then becomes a hard gate. CPU-only static fetch.""" + +from __future__ import annotations + +import re + +import pytest + +from tests.version_compat._fetch import fetch_text + + +ZOO_TAG = "main" + + +def _fetch_saving_utils() -> str: + src = fetch_text("unslothai/unsloth-zoo", ZOO_TAG, "unsloth_zoo/saving_utils.py") + if src is None: + pytest.skip("unsloth_zoo/saving_utils.py not fetchable") + return src + + +def _fetch_merge_tests() -> str: + src = fetch_text( + "unslothai/unsloth-zoo", + ZOO_TAG, + "tests/test_unsloth_zoo_lora_merge.py", + ) + if src is None: + pytest.skip("tests/test_unsloth_zoo_lora_merge.py not fetchable") + return src + + +def _skip_until_pr_647_lands(src: str) -> None: + if not any( + m in src + for m in ( + "_MOE_MERGE_STATE", + "_detect_moe_lora_layout", + "_resolve_num_experts_from_lora_stats", + ) + ): + pytest.skip( + "unslothai/unsloth-zoo#647 has not yet merged into main; " + "tests auto-promote to hard gates once it lands." + ) + + +def test_zoo_saving_utils_has_moe_merge_state(): + src = _fetch_saving_utils() + _skip_until_pr_647_lands(src) + for sym in ( + "_MOE_MERGE_STATE", + "_reset_moe_merge_state", + "_record_moe_merge_fallback", + ): + assert sym in src, f"{sym} missing from saving_utils.py (issue #5410 guard)." + # zoo#647 wraps the fallback guard's message onto a second line; + # allow the regex to span newlines via re.DOTALL. + assert re.search( + r"raise\s+RuntimeError\b.*?MoE", src, re.IGNORECASE | re.DOTALL + ), "no `raise RuntimeError(...MoE...)`; post-loop guard weakened." + + +def test_zoo_saving_utils_has_layout_detector(): + src = _fetch_saving_utils() + _skip_until_pr_647_lands(src) + assert ( + "_detect_moe_lora_layout" in src + ), "_detect_moe_lora_layout removed (issue #5410)." + assert ( + '"swapped"' in src and '"standard"' in src + ), "one of the layout labels removed." + + +def test_zoo_saving_utils_has_num_experts_resolver(): + src = _fetch_saving_utils() + _skip_until_pr_647_lands(src) + assert "_resolve_num_experts_from_lora_stats" in src, "resolver removed (#5410)." + assert re.search( + r"for\s+_\s+in\s+range\s*\(\s*\d+\s*\)", src + ), "resolver walk no longer bounded by `for _ in range(N):`." + + +def test_zoo_saving_utils_writes_generation_config(): + src = _fetch_saving_utils() + _skip_until_pr_647_lands(src) + # zoo#647 binds the generation_config attr to a local var + # (`gen_cfg = getattr(model, "generation_config", ...); ... + # gen_cfg.save_pretrained(save_directory)`) so an exact + # `generation_config.save_pretrained(` substring no longer + # matches. Anchor on the conceptual operation: a `generation_config` + # mention plus a `.save_pretrained(` call nearby, which is what + # the canary actually cares about. + assert re.search( + r"generation_config[\s\S]{0,400}?\.save_pretrained\s*\(", src + ), "generation_config.json no longer saved (#5410)." + + +def test_zoo_lora_merge_tests_have_standard_layout_coverage(): + src = _fetch_merge_tests() + if "test_merge_moe_gate_expert_standard_layout" not in src: + pytest.skip("unslothai/unsloth-zoo#647 not yet merged; coverage appears later.") + for name in ( + "test_merge_moe_gate_expert_standard_layout", + "test_merge_moe_up_expert_standard_layout", + "test_merge_moe_down_proj_expert_standard_layout", + "test_detect_moe_lora_layout_classifies_both_conventions", + "test_moe_merge_fallback_counter_records_bad_layout", + "test_resolve_num_experts_walks_base_layer_chain", + ): + assert name in src, f"regression test `{name}` removed." + + +def test_unsloth_save_pretrained_merged_entry_point_exists(): + import pathlib + + save_py = pathlib.Path(__file__).resolve().parents[2] / "unsloth" / "save.py" + if not save_py.is_file(): + pytest.skip(f"{save_py} not present") + text = save_py.read_text(encoding = "utf-8", errors = "replace") + assert "save_pretrained_merged" in text, "entry point removed from unsloth/save.py." + assert ( + "merge_and_overwrite_lora" in text + ), "no dispatch into unsloth_zoo merge; #647 bypassed." From e20bbeff9afdb0cc1f96093e558aac7e434278c4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 17 May 2026 01:40:20 -0700 Subject: [PATCH 2/8] intel-gpu: pin unsloth_zoo>=2026.5.2 via huggingfacenotorch (#5499) Fixes unslothai/unsloth#5494: installing any intel-gpu-torch* extra without also pulling `huggingface` or `colab-new` lets the resolver silently fall back to a stale unsloth_zoo (2026.3.6 in the original report) because no version floor is enforced on `unsloth_zoo[intelgpu]` in those blocks. unsloth_zoo 2026.5.2 (just released) also relaxes its own torch upper bound from <2.11.0 to <2.13.0, which is what unblocks the resolver for the intel-gpu-torch2110 and intel-gpu-torch2120 extras shipped in #5484. Adding the floor in `huggingfacenotorch` propagates it to every extra that includes the HF-without-torch base: amd, huggingface, and all nine intelgputorch* blocks. Single line, single source of truth. Requires unsloth_zoo 2026.5.2 to be on PyPI for end-to-end resolution. --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 3468f7f8a7..81cf5ac215 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,7 @@ triton = [ ] huggingfacenotorch = [ + "unsloth_zoo>=2026.5.2", "wheel>=0.42.0", "packaging", "numpy", From fb13275787451bd67bba88e546c94e4636f353e2 Mon Sep 17 00:00:00 2001 From: Etherll <61019402+Etherll@users.noreply.github.com> Date: Sun, 17 May 2026 14:02:36 +0300 Subject: [PATCH 3/8] fix(sentence_transformer): resume PEFT checkpoints under sentence-transformers >= 5.4 (#5454) Saves the base config.json next to adapter_config.json when checkpointing PEFT-wrapped sentence-transformer models, and overrides SentenceTransformerTrainer._load_from_checkpoint to load adapter weights via set_peft_model_state_dict and rebuild aux modules (Pooling, Normalize, Dense) from modules.json with strict type and path validation. Patches only activate on Unsloth-managed Transformer modules so non-Unsloth pipelines fall through to upstream behaviour. Fixes https://github.com/unslothai/unsloth/issues/5373 --- unsloth/models/sentence_transformer.py | 181 ++++++++++++++++++++++++- 1 file changed, 180 insertions(+), 1 deletion(-) diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index c53e3a7a81..aafc416221 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -527,6 +527,45 @@ This sentence-transformers model was finetuned and converted to GGUF format usin class FastSentenceTransformer(FastModel): + @staticmethod + def _save_base_config_for_processor_resume(config, output_path): + """ + sentence-transformers >= 5.4 reloads Transformer modules through + AutoProcessor. Tokenizer-only checkpoint roots make AutoProcessor fall + back to AutoConfig, so PEFT adapter checkpoints still need the base + config.json next to adapter_config.json. + """ + if config is None or not getattr(config, "model_type", None): + return + if hasattr(config, "save_pretrained"): + config.save_pretrained(output_path) + elif hasattr(config, "to_json_file"): + config_path = os.path.join(output_path, "config.json") + config.to_json_file(config_path) + + @staticmethod + def _patch_transformer_module_save_config(transformer_module, base_config = None): + transformer_module._unsloth_st_managed = True + if base_config is not None and getattr(base_config, "model_type", None): + transformer_module._unsloth_base_config = base_config + + if getattr(transformer_module, "_unsloth_save_config_patched", False): + return transformer_module + + original_save = transformer_module.save + + def _save_with_base_config(self, output_path, *args, **kwargs): + original_save(output_path, *args, **kwargs) + FastSentenceTransformer._save_base_config_for_processor_resume( + getattr(self, "_unsloth_base_config", None), output_path + ) + + transformer_module.save = types.MethodType( + _save_with_base_config, transformer_module + ) + transformer_module._unsloth_save_config_patched = True + return transformer_module + @staticmethod def _read_pooling_mode(model_name, token): """ @@ -1157,6 +1196,9 @@ class FastSentenceTransformer(FastModel): config_keys.append(config_key) transformer_module.config_keys = config_keys transformer_module.save_in_root = True + FastSentenceTransformer._patch_transformer_module_save_config( + transformer_module, getattr(model, "config", None) + ) if hasattr(model, "config"): model.config.tokenizer_class = tokenizer.__class__.__name__ @@ -1644,6 +1686,9 @@ class FastSentenceTransformer(FastModel): st_model._dtype = dtype st_model._load_in_4bit = load_in_4bit st_model.no_modules = False + FastSentenceTransformer._patch_transformer_module_save_config( + st_model[0], getattr(st_model[0].auto_model, "config", None) + ) # Add save methods def _save_pretrained_merged(self, save_directory, **save_kwargs): @@ -2067,6 +2112,9 @@ class FastSentenceTransformer(FastModel): transformer_module.model = peft_model else: transformer_module.auto_model = peft_model + FastSentenceTransformer._patch_transformer_module_save_config( + transformer_module, getattr(inner_model, "config", None) + ) # Store compile info for auto-compile at trainer time # torch.compile is deferred until training starts so we can check max_steps @@ -2121,6 +2169,9 @@ class FastSentenceTransformer(FastModel): transformer_module.model = peft_model else: transformer_module.auto_model = peft_model + FastSentenceTransformer._patch_transformer_module_save_config( + transformer_module, getattr(inner_model, "config", None) + ) return model else: return FastModel.get_peft_model( @@ -2235,5 +2286,133 @@ def _patch_sentence_transformer_trainer(): SentenceTransformerTrainer._unsloth_auto_compile_patched = True -# Auto-patch trainer on module import +def _patch_st_trainer_load_from_checkpoint(): + try: + from sentence_transformers import SentenceTransformerTrainer + except ImportError: + return + if getattr( + SentenceTransformerTrainer, "_unsloth_load_from_checkpoint_patched", False + ): + return + if not hasattr(SentenceTransformerTrainer, "_load_from_checkpoint"): + return + + _original = SentenceTransformerTrainer._load_from_checkpoint + + def _unsloth_load_from_checkpoint(self, checkpoint_path): + try: + from peft import PeftModel, load_peft_weights, set_peft_model_state_dict + except ImportError: + return _original(self, checkpoint_path) + + try: + mod0 = self.model[0] + except (IndexError, TypeError): + return _original(self, checkpoint_path) + + if isinstance(getattr(type(mod0), "auto_model", None), property): + inner = getattr(mod0, "model", None) + else: + inner = getattr(mod0, "auto_model", None) + inner = getattr(inner, "_orig_mod", inner) + + if not isinstance(inner, PeftModel): + return _original(self, checkpoint_path) + if not getattr(mod0, "_unsloth_st_managed", False): + return _original(self, checkpoint_path) + + if not any( + os.path.isfile(os.path.join(checkpoint_path, fn)) + for fn in ("adapter_model.safetensors", "adapter_model.bin") + ): + return _original(self, checkpoint_path) + + adapter_name = getattr(inner, "active_adapter", None) + if adapter_name is None and callable(getattr(inner, "active_adapters", None)): + adapter_name = inner.active_adapters() + if isinstance(adapter_name, (list, tuple, set)): + if len(adapter_name) != 1: + raise RuntimeError( + "Unsloth: Cannot resume multiple active PEFT adapters." + ) + adapter_name = next(iter(adapter_name)) + adapter_name = adapter_name or "default" + if adapter_name not in getattr(inner, "peft_config", {}): + raise RuntimeError(f"Unsloth: PEFT adapter {adapter_name!r} is not loaded.") + + load_result = set_peft_model_state_dict( + inner, load_peft_weights(checkpoint_path), adapter_name = adapter_name + ) + unexpected = getattr(load_result, "unexpected_keys", []) or [] + missing = [ + x + for x in (getattr(load_result, "missing_keys", []) or []) + if f".{adapter_name}." in x or x.endswith(f".{adapter_name}") + ] + if unexpected or missing: + raise RuntimeError( + "Unsloth: PEFT checkpoint does not match the active adapter " + f"(missing={missing[:8]}, unexpected={unexpected[:8]})." + ) + + modules_json = os.path.join(checkpoint_path, "modules.json") + if not os.path.isfile(modules_json): + raise RuntimeError("Unsloth: PEFT checkpoint is missing modules.json.") + try: + with open(modules_json, "r") as f: + module_configs = json.load(f) + except Exception as e: + raise RuntimeError("Unsloth: Cannot parse checkpoint modules.json.") from e + + root = os.path.abspath(os.fspath(checkpoint_path)) + restored = set() + for entry in module_configs: + idx = int(entry.get("idx", -1)) + if idx == 0: + continue + if idx < 0 or idx >= len(self.model): + raise RuntimeError(f"Unsloth: Bad module index in modules.json: {idx}.") + module = self.model[idx] + module_cls = type(module) + saved_type = entry.get("type", "") + if saved_type and not saved_type.endswith(f".{module_cls.__name__}"): + raise RuntimeError(f"Unsloth: Checkpoint module {idx} type mismatch.") + module_path = entry.get("path") + module_dir = os.path.abspath( + os.path.join(root, os.fspath(module_path or "")) + ) + try: + inside_root = os.path.commonpath([root, module_dir]) == root + except ValueError: + inside_root = False + if not module_path or not inside_root or not os.path.isdir(module_dir): + raise RuntimeError( + f"Unsloth: Bad checkpoint module path for index {idx}." + ) + if not hasattr(module_cls, "load"): + raise RuntimeError(f"Unsloth: Module {idx} cannot be reloaded.") + fresh = module_cls.load(module_dir) + if not isinstance(fresh, module_cls): + raise RuntimeError(f"Unsloth: Module {idx} reload returned wrong type.") + # Parameterless modules (Pooling, Normalize) make + # next(module.parameters()) raise StopIteration; route through + # the SentenceTransformer's device property instead. + try: + fresh.to(self.model.device) + except AttributeError: + pass + self.model[idx] = fresh + restored.add(idx) + missing_idx = sorted(set(range(1, len(self.model))) - restored) + if missing_idx: + raise RuntimeError( + f"Unsloth: Checkpoint modules.json is incomplete (missing idx={missing_idx[:8]})." + ) + + SentenceTransformerTrainer._load_from_checkpoint = _unsloth_load_from_checkpoint + SentenceTransformerTrainer._unsloth_load_from_checkpoint_patched = True + + _patch_sentence_transformer_trainer() +_patch_st_trainer_load_from_checkpoint() From ab56e4a9edaeae839d6b229429b30e36605b26a7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 17 May 2026 04:16:09 -0700 Subject: [PATCH 4/8] Studio: serialise GGUF reload and inherit unsloth-run extra args (#5427) * Studio: serialise GGUF reload and inherit unsloth-run extra args Closes #5401. Three related GGUF reload bugs reproduced against `unsloth studio run -m unsloth/Qwen3-0.6B-GGUF --gguf-variant Q4_K_M --top-k 20 --seed 42`: 1. The `POST /api/inference/load` already-loaded short-circuit only compared `model_identifier` and `hf_variant`. A same-(model, variant) Apply that flipped `cache_type_kv` / `speculative_type` / `chat_template_override` / `max_seq_length` / `llama_extra_args` returned `status="already_loaded"` and the new setting silently never reached llama-server. 2. The frontend chat-settings Apply path POSTs `/unload` then `/load` without round-tripping `llama_extra_args`. Every reload after `unsloth run --some-flag X` quietly dropped `--some-flag X` from the spawned `llama-server` command line. 3. `LlamaCppBackend.load_model` released `_lock` between Phase 1 (kill) and Phase 3 (spawn) so two concurrent loads each passed Phase 1 with `self._process is None`. Both ran Phase 2 (download), both reached Phase 3, and the Phase 3 defensive `_kill_process()` from #5171 collapsed them to one survivor only after both `subprocess.Popen` calls had landed. For the 86 GB MoE in #5161 / the model in #5401 the overlap window was tens of seconds, long enough to OOM the host. With a 0.6B model the pgrep timeline showed two simultaneous PIDs for 3.3 s on `main`. Fix: `studio/backend/core/inference/llama_cpp.py` * Add `self._serial_load_lock = threading.Lock()`. The whole body of `load_model` runs under this lock so two concurrent `/api/inference/load` requests are strictly sequential. The fine-grained `_lock` and the Phase 3 defensive `_kill_process()` from #5171 are kept as a second layer. `/unload`, `/status`, and `/load-progress` are unaffected because they only touch the fine-grained lock or read properties. * Add `self._extra_args` plus an `extra_args` property, written inside `load_model` whenever the caller supplies a non-`None` value. `unload_model()` deliberately does not reset it so the route layer can inherit the args across the frontend's `/unload` + `/load` gap. `studio/backend/routes/inference.py` * Add `_request_matches_loaded_settings(request, llama_backend)` that compares `max_seq_length`, `cache_type_kv`, `speculative_type`, `chat_template_override`, and `llama_extra_args` between the incoming request and the live backend. Same-(model, variant) requests whose runtime settings differ now fall through to a real reload instead of returning `already_loaded`. A missing `llama_extra_args` field on the request is treated as "inherit current", so the short-circuit still fires when the only difference is the frontend not echoing the CLI flags back. * GGUF load branch inherits `llama_extra_args` from `llama_backend.extra_args` when the request omits the field, re-validates through `validate_extra_args`, and forwards the result to `load_model(...)`. An explicit `[]` from the caller is still honoured as "clear". Verified end to end against a live `unsloth studio run` instance: | Scenario | Before | After | | --------------------------------------------------------------- | --------- | ------------------------------------------------------------------------ | | `/load` same (model, variant, settings) | 1 PID, `already_loaded` | unchanged | | `/load` same model, variant, new `cache_type_kv=q8_0` ctx=8192 | `already_loaded`, settings dropped | `loaded`, `/status` reports the new settings, new server has `-c 8192 --cache-type-k q8_0 --top-k 20 --seed 42` | | Frontend Apply `/unload` + `/load`, new settings, no `llama_extra_args` field | Drops `--top-k 20 --seed 42` | Preserves `--top-k 20 --seed 42` | | `/unload` + two parallel `/load` | Two PIDs for 3.3 s | Max simultaneous count = 1 across the full pgrep timeline | | `/load` with `llama_extra_args=[]` (explicit clear) | n/a | `loaded`, new server has no `--top-k` / `--seed` | | `/load` with `llama_extra_args=["--top-k","30","--seed","7"]` (override) | n/a | `loaded`, new server has the supplied flags | `pytest studio/backend/tests` is green except for one pre-existing terminal-width-sensitive assertion (`test_studio_api.py::test_help_output`) and the pre-existing `test_studio_api.py` fixture errors that fail on unmodified main too. No new regressions. * Studio: track requested n_ctx so Auto-slider flips trigger a reload Review feedback on PR #5427 from gemini-code-assist. The original short-circuit compared ``request.max_seq_length`` against ``llama_backend.context_length`` (the effective context). VRAM-fit logic can cap the running server below what the caller asked for, so this comparison incorrectly returns ``already_loaded`` when the user flips the slider from an explicit length (e.g. 8192) back to "Auto" (0): the explicit request was capped to, say, 4096, and the new "Auto" request reads ``backend.context_length == 4096`` and decides nothing changed. Track the originally requested ``n_ctx`` on the backend instead and compare against that. ``requested_n_ctx == 0`` means the last load asked for the model's native length; ``request.max_seq_length == 0`` matches it. Verified in the sandbox suite (now 90 tests): - ``test_explicit_to_auto_triggers_reload`` -- loaded with explicit 8192, then Apply with ``max_seq_length=0`` falls through to a real reload and the new server runs at the native 40960. - ``test_auto_to_explicit_triggers_reload`` -- inverse direction. - ``test_explicit_to_same_explicit_short_circuits`` -- re-Apply with the same explicit value still short-circuits (no needless reload). - Existing scenarios (kv change, spec change, template change, extra args inherit, parallel-load stress, frontend Apply flow) unchanged. ``pytest studio/backend/tests`` still green on the same set of tests; the pre-existing ``test_help_output`` failure and ``test_studio_api`` fixture errors are unaffected. * Studio: tighten comments in the 5401 fix Trim the verbose explanatory comments and docstrings introduced in f9cbec3b and dd0b1d58 down to one-line summaries. The "why" still points at issue #5401; the multi-paragraph rationale belonged in the PR body, not the source. No behaviour change. * ci: retrigger after zoo drift + IPython fixes landed in main * ci: retrigger Mac Studio UI CI after transient fetch flake * Studio: address six P2 followups on the 5401 reload PR Tightens the inheritance and serial-load paths to close the six P2 findings raised by codex-connector on PR #5427 against `f9cbec3b` / `dd0b1d58`. 1. Re-check loaded state before killing queued loads. Two duplicate `/api/inference/load` requests both pass the route-level `is_loaded` gate before the first publishes `_healthy = True`. The second waits on `_serial_load_lock`, enters Phase 1, and tears down the just-spawned llama-server for a redundant full reload. Added `LlamaCppBackend._already_in_target_state(...)` and a short-circuit at the top of the serial-lock block: if the live server already satisfies the kwargs, return True without killing. 2. Don't inherit CLI overrides that shadow new first-class settings. `unsloth run -c 4096` is a permitted pass-through; the validator docs explicitly call out `-c`/`--ctx-size`. Stored in `_extra_args` and appended after Studio's own flags, the inherited `-c 4096` silently won the last-wins parse against a new `max_seq_length=8192`. Added `strip_shadowing_flags` in `llama_server_args.py` (covers `-c`, `--cache-type-k/v`, `--spec-*`, `--chat-template*`, `--jinja`/`--no-jinja`) and the route runs the inherited list through it before validate + forward. 3. Restrict inherited llama args to the same GGUF model. `_extra_args` is deliberately preserved across `unload_model()` for the chat- settings Apply flow (`/unload` + `/load` with no `llama_extra_args` field). Now also track `_extra_args_source = (model_identifier, hf_variant)` so the route can refuse cross-model inheritance. `LlamaCppBackend.extra_args_source` exposes the tuple. 4. Persist extras only after a successful load. `_extra_args` was written at the top of `load_model` before Popen + health check, so a failed startup left bad args in place to poison the next UI retry. The write (along with `_requested_n_ctx`) is now deferred until after `_healthy = True`. 5. Ignore speculative diffs for vision loads. `load_model` silently gates speculative decoding on `not is_vision`, so the backend's `_speculative_type` stays `None` for vision models. The route's comparator now normalises the request's value to `"off"` when `llama_backend.is_vision` to avoid a no-op reload of a vision server every time the dropdown defaults to `default`. The `_already_in_target_state` helper applies the same rule. 6. Wait for the replacement server before short-circuiting. `_kill_process` did not clear `_healthy`; the new first-class settings (`_cache_type_kv`, `_speculative_type`, `_chat_template_override`) are written under `_lock` BEFORE Popen + `_wait_for_health`. A duplicate `/load` arriving during the new server's warm-up window could short-circuit against the not-yet-healthy replacement and the caller would start inference against a server that was still loading. `_kill_process` now sets `_healthy = False` in its `finally` block so `is_loaded` returns False from the moment the old server is killed until the new one finishes warm-up. Tests: - Sandbox suite under `./temp/sim_5401/` extended to 136 tests (was 90): new unit coverage for `strip_shadowing_flags` (12 cases), `_kill_process` clears `_healthy`, `extra_args_source` lifecycle and cross-model behaviour, failed-load preserving prior extras, and the duplicate-load short-circuit at `load_model` level. New live integration cases verify shadow-strip via `pgrep` on the live llama-server cmdline, cross-model refusal, and PID stability across a duplicate-load race. All 136 pass. - `pytest studio/backend/tests --deselect test_studio_api.py`: 1079 passed, 46 skipped, identical to the pre-change count. The pre-existing `test_studio_api.py` fixture errors and the terminal-width-sensitive `test_help_output` are unaffected. - Ruff: clean on the three modified files. * Studio: tighten GGUF reload inheritance and duplicate-load guard Re-narrow llama_extra_args to None after validate_extra_args when the incoming request omitted the field, so the backend can distinguish "caller omitted, inherit prior load" from "caller explicitly cleared to []". Without this a queued duplicate /load reaches the backend as [] and fails _already_in_target_state's exact-equality check, killing the just-started llama-server. The pass-through validate call from the original "forward llama-server args from unsloth studio run / unsloth run" change is preserved as-is; only the post-pass narrowing is new. Cross-source loads now explicitly clear extras so a model switch can't accidentally inherit via the backend's "no opinion" semantics. Store the caller's hf_variant kwarg (None for local GGUF files) in _extra_args_source instead of the derived self._hf_variant (an extracted filename quant label like "Q4_K_M"). Same-source check in the route is now symmetric for HF and direct-file loads. Add gguf_path to _already_in_target_state and prefer on-disk path identity when both backend and caller have a path. This stops the duplicate-load guard from killing a healthy server on repeat local loads (where hf_variant is None on the caller side but extracted on the backend side). Split shadow-flag stripping into per-group toggles (context / cache / spec / template). The route now opts into stripping only the groups whose first-class field was actually set on the incoming request, so an inherited --chat-template-file survives an Apply that omits chat_template_override. _request_matches_loaded_settings detects shadowing extras on the inherit path and falls through to a real reload so the strip can run. Mark --spec-default, --jinja, --no-jinja as boolean inside the shadow stripper so the value-consuming heuristic no longer eats the following positional token. * Studio: trim comments around GGUF reload inheritance * Studio: cover GGUF reload inheritance and shadow-flag stripping * Studio: drop redundant issue refs from inheritance comments * Studio: drop redundant issue refs from inheritance comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: key inheritance source off resolved gguf_variant codex-connector P2 on PR #5427 cd14cae1: the inheritance gate at ``routes/inference.py:696`` compared the stored ``source[1]`` against ``request.gguf_variant``, but the HF branch loaded with ``hf_variant = config.gguf_variant`` (the *resolved* variant after ModelConfig auto-pick). When the caller omitted ``gguf_variant`` on a follow-up Apply, ``source[1] == "Q4_K_M"`` but ``(request.gguf_variant or "") == ""``, ``same_source`` returned False, and the chat-settings Apply silently dropped CLI pass-through flags for every auto-pick / local-file load. Fix both sides of the comparison to key off ``config.gguf_variant``: * The route compares ``source[1]`` to ``config.gguf_variant`` (the resolved label) rather than the request field. * The local-mode load_model call now passes ``hf_variant = config.gguf_variant`` so ``_extra_args_source`` stores the same string the route reads back. The HF branch already did this. Sandbox: added test_source_records_caller_variant_not_extracted_label to lock the storage key contract. ``pytest studio/backend/tests --deselect test_studio_api.py``: 1100 passed, identical to pre-change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: deny upstream --ui family on llama-server pass-through The validator's web-UI block named only ``--webui`` / ``--no-webui``, which is llama.cpp's pre-rename spelling. Current upstream (``tools/server/README.md``) uses ``--ui`` / ``--no-ui`` plus ``--ui-config``, ``--ui-config-file``, and ``--ui-mcp-proxy`` / ``--no-ui-mcp-proxy``. Without these in the denylist a user could ``unsloth run --ui`` and enable llama-server's built-in web UI on the port Studio's reverse proxy targets, breaking the UI surface. Keep the legacy ``--webui`` group so the validator still rejects old binaries that haven't been re-spelled. Cross-referenced against the README's full flag list; this was the only gap for the post-#5401 inheritance / shadow-strip work. Pass- through flags from every other README category (sampling, jinja, ctx, cache, threads, GPU, reasoning, grammar, chat-template-kwargs) already validate cleanly; sandbox suite exercises ~60 of them in the new ``test_08_llama_server_pass_through.py``. ``pytest studio/backend/tests --deselect test_studio_api.py``: 1100 passed, identical to pre-change. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 1357 +++++++++-------- .../core/inference/llama_server_args.py | 100 ++ studio/backend/routes/inference.py | 137 +- .../tests/test_gguf_reload_inheritance.py | 237 +++ .../backend/tests/test_llama_server_args.py | 118 ++ 5 files changed, 1348 insertions(+), 601 deletions(-) create mode 100644 studio/backend/tests/test_gguf_reload_inheritance.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7ef687035c..3682f1dbbb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -470,6 +470,17 @@ class LlamaCppBackend: # their own cache (Gemma 3n / Gemma 4: .attention.shared_kv_layers). self._shared_kv_layers: Optional[int] = None self._lock = threading.Lock() + # Wraps load_model() end-to-end so concurrent loads serialise + # and never coexist as two llama-server processes (#5401). + self._serial_load_lock = threading.Lock() + # Last extra_args / requested n_ctx, preserved across unload so + # the chat UI's /unload+/load Apply path can inherit them (#5401). + # ``_extra_args_source`` records the (model_identifier, hf_variant) + # the stored args came from so the route can refuse cross-model + # inheritance. + self._extra_args: Optional[List[str]] = None + self._extra_args_source: Optional[tuple[str, Optional[str]]] = None + self._requested_n_ctx: int = 0 self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None self._cancel_event = threading.Event() @@ -505,6 +516,25 @@ class LlamaCppBackend: def hf_variant(self) -> Optional[str]: return self._hf_variant + @property + def extra_args(self) -> Optional[List[str]]: + """Extra llama-server flags from the last load. Copy; None = never + set, [] = explicitly cleared. Used by the route for inheritance.""" + return list(self._extra_args) if self._extra_args is not None else None + + @property + def requested_n_ctx(self) -> int: + """n_ctx the last load was invoked with (not the effective cap). + 0 means Auto. Used by the route to detect Auto-vs-explicit flips.""" + return self._requested_n_ctx + + @property + def extra_args_source(self) -> Optional[tuple[str, Optional[str]]]: + """(model_identifier, hf_variant) the stored extra_args came from. + ``None`` if no extras have ever been recorded. Used by the route + to refuse cross-model inheritance (#5401).""" + return self._extra_args_source + @property def context_length(self) -> Optional[int]: """Return the effective context length the server is running at.""" @@ -1983,653 +2013,778 @@ class LlamaCppBackend: Returns True if server started and health check passed. """ - self._cancel_event.clear() - - # ── Phase 1: kill old process (under lock, fast) ────────── - with self._lock: - self._kill_process() - - binary = self._find_llama_server_binary() - if not binary: - raise RuntimeError( - "llama-server binary not found. " - "Run setup.sh to build it, install llama.cpp, " - "or set LLAMA_SERVER_PATH environment variable." - ) - - # ── Phase 2: download (NO lock held, so cancel can proceed) ── - if hf_repo: - model_path = self._download_gguf( - hf_repo = hf_repo, + # Serialise the whole load so concurrent /load calls never + # leave two llama-server processes alive (#5401 / #5161). Does + # not block /unload, /status, /load-progress. + with self._serial_load_lock: + # Duplicate /load that raced past the route-level check + # (the first one hadn't published _healthy=True yet). If the + # live server already satisfies this request, do nothing. + if self._already_in_target_state( + gguf_path = gguf_path, + model_identifier = model_identifier, hf_variant = hf_variant, - hf_token = hf_token, - ) - # Auto-download mmproj for vision models - if is_vision and not mmproj_path: - mmproj_path = self._download_mmproj( + n_ctx = n_ctx, + cache_type_kv = cache_type_kv, + speculative_type = speculative_type, + chat_template_override = chat_template_override, + extra_args = extra_args, + is_vision = is_vision, + ): + logger.info( + f"load_model: backend already in target state for " + f"'{model_identifier}', skipping reload" + ) + return True + + self._cancel_event.clear() + + # ── Phase 1: kill old process (under lock, fast) ────────── + with self._lock: + self._kill_process() + + binary = self._find_llama_server_binary() + if not binary: + raise RuntimeError( + "llama-server binary not found. " + "Run setup.sh to build it, install llama.cpp, " + "or set LLAMA_SERVER_PATH environment variable." + ) + + # ── Phase 2: download (NO lock held, so cancel can proceed) ── + if hf_repo: + model_path = self._download_gguf( hf_repo = hf_repo, + hf_variant = hf_variant, hf_token = hf_token, ) - elif gguf_path: - if not Path(gguf_path).is_file(): - raise FileNotFoundError(f"GGUF file not found: {gguf_path}") - model_path = gguf_path - else: - raise ValueError("Either gguf_path or hf_repo must be provided") + # Auto-download mmproj for vision models + if is_vision and not mmproj_path: + mmproj_path = self._download_mmproj( + hf_repo = hf_repo, + hf_token = hf_token, + ) + elif gguf_path: + if not Path(gguf_path).is_file(): + raise FileNotFoundError(f"GGUF file not found: {gguf_path}") + model_path = gguf_path + else: + raise ValueError("Either gguf_path or hf_repo must be provided") - # Set identifier early so _read_gguf_metadata can use it for DeepSeek detection - self._model_identifier = model_identifier + # Set identifier early so _read_gguf_metadata can use it for DeepSeek detection + self._model_identifier = model_identifier - # Read GGUF metadata (context_length, chat_template) -- fast, header only - self._read_gguf_metadata(model_path) + # Read GGUF metadata (context_length, chat_template) -- fast, header only + self._read_gguf_metadata(model_path) - # Check cancel after download - if self._cancel_event.is_set(): - logger.info("Load cancelled after download phase") - return False - - # ── Phase 3: start llama-server (under lock) ────────────── - with self._lock: - # Re-check cancel inside lock + # Check cancel after download if self._cancel_event.is_set(): - logger.info("Load cancelled before server start") + logger.info("Load cancelled after download phase") return False - self._port = self._find_free_port() + # ── Phase 3: start llama-server (under lock) ────────────── + with self._lock: + # Re-check cancel inside lock + if self._cancel_event.is_set(): + logger.info("Load cancelled before server start") + return False - # Select GPU(s) based on model size + estimated KV cache. - # Seed safe defaults before GPU probing so the except path - # still has valid state to publish. - effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0) - max_available_ctx = self._context_length or effective_ctx - gpus: list[tuple[int, int]] = [] - try: - model_size = self._get_gguf_size_bytes(model_path) - gpus = self._get_gpu_free_memory() + self._port = self._find_free_port() - # Resolve effective context: 0 means let llama-server use the - # model's native length. Only expand to a known native length - # if metadata is available; otherwise preserve 0 as a sentinel. - if n_ctx > 0: - effective_ctx = n_ctx - elif self._context_length is not None: - effective_ctx = self._context_length - else: - effective_ctx = 0 - original_ctx = effective_ctx - # Default UI ceiling to the model's native context length. - # GPU/VRAM-fit logic below may shrink this if hardware is limited. + # Select GPU(s) based on model size + estimated KV cache. + # Seed safe defaults before GPU probing so the except path + # still has valid state to publish. + effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0) max_available_ctx = self._context_length or effective_ctx + gpus: list[tuple[int, int]] = [] + try: + model_size = self._get_gguf_size_bytes(model_path) + gpus = self._get_gpu_free_memory() - # Auto-cap context to fit in GPU VRAM and select GPUs. - # - # Two policies depending on whether the user set n_ctx: - # - # Explicit n_ctx (user chose a context length): - # Honor it. Try the full requested context with _select_gpus - # (which uses as many GPUs as needed). Only cap if it doesn't - # fit on any GPU combination. - # - # Auto n_ctx=0 (model's native context): - # Prefer fewer GPUs with reduced context over more GPUs, - # since multi-GPU is slower and the user didn't ask for a - # specific context length. - gpu_indices, use_fit = None, True - explicit_ctx = n_ctx > 0 + # Resolve effective context: 0 means let llama-server use the + # model's native length. Only expand to a known native length + # if metadata is available; otherwise preserve 0 as a sentinel. + if n_ctx > 0: + effective_ctx = n_ctx + elif self._context_length is not None: + effective_ctx = self._context_length + else: + effective_ctx = 0 + original_ctx = effective_ctx + # Default UI ceiling to the model's native context length. + # GPU/VRAM-fit logic below may shrink this if hardware is limited. + max_available_ctx = self._context_length or effective_ctx - if gpus and self._can_estimate_kv() and effective_ctx > 0: - # Compute the largest hardware-aware cap from the model's - # native context across all usable GPU subsets (for UI - # bounds), independent of the currently requested context. - native_ctx_for_cap = self._context_length or effective_ctx - if native_ctx_for_cap > 0: - ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) - best_cap = 0 - for n_gpus in range(1, len(ranked_for_cap) + 1): - subset = ranked_for_cap[:n_gpus] - pool_mib = sum(free for _, free in subset) - capped = self._fit_context_to_vram( - native_ctx_for_cap, - pool_mib, - model_size, - cache_type_kv, - n_parallel = n_parallel, + # Auto-cap context to fit in GPU VRAM and select GPUs. + # + # Two policies depending on whether the user set n_ctx: + # + # Explicit n_ctx (user chose a context length): + # Honor it. Try the full requested context with _select_gpus + # (which uses as many GPUs as needed). Only cap if it doesn't + # fit on any GPU combination. + # + # Auto n_ctx=0 (model's native context): + # Prefer fewer GPUs with reduced context over more GPUs, + # since multi-GPU is slower and the user didn't ask for a + # specific context length. + gpu_indices, use_fit = None, True + explicit_ctx = n_ctx > 0 + + if gpus and self._can_estimate_kv() and effective_ctx > 0: + # Compute the largest hardware-aware cap from the model's + # native context across all usable GPU subsets (for UI + # bounds), independent of the currently requested context. + native_ctx_for_cap = self._context_length or effective_ctx + if native_ctx_for_cap > 0: + ranked_for_cap = sorted( + gpus, key = lambda g: g[1], reverse = True ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv, n_parallel = n_parallel + best_cap = 0 + for n_gpus in range(1, len(ranked_for_cap) + 1): + subset = ranked_for_cap[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = self._fit_context_to_vram( + native_ctx_for_cap, + pool_mib, + model_size, + cache_type_kv, + n_parallel = n_parallel, + ) + kv = self._estimate_kv_cache_bytes( + capped, cache_type_kv, n_parallel = n_parallel + ) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * 0.90: + best_cap = max(best_cap, capped) + if best_cap > 0: + max_available_ctx = best_cap + else: + # Weights exceed 90% of every GPU subset's free + # memory, so there is no fitting context. Anchor + # the UI's "safe zone" threshold at 4096 (the + # spec's default when the model cannot fit) so + # the ctx slider shows the "might be slower" + # warning as soon as the user drags above the + # fallback default instead of never. + max_available_ctx = min(4096, native_ctx_for_cap) + + if explicit_ctx: + # Honor the user's requested context verbatim. If it + # fits, pin GPUs and skip --fit; if it doesn't, ship + # -c --fit on and let llama-server flex + # -ngl (CPU layer offload). The UI is expected to + # have surfaced the "might be slower" warning before + # the user submitted a ctx above the fit ceiling. + requested_total = ( + model_size + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * 0.90: - best_cap = max(best_cap, capped) - if best_cap > 0: - max_available_ctx = best_cap + gpu_indices, use_fit = self._select_gpus( + requested_total, gpus + ) + # No silent shrink: effective_ctx stays == n_ctx. else: - # Weights exceed 90% of every GPU subset's free - # memory, so there is no fitting context. Anchor - # the UI's "safe zone" threshold at 4096 (the - # spec's default when the model cannot fit) so - # the ctx slider shows the "might be slower" - # warning as soon as the user drags above the - # fallback default instead of never. - max_available_ctx = min(4096, native_ctx_for_cap) + # Auto context: prefer fewer GPUs, cap context + # to fit. Same headroom threshold as + # _select_gpus (#5106). + ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + pin_fraction = self._GPU_PIN_VRAM_FRACTION + for n_gpus in range(1, len(ranked) + 1): + subset = ranked[:n_gpus] + pool_mib = sum(free for _, free in subset) + capped = self._fit_context_to_vram( + effective_ctx, + pool_mib, + model_size, + cache_type_kv, + n_parallel = n_parallel, + ) + kv = self._estimate_kv_cache_bytes( + capped, cache_type_kv, n_parallel = n_parallel + ) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * pin_fraction: + effective_ctx = capped + gpu_indices = sorted(idx for idx, _ in subset) + use_fit = False + break + else: + # Native ctx doesn't fit. Drop to 4096 and + # re-check before deferring to --fit on: + # a model that overflows at 131k may pin + # comfortably with a 4096 KV cache (#5106). + effective_ctx = min(4096, effective_ctx) + if effective_ctx > 0: + for n_gpus in range(1, len(ranked) + 1): + subset = ranked[:n_gpus] + pool_mib = sum(free for _, free in subset) + kv = self._estimate_kv_cache_bytes( + effective_ctx, + cache_type_kv, + n_parallel = n_parallel, + ) + total_mib = (model_size + kv) / (1024 * 1024) + if total_mib <= pool_mib * pin_fraction: + gpu_indices = sorted( + idx for idx, _ in subset + ) + use_fit = False + break - if explicit_ctx: - # Honor the user's requested context verbatim. If it - # fits, pin GPUs and skip --fit; if it doesn't, ship - # -c --fit on and let llama-server flex - # -ngl (CPU layer offload). The UI is expected to - # have surfaced the "might be slower" warning before - # the user submitted a ctx above the fit ceiling. - requested_total = model_size + self._estimate_kv_cache_bytes( + elif gpus: + # Can't estimate KV -- fall back to file-size-only check. + # Without KV estimation we cannot prove a hardware cap, so + # keep the ceiling at the native context (already the default). + logger.debug( + "Falling back to file-size-only GPU selection", + model_size_gb = round(model_size / (1024**3), 2), + ) + gpu_indices, use_fit = self._select_gpus(model_size, gpus) + if use_fit and not explicit_ctx: + # Weights don't fit on any subset. Default the UI to + # 4096 so the slider doesn't land on an unusable native + # context. --fit on will flex -ngl at runtime. + effective_ctx = ( + min(4096, effective_ctx) if effective_ctx > 0 else 4096 + ) + + if effective_ctx < original_ctx: + kv_est = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel ) - gpu_indices, use_fit = self._select_gpus(requested_total, gpus) - # No silent shrink: effective_ctx stays == n_ctx. - else: - # Auto context: prefer fewer GPUs, cap context - # to fit. Same headroom threshold as - # _select_gpus (#5106). - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) - pin_fraction = self._GPU_PIN_VRAM_FRACTION - for n_gpus in range(1, len(ranked) + 1): - subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) - capped = self._fit_context_to_vram( - effective_ctx, - pool_mib, - model_size, - cache_type_kv, - n_parallel = n_parallel, - ) - kv = self._estimate_kv_cache_bytes( - capped, cache_type_kv, n_parallel = n_parallel - ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: - effective_ctx = capped - gpu_indices = sorted(idx for idx, _ in subset) - use_fit = False - break - else: - # Native ctx doesn't fit. Drop to 4096 and - # re-check before deferring to --fit on: - # a model that overflows at 131k may pin - # comfortably with a 4096 KV cache (#5106). - effective_ctx = min(4096, effective_ctx) - if effective_ctx > 0: - for n_gpus in range(1, len(ranked) + 1): - subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) - kv = self._estimate_kv_cache_bytes( - effective_ctx, - cache_type_kv, - n_parallel = n_parallel, - ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: - gpu_indices = sorted(idx for idx, _ in subset) - use_fit = False - break - - elif gpus: - # Can't estimate KV -- fall back to file-size-only check. - # Without KV estimation we cannot prove a hardware cap, so - # keep the ceiling at the native context (already the default). - logger.debug( - "Falling back to file-size-only GPU selection", - model_size_gb = round(model_size / (1024**3), 2), - ) - gpu_indices, use_fit = self._select_gpus(model_size, gpus) - if use_fit and not explicit_ctx: - # Weights don't fit on any subset. Default the UI to - # 4096 so the slider doesn't land on an unusable native - # context. --fit on will flex -ngl at runtime. - effective_ctx = ( - min(4096, effective_ctx) if effective_ctx > 0 else 4096 + logger.info( + f"Context auto-reduced: {original_ctx} -> {effective_ctx} " + f"(model: {model_size / (1024**3):.1f} GB, " + f"est. KV cache: {kv_est / (1024**3):.1f} GB)" ) - if effective_ctx < original_ctx: - kv_est = self._estimate_kv_cache_bytes( + kv_cache_bytes = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel ) logger.info( - f"Context auto-reduced: {original_ctx} -> {effective_ctx} " - f"(model: {model_size / (1024**3):.1f} GB, " - f"est. KV cache: {kv_est / (1024**3):.1f} GB)" + f"GGUF size: {model_size / (1024**3):.1f} GB, " + f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " + f"context: {effective_ctx}, " + f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" ) + except Exception as e: + logger.warning(f"GPU selection failed ({e}), using --fit on") + gpu_indices, use_fit = None, True + effective_ctx = n_ctx # fall back to original - kv_cache_bytes = self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel - ) - logger.info( - f"GGUF size: {model_size / (1024**3):.1f} GB, " - f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " - f"context: {effective_ctx}, " - f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" - ) - except Exception as e: - logger.warning(f"GPU selection failed ({e}), using --fit on") - gpu_indices, use_fit = None, True - effective_ctx = n_ctx # fall back to original + cmd = [ + binary, + "-m", + model_path, + "--port", + str(self._port), + "-c", + str(effective_ctx) if effective_ctx > 0 else "0", + "--parallel", + str(n_parallel), + "--flash-attn", + "on", # Force flash attention for speed + # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". + "--no-context-shift", + ] - cmd = [ - binary, - "-m", - model_path, - "--port", - str(self._port), - "-c", - str(effective_ctx) if effective_ctx > 0 else "0", - "--parallel", - str(n_parallel), - "--flash-attn", - "on", # Force flash attention for speed - # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". - "--no-context-shift", - ] + if use_fit: + cmd.extend(["--fit", "on"]) + elif gpu_indices is not None: + # Model fits on selected GPU(s) -- offload all layers + cmd.extend(["-ngl", "-1"]) - if use_fit: - cmd.extend(["--fit", "on"]) - elif gpu_indices is not None: - # Model fits on selected GPU(s) -- offload all layers - cmd.extend(["-ngl", "-1"]) - - # -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we - # do not inherit llama-server's internal default, which has historically - # varied (hardware concurrency incl. hyperthreads on some builds). - cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)]) - - # Always enable Jinja chat template rendering for proper template support - cmd.extend(["--jinja"]) - - # KV cache data type - _valid_cache_types = { - "f16", - "bf16", - "q8_0", - "q4_0", - "q4_1", - "q5_0", - "q5_1", - "iq4_nl", - "f32", - } - if cache_type_kv and cache_type_kv in _valid_cache_types: + # -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we + # do not inherit llama-server's internal default, which has historically + # varied (hardware concurrency incl. hyperthreads on some builds). cmd.extend( - ["--cache-type-k", cache_type_kv, "--cache-type-v", cache_type_kv] + ["--threads", str(n_threads if n_threads is not None else -1)] ) - self._cache_type_kv = cache_type_kv - logger.info(f"KV cache type: {cache_type_kv}") - else: - self._cache_type_kv = None - # Speculative decoding (n-gram self-speculation, zero VRAM cost) - # ngram-mod: ~16 MB shared hash pool, constant memory/complexity, - # variable draft lengths. Helps most when the model repeats - # existing text (code refactoring, summarization, reasoning). - # For general chat with low repetition, overhead is ~5 ms. - # - # Benchmarks from upstream llama.cpp speculative-decoding PRs: - # Scenario | Without | With | Speedup - # gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x - # Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x - # gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x - # - # Params from llama.cpp docs (docs/speculative.md): - # --spec-ngram-size-n 24 (small n not recommended) - # --draft-min 48 --draft-max 64 (MoEs need long drafts; - # dense models can reduce these) - # ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md - # ref: https://github.com/ggml-org/llama.cpp/pull/19164 - # ref: https://github.com/ggml-org/llama.cpp/pull/18471 - # ``"default"`` -> let llama-server pick a sensible spec - # config via ``--spec-default``. Explicit type names are - # passed through with the manual draft tuning we've shipped - # historically so power users keep their overrides. - _valid_spec_types = {"ngram-simple", "ngram-mod"} - normalized_spec = ( - speculative_type.lower().strip() if speculative_type else None - ) - if normalized_spec and normalized_spec != "off" and not is_vision: - if normalized_spec == "default": - cmd.append("--spec-default") - self._speculative_type = "default" - elif normalized_spec in _valid_spec_types: - cmd.extend(["--spec-type", normalized_spec]) - if normalized_spec == "ngram-mod": - cmd.extend( - [ - "--spec-ngram-size-n", - "24", - "--draft-min", - "48", - "--draft-max", - "64", - ] - ) - self._speculative_type = normalized_spec + # Always enable Jinja chat template rendering for proper template support + cmd.extend(["--jinja"]) + + # KV cache data type + _valid_cache_types = { + "f16", + "bf16", + "q8_0", + "q4_0", + "q4_1", + "q5_0", + "q5_1", + "iq4_nl", + "f32", + } + if cache_type_kv and cache_type_kv in _valid_cache_types: + cmd.extend( + [ + "--cache-type-k", + cache_type_kv, + "--cache-type-v", + cache_type_kv, + ] + ) + self._cache_type_kv = cache_type_kv + logger.info(f"KV cache type: {cache_type_kv}") + else: + self._cache_type_kv = None + + # Speculative decoding (n-gram self-speculation, zero VRAM cost) + # ngram-mod: ~16 MB shared hash pool, constant memory/complexity, + # variable draft lengths. Helps most when the model repeats + # existing text (code refactoring, summarization, reasoning). + # For general chat with low repetition, overhead is ~5 ms. + # + # Benchmarks from upstream llama.cpp speculative-decoding PRs: + # Scenario | Without | With | Speedup + # gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x + # Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x + # gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x + # + # Params from llama.cpp docs (docs/speculative.md): + # --spec-ngram-size-n 24 (small n not recommended) + # --draft-min 48 --draft-max 64 (MoEs need long drafts; + # dense models can reduce these) + # ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md + # ref: https://github.com/ggml-org/llama.cpp/pull/19164 + # ref: https://github.com/ggml-org/llama.cpp/pull/18471 + # ``"default"`` -> let llama-server pick a sensible spec + # config via ``--spec-default``. Explicit type names are + # passed through with the manual draft tuning we've shipped + # historically so power users keep their overrides. + _valid_spec_types = {"ngram-simple", "ngram-mod"} + normalized_spec = ( + speculative_type.lower().strip() if speculative_type else None + ) + if normalized_spec and normalized_spec != "off" and not is_vision: + if normalized_spec == "default": + cmd.append("--spec-default") + self._speculative_type = "default" + elif normalized_spec in _valid_spec_types: + cmd.extend(["--spec-type", normalized_spec]) + if normalized_spec == "ngram-mod": + cmd.extend( + [ + "--spec-ngram-size-n", + "24", + "--draft-min", + "48", + "--draft-max", + "64", + ] + ) + self._speculative_type = normalized_spec + else: + self._speculative_type = None else: self._speculative_type = None - else: - self._speculative_type = None - # Apply custom chat template override if provided - self._chat_template_override = chat_template_override - if chat_template_override: - import tempfile + # Apply custom chat template override if provided + self._chat_template_override = chat_template_override + if chat_template_override: + import tempfile - flags = detect_reasoning_flags( - chat_template_override, - self._model_identifier, - log_source = "GGUF chat template override", - ) - self._supports_reasoning = flags["supports_reasoning"] - self._reasoning_style = flags["reasoning_style"] - self._reasoning_always_on = flags["reasoning_always_on"] - self._supports_preserve_thinking = flags["supports_preserve_thinking"] - self._supports_tools = flags["supports_tools"] - - self._chat_template_file = tempfile.NamedTemporaryFile( - mode = "w", - suffix = ".jinja", - delete = False, - prefix = "unsloth_chat_template_", - ) - self._chat_template_file.write(chat_template_override) - self._chat_template_file.close() - cmd.extend(["--chat-template-file", self._chat_template_file.name]) - logger.info( - f"Using custom chat template file: {self._chat_template_file.name}" - ) - - # For reasoning models, set default thinking mode. - # Qwen3.5/3.6 models below 9B (0.8B, 2B, 4B) disable thinking by default. - # Only 9B and larger enable thinking. - # Always-on templates ignore the kwarg entirely, so skip. - if self._supports_reasoning and not self._reasoning_always_on: - thinking_default = True - mid = (model_identifier or "").lower() - if "qwen3.5" in mid or "qwen3.6" in mid: - size_val = _extract_model_size_b(mid) - if size_val is not None and size_val < 9: - thinking_default = False - self._reasoning_default = thinking_default - reasoning_kw = self._reasoning_kwargs(thinking_default) - cmd.extend( - [ - "--chat-template-kwargs", - json.dumps(reasoning_kw), + flags = detect_reasoning_flags( + chat_template_override, + self._model_identifier, + log_source = "GGUF chat template override", + ) + self._supports_reasoning = flags["supports_reasoning"] + self._reasoning_style = flags["reasoning_style"] + self._reasoning_always_on = flags["reasoning_always_on"] + self._supports_preserve_thinking = flags[ + "supports_preserve_thinking" ] - ) - logger.info(f"Reasoning model: {reasoning_kw} by default") + self._supports_tools = flags["supports_tools"] - if mmproj_path: - if not Path(mmproj_path).is_file(): - logger.warning(f"mmproj file not found: {mmproj_path}") - else: - # #5347 guard for paths that bypass detect_mmproj_file. - from utils.models.model_config import ( - mmproj_matches_model_family, + self._chat_template_file = tempfile.NamedTemporaryFile( + mode = "w", + suffix = ".jinja", + delete = False, + prefix = "unsloth_chat_template_", + ) + self._chat_template_file.write(chat_template_override) + self._chat_template_file.close() + cmd.extend(["--chat-template-file", self._chat_template_file.name]) + logger.info( + f"Using custom chat template file: {self._chat_template_file.name}" ) - if not mmproj_matches_model_family(model_path, mmproj_path): - logger.warning( - f"Skipping mmproj with mismatched family: " - f"model={Path(model_path).name}, " - f"mmproj={Path(mmproj_path).name}" - ) + # For reasoning models, set default thinking mode. + # Qwen3.5/3.6 models below 9B (0.8B, 2B, 4B) disable thinking by default. + # Only 9B and larger enable thinking. + # Always-on templates ignore the kwarg entirely, so skip. + if self._supports_reasoning and not self._reasoning_always_on: + thinking_default = True + mid = (model_identifier or "").lower() + if "qwen3.5" in mid or "qwen3.6" in mid: + size_val = _extract_model_size_b(mid) + if size_val is not None and size_val < 9: + thinking_default = False + self._reasoning_default = thinking_default + reasoning_kw = self._reasoning_kwargs(thinking_default) + cmd.extend( + [ + "--chat-template-kwargs", + json.dumps(reasoning_kw), + ] + ) + logger.info(f"Reasoning model: {reasoning_kw} by default") + + if mmproj_path: + if not Path(mmproj_path).is_file(): + logger.warning(f"mmproj file not found: {mmproj_path}") else: - cmd.extend(["--mmproj", mmproj_path]) - logger.info(f"Using mmproj for vision: {mmproj_path}") - - # Option C: add --api-key for direct client access when enabled - import os as _os - import secrets as _secrets - - if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": - self._api_key = _secrets.token_urlsafe(32) - cmd.extend(["--api-key", self._api_key]) - logger.info("llama-server started with --api-key for direct streaming") - else: - self._api_key = None - - # User-supplied pass-through args go last so llama.cpp's - # last-wins flag parsing lets the user override Studio's - # auto-set tier-2 flags (e.g. --cache-type-k, --spec-type). - # The route layer has already validated this list against - # the managed-flag denylist via validate_extra_args(). - if extra_args: - cmd.extend(str(a) for a in extra_args) - logger.info( - f"Appending user extra args to llama-server: {list(extra_args)}" - ) - - _log_cmd = list(cmd) - if "--api-key" in _log_cmd: - _ki = _log_cmd.index("--api-key") + 1 - if _ki < len(_log_cmd): - _log_cmd[_ki] = "" - logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") - - # Set library paths so llama-server can find its shared libs and CUDA DLLs - import os - import sys - - env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) - - if sys.platform == "win32": - # CUDA DLLs (cudart64_X.dll, cublas64_X.dll, etc.) must - # be on PATH. Order: binary_dir, torch's pip-installed - # nvidia wheels, then a system CUDA toolkit. Pip wheels - # are the canonical source per Studio's install design - # (mirrors the Linux LD_LIBRARY_PATH block below) and - # CUDA_PATH covers users with a system toolkit. #5106. - path_dirs = [binary_dir] - path_dirs.extend(self._windows_pip_nvidia_dll_dirs(sys.prefix)) - cuda_path = os.environ.get("CUDA_PATH", "") - if cuda_path: - cuda_bin = os.path.join(cuda_path, "bin") - if os.path.isdir(cuda_bin): - path_dirs.append(cuda_bin) - # Some CUDA installs put DLLs in bin\x64 - cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64") - if os.path.isdir(cuda_bin_x64): - path_dirs.append(cuda_bin_x64) - existing_path = env.get("PATH", "") - env["PATH"] = ";".join(path_dirs) + ";" + existing_path - else: - # Linux: set LD_LIBRARY_PATH for shared libs next to the binary - # and CUDA runtime libs (libcudart, libcublas, etc.) - import platform - - lib_dirs = [binary_dir] - _arch = platform.machine() # x86_64, aarch64, etc. - - # Pip-installed nvidia CUDA runtime libs (e.g. torch's - # bundled cuda-bindings). The prebuilt llama.cpp binary - # links against libcudart.so.13 / libcublas.so.13 which - # live here, not in /usr/local/cuda. - import glob as _glob - - for _nv_pattern in [ - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cu*", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cudnn", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "nvjitlink", - "lib", - ), - ]: - for _nv_dir in _glob.glob(_nv_pattern): - if os.path.isdir(_nv_dir): - lib_dirs.append(_nv_dir) - - for cuda_lib in [ - "/usr/local/cuda/lib64", - f"/usr/local/cuda/targets/{_arch}-linux/lib", - # Fallback CUDA compat paths (e.g. binary built with - # CUDA 12 on a system where default /usr/local/cuda - # points to CUDA 13+). - "/usr/local/cuda-12/lib64", - "/usr/local/cuda-12.8/lib64", - f"/usr/local/cuda-12/targets/{_arch}-linux/lib", - f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", - ]: - if os.path.isdir(cuda_lib): - lib_dirs.append(cuda_lib) - existing_ld = env.get("LD_LIBRARY_PATH", "") - new_ld = ":".join(lib_dirs) - env["LD_LIBRARY_PATH"] = ( - f"{new_ld}:{existing_ld}" if existing_ld else new_ld - ) - - # Pin to selected GPU(s). On ROCm, llama-server (and any torch - # in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES; - # narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing - # the full HIP/ROCR set the parent inherited. - if gpu_indices is not None: - pinned = ",".join(str(i) for i in gpu_indices) - env["CUDA_VISIBLE_DEVICES"] = pinned - try: - import torch as _torch - - if getattr(_torch.version, "hip", None) is not None: - env["HIP_VISIBLE_DEVICES"] = pinned - env["ROCR_VISIBLE_DEVICES"] = pinned - except Exception as e: - logger.debug( - "Failed to set ROCm visibility env vars for child: %s", e - ) - - # Defensive kill: if a concurrent load slipped past Phase 1 - # (because its `self._process` was None at the time) and - # already stored a Popen handle here, drop that orphan - # before we overwrite the reference. See issue #5161. - self._kill_process() - - self._stdout_lines = [] - self._process = subprocess.Popen( - cmd, - stdout = subprocess.PIPE, - stderr = subprocess.STDOUT, - text = True, - env = env, - **_windows_hidden_subprocess_kwargs(), - ) - - # Start background thread to drain stdout and prevent pipe deadlock - self._stdout_thread = threading.Thread( - target = self._drain_stdout, daemon = True, name = "llama-stdout" - ) - self._stdout_thread.start() - - # Store the resolved on-disk path, not the caller's kwarg. In - # HF mode the caller passes gguf_path=None and the real path - # (``model_path``) is what llama-server is actually mmap'ing. - # Downstream consumers (load_progress, log lines, etc.) need - # the path that exists on disk. - self._gguf_path = model_path - self._hf_repo = hf_repo - # For local GGUF files, extract variant from filename if not provided - if hf_variant: - self._hf_variant = hf_variant - elif gguf_path: - try: - from utils.models.model_config import _extract_quant_label - - self._hf_variant = _extract_quant_label(gguf_path) - except Exception: - self._hf_variant = None - else: - self._hf_variant = None - self._is_vision = is_vision - self._model_identifier = model_identifier - - # Store the effective (possibly capped) context separately. - # Do NOT overwrite _context_length -- it holds the model's native - # context length from GGUF metadata and is used for display/info. - self._effective_context_length = ( - effective_ctx if effective_ctx > 0 else self._context_length - ) - self._max_context_length = ( - max_available_ctx - if max_available_ctx > 0 - else self._effective_context_length - ) - - # Wait for llama-server to become healthy - if not self._wait_for_health(timeout = 600.0): - self._kill_process() - _gguf = gguf_path or "" - _is_ollama = ( - ".studio_links" in _gguf - or os.sep + "ollama_links" + os.sep in _gguf - or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf - or (self._model_identifier or "").startswith("ollama/") - ) - # Only show the Ollama-specific message when the server - # output indicates a GGUF compatibility issue, not for - # unrelated failures like OOM or missing binaries. - if _is_ollama: - _output = "\n".join(self._stdout_lines[-50:]).lower() - _gguf_compat_hints = ( - "key not found", - "unknown model architecture", - "failed to load model", - ) - if any(h in _output for h in _gguf_compat_hints): - raise RuntimeError( - "Some Ollama models do not work with llama.cpp. " - "Try a different model, or use this model directly through Ollama instead." + # #5347 guard for paths that bypass detect_mmproj_file. + from utils.models.model_config import ( + mmproj_matches_model_family, ) - raise RuntimeError( - "llama-server failed to start. " - "Check that the GGUF file is valid and you have enough memory." + + if not mmproj_matches_model_family(model_path, mmproj_path): + logger.warning( + f"Skipping mmproj with mismatched family: " + f"model={Path(model_path).name}, " + f"mmproj={Path(mmproj_path).name}" + ) + else: + cmd.extend(["--mmproj", mmproj_path]) + logger.info(f"Using mmproj for vision: {mmproj_path}") + + # Option C: add --api-key for direct client access when enabled + import os as _os + import secrets as _secrets + + if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1": + self._api_key = _secrets.token_urlsafe(32) + cmd.extend(["--api-key", self._api_key]) + logger.info( + "llama-server started with --api-key for direct streaming" + ) + else: + self._api_key = None + + # User-supplied pass-through args go last so llama.cpp's + # last-wins flag parsing lets the user override Studio's + # auto-set tier-2 flags (e.g. --cache-type-k, --spec-type). + # The route layer has already validated this list against + # the managed-flag denylist via validate_extra_args(). + if extra_args: + cmd.extend(str(a) for a in extra_args) + logger.info( + f"Appending user extra args to llama-server: {list(extra_args)}" + ) + + _log_cmd = list(cmd) + if "--api-key" in _log_cmd: + _ki = _log_cmd.index("--api-key") + 1 + if _ki < len(_log_cmd): + _log_cmd[_ki] = "" + logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") + + # Set library paths so llama-server can find its shared libs and CUDA DLLs + import os + import sys + + env = child_env_without_native_path_secret() + binary_dir = str(Path(binary).parent) + + if sys.platform == "win32": + # CUDA DLLs (cudart64_X.dll, cublas64_X.dll, etc.) must + # be on PATH. Order: binary_dir, torch's pip-installed + # nvidia wheels, then a system CUDA toolkit. Pip wheels + # are the canonical source per Studio's install design + # (mirrors the Linux LD_LIBRARY_PATH block below) and + # CUDA_PATH covers users with a system toolkit. #5106. + path_dirs = [binary_dir] + path_dirs.extend(self._windows_pip_nvidia_dll_dirs(sys.prefix)) + cuda_path = os.environ.get("CUDA_PATH", "") + if cuda_path: + cuda_bin = os.path.join(cuda_path, "bin") + if os.path.isdir(cuda_bin): + path_dirs.append(cuda_bin) + # Some CUDA installs put DLLs in bin\x64 + cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64") + if os.path.isdir(cuda_bin_x64): + path_dirs.append(cuda_bin_x64) + existing_path = env.get("PATH", "") + env["PATH"] = ";".join(path_dirs) + ";" + existing_path + else: + # Linux: set LD_LIBRARY_PATH for shared libs next to the binary + # and CUDA runtime libs (libcudart, libcublas, etc.) + import platform + + lib_dirs = [binary_dir] + _arch = platform.machine() # x86_64, aarch64, etc. + + # Pip-installed nvidia CUDA runtime libs (e.g. torch's + # bundled cuda-bindings). The prebuilt llama.cpp binary + # links against libcudart.so.13 / libcublas.so.13 which + # live here, not in /usr/local/cuda. + import glob as _glob + + for _nv_pattern in [ + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cu*", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cudnn", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "nvjitlink", + "lib", + ), + ]: + for _nv_dir in _glob.glob(_nv_pattern): + if os.path.isdir(_nv_dir): + lib_dirs.append(_nv_dir) + + for cuda_lib in [ + "/usr/local/cuda/lib64", + f"/usr/local/cuda/targets/{_arch}-linux/lib", + # Fallback CUDA compat paths (e.g. binary built with + # CUDA 12 on a system where default /usr/local/cuda + # points to CUDA 13+). + "/usr/local/cuda-12/lib64", + "/usr/local/cuda-12.8/lib64", + f"/usr/local/cuda-12/targets/{_arch}-linux/lib", + f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", + ]: + if os.path.isdir(cuda_lib): + lib_dirs.append(cuda_lib) + existing_ld = env.get("LD_LIBRARY_PATH", "") + new_ld = ":".join(lib_dirs) + env["LD_LIBRARY_PATH"] = ( + f"{new_ld}:{existing_ld}" if existing_ld else new_ld + ) + + # Pin to selected GPU(s). On ROCm, llama-server (and any torch + # in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES; + # narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing + # the full HIP/ROCR set the parent inherited. + if gpu_indices is not None: + pinned = ",".join(str(i) for i in gpu_indices) + env["CUDA_VISIBLE_DEVICES"] = pinned + try: + import torch as _torch + + if getattr(_torch.version, "hip", None) is not None: + env["HIP_VISIBLE_DEVICES"] = pinned + env["ROCR_VISIBLE_DEVICES"] = pinned + except Exception as e: + logger.debug( + "Failed to set ROCm visibility env vars for child: %s", e + ) + + # Defensive kill: if a concurrent load slipped past Phase 1 + # (because its `self._process` was None at the time) and + # already stored a Popen handle here, drop that orphan + # before we overwrite the reference. See issue #5161. + self._kill_process() + + self._stdout_lines = [] + self._process = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = env, + **_windows_hidden_subprocess_kwargs(), ) - self._healthy = True + # Start background thread to drain stdout and prevent pipe deadlock + self._stdout_thread = threading.Thread( + target = self._drain_stdout, daemon = True, name = "llama-stdout" + ) + self._stdout_thread.start() - # Catch silent CPU fallback when GPU was intended (#5106). - self._gpu_offload_active = self._classify_gpu_offload( - gpu_indices is not None or use_fit, gpus or [] - ) - if self._gpu_offload_active is False: - logger.warning( - "llama-server appears to have loaded the model entirely " - "on CPU even though Studio detected at least one GPU. " - "This usually means the prebuilt binary's GPU backend " - "failed to load -- on Windows, cudart64_X.dll / " - "cublas64_X.dll could not be resolved. Reinstall the " - "Studio llama.cpp prebuilt or install a matching CUDA " - "toolkit (issue unslothai/unsloth#5106).", + # Store the resolved on-disk path, not the caller's kwarg. In + # HF mode the caller passes gguf_path=None and the real path + # (``model_path``) is what llama-server is actually mmap'ing. + # Downstream consumers (load_progress, log lines, etc.) need + # the path that exists on disk. + self._gguf_path = model_path + self._hf_repo = hf_repo + # For local GGUF files, extract variant from filename if not provided + if hf_variant: + self._hf_variant = hf_variant + elif gguf_path: + try: + from utils.models.model_config import _extract_quant_label + + self._hf_variant = _extract_quant_label(gguf_path) + except Exception: + self._hf_variant = None + else: + self._hf_variant = None + self._is_vision = is_vision + self._model_identifier = model_identifier + + # Store the effective (possibly capped) context separately. + # Do NOT overwrite _context_length -- it holds the model's native + # context length from GGUF metadata and is used for display/info. + self._effective_context_length = ( + effective_ctx if effective_ctx > 0 else self._context_length + ) + self._max_context_length = ( + max_available_ctx + if max_available_ctx > 0 + else self._effective_context_length ) - logger.info( - f"llama-server ready on port {self._port} " - f"for model '{model_identifier}'" - ) - return True + # Wait for llama-server to become healthy + if not self._wait_for_health(timeout = 600.0): + self._kill_process() + _gguf = gguf_path or "" + _is_ollama = ( + ".studio_links" in _gguf + or os.sep + "ollama_links" + os.sep in _gguf + or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf + or (self._model_identifier or "").startswith("ollama/") + ) + # Only show the Ollama-specific message when the server + # output indicates a GGUF compatibility issue, not for + # unrelated failures like OOM or missing binaries. + if _is_ollama: + _output = "\n".join(self._stdout_lines[-50:]).lower() + _gguf_compat_hints = ( + "key not found", + "unknown model architecture", + "failed to load model", + ) + if any(h in _output for h in _gguf_compat_hints): + raise RuntimeError( + "Some Ollama models do not work with llama.cpp. " + "Try a different model, or use this model directly through Ollama instead." + ) + raise RuntimeError( + "llama-server failed to start. " + "Check that the GGUF file is valid and you have enough memory." + ) + + self._healthy = True + + # Commit caller intent only after _healthy=True so a + # failed startup can't poison the next inheritance check. + # None keeps prior, [] clears, list sets. Source records + # the caller's hf_variant (None for local files) so the + # route's same_source check stays symmetric. + if extra_args is not None: + self._extra_args = list(extra_args) + self._extra_args_source = (model_identifier, hf_variant) + self._requested_n_ctx = int(n_ctx) + + # Catch silent CPU fallback when GPU was intended (#5106). + self._gpu_offload_active = self._classify_gpu_offload( + gpu_indices is not None or use_fit, gpus or [] + ) + if self._gpu_offload_active is False: + logger.warning( + "llama-server appears to have loaded the model entirely " + "on CPU even though Studio detected at least one GPU. " + "This usually means the prebuilt binary's GPU backend " + "failed to load -- on Windows, cudart64_X.dll / " + "cublas64_X.dll could not be resolved. Reinstall the " + "Studio llama.cpp prebuilt or install a matching CUDA " + "toolkit (issue unslothai/unsloth#5106).", + ) + + logger.info( + f"llama-server ready on port {self._port} " + f"for model '{model_identifier}'" + ) + return True + + def _already_in_target_state( + self, + *, + model_identifier: str, + hf_variant: Optional[str], + n_ctx: int, + cache_type_kv: Optional[str], + speculative_type: Optional[str], + chat_template_override: Optional[str], + extra_args: Optional[List[str]], + is_vision: bool, + gguf_path: Optional[str] = None, + ) -> bool: + """True iff the live server already satisfies these load kwargs. + + Mirrors ``routes/inference.py:_request_matches_loaded_settings`` + but compares raw kwargs so ``load_model`` can short-circuit a + duplicate /load that raced past the route-level check (#5401). + """ + if not self.is_loaded: + return False + if (self._model_identifier or "").lower() != (model_identifier or "").lower(): + return False + # Direct-file loads pass hf_variant=None while the backend + # stores an extracted filename label; compare paths instead + # to keep the guard symmetric. + if gguf_path is not None and self._gguf_path: + try: + if Path(self._gguf_path).resolve() != Path(gguf_path).resolve(): + return False + except OSError: + return False + elif (self._hf_variant or "").lower() != (hf_variant or "").lower(): + return False + if self._requested_n_ctx != int(n_ctx): + return False + + def _norm(value): + if value is None: + return None + if isinstance(value, str): + stripped = value.strip().lower() + return stripped or None + return value + + if _norm(self._cache_type_kv) != _norm(cache_type_kv): + return False + + # Vision GGUFs silently drop speculative decoding in + # load_model (the spec gate is "not is_vision"); treat the + # request's value as "off" so a vision load with + # speculative_type="default" still matches. + if self._is_vision or is_vision: + req_spec = "off" + else: + req_spec = _norm(speculative_type) or "off" + backend_spec = _norm(self._speculative_type) or "off" + if req_spec != backend_spec: + return False + + if (self._chat_template_override or None) != (chat_template_override or None): + return False + + # extra_args=None means "no opinion" (inherit semantics handled + # at the route layer); only an explicit list forces equality. + if extra_args is not None: + current = list(self._extra_args) if self._extra_args is not None else [] + if list(extra_args) != current: + return False + return True def _classify_gpu_offload( self, @@ -2737,6 +2892,10 @@ class LlamaCppBackend: logger.warning(f"Error killing llama-server process: {e}") finally: self._process = None + # Clear healthy so a /load arriving during the replacement + # server's warm-up window cannot short-circuit against the + # previous server's health (#5401). + self._healthy = False if self._stdout_thread is not None: self._stdout_thread.join(timeout = 2) self._stdout_thread = None diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 44c7d542c7..0f6927fc5a 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -69,7 +69,15 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( # Single-model server -- Studio runs one model per llama-server # process and serves its own UI. Enabling multi-model loading or # llama-server's built-in web UI changes the surface clients see. + # ``--webui``/``--no-webui`` are the legacy spelling; current + # upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions. + # Keep both so the denylist matches old and new llama-server + # binaries (Studio's prebuilt vs system-llama.cpp). frozenset({"--webui", "--no-webui"}), + frozenset({"--ui", "--no-ui"}), + frozenset({"--ui-config"}), + frozenset({"--ui-config-file"}), + frozenset({"--ui-mcp-proxy", "--no-ui-mcp-proxy"}), frozenset({"--models-dir"}), frozenset({"--models-preset"}), frozenset({"--models-max"}), @@ -118,3 +126,95 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]: def is_managed_flag(flag: str) -> bool: """True if ``flag`` is a Studio-managed llama-server flag.""" return flag in _DENYLIST + + +# Pass-through flags that shadow first-class ``LoadRequest`` fields +# (max_seq_length, cache_type_kv, speculative_type, +# chat_template_override). Stripped from inherited extras so they +# can't last-wins-override an Apply that re-sets the same first-class +# field. +_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) +_CACHE_FLAGS: frozenset[str] = frozenset( + {"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"} +) +_SPEC_FLAGS: frozenset[str] = frozenset( + { + "--spec-default", + "--spec-type", + "--spec-ngram-size-n", + "--spec-ngram-size", + "--draft-min", + "--draft-max", + } +) +_TEMPLATE_FLAGS: frozenset[str] = frozenset( + { + "--chat-template", + "--chat-template-file", + "--chat-template-kwargs", + "--jinja", + "--no-jinja", + } +) + +_SHADOWING_FLAGS: frozenset[str] = ( + _CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS +) + +# Boolean flags inside _SHADOWING_FLAGS that take no value. The +# value-consuming heuristic in strip_shadowing_flags must skip just the +# flag for these, never the following token. +_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset( + {"--spec-default", "--jinja", "--no-jinja"} +) + + +def strip_shadowing_flags( + args: Iterable[str], + *, + strip_context: bool = True, + strip_cache: bool = True, + strip_spec: bool = True, + strip_template: bool = True, +) -> list[str]: + """Strip flags that shadow first-class Studio settings. + + Used when the route inherits a previous load's ``llama_extra_args`` + so that an inherited ``-c 4096`` cannot override the current + request's ``max_seq_length`` (and equivalents for cache / + speculative / chat template). Each ``strip_*`` flag controls one + group; the route only strips groups whose corresponding first-class + field was actually supplied by the caller, so an inherited + ``--chat-template-file`` survives an Apply that omits both + ``llama_extra_args`` and ``chat_template_override``. + """ + shadowing: set[str] = set() + if strip_context: + shadowing |= _CONTEXT_FLAGS + if strip_cache: + shadowing |= _CACHE_FLAGS + if strip_spec: + shadowing |= _SPEC_FLAGS + if strip_template: + shadowing |= _TEMPLATE_FLAGS + + tokens = [str(a) for a in (args or [])] + out: list[str] = [] + i, n = 0, len(tokens) + while i < n: + tok = tokens[i] + flag = _flag_name(tok) + if flag is None or flag not in shadowing: + out.append(tok) + i += 1 + continue + # Drop this token. Boolean shadowing flags never carry a value; + # other shadowing flags consume the next token when it isn't a + # flag and the value isn't already packed as ``--key=value``. + if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok: + i += 1 + elif i + 1 < n and _flag_name(tokens[i + 1]) is None: + i += 2 + else: + i += 1 + return out diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 76bbb59c94..60078ecc9b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -119,7 +119,10 @@ try: _DEFAULT_T_MAX_PREDICT_MS, detect_reasoning_flags, ) - from core.inference.llama_server_args import validate_extra_args + from core.inference.llama_server_args import ( + strip_shadowing_flags, + validate_extra_args, + ) from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import load_model_defaults @@ -141,7 +144,10 @@ except ImportError: _DEFAULT_T_MAX_PREDICT_MS, detect_reasoning_flags, ) - from core.inference.llama_server_args import validate_extra_args + from core.inference.llama_server_args import ( + strip_shadowing_flags, + validate_extra_args, + ) from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import load_model_defaults @@ -406,6 +412,57 @@ def _validate_native_mmproj_companion( ) from exc +def _normalise_settings_str(value: Optional[str]) -> Optional[str]: + """Lowercase + strip a settings string, mapping blank/None to None.""" + if value is None: + return None + if isinstance(value, str): + stripped = value.strip().lower() + return stripped or None + return value + + +def _request_matches_loaded_settings( + request: LoadRequest, llama_backend: LlamaCppBackend +) -> bool: + """True iff every runtime setting on the request matches the loaded + server. Caller has already checked model+variant+is_loaded. See #5401.""" + # Compare requested n_ctx (not effective) so VRAM-cap doesn't mask + # an Auto-vs-explicit slider flip. + if request.max_seq_length != llama_backend.requested_n_ctx: + return False + if _normalise_settings_str(request.cache_type_kv) != _normalise_settings_str( + llama_backend.cache_type_kv + ): + return False + # Vision loads silently drop speculative decoding (llama_cpp.py gates + # spec on ``not is_vision``), so treat the request as ``off`` against + # the backend's ``None`` to avoid forcing a redundant reload. + if llama_backend.is_vision: + req_spec = "off" + else: + req_spec = _normalise_settings_str(request.speculative_type) or "off" + backend_spec = _normalise_settings_str(llama_backend.speculative_type) or "off" + if req_spec != backend_spec: + return False + if (request.chat_template_override or None) != ( + llama_backend.chat_template_override or None + ): + return False + # llama_extra_args=None means "inherit"; only an explicit list that + # differs forces a reload. On the inherit path, refuse to match if + # stored extras contain any shadow flag, so the reload path can + # strip them instead of leaving a stale override in effect. + backend_extra = list(llama_backend.extra_args) if llama_backend.extra_args else [] + if request.llama_extra_args is None: + if backend_extra and strip_shadowing_flags(backend_extra) != backend_extra: + return False + else: + if list(request.llama_extra_args) != backend_extra: + return False + return True + + def _resolve_model_identifier_for_request( request: LoadRequest | ValidateModelRequest, *, @@ -461,6 +518,11 @@ async def load_model( extra_llama_args = validate_extra_args(request.llama_extra_args) except ValueError as exc: raise HTTPException(status_code = 400, detail = str(exc)) + # Re-narrow []-from-None back to None so the inheritance path + # below can tell "caller omitted" from "caller explicit []". + extra_llama_args: Optional[list[str]] = ( + None if request.llama_extra_args is None else extra_llama_args + ) model_identifier, model_log_label, native_grant_backed = ( _resolve_model_identifier_for_request(request, operation = "load-model") @@ -479,6 +541,9 @@ async def load_model( and llama_backend.hf_variant.lower() == request.gguf_variant.lower() and llama_backend.model_identifier and llama_backend.model_identifier.lower() == model_identifier.lower() + # Also require runtime settings to match so Apply changes + # aren't silently dropped (#5401). + and _request_matches_loaded_settings(request, llama_backend) ): logger.info( f"Model already loaded (GGUF): {model_log_label} variant={request.gguf_variant}, skipping reload" @@ -613,6 +678,70 @@ async def load_model( ) unsloth_backend.unload_model(unsloth_backend.active_model_name) + # Inherit llama_extra_args from the previous load when the + # request omits the field (the chat-settings Apply path + # does not round-trip them; explicit [] still clears). + # Inheritance is gated on (model_identifier, hf_variant) + # to refuse cross-model pickup, and shadowing flags are + # stripped so an inherited override can't win the last-wins + # CLI parse against a freshly-supplied first-class field. + if request.llama_extra_args is None and llama_backend.extra_args: + source = llama_backend.extra_args_source + # Compare against the resolved variant, not the request + # field: callers commonly omit gguf_variant for local + # ``.gguf`` paths and HF auto-pick flows. ``config.gguf_ + # variant`` is the variant load_model was actually + # invoked with (see the HF / local branches below), so + # both sides of the comparison key off the same string. + resolved_variant = config.gguf_variant + same_source = bool( + source + and source[0] + and source[0].lower() == model_identifier.lower() + and (source[1] or "").lower() == (resolved_variant or "").lower() + ) + if not same_source: + logger.info( + "Not inheriting llama_extra_args: stored args came " + "from %s, loading %s", + source, + (model_identifier, resolved_variant), + ) + # Cross-model: clear explicitly so the backend + # doesn't inherit via "no opinion" semantics. + extra_llama_args = [] + else: + # Strip only the groups whose first-class field + # was actually set by the caller, so an inherited + # --chat-template-file survives an Apply that omits + # chat_template_override. + fields_set = getattr(request, "model_fields_set", set()) + stripped = strip_shadowing_flags( + llama_backend.extra_args, + strip_context = "max_seq_length" in fields_set, + strip_cache = "cache_type_kv" in fields_set, + strip_spec = "speculative_type" in fields_set, + strip_template = "chat_template_override" in fields_set, + ) + try: + extra_llama_args = validate_extra_args(stripped) + except ValueError: + # Should not happen on already-validated args; degrade + # to no-extras rather than 400 if managed flags changed. + logger.warning( + "Stored llama_extra_args failed revalidation; " + "loading without them: %s", + stripped, + ) + extra_llama_args = [] + else: + if extra_llama_args: + logger.info( + "Inheriting llama_extra_args from previous " + "load (same model, shadow-stripped): %s", + extra_llama_args, + ) + # Route to HF mode or local mode based on config # Run in a thread so the event loop stays free for progress # polling and other requests during the (potentially long) @@ -645,6 +774,10 @@ async def load_model( llama_backend.load_model, gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, + # Pass the resolved variant so _extra_args_source + # is keyed off the same string the inheritance + # check at the top of /load uses (#5401 followup). + hf_variant = config.gguf_variant, model_identifier = config.identifier, is_vision = config.is_vision, n_ctx = request.max_seq_length, diff --git a/studio/backend/tests/test_gguf_reload_inheritance.py b/studio/backend/tests/test_gguf_reload_inheritance.py new file mode 100644 index 0000000000..4b0b450cb0 --- /dev/null +++ b/studio/backend/tests/test_gguf_reload_inheritance.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Backend contract for the GGUF reload duplicate-load guard. + +``LlamaCppBackend._already_in_target_state`` is the in-process +short-circuit that prevents a serialised duplicate /load from killing +the just-spawned llama-server. These tests pin the local-file +identity, the HF-mode hf_variant fallback, and the ``extra_args`` +None-vs-[] inherit semantics so the guard cannot silently regress. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +_httpx_stub = _types.ModuleType("httpx") +for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) +_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) +_httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend + + +class _FakeProcess: + """Stand-in for subprocess.Popen so atexit cleanup doesn't crash.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def _loaded_backend(**overrides): + backend = LlamaCppBackend() + backend._process = _FakeProcess() # is_loaded only checks "is not None" + backend._healthy = True + backend._model_identifier = "owner/repo" + backend._hf_variant = "Q4_K_M" + backend._requested_n_ctx = 8192 + backend._cache_type_kv = None + backend._speculative_type = None + backend._chat_template_override = None + backend._is_vision = False + backend._extra_args = None + backend._extra_args_source = None + backend._gguf_path = None + for key, value in overrides.items(): + setattr(backend, key, value) + return backend + + +# ── Local-file identity via gguf_path ──────────────────────────────── + + +def test_already_in_target_state_uses_gguf_path_when_present(tmp_path): + gguf_file = tmp_path / "model.Q4_K_M.gguf" + gguf_file.write_bytes(b"") + backend = _loaded_backend( + _hf_variant = "Q4_K_M", + _gguf_path = str(gguf_file), + ) + assert ( + backend._already_in_target_state( + gguf_path = str(gguf_file), + model_identifier = "owner/repo", + hf_variant = None, + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_rejects_different_gguf_path(tmp_path): + a = tmp_path / "a.gguf" + a.write_bytes(b"") + b = tmp_path / "b.gguf" + b.write_bytes(b"") + backend = _loaded_backend(_gguf_path = str(a)) + assert ( + backend._already_in_target_state( + gguf_path = str(b), + model_identifier = "owner/repo", + hf_variant = None, + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is False + ) + + +# ── HF mode falls back to hf_variant comparison ────────────────────── + + +def test_already_in_target_state_falls_back_to_hf_variant_for_hf_loads(): + backend = _loaded_backend(_hf_variant = "Q4_K_M", _gguf_path = None) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q8_0", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is False + ) + + +def test_already_in_target_state_hf_same_variant_matches(): + backend = _loaded_backend(_hf_variant = "Q4_K_M", _gguf_path = None) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +# ── extra_args: None inherits, [] forces reload, list enforces ─────── + + +def test_already_in_target_state_none_extras_inherits_stored(): + backend = _loaded_backend(_extra_args = ["--top-k", "20"]) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = None, + is_vision = False, + ) + is True + ) + + +def test_already_in_target_state_empty_extras_forces_reload_when_stored(): + backend = _loaded_backend(_extra_args = ["--top-k", "20"]) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = [], + is_vision = False, + ) + is False + ) + + +def test_already_in_target_state_explicit_extras_match(): + backend = _loaded_backend(_extra_args = ["--top-k", "20"]) + assert ( + backend._already_in_target_state( + gguf_path = None, + model_identifier = "owner/repo", + hf_variant = "Q4_K_M", + n_ctx = 8192, + cache_type_kv = None, + speculative_type = None, + chat_template_override = None, + extra_args = ["--top-k", "20"], + is_vision = False, + ) + is True + ) + + +def test_extra_args_source_default_is_none(): + backend = LlamaCppBackend() + assert backend.extra_args_source is None diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py index 351fbd014d..3013acfdb8 100644 --- a/studio/backend/tests/test_llama_server_args.py +++ b/studio/backend/tests/test_llama_server_args.py @@ -15,6 +15,7 @@ import pytest from core.inference.llama_server_args import ( is_managed_flag, + strip_shadowing_flags, validate_extra_args, ) @@ -187,3 +188,120 @@ def test_is_managed_flag_false_for_pass_through(): assert is_managed_flag("--flash-attn") is False assert is_managed_flag("-ngl") is False assert is_managed_flag("--threads") is False + + +# ── strip_shadowing_flags ───────────────────────────────────────────── + + +def test_strip_shadowing_flags_drops_context_when_requested(): + out = strip_shadowing_flags( + ["-c", "4096", "--top-k", "20"], + strip_context = True, + strip_cache = False, + strip_spec = False, + strip_template = False, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_context_when_not_requested(): + out = strip_shadowing_flags( + ["-c", "4096", "--top-k", "20"], + strip_context = False, + strip_cache = False, + strip_spec = False, + strip_template = False, + ) + assert out == ["-c", "4096", "--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_chat_template_when_template_disabled(): + # Caller did not supply chat_template_override; the inherited + # --chat-template-file must survive the strip. + out = strip_shadowing_flags( + ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"], + strip_context = True, + strip_cache = True, + strip_spec = True, + strip_template = False, + ) + assert out == ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"] + + +def test_strip_shadowing_flags_drops_template_when_requested(): + out = strip_shadowing_flags( + ["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"], + strip_template = True, + ) + assert out == ["--top-k", "20"] + + +def test_strip_shadowing_flags_keeps_cache_when_cache_disabled(): + out = strip_shadowing_flags( + ["--cache-type-k", "q8_0", "--cache-type-v", "q8_0", "--top-k", "20"], + strip_cache = False, + ) + assert out == [ + "--cache-type-k", + "q8_0", + "--cache-type-v", + "q8_0", + "--top-k", + "20", + ] + + +def test_strip_shadowing_flags_keeps_spec_when_spec_disabled(): + out = strip_shadowing_flags( + ["--spec-type", "ngram-mod", "--draft-min", "48", "--top-k", "20"], + strip_spec = False, + ) + assert out == [ + "--spec-type", + "ngram-mod", + "--draft-min", + "48", + "--top-k", + "20", + ] + + +def test_strip_shadowing_flags_boolean_does_not_consume_next_token(): + # --spec-default is a boolean shadowing flag; the value-skipping + # heuristic must skip just the flag, not the following positional. + out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True) + assert out == ["ngram-mod"] + + +def test_strip_shadowing_flags_jinja_boolean_preserves_positional(): + out = strip_shadowing_flags(["--jinja", "trailing-positional"], strip_template = True) + assert out == ["trailing-positional"] + + +def test_strip_shadowing_flags_no_jinja_boolean_preserves_positional(): + out = strip_shadowing_flags( + ["--no-jinja", "trailing-positional"], strip_template = True + ) + assert out == ["trailing-positional"] + + +def test_strip_shadowing_flags_equals_form_drops_only_the_flag(): + out = strip_shadowing_flags(["--ctx-size=4096", "--seed", "-1"], strip_context = True) + assert out == ["--seed", "-1"] + + +def test_strip_shadowing_flags_handles_none_input(): + assert strip_shadowing_flags(None) == [] + + +def test_strip_shadowing_flags_handles_empty_input(): + assert strip_shadowing_flags([]) == [] + + +def test_strip_shadowing_flags_defaults_strip_everything(): + # The route's already-loaded comparator calls strip_shadowing_flags + # with no kwargs to detect ANY shadowing flag in stored extras. + out = strip_shadowing_flags( + ["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"] + ) + assert out == [] From 0542dc07259c693fb98c32805416bf564e437aec Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 17 May 2026 04:20:46 -0700 Subject: [PATCH 5/8] Studio: IME / multilingual composer regression test + RTL dir="auto" (#5485) Adds dir="auto" to the main, edit, and compare chat composers so RTL scripts (Arabic, Hebrew, Persian, Urdu) flow right to left without forcing the rest of the UI into RTL. Wires a model-free Playwright smoke (multilingual paste round trip across 31 scripts + a stuck-IME composition repro for issue #5318 / PR #5327) into the Studio UI CI job as a third Studio boot, plus a pure-Python static-guard test that locks down dir="auto" on all three composers and the minimal env contract for the smoke. --- .github/workflows/studio-ui-smoke.yml | 55 ++- .../src/components/assistant-ui/thread.tsx | 5 + .../src/features/chat/shared-composer.tsx | 3 + tests/studio/playwright_chat_ime_i18n.py | 457 ++++++++++++++++++ .../test_composer_rtl_bidi_attribute.py | 73 +++ 5 files changed, 588 insertions(+), 5 deletions(-) create mode 100644 tests/studio/playwright_chat_ime_i18n.py create mode 100644 tests/studio/test_composer_rtl_bidi_attribute.py diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index 79476a62ea..455fe4b7e1 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -229,12 +229,55 @@ jobs: kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true sleep 2 + # IME + multilingual paste regression (issue #5318 / PR #5327). + # Third Studio on its own port so a hang here cannot poison the + # earlier UI tests. No GGUF -- the bug surface is the composer. + - name: Reset auth + boot Studio for IME / i18n tests (port 18896) + run: | + unsloth studio reset-password + mkdir -p logs + UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18896 \ + > logs/studio_ime.log 2>&1 & + echo "STUDIO_IME_PID=$!" >> "$GITHUB_ENV" + + - name: Wait for /api/health on 18896 + run: | + for i in $(seq 1 180); do + if curl -fs "http://127.0.0.1:18896/api/health" > /tmp/health3.json; then + jq -e '.status == "healthy"' /tmp/health3.json && break + fi + sleep 1 + done + jq -e '.status == "healthy"' /tmp/health3.json + + - name: Pass bootstrap pw for IME / i18n test + # IME smoke does the change-password against the bootstrap that + # Studio's frontend injects into the page, so it only needs the + # NEW password. + run: | + NEW="CIIme-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')" + echo "::add-mask::$NEW" + echo "STUDIO_IME_NEW_PW=$NEW" >> "$GITHUB_ENV" + + - name: Drive IME + multilingual paste regression with Playwright + env: + BASE_URL: http://127.0.0.1:18896 + STUDIO_NEW_PW: ${{ env.STUDIO_IME_NEW_PW }} + PW_ART_DIR: logs/playwright_ime + STUDIO_UI_STRICT: '1' + run: | + mkdir -p logs/playwright_ime + python tests/studio/playwright_chat_ime_i18n.py + + - name: Stop third Studio + if: always() + run: | + kill "${STUDIO_IME_PID}" 2>/dev/null || true + sleep 2 + - name: Upload Playwright artifacts - # Always upload (not just failure) so a green run's screenshots - # are reviewable in the Actions UI -- catches "passed but the - # UI is silently broken" regressions that would be invisible - # otherwise. Both Studio's logs (chat + extra) and BOTH - # Playwright artifact dirs are bundled. + # Always upload so a green run's screenshots stay reviewable -- + # catches "passed but the UI is silently broken" regressions. if: always() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: @@ -242,7 +285,9 @@ jobs: path: | logs/studio.log logs/studio_extra.log + logs/studio_ime.log logs/install.log logs/playwright logs/playwright_extra + logs/playwright_ime retention-days: 7 diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index fb63748bf1..c8fea9c704 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -328,6 +328,9 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => { autoFocus={!disabled} disabled={disabled} aria-label="Message input" + // dir="auto": browser picks LTR/RTL from the first strong char; + // no effect on Latin / CJK / Devanagari. + dir="auto" {...inputProps} /> {
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index c4ffa98467..cd31d37cd7 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -690,6 +690,9 @@ export function SharedComposer({ placeholder="Send to both models..." className="composer-input" rows={1} + // dir="auto" auto-detects RTL (Arabic / Hebrew / Persian / Urdu) + // from the first strong character; no effect on LTR scripts. + dir="auto" />
diff --git a/tests/studio/playwright_chat_ime_i18n.py b/tests/studio/playwright_chat_ime_i18n.py new file mode 100644 index 0000000000..c882d88cbd --- /dev/null +++ b/tests/studio/playwright_chat_ime_i18n.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Studio chat composer IME + multilingual regression smoke. + +Covers two surfaces: + A. Stuck IME composition (issue #5318 / PR #5327): duplicate + compositionstart with no compositionend left isComposing=true, + dropping all subsequent keystrokes including ASCII. + B. Multilingual paste round-trip across 31 scripts -- guards the + controlled-textarea / React state plumbing against Unicode mangling. + +Model-free; the bug surface is the composer, not inference. + +Env contract matches playwright_chat_ui.py: + BASE_URL, STUDIO_NEW_PW, PW_ART_DIR, STUDIO_UI_STRICT. +""" + +import os +import sys +from pathlib import Path + +from playwright.sync_api import expect, sync_playwright + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from _playwright_robust import ( # noqa: E402 + chromium_launch_args, + click_and_wait_for_response, + install_view_transition_killer, + install_wall_clock_watchdog, + is_benign_console_error, + is_benign_page_error, + recover_or_replace_page, + wait_for_health, +) + +BASE = os.environ["BASE_URL"] +NEW = os.environ["STUDIO_NEW_PW"] +ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright_ime") +ART = Path(ART_DIR) +ART.mkdir(parents = True, exist_ok = True) +STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1" + +# Wall-clock cap. Realistic run is 30-60s; 5 min leaves cold-launch headroom. +WALL_TIMEOUT_S = float(os.environ.get("STUDIO_IME_WALL_TIMEOUT_S", "300")) + + +# One short greeting + arithmetic per script (ordered by speaker count) -- +# each entry catches a distinct class of Unicode regression. +I18N_SAMPLES = [ + ("en", "English", "Hello, 1+1=2"), + ("zh-CN", "Chinese (Simplified)", "你好，1+1=2"), + ("es", "Spanish", "Hola, 1+1=2"), + ("hi", "Hindi (Devanagari)", "नमस्ते, 1+1=2"), + ("ar", "Arabic (RTL)", "مرحبا، ١+١=٢"), + ("bn", "Bengali", "নমস্কার, ১+১=২"), + ("pt", "Portuguese", "Olá, 1+1=2"), + ("ru", "Russian (Cyrillic)", "Привет, 1+1=2"), + ("ja", "Japanese", "こんにちは、1+1=2"), + ("pa", "Punjabi (Gurmukhi)", "ਸਤ ਸ੍ਰੀ ਅਕਾਲ, 1+1=2"), + ("de", "German", "Hallo, 1+1=2"), + ("jv", "Javanese", "Halo, 1+1=2"), + ("ko", "Korean (Hangul)", "안녕하세요, 1+1=2"), + ("fr", "French", "Bonjour, 1+1=2"), + ("tr", "Turkish", "Merhaba, 1+1=2"), + ("vi", "Vietnamese (diacritics)", "Xin chào, 1+1=2"), + ("ur", "Urdu (Arabic-Naskh)", "ہیلو، 1+1=2"), + ("ta", "Tamil", "வணக்கம், 1+1=2"), + ("te", "Telugu", "నమస్తే, 1+1=2"), + ("mr", "Marathi (Devanagari)", "नमस्कार, 1+1=2"), + ("it", "Italian", "Ciao, 1+1=2"), + ("th", "Thai", "สวัสดี, ๑+๑=๒"), + ("pl", "Polish", "Cześć, 1+1=2"), + ("uk", "Ukrainian (Cyrillic)", "Привіт, 1+1=2"), + ("fa", "Persian (RTL)", "سلام، ۱+۱=۲"), + ("nl", "Dutch", "Hallo, 1+1=2"), + ("he", "Hebrew (RTL)", "שלום, 1+1=2"), + ("el", "Greek", "Γειά, 1+1=2"), + ("id", "Indonesian", "Halo, 1+1=2"), + ("sw", "Swahili", "Habari, 1+1=2"), + ("emoji", "Emoji + ZWJ + flag", "👋 🇺🇳 👨‍👩‍👧‍👦 1+1=2"), +] + + +_n = [0] + + +def step(s): + print(f"[ime] STEP {s}", flush = True) + + +def info(s): + print(f"[ime] {s}", flush = True) + + +def fail(m): + raise AssertionError(f"[ime] FAIL: {m}") + + +def soft_fail(m): + """Hard fail in STRICT mode, info-warn otherwise. Mirrors playwright_chat_ui.py.""" + if STRICT: + fail(m) + info(f"WARN (strict-off): {m}") + + +with sync_playwright() as p: + _watchdog = install_wall_clock_watchdog( + WALL_TIMEOUT_S, + label = "ime", + info = info, + ) + wait_for_health(BASE, timeout = 30.0, info = info) + browser = p.chromium.launch( + headless = True, + args = chromium_launch_args(), + ) + ctx = browser.new_context( + viewport = {"width": 1280, "height": 900}, + reduced_motion = "reduce", + ) + install_view_transition_killer(ctx) + page = ctx.new_page() + page.set_default_timeout(60_000) + + page_errors: list[str] = [] + console_errors: list[str] = [] + + def _on_console(m): + if m.type != "error": + return + try: + console_errors.append(m.text) + except Exception: + return + + def _attach_listeners(target): + target.on("pageerror", lambda e: page_errors.append(str(e))) + target.on("console", _on_console) + + _attach_listeners(page) + + def shoot(name): + _n[0] += 1 + try: + page.screenshot( + path = str(ART / f"{_n[0]:02d}-{name}.png"), + full_page = True, + timeout = 90_000, + animations = "disabled", + ) + except Exception as _shoot_err: + info(f"WARN: screenshot {name} failed: {_shoot_err}") + + # 1. Bootstrap auth via /change-password (mirrors playwright_chat_ui.py + # retry-on-rerender to absorb React form-detach races). + step("change-password through UI (Setup your account)") + form_err: Exception | None = None + for _form_attempt in range(3): + try: + page.goto( + f"{BASE}/change-password", + wait_until = "domcontentloaded", + timeout = 60_000, + ) + try: + page.wait_for_load_state("networkidle", timeout = 30_000) + except Exception: + pass + pw_field = page.locator("#new-password") + pw_field.wait_for(state = "visible", timeout = 60_000) + pw_field.fill(NEW, timeout = 60_000) + page.fill("#confirm-password", NEW, timeout = 60_000) + shoot("01-change-password-filled") + status, _ = click_and_wait_for_response( + page, + url_substr = "/api/auth/change-password", + method = "POST", + do_click = lambda: page.locator('button[type="submit"]').click(), + timeout_ms = 30_000, + info = lambda m: print(f"[ime] {m}", flush = True), + ) + if status is not None and status >= 400: + raise AssertionError(f"change-password POST returned {status}") + form_err = None + break + except Exception as e: + form_err = e + info( + f"change-password attempt {_form_attempt + 1} failed: " + f"{type(e).__name__}: {str(e)[:200]}" + ) + if _form_attempt < 2: + page = recover_or_replace_page( + page, + ctx, + default_timeout_ms = 60_000, + info = lambda m: print(f"[ime] recovery: {m}", flush = True), + ) + _attach_listeners(page) + if form_err is not None: + raise form_err + + # 2. Wait for composer mount. No GGUF: the bug surface is React state, not inference. + step("wait for composer to mount") + try: + page.wait_for_load_state("networkidle", timeout = 30_000) + except Exception: + pass + composer = page.locator('textarea[aria-label="Message input"]') + _mount_err: Exception | None = None + for _mount_attempt in range(2): + try: + composer.wait_for(state = "visible", timeout = 60_000) + _mount_err = None + break + except Exception as e: + _mount_err = e + info( + f"composer.wait_for attempt {_mount_attempt + 1} failed: " + f"{type(e).__name__}: {str(e)[:200]}" + ) + try: + shoot(f"02-composer-wait-attempt-{_mount_attempt + 1}-fail") + except Exception: + pass + if _mount_attempt == 0: + page = recover_or_replace_page( + page, + ctx, + default_timeout_ms = 60_000, + info = lambda m: print(f"[ime] recovery: {m}", flush = True), + ) + _attach_listeners(page) + composer = page.locator('textarea[aria-label="Message input"]') + if _mount_err is not None: + raise _mount_err + composer.click() + shoot("02-composer-focused") + + # Main composer must carry dir="auto" so RTL flows right-to-left. + dir_attr = composer.evaluate("(el) => el.getAttribute('dir')") + if dir_attr != "auto": + soft_fail( + f'composer is missing dir="auto" (got {dir_attr!r}); RTL ' + "languages will render LTR." + ) + else: + info('composer dir="auto" present') + + # Source-level guard for the edit and compare composers (neither + # is mounted here): grep the JSX for dir="auto" inside each block. + _repo_root = Path(__file__).resolve().parents[2] + _thread_src = ( + _repo_root / "studio/frontend/src/components/assistant-ui/thread.tsx" + ).read_text() + _shared_src = ( + _repo_root / "studio/frontend/src/features/chat/shared-composer.tsx" + ).read_text() + _edit_idx = _thread_src.find("aui-edit-composer-input") + if _edit_idx == -1 or 'dir="auto"' not in _thread_src[_edit_idx : _edit_idx + 600]: + soft_fail('edit composer source is missing dir="auto"') + else: + info('edit composer dir="auto" present (source)') + _compare_idx = _shared_src.find("Send to both models") + if ( + _compare_idx == -1 + or 'dir="auto"' + not in _shared_src[max(_compare_idx - 400, 0) : _compare_idx + 400] + ): + soft_fail('compare composer source is missing dir="auto"') + else: + info('compare composer dir="auto" present (source)') + + def read_value() -> str: + return composer.evaluate("(el) => el.value") + + def set_value_via_setter(s: str) -> str: + """Write via React's monkey-patched setter + paste input event, + then await two rAFs so the controlled value is committed before + readback (plain `.value=s` would be overwritten on next render).""" + return composer.evaluate( + """async (el, v) => { + const setter = Object.getOwnPropertyDescriptor( + window.HTMLTextAreaElement.prototype, 'value' + ).set; + setter.call(el, v); + el.dispatchEvent(new InputEvent('input', { + bubbles: true, + inputType: 'insertFromPaste', + data: v, + })); + await new Promise((r) => requestAnimationFrame(r)); + await new Promise((r) => requestAnimationFrame(r)); + return el.value; + }""", + s, + ) + + def clear() -> None: + set_value_via_setter("") + + # 3. Baseline: ASCII keyboard typing works. Bail fast if not. + step("baseline ASCII keyboard typing") + clear() + composer.click() + for ch in "hello world": + page.keyboard.type(ch) + got = read_value() + if got != "hello world": + fail(f"ASCII typing readback {got!r} != 'hello world'") + info("baseline ASCII OK") + shoot("03-baseline-ascii") + clear() + + # 4. Multilingual paste round-trip; byte-for-byte readback required. + step(f"multilingual paste round-trip ({len(I18N_SAMPLES)} samples)") + paste_failures: list[tuple[str, str, str, str]] = [] + for code, label, text in I18N_SAMPLES: + got = set_value_via_setter(text) + if got != text: + paste_failures.append((code, label, text, got)) + info(f" {code:>6} ({label}): FAIL -- got {got!r}") + else: + info(f" {code:>6} ({label}): OK") + clear() + if paste_failures: + shoot("04-paste-failures") + lines = [ + f" {code} ({label}): want={want!r} got={got!r}" + for code, label, want, got in paste_failures + ] + fail( + f"{len(paste_failures)}/{len(I18N_SAMPLES)} languages failed paste round-trip:\n" + + "\n".join(lines) + ) + info(f"all {len(I18N_SAMPLES)} multilingual paste samples OK") + shoot("04-paste-all-ok") + + # 5. Healthy IME composition (compositionstart/update/end + insert events). + step("normal IME composition (compose 你好)") + clear() + composer.click() + composer.evaluate( + """(el) => { + el.focus(); + el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''})); + el.dispatchEvent(new CompositionEvent('compositionupdate', {bubbles:true, data:'你'})); + el.dispatchEvent(new CompositionEvent('compositionupdate', {bubbles:true, data:'你好'})); + const setter = Object.getOwnPropertyDescriptor( + window.HTMLTextAreaElement.prototype, 'value' + ).set; + setter.call(el, el.value + '你好'); + el.dispatchEvent(new InputEvent('input', { + bubbles:true, inputType:'insertCompositionText', + data:'你好', isComposing:true, + })); + el.dispatchEvent(new CompositionEvent('compositionend', {bubbles:true, data:'你好'})); + el.dispatchEvent(new InputEvent('input', { + bubbles:true, inputType:'insertFromComposition', data:'你好', + })); + }""" + ) + got = read_value() + if "你好" not in got: + shoot("05-normal-composition-FAIL") + fail(f"normal composition readback {got!r} missing '你好'") + info(f"normal composition OK: ta.value={got!r}") + shoot("05-normal-composition") + clear() + + # 6. Stuck IME repro for issue #5318: duplicate compositionstart with + # no compositionend wedged isComposing=true and dropped ASCII keys. + # PR #5327 cleared the stale state on non-composing input. + step("BUG REPRO: stuck IME composition recovery (issue #5318)") + clear() + composer.click() + composer.evaluate( + """(el) => { + el.focus(); + el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''})); + // Duplicate compositionstart with NO matching compositionend. + // This is exactly the event sequence observed from the IMEs + // in issue #5318 (kei-yamazaki / langxiaopiao030 / PapyrusNotes). + el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''})); + }""" + ) + # Drive the real keyboard path; on the broken build React drops + # 'abcd' and reconciles el.value back to ''. wait_for_function + # crosses the microtask boundary so we see committed React state. + page.keyboard.type("abcd") + try: + page.wait_for_function( + """(el) => el.value === 'abcd'""", + composer.element_handle(), + timeout = 5_000, + ) + except Exception: + pass + after_key = read_value() + info(f"after_key='abcd' readback={after_key!r}") + shoot("06-stuck-composition-recovery") + if after_key != "abcd": + fail( + "stuck-composition repro: keyboard 'abcd' was not preserved after " + f"duplicate compositionstart; readback {after_key!r}. React state " + "likely still stuck in isComposing=true (issue #5318 / before " + "PR #5327)." + ) + # Cross-check React's view of isComposing via the Send button: + # ComposerAction stays disabled while isComposing is true (PR #5327). + send_btn = page.locator('button[aria-label="Send message"]') + if send_btn.count() == 0: + soft_fail("Send button not found after stuck-composition recovery") + else: + try: + expect(send_btn).not_to_be_disabled(timeout = 5_000) + info("Send button correctly enabled after stuck-composition recovery") + except Exception: + soft_fail( + "Send button still disabled after stuck-composition recovery -- " + "React isComposing state likely never cleared" + ) + info("stuck-composition recovery PASS") + clear() + + # 7. Final state. The change-password redirect emits benign 401 noise, + # so we filter via is_benign_* and only fail on real errors. + shoot("07-final") + real_page_errors = [e for e in page_errors if not is_benign_page_error(e)] + real_console_errors = [e for e in console_errors if not is_benign_console_error(e)] + info( + f"page_errors={len(page_errors)} ({len(real_page_errors)} non-benign); " + f"console_errors={len(console_errors)} " + f"({len(real_console_errors)} non-benign)" + ) + if page_errors: + info(f"first page error: {page_errors[0][:200]!r}") + if console_errors: + info(f"first console error: {console_errors[0][:200]!r}") + if real_page_errors: + fail( + f"{len(real_page_errors)} non-benign pageerror events; " + f"first={real_page_errors[0][:200]!r}" + ) + if real_console_errors: + fail( + f"{len(real_console_errors)} non-benign console.error events; " + f"first={real_console_errors[0][:200]!r}" + ) + + info( + f"DONE: ascii=OK paste={len(I18N_SAMPLES)}/{len(I18N_SAMPLES)} " + f"normal_composition=OK stuck_recovery=OK" + ) + _watchdog.cancel() + browser.close() diff --git a/tests/studio/test_composer_rtl_bidi_attribute.py b/tests/studio/test_composer_rtl_bidi_attribute.py new file mode 100644 index 0000000000..5b1437b4fc --- /dev/null +++ b/tests/studio/test_composer_rtl_bidi_attribute.py @@ -0,0 +1,73 @@ +"""Lock down the RTL bidi auto-detection contract on the chat composers. + +The browser's Unicode bidi algorithm only flows Arabic / Hebrew / Persian / +Urdu right-to-left when the textarea carries `dir="auto"`. The three +composer surfaces (main chat, inline edit, compare mode) each need the +attribute, and the IME / i18n Playwright smoke must keep its env contract +minimal (no dead `STUDIO_OLD_PW`). +""" + +from __future__ import annotations + +import re +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +THREAD_TSX = REPO / "studio/frontend/src/components/assistant-ui/thread.tsx" +SHARED_TSX = REPO / "studio/frontend/src/features/chat/shared-composer.tsx" +WORKFLOW_YML = REPO / ".github/workflows/studio-ui-smoke.yml" +IME_PY = REPO / "tests/studio/playwright_chat_ime_i18n.py" + + +def _block_around(src: str, anchor: str, radius: int = 600) -> str: + idx = src.find(anchor) + assert idx != -1, f"anchor {anchor!r} not found" + return src[max(idx - radius, 0) : idx + radius] + + +def test_main_composer_has_dir_auto(): + block = _block_around(THREAD_TSX.read_text(), 'aria-label="Message input"') + assert 'dir="auto"' in block, 'main composer is missing dir="auto"' + + +def test_edit_composer_has_dir_auto(): + block = _block_around(THREAD_TSX.read_text(), "aui-edit-composer-input") + assert 'dir="auto"' in block, 'edit composer is missing dir="auto"' + + +def test_compare_composer_has_dir_auto(): + block = _block_around(SHARED_TSX.read_text(), "Send to both models") + assert 'dir="auto"' in block, 'compare composer is missing dir="auto"' + + +def test_ime_workflow_step_does_not_set_studio_old_pw(): + yml = WORKFLOW_YML.read_text() + drive_idx = yml.find("Drive IME + multilingual paste regression") + assert drive_idx != -1, "IME drive step not found in workflow" + next_step_idx = yml.find("- name:", drive_idx + 1) + drive_block = yml[drive_idx : next_step_idx if next_step_idx != -1 else None] + assert ( + "STUDIO_OLD_PW" not in drive_block + ), "IME drive step still passes dead STUDIO_OLD_PW env var" + assert "STUDIO_NEW_PW" in drive_block, "IME drive step missing STUDIO_NEW_PW" + + +def test_ime_pass_password_step_does_not_export_old_pw(): + yml = WORKFLOW_YML.read_text() + pass_idx = yml.find("Pass bootstrap pw for IME / i18n test") + assert pass_idx != -1, "IME password setup step not found" + next_step_idx = yml.find("- name:", pass_idx + 1) + pass_block = yml[pass_idx : next_step_idx if next_step_idx != -1 else None] + assert ( + "STUDIO_IME_OLD_PW" not in pass_block + ), "IME password setup still exports dead STUDIO_IME_OLD_PW" + assert "STUDIO_IME_NEW_PW" in pass_block + + +def test_ime_playwright_script_does_not_read_studio_old_pw(): + src = IME_PY.read_text() + code_only = re.sub(r'""".*?"""', "", src, flags = re.DOTALL) + assert ( + "STUDIO_OLD_PW" not in code_only + ), "IME Playwright script still references dead STUDIO_OLD_PW env var" + assert 'os.environ["STUDIO_NEW_PW"]' in code_only From f0270bcb17287fb274d55ce60054d50f983a41e4 Mon Sep 17 00:00:00 2001 From: h34v3nzc0dex Date: Sun, 17 May 2026 06:04:00 -0600 Subject: [PATCH 6/8] fix(studio/worker): inject --gcc-install-dir for HIP source builds on Ubuntu 24.04 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On Ubuntu 24.04 + ROCm clang-20, the HIP source-build fallback in `_install_package_wheel_first` (causal-conv1d, mamba-ssm source fallback, flash-attn source fallback) dies at: /opt/rocm-X.Y/lib/llvm/lib/clang/20/include/__clang_hip_runtime_wrapper.h:112:10: fatal error: 'cstdlib' file not found Root cause: clang-20 picks the highest-numbered /usr/lib/gcc/x86_64-linux-gnu/ runtime dir by default. On 24.04 that's gcc-14, whose runtime objects ship in the gcc-14 package but whose C++ headers (/usr/include/c++/14) come from libstdc++-14-dev — NOT in the default apt set. libstdc++-13-dev IS in the default set, so /usr/include/c++/13 exists. clang has no way to discover that asymmetry and the build fails. Fix: new `_hipcc_gcc_install_dir()` helper iterates gcc 14 → 11 and returns the first /usr/lib/gcc/x86_64-linux-gnu/ dir where BOTH the runtime AND /usr/include/c++/ exist. The HIP branch of `_install_package_wheel_first` appends `--gcc-install-dir=` to HIPCC_COMPILE_FLAGS_APPEND before invoking pip. Respects an existing `--gcc-install-dir` in the env var (user-set takes precedence); preserves any other flags the user has set (appends to the end rather than overwriting). No-op on non-HIP, non-Linux, non-x86_64. Mirrors the same fix bbf004c added to studio/setup.sh for the llama.cpp HIP build branch (#5301), but via env var since pip-driven source builds can't take CMake flags directly. Verified on Ryzen AI MAX+ 395 / Radeon 8060S (gfx1151) / Ubuntu 24.04 / ROCm 7.13 nightly: `_hipcc_gcc_install_dir()` returns `/usr/lib/gcc/x86_64-linux-gnu/13`, which matches the manual workaround that already lets `pip install causal-conv1d` succeed on this hardware. Tests added (8 new in test_training_worker_flash_attn.py): - test_hipcc_gcc_install_dir_picks_highest_with_headers - test_hipcc_gcc_install_dir_picks_14_when_headers_exist - test_hipcc_gcc_install_dir_returns_none_when_no_match - test_hipcc_gcc_install_dir_returns_none_on_non_linux - test_hipcc_gcc_install_dir_returns_none_on_non_x86_64 - test_install_injects_gcc_install_dir_on_hip_source_build - test_install_appends_to_existing_hipcc_compile_flags - test_install_respects_user_gcc_install_dir - test_install_does_not_inject_env_on_cuda Per @danielhanchen's suggestion in https://github.com/unslothai/unsloth/pull/5434#issuecomment-4469980122 --- studio/backend/core/training/worker.py | 59 ++++ .../tests/test_training_worker_flash_attn.py | 291 ++++++++++++++++++ 2 files changed, 350 insertions(+) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 4434436ca3..bca551fbd2 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -77,6 +77,38 @@ def _model_wants_causal_conv1d(model_name: str) -> bool: ) +def _hipcc_gcc_install_dir() -> str | None: + """Return the highest-numbered ``/usr/lib/gcc/x86_64-linux-gnu/`` that has + BOTH the gcc runtime dir AND the corresponding ``/usr/include/c++/`` C++ + headers, or ``None`` if no match (or non-Linux / non-x86_64). + + Ubuntu 24.04 ships ``/usr/lib/gcc/x86_64-linux-gnu/14/`` (gcc-14 runtime + objects) but does NOT ship ``/usr/include/c++/14`` in its default apt set; + libstdc++ headers come from ``libstdc++-13-dev``. ROCm clang-20 picks the + highest-numbered runtime dir by default, finds no ````, and the + HIP source build fails with:: + + /opt/rocm-X.Y/lib/llvm/lib/clang/20/include/__clang_hip_runtime_wrapper.h:112:10: + fatal error: 'cstdlib' file not found + + Returning a path lets the caller pass ``--gcc-install-dir=`` to clang + via ``HIPCC_COMPILE_FLAGS_APPEND``. Mirrors the same loop ``bbf004c`` added + to ``studio/setup.sh`` for the llama.cpp HIP build branch (PR #5301). + """ + if not sys.platform.startswith("linux"): + return None + import platform as _platform + + if _platform.machine().lower() != "x86_64": + return None + for _ver in (14, 13, 12, 11): + _runtime = f"/usr/lib/gcc/x86_64-linux-gnu/{_ver}/include" + _headers = f"/usr/include/c++/{_ver}" + if os.path.isdir(_runtime) and os.path.isdir(_headers): + return f"/usr/lib/gcc/x86_64-linux-gnu/{_ver}" + return None + + def _install_package_wheel_first( *, event_queue: Any, @@ -212,6 +244,33 @@ def _install_package_wheel_first( } if is_hip: _run_kwargs["timeout"] = 1800 + # On Ubuntu 24.04 + ROCm clang-20, the HIP source build (causal-conv1d, + # mamba-ssm source fallback, flash-attn source fallback) defaults to + # /usr/lib/gcc/x86_64-linux-gnu/14/ which has the runtime dir but no + # /usr/include/c++/14 headers, and dies at: + # __clang_hip_runtime_wrapper.h:112:10: + # fatal error: 'cstdlib' file not found + # Inject --gcc-install-dir for a gcc whose C++ headers actually exist. + # Respect any pre-existing --gcc-install-dir in HIPCC_COMPILE_FLAGS_APPEND + # (user knows best); otherwise append. Mirrors the same fix bbf004c + # added to studio/setup.sh for the llama.cpp HIP build (PR #5301). + _existing_flags = os.environ.get("HIPCC_COMPILE_FLAGS_APPEND", "") + if "--gcc-install-dir" not in _existing_flags: + _gcc_dir = _hipcc_gcc_install_dir() + if _gcc_dir is not None: + _appended = ( + f"{_existing_flags} --gcc-install-dir={_gcc_dir}" + ).strip() + _run_kwargs["env"] = { + **os.environ, + "HIPCC_COMPILE_FLAGS_APPEND": _appended, + } + logger.info( + "HIP source build for %s: appended " + "--gcc-install-dir=%s to HIPCC_COMPILE_FLAGS_APPEND", + display_name, + _gcc_dir, + ) try: result = _sp.run(pypi_cmd, **_run_kwargs) diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 0737bdc82f..d99c6b1f4a 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -6,6 +6,7 @@ from __future__ import annotations import builtins import subprocess import sys +from typing import Any from unittest import mock from core.training import worker @@ -22,6 +23,17 @@ def _missing_flash_attn_import(): return fake_import +def _missing_module_import(missing: str): + real_import = builtins.__import__ + + def fake_import(name, globals = None, locals = None, fromlist = (), level = 0): + if name == missing: + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + return fake_import + + def test_should_try_runtime_flash_attn_install_threshold_and_skip(monkeypatch): monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) assert worker._should_try_runtime_flash_attn_install(32767) is False @@ -193,3 +205,282 @@ def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch): release_tag = worker._MAMBA_SSM_RELEASE_TAG, release_base_url = "https://github.com/state-spaces/mamba/releases/download", ) + + +# ──────────────────────────────────────────────────────────────────── +# HIP source-build gcc-install-dir coverage (h34v3nzc0dex Strix Halo). +# Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14, +# so ROCm clang-20 picks it and fails with 'cstdlib' file not found +# when building causal-conv1d (or any other HIP source fallback). +# _hipcc_gcc_install_dir() finds a gcc dir that has both halves; the +# _install_package_wheel_first HIP branch passes it to clang via +# HIPCC_COMPILE_FLAGS_APPEND. Parallel to bbf004c's setup.sh fix for +# the llama.cpp HIP build (PR #5301). +# ──────────────────────────────────────────────────────────────────── + + +def _isdir_for_layout(*existing: str): + """Return an os.path.isdir replacement that only treats the given + absolute paths as directories. Lets a test simulate exactly which + gcc runtime dirs and C++ header dirs exist on the host.""" + valid = set(existing) + + def fake_isdir(path: str) -> bool: + return path in valid + + return fake_isdir + + +def test_hipcc_gcc_install_dir_picks_highest_with_headers(monkeypatch): + """gcc-14 has runtime but no /usr/include/c++/14; loop falls through + to gcc-13 which has both. This is the exact Ubuntu 24.04 layout.""" + monkeypatch.setattr(sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "x86_64") + monkeypatch.setattr( + worker.os.path, + "isdir", + _isdir_for_layout( + "/usr/lib/gcc/x86_64-linux-gnu/14/include", # runtime present + # but no /usr/include/c++/14 — typical Ubuntu 24.04 default + "/usr/lib/gcc/x86_64-linux-gnu/13/include", + "/usr/include/c++/13", # libstdc++-13-dev installed + ), + ) + assert worker._hipcc_gcc_install_dir() == "/usr/lib/gcc/x86_64-linux-gnu/13" + + +def test_hipcc_gcc_install_dir_picks_14_when_headers_exist(monkeypatch): + """If the user has libstdc++-14-dev installed, prefer gcc-14.""" + monkeypatch.setattr(sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "x86_64") + monkeypatch.setattr( + worker.os.path, + "isdir", + _isdir_for_layout( + "/usr/lib/gcc/x86_64-linux-gnu/14/include", + "/usr/include/c++/14", + ), + ) + assert worker._hipcc_gcc_install_dir() == "/usr/lib/gcc/x86_64-linux-gnu/14" + + +def test_hipcc_gcc_install_dir_returns_none_when_no_match(monkeypatch): + """No gcc dir has both halves → return None and skip the env injection + rather than guessing wrong and surfacing a confusing build failure.""" + monkeypatch.setattr(sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "x86_64") + monkeypatch.setattr(worker.os.path, "isdir", lambda path: False) + assert worker._hipcc_gcc_install_dir() is None + + +def test_hipcc_gcc_install_dir_returns_none_on_non_linux(monkeypatch): + """Don't probe gcc layout on macOS / Windows — early-return.""" + monkeypatch.setattr(sys, "platform", "darwin") + + def _isdir_should_not_be_called(_path): + raise AssertionError("isdir should not be called on non-Linux") + + monkeypatch.setattr(worker.os.path, "isdir", _isdir_should_not_be_called) + assert worker._hipcc_gcc_install_dir() is None + + +def test_hipcc_gcc_install_dir_returns_none_on_non_x86_64(monkeypatch): + """ROCm clang-20 on aarch64 has a different libstdc++ layout.""" + monkeypatch.setattr(sys, "platform", "linux") + import platform as _platform + + monkeypatch.setattr(_platform, "machine", lambda: "aarch64") + assert worker._hipcc_gcc_install_dir() is None + + +def _make_hip_install_env(monkeypatch, *, gcc_dir: str | None): + """Common scaffolding for tests that exercise the HIP source-build + branch of _install_package_wheel_first end-to-end. The package isn't + installed yet, no prebuilt wheel exists, hipcc is on PATH, and the + fake env reports an HIP torch.""" + monkeypatch.setattr( + builtins, "__import__", _missing_module_import("causal_conv1d") + ) + monkeypatch.setattr( + worker, + "probe_torch_wheel_env", + lambda timeout = 30: { + "hip_version": "7.13.26176", + "python_tag": "cp312", + "torch_mm": "2.11", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + }, + ) + monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None) + monkeypatch.setattr( + worker.shutil, + "which", + lambda name: "/opt/rocm/bin/hipcc" if name == "hipcc" else None, + ) + monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) + monkeypatch.setattr(worker, "_hipcc_gcc_install_dir", lambda: gcc_dir) + + +def test_install_injects_gcc_install_dir_on_hip_source_build(monkeypatch): + """HIP source-build with no user-set HIPCC_COMPILE_FLAGS_APPEND → + subprocess env carries --gcc-install-dir=.""" + monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False) + _make_hip_install_env( + monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13" + ) + + captured: dict[str, str] = {} + + def fake_run(cmd, **kwargs): + captured.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._install_package_wheel_first( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = "1.6.2.post1", + filename_prefix = "causal_conv1d", + release_tag = "v1.6.2.post1", + release_base_url = "https://example.com", + ) + + assert ( + captured.get("HIPCC_COMPILE_FLAGS_APPEND") + == "--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13" + ) + + +def test_install_appends_to_existing_hipcc_compile_flags(monkeypatch): + """User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' set → final value + keeps the user's flags AND adds --gcc-install-dir at the end.""" + monkeypatch.setenv("HIPCC_COMPILE_FLAGS_APPEND", "-O3 -DFOO") + _make_hip_install_env( + monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13" + ) + + captured: dict[str, str] = {} + + def fake_run(cmd, **kwargs): + captured.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._install_package_wheel_first( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = "1.6.2.post1", + filename_prefix = "causal_conv1d", + release_tag = "v1.6.2.post1", + release_base_url = "https://example.com", + ) + + assert captured.get("HIPCC_COMPILE_FLAGS_APPEND") == ( + "-O3 -DFOO --gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13" + ) + + +def test_install_respects_user_gcc_install_dir(monkeypatch): + """User explicitly set --gcc-install-dir=… already → don't touch it. + Avoids two competing --gcc-install-dir flags on the clang command line.""" + monkeypatch.setenv( + "HIPCC_COMPILE_FLAGS_APPEND", + "--gcc-install-dir=/opt/custom/gcc-13", + ) + _make_hip_install_env( + monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13" + ) + + captured: dict[str, str] | None = {"_called": "no"} + + def fake_run(cmd, **kwargs): + env = kwargs.get("env") + if env is not None: + captured.clear() + captured.update(env) + else: + captured["_called"] = "yes_no_env" + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._install_package_wheel_first( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = "1.6.2.post1", + filename_prefix = "causal_conv1d", + release_tag = "v1.6.2.post1", + release_base_url = "https://example.com", + ) + + # subprocess.run was invoked without env override (the user already + # set HIPCC_COMPILE_FLAGS_APPEND with --gcc-install-dir, so we left + # the env alone — the existing value is inherited normally). + assert captured == {"_called": "yes_no_env"} + + +def test_install_does_not_inject_env_on_cuda(monkeypatch): + """CUDA path (no hip_version in env) → no env override at all.""" + monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False) + monkeypatch.setattr( + builtins, "__import__", _missing_module_import("causal_conv1d") + ) + monkeypatch.setattr( + worker, + "probe_torch_wheel_env", + lambda timeout = 30: { + "python_tag": "cp312", + "torch_mm": "2.11", + "cuda_major": "12", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + }, + ) + monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None) + monkeypatch.setattr(worker.shutil, "which", lambda name: None) + monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) + # If _hipcc_gcc_install_dir were called on CUDA we'd want to know. + monkeypatch.setattr( + worker, + "_hipcc_gcc_install_dir", + lambda: (_ for _ in ()).throw( + AssertionError("must not run on CUDA") + ), + ) + + captured: dict[str, Any] = {} + + def fake_run(cmd, **kwargs): + captured["env_in_kwargs"] = "env" in kwargs + return subprocess.CompletedProcess(cmd, 0, "") + + monkeypatch.setattr(worker._sp, "run", fake_run) + + worker._install_package_wheel_first( + event_queue = [], + import_name = "causal_conv1d", + display_name = "causal-conv1d", + pypi_name = "causal-conv1d", + pypi_version = "1.6.2.post1", + filename_prefix = "causal_conv1d", + release_tag = "v1.6.2.post1", + release_base_url = "https://example.com", + ) + + # CUDA branch never sets the env, never invokes the gcc helper. + assert captured.get("env_in_kwargs") is False From 81ae3583e77f41a5d513b27b257ca5d5e9244e15 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 17 May 2026 12:05:47 +0000 Subject: [PATCH 7/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/training/worker.py | 4 +--- .../tests/test_training_worker_flash_attn.py | 24 +++++-------------- 2 files changed, 7 insertions(+), 21 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index bca551fbd2..84be40d7f1 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -258,9 +258,7 @@ def _install_package_wheel_first( if "--gcc-install-dir" not in _existing_flags: _gcc_dir = _hipcc_gcc_install_dir() if _gcc_dir is not None: - _appended = ( - f"{_existing_flags} --gcc-install-dir={_gcc_dir}" - ).strip() + _appended = (f"{_existing_flags} --gcc-install-dir={_gcc_dir}").strip() _run_kwargs["env"] = { **os.environ, "HIPCC_COMPILE_FLAGS_APPEND": _appended, diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index d99c6b1f4a..733b726656 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -304,9 +304,7 @@ def _make_hip_install_env(monkeypatch, *, gcc_dir: str | None): branch of _install_package_wheel_first end-to-end. The package isn't installed yet, no prebuilt wheel exists, hipcc is on PATH, and the fake env reports an HIP torch.""" - monkeypatch.setattr( - builtins, "__import__", _missing_module_import("causal_conv1d") - ) + monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d")) monkeypatch.setattr( worker, "probe_torch_wheel_env", @@ -332,9 +330,7 @@ def test_install_injects_gcc_install_dir_on_hip_source_build(monkeypatch): """HIP source-build with no user-set HIPCC_COMPILE_FLAGS_APPEND → subprocess env carries --gcc-install-dir=.""" monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False) - _make_hip_install_env( - monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13" - ) + _make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13") captured: dict[str, str] = {} @@ -365,9 +361,7 @@ def test_install_appends_to_existing_hipcc_compile_flags(monkeypatch): """User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' set → final value keeps the user's flags AND adds --gcc-install-dir at the end.""" monkeypatch.setenv("HIPCC_COMPILE_FLAGS_APPEND", "-O3 -DFOO") - _make_hip_install_env( - monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13" - ) + _make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13") captured: dict[str, str] = {} @@ -400,9 +394,7 @@ def test_install_respects_user_gcc_install_dir(monkeypatch): "HIPCC_COMPILE_FLAGS_APPEND", "--gcc-install-dir=/opt/custom/gcc-13", ) - _make_hip_install_env( - monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13" - ) + _make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13") captured: dict[str, str] | None = {"_called": "no"} @@ -437,9 +429,7 @@ def test_install_respects_user_gcc_install_dir(monkeypatch): def test_install_does_not_inject_env_on_cuda(monkeypatch): """CUDA path (no hip_version in env) → no env override at all.""" monkeypatch.delenv("HIPCC_COMPILE_FLAGS_APPEND", raising = False) - monkeypatch.setattr( - builtins, "__import__", _missing_module_import("causal_conv1d") - ) + monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d")) monkeypatch.setattr( worker, "probe_torch_wheel_env", @@ -458,9 +448,7 @@ def test_install_does_not_inject_env_on_cuda(monkeypatch): monkeypatch.setattr( worker, "_hipcc_gcc_install_dir", - lambda: (_ for _ in ()).throw( - AssertionError("must not run on CUDA") - ), + lambda: (_ for _ in ()).throw(AssertionError("must not run on CUDA")), ) captured: dict[str, Any] = {} From aa30ae5df1cfb098ec5017cae0353942d2b057e5 Mon Sep 17 00:00:00 2001 From: h34v3nzc0dex Date: Sun, 17 May 2026 06:17:46 -0600 Subject: [PATCH 8/8] review: apply gemini-code-assist suggestion on _run_kwargs env handling Use _run_kwargs.get("env", os.environ).copy() + key-mutation instead of rebuilding env from os.environ directly. Today both forms are equivalent (no earlier code in _install_package_wheel_first sets _run_kwargs["env"]), but the .get().copy() pattern survives any future env modification added upstream of this block without silently throwing it away. No behavioural change; tests already assert the final HIPCC_COMPILE_FLAGS_APPEND value, not the env-construction pattern. Per https://github.com/unslothai/unsloth/pull/5517#discussion_r... (gemini-code-assist[bot]) --- studio/backend/core/training/worker.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 84be40d7f1..4b4f3af58f 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -259,10 +259,9 @@ def _install_package_wheel_first( _gcc_dir = _hipcc_gcc_install_dir() if _gcc_dir is not None: _appended = (f"{_existing_flags} --gcc-install-dir={_gcc_dir}").strip() - _run_kwargs["env"] = { - **os.environ, - "HIPCC_COMPILE_FLAGS_APPEND": _appended, - } + _env = _run_kwargs.get("env", os.environ).copy() + _env["HIPCC_COMPILE_FLAGS_APPEND"] = _appended + _run_kwargs["env"] = _env logger.info( "HIP source build for %s: appended " "--gcc-install-dir=%s to HIPCC_COMPILE_FLAGS_APPEND",