From 53a071cb14142c4d457aab3df2f1e2b287c790a1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 20:25:41 -0700 Subject: [PATCH 001/113] Harden Windows Pester install against missing PSGallery (#6892) The setup.ps1 unit-tests job intermittently fails on the windows-latest runner with 'No repository with the name PSGallery was found.' when the default PowerShell Gallery is not registered, so Set-PSRepository throws before Pester can be installed. Register the default gallery first when it is missing, then set its policy and install Pester as before. --- .github/workflows/studio-windows-inference-smoke.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 8186c07211..0bc216d65a 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -1610,6 +1610,13 @@ jobs: - name: Install Pester v5 shell: pwsh run: | + # PSGallery is intermittently absent from the repository list on GitHub's Windows + # runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the + # name 'PSGallery' was found." Re-register the default gallery first so the policy + # change and module install below always have a repository to target. + if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) { + Register-PSRepository -Default -ErrorAction SilentlyContinue + } Set-PSRepository PSGallery -InstallationPolicy Trusted Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser Import-Module Pester -MinimumVersion 5.5.0 From cb6737cbb8e2a8bf82985b19fcb126319e67cc5e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:13:25 -0700 Subject: [PATCH 002/113] Auto Xet to HTTP download fallback in from_pretrained; share Studio's fallback via unsloth_zoo (#6638) --- .github/workflows/consolidated-tests-ci.yml | 1 + studio/backend/tests/test_hf_xet_fallback.py | 550 +++++------ .../tests/test_model_update_robustness.py | 44 +- studio/backend/utils/hf_xet_fallback.py | 554 ++++------- .../model-selector/model-update-action.tsx | 11 +- .../assistant-ui/model-selector/pickers.tsx | 14 +- tests/test_prefetch_snapshot_scope.py | 916 ++++++++++++++++++ unsloth/models/_utils.py | 406 ++++++++ unsloth/models/diffusion.py | 41 +- unsloth/models/llama.py | 73 ++ unsloth/models/loader.py | 47 + unsloth/models/sentence_transformer.py | 144 ++- unsloth/models/vision.py | 74 ++ unsloth/tokenizer_utils.py | 18 +- 14 files changed, 2155 insertions(+), 738 deletions(-) create mode 100644 tests/test_prefetch_snapshot_scope.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 7978a200c0..ae4b386589 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -364,6 +364,7 @@ jobs: tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_prefetch_snapshot_scope.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 9e40fbf508..4d73213d15 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -1,18 +1,16 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP -transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on. -CPU-only, no network, no real subprocess (the per-attempt download seam is -monkeypatched). +"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback. + +The transport-policy matrix is tested once in unsloth_zoo; here we assert only the +Studio seam: re-exporting the shared API and injecting the marker-aware +prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess. """ from __future__ import annotations -import subprocess import sys -import threading -import time import types as _types from pathlib import Path @@ -22,9 +20,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Stub heavy/unavailable deps before importing the module under test. Use the -# real structlog when present; a bare stub left in sys.modules would break later -# modules that log at import time. +# Stub heavy/unavailable deps before importing the module under test. Use real structlog when present; +# a bare stub would break later modules that log at import time. _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) @@ -34,171 +31,59 @@ except ImportError: sys.modules["structlog"] = _types.ModuleType("structlog") import huggingface_hub -from huggingface_hub import constants as hf_constants + +try: + import unsloth_zoo.hf_xet_fallback as _shared_mod + shared = _shared_mod +except Exception: # noqa: BLE001 - still collect degraded-path tests when unsloth_zoo is unavailable + shared = None import utils.hf_xet_fallback as xf -# --------------------------------------------------------------------------- # -# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total. -# --------------------------------------------------------------------------- # -REPO = "ztest/xet-watchdog" - - -@pytest.fixture -def hf_cache(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) - return tmp_path - - -def _blobs_dir(root: Path, repo_id: str = REPO) -> Path: - d = root / f"models--{repo_id.replace('/', '--')}" / "blobs" - d.mkdir(parents = True, exist_ok = True) - return d - - -def _wait( - predicate, - timeout: float = 2.0, - step: float = 0.02, -) -> bool: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if predicate(): - return True - time.sleep(step) - return predicate() - - -def test_constant_incomplete_fires_stall(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - assert _wait( - lambda: len(calls) >= 1, timeout = 3.0 - ), "watchdog never fired on a constant-size .incomplete" - finally: - stop.set() - assert "stalled" in calls[0].lower() - - -def test_growing_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - part = blobs / "growing.incomplete" - part.write_bytes(b"\0" * 1024) - - grow_stop = threading.Event() - - def _grow(): - size = 1024 - while not grow_stop.wait(0.05): - size += 4096 - part.write_bytes(b"\0" * size) - - grower = threading.Thread(target = _grow, daemon = True) - grower.start() - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(1.0) # well past stall_timeout, but bytes keep growing - assert calls == [], "watchdog fired despite continuous progress" - finally: - stop.set() - grow_stop.set() - - -def test_no_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(0.8) - assert calls == [], "watchdog fired with no active .incomplete" - finally: - stop.set() - - -def test_stall_fires_at_most_once(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048) - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2 - ) - try: - assert _wait(lambda: len(calls) >= 1, timeout = 3.0) - time.sleep(0.6) # keep ticking; must not fire again - assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1" - finally: - stop.set() - - -def test_get_state_empty_cache(hf_cache): - assert xf.get_hf_download_state([REPO]) == (0, False) - - -def test_get_state_absent_cache_root(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache")) - assert xf.get_hf_download_state([REPO]) == (0, False) - - -def test_get_state_skips_local_paths(hf_cache): - # Filesystem paths are not HF repo IDs and must be ignored without error. - assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False) - - -def test_get_state_sparse_aware(hf_cache): - blobs = _blobs_dir(hf_cache) - sparse = blobs / "sparse.incomplete" - with open(sparse, "wb") as f: - f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks - st = sparse.stat() - if getattr(st, "st_blocks", 0) == 0: - pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable") - total, has_incomplete = xf.get_hf_download_state([REPO]) - assert has_incomplete is True - assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks" - - -# --------------------------------------------------------------------------- # -# Transport policy: cached short-circuit, cancel, error propagation, and the -# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn. -# --------------------------------------------------------------------------- # DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf" -@pytest.fixture(autouse = True) -def _no_real_cache_hit(monkeypatch): - """Default: the cached probe misses; tests override it to force a hit.""" +def _requires_shared(): + if shared is None: + pytest.skip("unsloth_zoo.hf_xet_fallback is not installed in this environment") + + +def test_shim_reexports_shared_api(): + _requires_shared() + assert xf.DownloadStallError is shared.DownloadStallError + for name in ( + "start_watchdog", + "get_hf_download_state", + "child_should_disable_xet", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", + ): + assert hasattr(xf, name), f"shim missing {name}" + + +def test_child_should_disable_xet_truth_table(): + assert xf.child_should_disable_xet({"disable_xet": True}) is True + assert xf.child_should_disable_xet({"disable_xet": False}) is False + assert xf.child_should_disable_xet({}) is False + + +def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): + """A Xet stall retries over HTTP and the shim runs Studio's marker-aware + ``prepare_cache_for_transport(..., 'http')`` before the retry.""" + _requires_shared() + for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"): + monkeypatch.delenv(var, raising = False) monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) + seen_disable_xet = [] -class _FakeAttempt: - """Records calls to the download seam and returns scripted results.""" - - def __init__(self, results): - self._results = list(results) - self.calls = [] - - def __call__( - self, + def fake_attempt( repo_id, - filename, - token, *, + kind, + params, + token, repo_type, disable_xet, cancel_event, @@ -208,146 +93,243 @@ class _FakeAttempt: on_status, force_download = False, ): - self.calls.append( - _types.SimpleNamespace( - repo_id = repo_id, - filename = filename, - disable_xet = disable_xet, - repo_type = repo_type, - ) - ) - return self._results[len(self.calls) - 1] + seen_disable_xet.append(disable_xet) + return ("ok", "/cache/model.gguf") if disable_xet else ("stall", None) + monkeypatch.setattr(shared, "_run_download_attempt", fake_attempt) -def _install(monkeypatch, results): - fake = _FakeAttempt(results) - monkeypatch.setattr(xf, "_run_download_attempt", fake) - return fake - - -def test_cached_file_short_circuits(monkeypatch, tmp_path): - cached = tmp_path / "cached.gguf" - cached.write_bytes(b"\0" * 8) - monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached)) - fake = _install(monkeypatch, []) # must not be called - - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == str(cached) - assert fake.calls == [], "spawned a download for an already-cached file" - - -def test_cancel_before_start_raises_no_attempt(monkeypatch): - fake = _install(monkeypatch, []) - ev = threading.Event() - ev.set() - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev) - assert fake.calls == [] - - -def test_nonstall_error_propagates_without_fallback(monkeypatch): - fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")]) - with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback" - assert fake.calls[0].disable_xet is False - - -def test_immediate_success_uses_xet_only(monkeypatch): - prepared = [] - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", - lambda *a, **k: prepared.append(a), - ) - fake = _install(monkeypatch, [("ok", "/cache/model.gguf")]) - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == "/cache/model.gguf" - assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False - assert prepared == [], "no cache prep should run when Xet succeeds first try" - - -def test_stall_then_http_fallback_succeeds(monkeypatch): prepared = [] monkeypatch.setattr( "hub.utils.download_registry.prepare_cache_for_transport", lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), ) - fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")]) out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) assert out == "/cache/model.gguf" - assert len(fake.calls) == 2 - assert fake.calls[0].disable_xet is False # Xet first - assert fake.calls[1].disable_xet is True # HTTP fallback - assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry" + assert seen_disable_xet == [False, True] # Xet first, then HTTP + assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep" -def test_second_stall_raises_download_stall_error(monkeypatch): - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("stall", None), ("stall", None)]) - with pytest.raises(xf.DownloadStallError): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 2 +def test_shim_snapshot_injects_studio_prepare(monkeypatch): + """The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper.""" + captured = {} + + def fake_snapshot(repo_id, **kwargs): + captured["repo_id"] = repo_id + captured["prepare_for_http_fn"] = kwargs.get("prepare_for_http_fn") + return "/tmp/snap-dir" + + monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot) + out = xf.snapshot_download_with_xet_fallback("org/model") + assert out == "/tmp/snap-dir" + assert captured["repo_id"] == "org/model" + assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http -def test_cancelled_midattempt_raises_no_fallback(monkeypatch): - fake = _install(monkeypatch, [("cancelled", None)]) - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1 +def test_degrades_gracefully_without_shared_helper(monkeypatch): + """On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio + boots) and exposes stub API doing plain HF downloads with the watchdog disabled.""" + import importlib + + class _BlockShared: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + raise ModuleNotFoundError(f"No module named '{name}'", name = name) + return None + + finder = _BlockShared() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + + # Boots without raising and mirrors the shared API surface. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.child_should_disable_xet({"disable_xet": True}) is True + assert degraded.get_hf_download_state(["x"]) is None # unmeasurable + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() # never fires + + # Degraded mode still emits heartbeats so the inactivity deadline is not tripped. + import time as _time + + beats = [] + hb_stop = degraded.start_watchdog( + repo_ids = ["x"], + on_stall = lambda m: None, + on_heartbeat = beats.append, + interval = 0.02, + ) + try: + deadline = _time.monotonic() + 2.0 + while not beats and _time.monotonic() < deadline: + _time.sleep(0.02) + assert beats, "degraded watchdog emitted no heartbeat" + finally: + hb_stop.set() + + # Downloads fall back to plain huggingface_hub (no watchdog, no crash). + called = {} + + def _fake_snapshot(repo_id, **kwargs): + called["repo_id"] = repo_id + return "/snap-dir" + + monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot) + assert degraded.snapshot_download_with_xet_fallback("org/model") == "/snap-dir" + assert called["repo_id"] == "org/model" + + # Cancellation still holds: an already-set cancel_event aborts before the HF download. + import threading as _threading + + cancelled = _threading.Event() + cancelled.set() + called.clear() + with pytest.raises(RuntimeError, match = "Cancelled"): + degraded.snapshot_download_with_xet_fallback("org/model", cancel_event = cancelled) + assert "repo_id" not in called, "degraded download ran despite cancellation" + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -def test_per_file_independent_fallback(monkeypatch): - """A stalled shard falls back; a sibling shard that succeeds does not.""" - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")]) - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a" - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b" - assert [c.disable_xet for c in fake.calls] == [False, False, True] +def test_degrades_when_unsloth_zoo_entirely_absent(): + """When unsloth_zoo is absent entirely, the import raises + ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still + degrades and does not re-raise, breaking every Studio import that pulls it in.""" + import importlib + + class _BlockZoo: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Whole package absent, so ModuleNotFoundError.name is the top-level 'unsloth_zoo'. + if name == "unsloth_zoo" or name.startswith("unsloth_zoo."): + raise ModuleNotFoundError("No module named 'unsloth_zoo'", name = "unsloth_zoo") + return None + + finder = _BlockZoo() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # Boots without raising and exposes the stub API. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -# --------------------------------------------------------------------------- # -# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect -# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it). -# --------------------------------------------------------------------------- # -def _safe_path() -> str: +def test_degrades_when_shared_helper_import_raises_importerror(): + """unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only + Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too.""" + import importlib + + class _BlockWithImportError: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + # Mirror a torch-less install: a plain ImportError with no .name. + raise ImportError("Unsloth: Pytorch is not installed.") + return None + + finder = _BlockWithImportError() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_zoo = sys.modules.pop("unsloth_zoo", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_zoo is not None: + sys.modules["unsloth_zoo"] = saved_zoo + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): + """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails.""" + import importlib import os - return os.environ.get("PATH", "") + monkeypatch.delenv("UNSLOTH_ZOO_DISABLE_GPU_INIT", raising = False) + seen_env = [] -def test_disable_xet_constant_set_in_fresh_interpreter(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()}, - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True " - f"(rc={proc.returncode}): {proc.stderr}" - ) + class _GpuGatedBlocker: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Crash is in unsloth_zoo's __init__, so intercept "unsloth_zoo" itself (the parent). + if name == "unsloth_zoo": + # Record the env each attempt sees; raise the no-GPU error both times so the shim + # degrades. + seen_env.append(os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT")) + raise NotImplementedError("Unsloth cannot find any torch accelerator") + return None - -def test_default_leaves_xet_enabled(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"without the env var, constants.HF_HUB_DISABLE_XET was not False " - f"(rc={proc.returncode}): {proc.stderr}" - ) + finder = _GpuGatedBlocker() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # First attempt without the light env, then a retry with it set. + assert seen_env == [None, "1"], seen_env + # Both attempts raised -> Studio still boots in degraded mode. + assert issubclass(degraded.DownloadStallError, RuntimeError) + # The env override must not leak past the import. + assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py index 9cf2a62c39..300eb587b3 100644 --- a/studio/backend/tests/test_model_update_robustness.py +++ b/studio/backend/tests/test_model_update_robustness.py @@ -5,8 +5,8 @@ Covers: * GGUF variant listing computes update_available from the already-fetched sibling metadata instead of a second Hub call. - * hf_hub_download_with_xet_fallback(force_download=True) bypasses the - try_to_load_from_cache cache-first early-return. + * hf_hub_download_with_xet_fallback forwards force_download through the shim to the + shared unsloth_zoo helper (which owns the cache-first early-return and its bypass). The cache "Update" action now runs through the download manager as a normal managed download (so it shows in the Downloads panel with progress + cancel), @@ -341,44 +341,26 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): # ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ─── -def test_force_download_bypasses_cache_first_early_return(monkeypatch): - """force_download=True skips the try_to_load_from_cache early-return and - proceeds to the real download path; force_download=False returns the cached - path without ever attempting a download (X2/F2).""" - import huggingface_hub as hf +def test_force_download_is_forwarded_through_the_shim(monkeypatch): + """The shim's contract is to forward force_download unchanged to the shared helper (which owns the + cache-first early-return and bypass). Verify both False and True reach it (X2/F2).""" import utils.hf_xet_fallback as X - cached_path = "/cache/blob/cached.gguf" + seen = [] - # Pretend the blob IS cached on disk (try_to_load_from_cache is imported - # inside the function from huggingface_hub, and os.path.exists must agree). - monkeypatch.setattr(hf, "try_to_load_from_cache", lambda *a, **k: cached_path, raising = False) - monkeypatch.setattr(X.os.path, "exists", lambda p: True, raising = False) + def fake_shared(repo_id, filename, token, **kwargs): + seen.append(kwargs.get("force_download")) + return "/downloaded/path" - attempts = [] + monkeypatch.setattr(X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True) - def fake_attempt(repo_id, filename, token, **kwargs): - attempts.append( - {"repo_id": repo_id, "filename": filename, "force": kwargs.get("force_download")} - ) - return ("ok", "/freshly/downloaded/path") - - monkeypatch.setattr(X, "_run_download_attempt", fake_attempt, raising = True) - - # force_download=False: cache-first early-return, no download attempt. - out = X.hf_hub_download_with_xet_fallback( + X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = False ) - assert out == cached_path - assert attempts == [] # never reached the real download - - # force_download=True: bypass the early-return, run the real download. - out2 = X.hf_hub_download_with_xet_fallback( + X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = True ) - assert out2 == "/freshly/downloaded/path" - assert len(attempts) == 1 - assert attempts[0]["force"] is True + assert seen == [False, True] # the shim forwards force_download to the shared helper unchanged # ── multi-revision GGUF blob comparison and update reclaim ── diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 15961ac03a..2dd2247396 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -1,341 +1,204 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall. +"""Studio shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback. -Xet (``hf_xet``) is the fast default but can hang with no progress and no -exception, and a blocked native thread cannot be killed. Keep Xet primary; fall -back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET`` -is read at import time, so the fallback runs in a fresh ``spawn`` child (not a -thread) that sets the env before importing ``huggingface_hub``. Cached files -short-circuit with no child; deterministic errors (401/403/404/disk-full) and -cancellation propagate without a fallback. Mirrors the safetensors inference -recovery in core/inference/{orchestrator,worker}.py. +Re-exports the shared API and injects Studio's marker-aware cache purge +(``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` +marker semantics on the HTTP retry. """ from __future__ import annotations -import multiprocessing as mp -import os -import queue -import signal -import sys import threading -import time from typing import Any, Callable, Optional -from loggers import get_logger +_shared_import_error = None +try: + import unsloth_zoo.hf_xet_fallback as _shared + _shared_available = True +except Exception as _exc: # noqa: BLE001 - any import failure must degrade, not crash + # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less Studio + # host. The download helper needs none of it, so retry via the light UNSLOTH_ZOO_DISABLE_GPU_INIT + # path before giving up. + _shared_import_error = _exc + import os as _os -logger = get_logger(__name__) - -_CTX = mp.get_context("spawn") - -# Defaults match the existing inference watchdog and hub shutdown deadline. -DEFAULT_HEARTBEAT_INTERVAL = 30.0 -DEFAULT_STALL_TIMEOUT = 180.0 -DEFAULT_GRACE_PERIOD = 10.0 -_POLL_INTERVAL = 0.5 - - -class DownloadStallError(RuntimeError): - """Raised when no download progress is observed for too long. - - Canonical home; orchestrator.py re-imports it so all paths share one type. - """ - - -def child_should_disable_xet(config: dict) -> bool: - """Single source of truth for the per-worker Xet env flip.""" - return bool(config.get("disable_xet")) - - -def get_hf_download_state( - repo_ids: Optional[list[str]] = None, *, repo_type: str = "model" -) -> Optional[tuple[int, bool]]: - """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache. - - Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete`` - is not mistaken for full-size progress. ``None`` means the state could not be - measured, so callers skip stall logic for that tick. - """ + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" try: - from hub.utils.hf_cache_state import ( - blob_bytes_present, - has_active_incomplete_blobs, - hf_cache_root, - iter_active_repo_cache_dirs, - ) + import unsloth_zoo.hf_xet_fallback as _shared + _shared_available = True + _shared_import_error = None + except Exception as _exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF downloads + _shared_import_error = _exc2 + _shared_available = False + finally: + if _prev_gpu_init is None: + _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) + else: + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init - if hf_cache_root() is None: - return (0, False) +if _shared_available: + # Bind by assignment so each public name shares one module-level binding with the degraded branch. + DEFAULT_GRACE_PERIOD = _shared.DEFAULT_GRACE_PERIOD + DEFAULT_HEARTBEAT_INTERVAL = _shared.DEFAULT_HEARTBEAT_INTERVAL + DEFAULT_STALL_TIMEOUT = _shared.DEFAULT_STALL_TIMEOUT + DownloadStallError = _shared.DownloadStallError + child_should_disable_xet = _shared.child_should_disable_xet + get_hf_download_state = _shared.get_hf_download_state + start_watchdog = _shared.start_watchdog + _shared_hf_hub_download_with_xet_fallback = _shared.hf_hub_download_with_xet_fallback + _shared_snapshot_download_with_xet_fallback = _shared.snapshot_download_with_xet_fallback +else: + # Degrade instead of crashing Studio: plain HF downloads, stall watchdog disabled. Thin stubs, + # not a second copy of the orchestration; recovery returns once unsloth_zoo is upgraded. + import logging as _logging - total = 0 - has_incomplete = False - for repo_id in repo_ids or []: - # Skip local paths: HF IDs never start with / . ~ or contain "\". - if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id: - continue - for entry in iter_active_repo_cache_dirs(repo_type, repo_id): - blobs_dir = entry / "blobs" - if not blobs_dir.is_dir(): - continue - for blob in blobs_dir.iterdir(): - try: - if blob.is_file(): - total += blob_bytes_present(blob) - except OSError: - pass - if has_active_incomplete_blobs(repo_type, repo_id): - has_incomplete = True - return (total, has_incomplete) - except Exception as e: - logger.debug("Failed to determine HF download state: %s", e) - return None + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " + "re-enable automatic Xet -> HTTP download recovery.", + _shared_import_error, + ) + DEFAULT_HEARTBEAT_INTERVAL = 30.0 + DEFAULT_STALL_TIMEOUT = 180.0 + DEFAULT_GRACE_PERIOD = 10.0 -def start_watchdog( - *, - repo_ids: list[str], - on_stall: Callable[[str], None], - repo_type: str = "model", - interval: float = DEFAULT_HEARTBEAT_INTERVAL, - stall_timeout: float = DEFAULT_STALL_TIMEOUT, - xet_disabled: bool = False, - on_heartbeat: Optional[Callable[[str], None]] = None, -) -> threading.Event: - """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a - ``*.incomplete`` is present AND the on-disk size is unchanged for - *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so - post-download init is never misread as a stall. Returns a stop event the - caller sets when the download phase ends. - """ - stop = threading.Event() - transport = "https" if xet_disabled else "xet" - fired = False + class DownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" - def _beat() -> None: - nonlocal fired - state = get_hf_download_state(repo_ids, repo_type = repo_type) - last_size = state[0] if state is not None else 0 - last_change = time.monotonic() + def child_should_disable_xet(config: dict) -> bool: + return bool(config.get("disable_xet")) - while not stop.wait(interval): - state = get_hf_download_state(repo_ids, repo_type = repo_type) - now = time.monotonic() + def get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires - if state is None: - if on_heartbeat is not None: + def start_watchdog( + *, + on_heartbeat: "Optional[Callable[[str], None]]" = None, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + xet_disabled: bool = False, + **kwargs: Any, + ) -> "threading.Event": + # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline + # is not tripped during a long download. + stop = threading.Event() + if on_heartbeat is None: + return stop + transport = "https" if xet_disabled else "xet" + + def _beat() -> None: + while not stop.wait(interval): + try: on_heartbeat(f"Downloading ({transport} transport)...") - continue + except Exception: + pass - current_size, has_incomplete = state - if current_size != last_size: - last_size = current_size - last_change = now + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() + return stop - # Reset unless .incomplete confirms an active download, so model init - # and lock waits are not counted as a stall. - if not has_incomplete: - last_change = now - elif now - last_change >= stall_timeout: - if not fired: - fired = True - on_stall( - f"Download appears stalled ({transport} transport) " - f"-- no progress for {int(now - last_change)}s" - ) - return + def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() - if on_heartbeat is not None: - on_heartbeat(f"Downloading ({transport} transport)...") + def _shared_hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + repo_type: str = "model", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, + ) -> str: + # Keep the cancellation contract: do not start or return a download once cancelled. + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") - threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start() - return stop - - -def _download_child_entry( - *, - repo_id: str, - filename: str, - token: Optional[str], - repo_type: str, - disable_xet: bool, - result_queue: Any, - force_download: bool = False, -) -> None: - """Spawn-child entrypoint: download one file and report the result. - - Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub, - forms its own process group so the parent can kill the whole transfer, and - never logs the token or signed URLs. - """ - # Die with Studio on Linux (this mp child gets no parent-set preexec_fn). - try: - from utils.process_lifetime import bind_current_process_to_parent_lifetime - bind_current_process_to_parent_lifetime() - except Exception: - pass - - if hasattr(os, "setsid"): - try: - os.setsid() - except OSError: - pass - - if disable_xet: - os.environ["HF_HUB_DISABLE_XET"] = "1" - # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse - # partials a sequential resume cannot safely continue). - os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" - os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") - - # Test-only fault injection (never set in production): stall the Xet attempt - # so the watchdog + HTTP fallback can be exercised against a real repo. - if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1": - import time as _t - try: - from huggingface_hub.constants import HF_HUB_CACHE - - blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs") - os.makedirs(blobs, exist_ok = True) - with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh: - fh.write(b"\0" * 4096) - except OSError: - pass - while True: - _t.sleep(3600) - - try: from huggingface_hub import hf_hub_download + path = hf_hub_download( repo_id = repo_id, filename = filename, - repo_type = repo_type, token = token, + repo_type = repo_type, + revision = revision, + cache_dir = cache_dir, force_download = force_download, ) - result_queue.put({"ok": True, "path": path}) - except BaseException as e: # noqa: BLE001 - report every failure to the parent - error = f"{type(e).__name__}: {e}" - try: - from hub.utils.download_registry import scrub_secrets - error = scrub_secrets(error, hf_token = token) - except Exception: - pass - result_queue.put({"ok": False, "error": error}) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + def _shared_snapshot_download_with_xet_fallback( + repo_id: str, + *, + revision: Optional[str] = None, + token: Optional[str] = None, + repo_type: str = "model", + cache_dir: Optional[str] = None, + allow_patterns: Optional[Any] = None, + ignore_patterns: Optional[Any] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, + ) -> str: + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") -def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None: - """Kill *proc* and its whole process group (Xet may spawn helper procs). + from huggingface_hub import snapshot_download - The child calls ``os.setsid()`` so its pgid equals its pid; signal via - ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a - group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*. - """ - pid = proc.pid - - def _signal_group(sig: int) -> None: - if pid is not None and hasattr(os, "killpg"): - try: - os.killpg(pid, sig) - return - except (ProcessLookupError, PermissionError, OSError): - pass - # Windows or pre-setsid: best effort on the single process. - try: - proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill() - except Exception: - pass - - _signal_group(getattr(signal, "SIGTERM", signal.SIGINT)) - proc.join(timeout = grace_period) - if proc.is_alive(): - _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM)) - proc.join(timeout = 5.0) - - -def _run_download_attempt( - repo_id: str, - filename: str, - token: Optional[str], - *, - repo_type: str, - disable_xet: bool, - cancel_event: Optional[threading.Event], - stall_timeout: float, - interval: float, - grace_period: float, - on_status: Optional[Callable[[str], None]], - force_download: bool = False, -) -> tuple[str, Optional[str]]: - """Run one download in a spawn child supervised by the no-progress watchdog. - - Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or - ``("error", message)``. This is the seam tests monkeypatch to avoid spawning. - """ - result_queue: Any = _CTX.Queue() - proc = _CTX.Process( - target = _download_child_entry, - kwargs = dict( + path = snapshot_download( repo_id = repo_id, - filename = filename, - token = token, repo_type = repo_type, - disable_xet = disable_xet, - result_queue = result_queue, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, force_download = force_download, - ), - daemon = True, - ) - proc.start() - from utils.process_lifetime import adopt_pid - - adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) - - stalled = threading.Event() - stop_watchdog = start_watchdog( - repo_ids = [repo_id], - on_stall = lambda msg: stalled.set(), - repo_type = repo_type, - interval = interval, - stall_timeout = stall_timeout, - xet_disabled = disable_xet, - on_heartbeat = on_status, - ) - - result: Optional[dict] = None - try: - while proc.is_alive(): - if cancel_event is not None and cancel_event.is_set(): - _terminate_process_group(proc, grace_period) - return ("cancelled", None) - if stalled.is_set(): - _terminate_process_group(proc, grace_period) - return ("stall", None) - try: - result = result_queue.get(timeout = _POLL_INTERVAL) - break - except queue.Empty: - continue - else: - # Process exited; drain any result it enqueued. - try: - result = result_queue.get_nowait() - except queue.Empty: - result = None - finally: - stop_watchdog.set() - proc.join(timeout = grace_period) - - if result is None: - return ( - "error", - f"download process for '{repo_id}/{filename}' exited " - f"(code={proc.exitcode}) without a result", ) - if result.get("ok"): - return ("ok", result["path"]) - return ("error", result.get("error") or "unknown download error") + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +__all__ = [ + "DEFAULT_GRACE_PERIOD", + "DEFAULT_HEARTBEAT_INTERVAL", + "DEFAULT_STALL_TIMEOUT", + "DownloadStallError", + "child_should_disable_xet", + "get_hf_download_state", + "start_watchdog", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", +] + + +def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None: + """Studio's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport`` + accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged, + not fatal to the retry.""" + try: + from hub.utils.download_registry import prepare_cache_for_transport + prepare_cache_for_transport(repo_type, repo_id, "http") + except Exception as exc: + try: + from loggers import get_logger + get_logger(__name__).debug( + "Studio prepare_cache_for_transport failed for %s: %s", repo_id, exc + ) + except ModuleNotFoundError as logger_exc: + if logger_exc.name != "loggers": + raise def hf_hub_download_with_xet_fallback( @@ -345,83 +208,32 @@ def hf_hub_download_with_xet_fallback( *, cancel_event: Optional[threading.Event] = None, repo_type: str = "model", + revision: Optional[str] = None, stall_timeout: float = DEFAULT_STALL_TIMEOUT, interval: float = DEFAULT_HEARTBEAT_INTERVAL, grace_period: float = DEFAULT_GRACE_PERIOD, on_status: Optional[Callable[[str], None]] = None, force_download: bool = False, ) -> str: - """Download a single file with Xet primary and HTTP as a stall-only fallback. + """Single-file download via the shared fallback with Studio's marker-aware HTTP-retry prep. + ``force_download`` re-fetches a newer blob over a cached one (Studio's model-update path).""" + return _shared_hf_hub_download_with_xet_fallback( + repo_id, + filename, + token, + cancel_event = cancel_event, + repo_type = repo_type, + revision = revision, + stall_timeout = stall_timeout, + interval = interval, + grace_period = grace_period, + on_status = on_status, + force_download = force_download, + prepare_for_http_fn = _studio_prepare_for_http, + ) - Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if - *cancel_event* is set, re-raises a deterministic child error unchanged (no - fallback), and raises ``DownloadStallError`` only if BOTH transports stall. - When *force_download* is True the cache-first early-return is skipped and the - flag is threaded to ``hf_hub_download`` so a newer remote blob is re-fetched - even if an older blob is already cached. - """ - # Finalized blob already cached: return it with no child and no network. - # Skipped when force_download is set so an update re-fetches a newer blob. - if not force_download: - try: - from huggingface_hub import try_to_load_from_cache - cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) - if isinstance(cached, str) and os.path.exists(cached): - return cached - except Exception as e: - logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) - - if cancel_event is not None and cancel_event.is_set(): - raise RuntimeError("Cancelled") - - disable_xet = False - for attempt in range(2): - if disable_xet: - # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume - # over a sparse Xet/hf_transfer partial silently corrupts the blob. - try: - from hub.utils.download_registry import prepare_cache_for_transport - prepare_cache_for_transport(repo_type, repo_id, "http") - except Exception as e: - logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e) - - kind, payload = _run_download_attempt( - repo_id, - filename, - token, - repo_type = repo_type, - disable_xet = disable_xet, - cancel_event = cancel_event, - stall_timeout = stall_timeout, - interval = interval, - grace_period = grace_period, - on_status = on_status, - force_download = force_download, - ) - - if kind == "ok": - return payload # type: ignore[return-value] - if kind == "cancelled": - raise RuntimeError("Cancelled") - if kind == "error": - # Deterministic failure: the other transport would fail identically. - raise RuntimeError(payload) - # kind == "stall" - if attempt == 0 and not disable_xet: - logger.warning( - "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1", - repo_id, - filename, - ) - if on_status is not None: - on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP") - disable_xet = True - continue - raise DownloadStallError( - f"Download stalled for '{repo_id}/{filename}' even with " - f"HF_HUB_DISABLE_XET=1 -- check your network connection" - ) - - # Unreachable: the loop either returns or raises on each attempt. - raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'") +def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str: + """Whole-repo download via the shared fallback with Studio's marker-aware HTTP-retry prep.""" + kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http) + return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx index d00c812325..db7628777a 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx @@ -42,10 +42,8 @@ export function ModelUpdateAction({ }: ModelUpdateActionProps) { const [open, setOpen] = useState(false); - // The update is a managed download (it surfaces in the global Downloads panel - // with progress + cancel). When this exact repo+variant finishes, refresh the - // caller so the "update available" cue clears once the new revision is on - // disk. A ref keeps the subscription stable across renders without resubscribing. + // Refresh the caller when this repo+variant's download finishes so the "update available" cue + // clears. A ref keeps the subscription stable across renders. const onUpdatedRef = useRef(onUpdated); onUpdatedRef.current = onUpdated; useEffect(() => { @@ -60,9 +58,8 @@ export function ModelUpdateAction({ }, [repoId, variant]); const handleConfirm = useCallback(() => { - // Start the background re-download and close the dialog immediately; the - // Downloads panel owns progress + cancel from here. Only a failure to START - // surfaces a toast — a failed download reports itself in the panel. + // Start the re-download and close the dialog; the Downloads panel owns progress + cancel. + // Only a failure to START toasts (a failed download shows in the panel). void Promise.resolve() .then(onConfirm) .catch((err) => { diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 20000c82ee..9890c5f574 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1247,11 +1247,8 @@ export function HubModelPicker({ onEject?: () => void; }) { const gpu = useGpuInfo(); - // The currently-loaded/running model id. We read params.checkpoint from the - // runtime store (backend-mirrored from /api/inference/status.active_model, see - // chat-runtime-store) rather than the dropdown `isSelected` highlight (which is - // just `value === repo_id` and can reflect a staged, not-yet-loaded pick). Used - // to disable the cached-row update action for the model that's live in memory. + // Live model id from the runtime store (backend-mirrored active_model), not the dropdown + // highlight which can be a staged pick. Disables the update action for it. const loadedModelId = useChatRuntimeStore((s) => s.params.checkpoint); // Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date). const loadTimes = useModelLoadTimes(value); @@ -1589,11 +1586,8 @@ export function HubModelPicker({ refreshLocalModelsList(); }, [hfToken, refreshLocalModelsList]); - // Updates run as MANAGED downloads (they show in the global Downloads panel - // with manifest-based progress + a working Cancel), instead of a blocking - // call. The worker re-resolves `main` and pulls only changed blobs, so the - // cached copy stays usable until the new revision lands. The row's - // ModelUpdateAction refreshes the list when this repo+variant completes. + // Updates run as managed downloads (Downloads panel: progress + Cancel), not a blocking + // call. The worker pulls only changed blobs, so the cached copy stays usable until done. const startManagedUpdate = useCallback((repoId: string, variant: string, expectedBytes: number) => { return downloadManager .requestStart({ diff --git a/tests/test_prefetch_snapshot_scope.py b/tests/test_prefetch_snapshot_scope.py new file mode 100644 index 0000000000..c7ec4f2c34 --- /dev/null +++ b/tests/test_prefetch_snapshot_scope.py @@ -0,0 +1,916 @@ +# Unsloth Zoo - Utilities for Unsloth +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""Pure-CPU, no-network unit tests for prefetch snapshot scoping in unsloth/models/_utils.py. + +maybe_prefetch_hf_snapshot warms the HF cache before the in-process load. The warm must cover at +least what the load reads (else the missing file falls to an unprotected in-process Xet fetch) but +not pull weights the load never reads. These tests lock the allow/ignore patterns each mode hands +snapshot_download_with_xet_fallback. The zoo downloader is monkeypatched to capture its kwargs. +""" + +import fnmatch +import sys +import types + +import pytest + +from unsloth.models import _utils as U + + +def _filter(names, allow_patterns, ignore_patterns): + """Mirror HF filter_repo_objects: keep on allow match (or None), drop on ignore match.""" + kept = [] + for name in names: + if allow_patterns is not None and not any(fnmatch.fnmatch(name, p) for p in allow_patterns): + continue + if ignore_patterns and any(fnmatch.fnmatch(name, p) for p in ignore_patterns): + continue + kept.append(name) + return kept + + +@pytest.fixture +def capture(monkeypatch): + """Run maybe_prefetch_hf_snapshot with a fake repo, capturing the patterns forwarded to a + fake injected zoo downloader (independent of the installed unsloth_zoo). Offline env cleared.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + + state = {} + + def fake_download(repo_id, **kw): + state["repo_id"] = repo_id + state["allow_patterns"] = kw.get("allow_patterns") + state["ignore_patterns"] = kw.get("ignore_patterns") + state["variant"] = kw.get("variant") + return "/tmp/fake-snapshot" + + fake_module = types.ModuleType("unsloth_zoo.hf_xet_fallback") + fake_module.snapshot_download_with_xet_fallback = fake_download + fake_module.DownloadStallError = type("DownloadStallError", (RuntimeError,), {}) + monkeypatch.setitem(sys.modules, "unsloth_zoo.hf_xet_fallback", fake_module) + + # Neutralize the model_info network call by default; tests exercising format selection + # install their own. + import huggingface_hub + + class _NoNetworkApi: + def model_info(self, *a, **k): + raise RuntimeError("no network in test") + + monkeypatch.setattr(huggingface_hub, "HfApi", _NoNetworkApi) + + def run(**call_kwargs): + state.clear() + ok = U.maybe_prefetch_hf_snapshot("some-org/some-repo", **call_kwargs) + return ok, state + + return run + + +# Representative repo listing: root weights + aux, subdir, adapter, checkpoint, merged weights. +_SAMPLE_FILES = [ + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + "model.safetensors.index.json", + "pytorch_model.bin", + "fp16/model.safetensors", + "experimental/model-00001-of-00002.safetensors", + "checkpoint-500/model.safetensors", + "adapter_config.json", + "adapter_model.safetensors", +] + + +def test_weights_at_root_excludes_subdir_weights(capture): + """A root load ignores subdir weights (fp16/, experimental/, checkpoint-500/) but keeps root weights.""" + ok, st = capture(weights_at_root = True, use_safetensors = True) + assert ok is True + assert st["allow_patterns"] is None + ig = st["ignore_patterns"] + assert "*/*.safetensors" in ig and "*/*.bin" in ig + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], ig) + assert "model-00001-of-00002.safetensors" in kept + assert "model.safetensors.index.json" in kept + assert "config.json" in kept + assert "fp16/model.safetensors" not in kept + assert "experimental/model-00001-of-00002.safetensors" not in kept + assert "checkpoint-500/model.safetensors" not in kept + + +def test_adapter_only_excludes_merged_weights(capture): + """An adapter warm keeps adapter files + root aux, not merged full-model weights.""" + ok, st = capture(adapter_only = True) + assert ok is True + assert st["ignore_patterns"] is None + allow = st["allow_patterns"] + assert "adapter_config.json" in allow and "adapter_model*" in allow + kept = _filter(_SAMPLE_FILES, allow, st["ignore_patterns"]) + assert "adapter_config.json" in kept + assert "adapter_model.safetensors" in kept + assert "config.json" in kept and "tokenizer.json" in kept + assert "model-00001-of-00002.safetensors" not in kept + assert "pytorch_model.bin" not in kept + assert "fp16/model.safetensors" not in kept + + +def test_adapter_only_warms_sharded_adapter(capture): + """A sharded adapter is still covered by the adapter_model* glob.""" + _, st = capture(adapter_only = True) + sharded = [ + "adapter_config.json", + "adapter_model-00001-of-00002.safetensors", + "adapter_model-00002-of-00002.safetensors", + "adapter_model.safetensors.index.json", + ] + kept = _filter(sharded, st["allow_patterns"], st["ignore_patterns"]) + assert set(kept) == set(sharded) + + +def test_tokenizer_only_warms_only_aux_files(capture): + """A tokenizer-only repo warms tokenizer/config/vocab files, never weights.""" + _, st = capture(tokenizer_only = True) + assert st["ignore_patterns"] is None + assert st["allow_patterns"] == list(U._ROOT_AUX_PREFETCH_PATTERNS) + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], st["ignore_patterns"]) + assert "tokenizer.json" in kept and "config.json" in kept + assert "model-00001-of-00002.safetensors" not in kept + assert "adapter_model.safetensors" not in kept + + +def test_aux_warm_covers_arbitrary_remote_code_modules(capture): + """The aux warm must cover any *.py, since trust_remote_code auto_map names modules freely.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + assert "*.py" in allow + remote_code = [ + "config.json", + "modeling.py", + "tokenization.py", + "my_custom_code.py", + "configuration_foo.py", + ] + kept = _filter(remote_code, allow, st["ignore_patterns"]) + for name in ("modeling.py", "tokenization.py", "my_custom_code.py", "configuration_foo.py"): + assert name in kept, name + + +def test_subfolder_warms_subfolder_plus_root_aux(capture): + """A subfolder load warms that subfolder's weights plus root aux; other subdirs/root weights skipped.""" + _, st = capture(subfolder = "fp16") + allow = st["allow_patterns"] + assert "fp16/*" in allow + assert all(p in allow for p in U._ROOT_AUX_PREFETCH_PATTERNS) + kept = _filter(_SAMPLE_FILES, allow, st["ignore_patterns"]) + assert "fp16/model.safetensors" in kept + assert "config.json" in kept + assert "experimental/model-00001-of-00002.safetensors" not in kept + + +def test_subfolder_takes_precedence_over_weights_at_root(capture): + """When a subfolder is requested the subfolder branch wins over weights_at_root.""" + _, st = capture(subfolder = "fp16", weights_at_root = True) + assert "fp16/*" in st["allow_patterns"] + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], st["ignore_patterns"]) + assert "fp16/model.safetensors" in kept + + +def test_local_dir_is_not_warmed(capture, tmp_path): + """A local directory path skips the warm (returns False).""" + d = tmp_path / "local-model" + d.mkdir() + ok = U.maybe_prefetch_hf_snapshot(str(d), weights_at_root = True) + assert ok is False + + +def _install_fake_model_info(monkeypatch, filenames): + """Make HfApi().model_info(...).siblings report filenames, with no network.""" + import huggingface_hub + + class _Sib: + def __init__(self, name): + self.rfilename = name + + class _Info: + def __init__(self, names): + self.siblings = [_Sib(n) for n in names] + + class _Api: + def model_info(self, *a, **k): + return _Info(filenames) + + monkeypatch.setattr(huggingface_hub, "HfApi", _Api) + + +# ----- Finding P: variant-aware weight-format selection ----- + + +def test_variant_keeps_bin_when_only_default_safetensors(monkeypatch): + """A default model.safetensors must not prove a variant .bin redundant; without a variant it does.""" + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.fp16.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" not in ig + ig_default = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig_default + + +def test_variant_drops_bin_when_variant_safetensors_present(monkeypatch): + """A variant-matching safetensors makes the variant .bin redundant, so .bin is dropped.""" + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.fp16.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig + + +def test_no_variant_keeps_bin_when_only_variant_safetensors(monkeypatch): + """For a no-variant load, only a canonical safetensors (not a lone variant) makes .bin redundant.""" + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" not in ig + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.bin"]) + ig2 = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig2 + + +def test_variant_keeps_bin_for_noncanonical_sidecar(monkeypatch): + """A non-canonical variant sidecar must not prove the variant .bin redundant; a canonical one does.""" + _install_fake_model_info( + monkeypatch, ["consolidated.fp16.safetensors", "pytorch_model.fp16.bin"] + ) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" not in ig + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.fp16.bin"]) + ig2 = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig2 + + +def test_is_canonical_model_weight_safetensors(): + """The canonical detector matches only non-variant model-weight safetensors names.""" + assert U._is_canonical_model_weight_safetensors("model.safetensors") is True + assert U._is_canonical_model_weight_safetensors("model-00001-of-00002.safetensors") is True + assert U._is_canonical_model_weight_safetensors("model.safetensors.index.json") is True + assert U._is_canonical_model_weight_safetensors("model.fp16.safetensors") is False + assert ( + U._is_canonical_model_weight_safetensors("model.fp16-00001-of-00002.safetensors") is False + ) + assert U._is_canonical_model_weight_safetensors("adapter_model.safetensors") is False + + +def test_st_prefetch_resolves_env_cache_and_runs_after_validation(): + """The ST prefetch must resolve SENTENCE_TRANSFORMERS_HOME and run after load-mode validation.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + prefetch_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ] + assert len(prefetch_calls) == 1, "expected exactly one ST prefetch call" + call = prefetch_calls[0] + # cache_dir kwarg resolves SENTENCE_TRANSFORMERS_HOME. + cache_dir_kw = next((kw for kw in call.keywords if kw.arg == "cache_dir"), None) + assert cache_dir_kw is not None, "ST prefetch must pass cache_dir" + assert "SENTENCE_TRANSFORMERS_HOME" in ast.dump( + cache_dir_kw.value + ), "ST prefetch cache_dir must resolve SENTENCE_TRANSFORMERS_HOME" + # Load-mode validation runs before the prefetch (fewer source lines = earlier). + val_lineno = src[: src.index("Can only load in 4bit or 8bit or 16bit")].count("\n") + assert val_lineno < call.lineno, "load-mode validation must precede the ST prefetch" + + +def test_st_cache_resolutions_honor_explicit_hf_cache_dir(): + """Every ST cache resolution falling back to SENTENCE_TRANSFORMERS_HOME must first honor an explicit HF cache_dir.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + resolutions = [ + kw + for kw in ast.walk(tree) + if isinstance(kw, ast.keyword) + and kw.arg == "cache_dir" + and "SENTENCE_TRANSFORMERS_HOME" in ast.dump(kw.value) + ] + assert resolutions, "expected cache_dir resolutions referencing SENTENCE_TRANSFORMERS_HOME" + for kw in resolutions: + assert "'cache_dir'" in ast.dump( + kw.value + ), "an ST cache_dir resolution must read an explicit kwargs.get('cache_dir') first" + + +def test_st_native_loads_map_hf_cache_dir_to_cache_folder(): + """Native SentenceTransformer loads take cache_folder, so an explicit HF cache_dir must be mapped onto it.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + # Every native SentenceTransformer(...) forwarding cache_folder must read cache_dir. + st_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "SentenceTransformer" + ] + cache_folder_kws = [kw for call in st_calls for kw in call.keywords if kw.arg == "cache_folder"] + assert cache_folder_kws, "expected a native SentenceTransformer call forwarding cache_folder" + for kw in cache_folder_kws: + assert "'cache_dir'" in ast.dump( + kw.value + ), "a native SentenceTransformer cache_folder must map the explicit HF cache_dir first" + # for_inference feeds cache_folder via st_kwargs; both native branches map cache_dir -> cache_folder. + normalized = "".join(src.split()) + assert ( + 'st_kwargs["cache_folder"]=' in normalized + ), "for_inference must set st_kwargs cache_folder" + assert ( + normalized.count('kwargs.get("cache_dir")orkwargs.get("cache_folder")') >= 2 + ), "both native ST branches (for_inference, fast-encoder) must map cache_dir -> cache_folder" + + +def test_vision_warms_vllm_tokenizer_after_remap(): + """On the vLLM path the tokenizer warm is deferred until after the fast_inference_setup remap.""" + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "vision.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + guard = "if _vllm_owns_weights and isinstance(tokenizer_name" + assert guard in src, "expected a vLLM-gated tokenizer warm" + assert src.index(guard) > src.index( + "fast_inference_setup(" + ), "the vLLM tokenizer warm must run after the fast_inference_setup remap" + + +def test_diffusion_forwards_variant_to_real_load(): + """FastDiffusionModel must forward variant to the real model_cls.from_pretrained load, not just the prefetch.""" + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "diffusion.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + assert ( + 'load_kwargs["variant"] = kwargs["variant"]' in src + ), "the diffusion load must forward variant to model_cls.from_pretrained" + + +def test_vision_prefetch_runs_after_load_mode_validation(): + """The FastBaseModel (vision) prefetch must run after the load-mode validation.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "vision.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + prefetch_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ] + assert prefetch_calls, "expected a vision prefetch call" + first_prefetch = min(call.lineno for call in prefetch_calls) + val_lineno = src[: src.index("Can only load in 4bit or 8bit or 16bit")].count("\n") + assert val_lineno < first_prefetch, "load-mode validation must precede the vision prefetch" + + +def test_llama_prefetch_skips_only_real_vllm_loads(): + """The llama prefetch's fast_inference skip must be gated on num_labels is None (a classification load still downloads).""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "llama.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + gated = False + for n in ast.walk(tree): + if not ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ): + continue + fi_kw = next((kw for kw in n.keywords if kw.arg == "fast_inference"), None) + if fi_kw is None: + continue + dumped = ast.dump(fi_kw.value) + if "fast_inference" in dumped and "num_labels" in dumped: + gated = True + assert gated, "llama prefetch fast_inference must be gated on num_labels is None" + + +def test_st_fallback_module_loads_resolve_env_cache(): + """Fallback module loads deriving cache_dir from cache_folder must also fall back to SENTENCE_TRANSFORMERS_HOME.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + + # Fallback sites (cache_dir derived from cache_folder) must resolve SENTENCE_TRANSFORMERS_HOME. + checked = 0 + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_load_modules"): + continue + cache_dir_kw = next((kw for kw in node.keywords if kw.arg == "cache_dir"), None) + if cache_dir_kw is None: + continue + dumped = ast.dump(cache_dir_kw.value) + if "cache_folder" not in dumped: + continue # internal pass-through, not a resolution site + checked += 1 + assert ( + "SENTENCE_TRANSFORMERS_HOME" in dumped + ), f"{node.func.attr} cache_dir resolves cache_folder but not SENTENCE_TRANSFORMERS_HOME" + assert ( + checked >= 2 + ), "expected the fallback _module_path and _load_modules calls to resolve the env cache" + + +def test_st_fallback_module_loads_forward_revision(): + """The fallback module loads must forward revision so module files match the revision-pinned weights. + Guards: (a) helpers accept revision, (b) every download primitive forwards it, (c) _load_modules + threads it into internal calls, (d) the from_pretrained fallback sites forward it.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + funcs = { + n.name: n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) + and n.name in ("_module_path", "_read_pooling_mode", "_load_modules") + } + assert set(funcs) == {"_module_path", "_read_pooling_mode", "_load_modules"} + + # (a) each helper takes a revision parameter. + for name, fn in funcs.items(): + arg_names = {a.arg for a in fn.args.args + fn.args.kwonlyargs} + assert "revision" in arg_names, f"{name} must accept a revision argument" + + # (b) every download primitive inside the helpers forwards revision. + downloads = 0 + for name, fn in funcs.items(): + for node in ast.walk(fn): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)): + continue + if node.func.id not in ("hf_hub_download", "load_dir_path"): + continue + downloads += 1 + assert any( + kw.arg == "revision" for kw in node.keywords + ), f"{node.func.id} in {name} must forward revision" + assert downloads >= 3, "expected the module-download primitives to be revision-guarded" + + # (c) _load_modules threads revision into its internal _module_path / _read_pooling_mode calls. + internal = 0 + for node in ast.walk(funcs["_load_modules"]): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_read_pooling_mode"): + continue + internal += 1 + assert any( + kw.arg == "revision" for kw in node.keywords + ), f"_load_modules must forward revision to {node.func.attr}" + assert internal >= 2, "expected _load_modules to call _module_path and _read_pooling_mode" + + # (d) the from_pretrained fallback _module_path / _load_modules sites forward revision. + checked = 0 + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_load_modules"): + continue + cache_dir_kw = next((kw for kw in node.keywords if kw.arg == "cache_dir"), None) + if cache_dir_kw is None or "cache_folder" not in ast.dump(cache_dir_kw.value): + continue # internal pass-through, not a fallback site + checked += 1 + rev_kw = next((kw for kw in node.keywords if kw.arg == "revision"), None) + assert rev_kw is not None and "revision" in ast.dump( + rev_kw.value + ), f"{node.func.attr} fallback call must forward revision" + assert ( + checked >= 2 + ), "expected the fallback _module_path and _load_modules calls to forward revision" + + +def test_st_fallback_model_load_resolves_env_cache(): + """from_pretrained must resolve the warmed ST cache into kwargs['cache_dir'] before the FastModel weight load.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + def _resolves_st_cache(value_node): + # Resolution may be inline or in the assignment to an intermediate variable the value references. + dumped = ast.dump(value_node) + if "cache_folder" in dumped and "SENTENCE_TRANSFORMERS_HOME" in dumped: + return True + if isinstance(value_node, ast.Name): + for n in ast.walk(tree): + if isinstance(n, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == value_node.id for t in n.targets + ): + d = ast.dump(n.value) + if "cache_folder" in d and "SENTENCE_TRANSFORMERS_HOME" in d: + return True + return False + + resolved_lines = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Assign): + continue + for tgt in node.targets: + if ( + isinstance(tgt, ast.Subscript) + and isinstance(tgt.value, ast.Name) + and tgt.value.id == "kwargs" + and isinstance(tgt.slice, ast.Constant) + and tgt.slice.value == "cache_dir" + and _resolves_st_cache(node.value) + ): + resolved_lines.append(node.lineno) + assert resolved_lines, "from_pretrained must resolve the ST cache into kwargs['cache_dir']" + + fastmodel_calls = [ + n.lineno + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and n.func.attr == "from_pretrained" + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "FastModel" + ] + assert fastmodel_calls, "expected a FastModel.from_pretrained call" + assert min(resolved_lines) < min( + fastmodel_calls + ), "kwargs['cache_dir'] must be resolved before the fallback FastModel weight load" + + +def test_canonical_variant_model_weight_matches_transformers_names(): + """The variant safetensors detector matches only canonical variant names, rejecting sidecars and wrong variants.""" + f = U._is_canonical_variant_model_weight_safetensors + assert f("model.fp16.safetensors", "fp16") is True + assert f("model.fp16-00001-of-00002.safetensors", "fp16") is True + assert f("model-00001-of-00002.fp16.safetensors", "fp16") is True + assert f("model.safetensors.index.fp16.json", "fp16") is True + assert f("consolidated.fp16.safetensors", "fp16") is False + assert f("model.safetensors", "fp16") is False + assert f("model-00001-of-00002.safetensors", "fp16") is False + assert f("model.bf16.safetensors", "fp16") is False + + +def test_variant_is_forwarded_to_downloader(capture): + """maybe_prefetch_hf_snapshot must forward variant to the downloader (absent a variant, nothing is forwarded).""" + _, st = capture(weights_at_root = True, use_safetensors = True, variant = "fp16") + assert st["variant"] == "fp16" + _, st = capture(weights_at_root = True, use_safetensors = True) + assert st["variant"] is None + + +def test_variant_drops_bin_for_sharded_variant_safetensors(monkeypatch): + """A sharded variant safetensors is recognized, so its redundant variant .bin is dropped.""" + _install_fake_model_info( + monkeypatch, + [ + "model.fp16-00001-of-00002.safetensors", + "model.fp16-00002-of-00002.safetensors", + "pytorch_model.fp16-00001-of-00002.bin", + ], + ) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig + + +def test_tokenizer_only_warms_extra_vocab_files(capture): + """tokenizer_only must warm SentencePiece / vocab / processor files, including a named jinja template.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + for name in ( + "spm.model", + "normalizer.json", + "video_preprocessor_config.json", + "tokenizer.model.v3", + ): + assert name in allow, name + sample = [ + "spm.model", + "normalizer.json", + "video_preprocessor_config.json", + "tokenizer.model.v3", + "additional_chat_templates/custom.jinja", + ] + kept = _filter(sample, allow, st["ignore_patterns"]) + assert set(kept) == set(sample) + + +def test_format_probe_runs_even_when_config_cached(capture, monkeypatch): + """A cached config.json must not skip the weight-format probe; model_info still drops the redundant .bin.""" + import huggingface_hub + + # Pretend config.json is cached (the AutoConfig side effect); this must not gate the probe. + monkeypatch.setattr( + huggingface_hub, "try_to_load_from_cache", lambda *a, **k: "/cache/config.json" + ) + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.bin"]) + _, st = capture(weights_at_root = True) + ig = st["ignore_patterns"] or [] + assert "*.bin" in ig + + +def test_optimizer_safetensors_does_not_drop_bin(monkeypatch): + """An optimizer.safetensors sidecar must not count as model safetensors, so the real .bin weights are kept.""" + _install_fake_model_info(monkeypatch, ["pytorch_model.bin", "optimizer.safetensors"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" not in ig + + +def test_model_safetensors_still_drops_bin(monkeypatch): + """Control for the optimizer case: a real model.safetensors next to pytorch_model.bin still drops the .bin.""" + _install_fake_model_info( + monkeypatch, ["model.safetensors", "pytorch_model.bin", "optimizer.safetensors"] + ) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig + + +def test_whole_multi_component_snapshot_keeps_subdir_bin(monkeypatch): + """A whole multi-component snapshot must not drop *.bin (it would strip a subdir module's weight); a root load still does.""" + _install_fake_model_info(monkeypatch, ["model.safetensors", "1_Dense/pytorch_model.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = False) + assert "*.bin" not in ig + ig_root = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig_root + + +def test_is_model_weight_safetensors_classification(): + """Real model weights count; adapter / trainer-state sidecars do not.""" + assert U._is_model_weight_safetensors("model.safetensors") is True + assert U._is_model_weight_safetensors("model-00001-of-00002.safetensors") is True + assert U._is_model_weight_safetensors("model.safetensors.index.json") is True + assert U._is_model_weight_safetensors("consolidated.safetensors") is True + assert U._is_model_weight_safetensors("adapter_model.safetensors") is False + assert U._is_model_weight_safetensors("optimizer.safetensors") is False + assert U._is_model_weight_safetensors("scheduler.safetensors") is False + assert U._is_model_weight_safetensors("rng_state_0.safetensors") is False + + +def test_tokenizer_only_warms_slow_sentencepiece_vocab(capture): + """tokenizer_only must warm the slow-tokenizer SentencePiece / BPE vocab files AutoTokenizer fetches first.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + for name in ( + "sentencepiece.bpe.model", + "source.spm", + "target.spm", + "bpe.codes", + "vocab.bpe", + "sentencepiece.model", + "vocab-src.json", + "vocab-tgt.json", + ): + assert name in allow, name + + +def test_adapter_safetensors_check_scoped_to_root(monkeypatch): + """_adapter_repo_has_safetensors must only count a root adapter_model*.safetensors, not a subdir one.""" + import huggingface_hub + + class _Sib: + def __init__(self, name): + self.rfilename = name + + class _Api: + def __init__(self, names): + self._names = names + + def model_info(self, *a, **k): + return type("MI", (), {"siblings": [_Sib(n) for n in self._names]})() + + # Subdir safetensors only -> not reported present. + monkeypatch.setattr( + huggingface_hub, + "HfApi", + lambda: _Api( + ["adapter_config.json", "adapter_model.bin", "checkpoint-5/adapter_model.safetensors"] + ), + ) + assert U._adapter_repo_has_safetensors("org/repo") is False + # Root safetensors -> reported present. + monkeypatch.setattr( + huggingface_hub, + "HfApi", + lambda: _Api(["adapter_config.json", "adapter_model.safetensors"]), + ) + assert U._adapter_repo_has_safetensors("org/repo") is True + + +def test_gguf_file_warm_keeps_gguf(capture): + """A gguf_file load allow-lists that GGUF while not pulling other quants the repo publishes.""" + _, st = capture(weights_at_root = True, gguf_file = "model-Q4_K_M.gguf") + allow = st["allow_patterns"] + ig = st["ignore_patterns"] + assert allow is not None and "model-Q4_K_M.gguf" in allow + sample = [ + "model-Q4_K_M.gguf", + "model-Q8_0.gguf", + "config.json", + "tokenizer.json", + ] + kept = _filter(sample, allow, ig) + assert "model-Q4_K_M.gguf" in kept + assert "config.json" in kept + assert "model-Q8_0.gguf" not in kept + + +# ----- Finding Q: adapter weight-format selection ----- + + +def test_adapter_only_prefers_safetensors_over_bin(capture, monkeypatch): + """A mixed-format adapter repo warms only the safetensors PeftModel reads, not both formats.""" + _install_fake_model_info( + monkeypatch, ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"] + ) + _, st = capture(adapter_only = True) + ig = st["ignore_patterns"] + assert ig is not None and "adapter_model*.bin" in ig + kept = _filter( + ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"], + st["allow_patterns"], + ig, + ) + assert "adapter_model.safetensors" in kept + assert "adapter_model.bin" not in kept + + +def test_adapter_only_bin_only_keeps_bin(capture, monkeypatch): + """A .bin-only adapter repo must keep adapter_model.bin (no safetensors found -> both formats eligible).""" + _install_fake_model_info(monkeypatch, ["adapter_config.json", "adapter_model.bin"]) + _, st = capture(adapter_only = True) + kept = _filter( + ["adapter_config.json", "adapter_model.bin"], st["allow_patterns"], st["ignore_patterns"] + ) + assert "adapter_model.bin" in kept + + +def test_adapter_only_explicit_use_safetensors_false_keeps_bin(capture): + """An explicit use_safetensors=False forces the .bin form without a model_info call.""" + _, st = capture(adapter_only = True, use_safetensors = False) + ig = st["ignore_patterns"] + assert ig is not None and "adapter_model*.safetensors" in ig + kept = _filter( + ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"], + st["allow_patterns"], + ig, + ) + assert "adapter_model.bin" in kept + assert "adapter_model.safetensors" not in kept + + +def test_gguf_file_with_subfolder_warms_subfolder_path(capture): + """gguf_file + subfolder: the warm allow-lists /, not the bare root name.""" + _, st = capture(weights_at_root = True, gguf_file = "model-Q4_K_M.gguf", subfolder = "gguf") + allow = st["allow_patterns"] + assert "gguf/model-Q4_K_M.gguf" in allow + kept = _filter(["gguf/model-Q4_K_M.gguf", "config.json"], allow, st["ignore_patterns"]) + assert "gguf/model-Q4_K_M.gguf" in kept and "config.json" in kept + + +def test_from_tf_root_load_ignores_nested_h5(capture): + """A from_tf root load keeps the root .h5 but drops nested .h5 / .msgpack checkpoints.""" + _, st = capture(weights_at_root = True, from_tf = True) + ig = st["ignore_patterns"] + assert "*/*.h5" in ig and "*/*.msgpack" in ig + kept = _filter(["model.h5", "checkpoint-1/model.h5", "config.json"], st["allow_patterns"], ig) + assert "model.h5" in kept + assert "checkpoint-1/model.h5" not in kept + + +def test_sentence_transformer_from_pretrained_is_prefetch_wired(): + """from_pretrained must call maybe_prefetch_hf_snapshot as an unconditional top-level statement before any return.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + cls = next( + n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "FastSentenceTransformer" + ) + fp = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "from_pretrained") + + def _prefetch_call(node): + # a bare call statement, or one whose return is captured (e.g. _st_prefetched = ...) + value = node.value if isinstance(node, (ast.Expr, ast.Assign)) else None + if ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Name) + and value.func.id == "maybe_prefetch_hf_snapshot" + ): + return value + return None + + prefetch_pos = next((i for i, n in enumerate(fp.body) if _prefetch_call(n)), None) + return_pos = next((i for i, n in enumerate(fp.body) if isinstance(n, ast.Return)), len(fp.body)) + assert ( + prefetch_pos is not None + ), "from_pretrained must call maybe_prefetch_hf_snapshot at top level" + assert prefetch_pos < return_pos, "prefetch must run before any top-level return" + # local_files_only must be forwarded so an offline load does not start a Hub download. + prefetch_call = _prefetch_call(fp.body[prefetch_pos]) + assert "local_files_only" in { + kw.arg for kw in prefetch_call.keywords + }, "prefetch must forward local_files_only" + + +def test_st_module_download_forwards_cache_folder(): + """_load_modules must forward the custom cache_folder into load_dir_path so per-module subdirs read the warmed cache.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "load_dir_path" + ] + assert calls, "expected a load_dir_path call in sentence_transformer.py" + assert all( + "cache_folder" in {kw.arg for kw in c.keywords} for c in calls + ), "every load_dir_path call must forward cache_folder" + + +def test_st_native_sentence_transformer_calls_forward_cache_folder(): + """Every native SentenceTransformer(model_name, ...) load must forward cache_folder; a modules-based build needs none.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + weight_loading_calls = [] + for n in ast.walk(tree): + if not ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "SentenceTransformer" + ): + continue + kw_names = {kw.arg for kw in n.keywords} + # A modules-based build downloads nothing; only a repo-name load reads the cache. + if "modules" in kw_names: + continue + weight_loading_calls.append(n) + assert ( + weight_loading_calls + ), "expected a repo-name SentenceTransformer load in sentence_transformer.py" + # cache_folder is forwarded explicitly or via a **kwargs unpacking (kw.arg == None). + for c in weight_loading_calls: + kw_names = {kw.arg for kw in c.keywords} + forwards = "cache_folder" in kw_names or None in kw_names + assert forwards, ( + "a repo-name SentenceTransformer load must forward cache_folder " + f"(explicitly or via **kwargs) at line {c.lineno}" + ) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 047783c35e..260fe36652 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -83,6 +83,7 @@ __all__ = [ "verify_fp8_support_if_applicable", "_get_inference_mode_context_manager", "hf_login", + "maybe_prefetch_hf_snapshot", "is_moe_model", "get_moe_target_parameters", "make_fast_generate_wrapper", @@ -905,6 +906,411 @@ logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITI TORCHAO_MSG = "Error: torchao not found, please install with `pip install torchao`" +# Artifacts a Transformers/PEFT load never reads (ONNX/TF/Flax/CoreML/GGUF/training state), skipped +# when prewarming so a mixed-format repo is not pulled in full. +_PREFETCH_IGNORE_PATTERNS = ( + "*.onnx", + "onnx/*", + "*.h5", + "*.msgpack", + "*.tflite", + "coreml/*", + "*.mlpackage/*", + "*.mlmodel", + "*.gguf", + # Training / checkpoint formats from_pretrained never reads. + "*.pt", + "*.pth", + "*.ckpt", + "optimizer.*", + "scheduler.*", + "rng_state*", + "trainer_state.json", + "events.out.tfevents*", + "checkpoint-*/*", +) + + +# Repo-root tokenizer / config / processor files from_pretrained reads from root even when weights +# load from a subfolder. Exact names (no wildcard) so they match only root-level files. +_ROOT_AUX_PREFETCH_PATTERNS = ( + "config.json", + "generation_config.json", + "tokenizer_config.json", + "tokenizer.json", + "tokenizer.model", + "special_tokens_map.json", + "added_tokens.json", + "vocab.json", + "vocab.txt", + "merges.txt", + "spiece.model", + # More VOCAB_FILES_NAMES the slow tokenizer may fetch (DeBERTa-v2, Whisper, Mistral, XLM-R/mBART, Marian, FSMT/XLM, GPT-2). + "spm.model", + "normalizer.json", + "tokenizer.model.v3", + "sentencepiece.bpe.model", + "source.spm", + "target.spm", + "bpe.codes", + "vocab.bpe", + # More VOCAB_FILES_NAMES (RemBERT, FSMT) a distinct-tokenizer-repo warm must cache too. + "sentencepiece.model", + "vocab-src.json", + "vocab-tgt.json", + "chat_template.jinja", + "chat_template.json", + # chat_template="" fetches additional_chat_templates/.jinja. + "additional_chat_templates/*.jinja", + "preprocessor_config.json", + "processor_config.json", + "video_preprocessor_config.json", # Qwen2.5-VL-style video processors + # trust_remote_code auto_map can name any module, so warm every *.py (tiny; none in a non-remote repo). + "*.py", + "*.tiktoken", # tiktoken vocab (e.g. Qwen's qwen.tiktoken) +) + + +# Files a PEFT adapter load reads: config + weights (glob covers sharded adapters). Any merged +# full-model weights the repo also ships match none of these. +_ADAPTER_PREFETCH_PATTERNS = ( + "adapter_config.json", + "adapter_model*", +) + + +# Weight files in a SUBDIRECTORY. A bare root load reads only root weights, so ignoring these drops +# alternate-precision/experimental dirs (fp16/, experimental/). "*/*" spans "/" (HF fnmatch), so nested +# weights match while root "model.safetensors" is kept. Only applied when weights_at_root (diffusion +# keeps weights in subfolders). +_SUBDIR_WEIGHT_IGNORE_PATTERNS = ( + "*/*.safetensors", + "*/*.bin", + "*/*.h5", + "*/*.msgpack", + "*/*.pt", + "*/*.pth", +) + + +def _in_requested_load_scope(filename, subfolder): + """True if *filename* is in the location being loaded (*subfolder*, else root). Scopes the ".bin is + redundant when safetensors exist" test so a .bin-only subfolder keeps its .bin.""" + filename = filename.replace("\\", "/") + if isinstance(subfolder, str) and subfolder.strip("/"): + return filename.startswith(subfolder.strip("/") + "/") + return "/" not in filename # root load: no directory component + + +# .safetensors training-state files that are NOT model weights (e.g. optimizer.safetensors next to a +# real pytorch_model.bin); counting them as "model safetensors present" would drop the needed .bin. +_NON_MODEL_WEIGHT_STEMS = frozenset( + { + "optimizer", + "scheduler", + "scaler", + "rng_state", + "training_args", + } +) + + +def _is_model_weight_safetensors(filename): + """True if *filename* is a model-weights safetensors, not a PEFT adapter/sidecar + (adapter_model.safetensors) or trainer-state (optimizer.safetensors). Only a real one proves the + .bin redundant; counting a sidecar would wrongly drop the needed .bin (fetched then without Xet fallback).""" + name = filename.replace("\\", "/").rsplit("/", 1)[-1] + if not name.endswith((".safetensors", ".safetensors.index.json")): + return False + if name.startswith("adapter_"): + return False + # Stem before first dot: "optimizer.safetensors" -> "optimizer" (real shards kept); rng_state via prefix. + stem = name.split(".", 1)[0].lower() + if stem in _NON_MODEL_WEIGHT_STEMS or stem.startswith("rng_state"): + return False + return True + + +def _is_canonical_variant_model_weight_safetensors(filename, variant): + """True for a canonical model-weights safetensors carrying the requested *variant*, in the forms + transformers reads (single, either numbered-shard layout, or the index). Strict (base must be + "model"): a sidecar like consolidated..safetensors does not prove the variant .bin redundant.""" + base = filename.replace("\\", "/").rsplit("/", 1)[-1] + v = re.escape(variant) + return bool( + re.match( + rf"^(?:model\.{v}\.safetensors" + rf"|model\.{v}-\d{{5}}-of-\d{{5}}\.safetensors" + rf"|model-\d{{5}}-of-\d{{5}}\.{v}\.safetensors" + rf"|model\.safetensors\.index\.{v}\.json)$", + base, + ) + ) + + +_CANONICAL_MODEL_WEIGHT_SAFETENSORS_RE = re.compile( + r"^(?:model\.safetensors|model-\d{5}-of-\d{5}\.safetensors|model\.safetensors\.index\.json)$" +) + + +def _is_canonical_model_weight_safetensors(filename): + """True for a canonical (non-variant) model-weights safetensors a default load reads (model.safetensors, + a numbered shard, or the index). Strict: an unrecognized name keeps both formats, so a variant-only + safetensors + pytorch_model.bin repo never has its .bin dropped for a no-variant load.""" + name = filename.replace("\\", "/").rsplit("/", 1)[-1] + return bool(_CANONICAL_MODEL_WEIGHT_SAFETENSORS_RE.match(name)) + + +def _adapter_repo_has_safetensors( + model_name, + *, + token = None, + revision = None, +): + """Best-effort: does the adapter repo ship a root safetensors adapter weight (making the .bin + redundant)? Scoped to root adapter_model* files; any failure returns False.""" + try: + from huggingface_hub import HfApi + siblings = HfApi().model_info(model_name, revision = revision, token = token).siblings or [] + return any( + "/" not in sibling.rfilename.replace("\\", "/") # root only + and sibling.rfilename.startswith("adapter_model") + and sibling.rfilename.endswith(".safetensors") + for sibling in siblings + ) + except Exception: + return False + + +def _prefetch_ignore_patterns( + model_name, + *, + token = None, + revision = None, + subfolder = None, + use_safetensors = None, + from_tf = False, + from_flax = False, + variant = None, + weights_at_root = False, +): + """ignore_patterns for the prewarm snapshot: the static skip list, minus the checkpoint guard when + loading from a checkpoint-* subfolder, minus the weight format the load will not read. use_safetensors + is a format allowlist (True -> skip *.bin, False -> skip *.safetensors); auto (None) skips *.bin only + when in-scope safetensors are shipped. from_tf/from_flax keep *.h5/*.msgpack. + + Suppressed for a whole multi-component snapshot (weights_at_root=False, no subfolder: ST/diffusers + repos with per-subfolder weights, each in its own format), since "*" spans "/" so dropping "*.bin" + would strip a module's only weight.""" + # Keep checkpoint-*/* under a checkpoint-* subfolder; keep *.h5 / *.msgpack under from_tf/flax. + ignore_patterns = [ + pattern + for pattern in _PREFETCH_IGNORE_PATTERNS + if not ( + ( + pattern == "checkpoint-*/*" + and isinstance(subfolder, str) + and subfolder.startswith("checkpoint-") + ) + or (from_tf and pattern == "*.h5") + or (from_flax and pattern == "*.msgpack") + ) + ] + # Drop the format the load will not read (the other doubles the download); skipped for a whole + # multi-component snapshot (see docstring). + whole_multi_component = not weights_at_root and not ( + isinstance(subfolder, str) and subfolder.strip("/") + ) + if whole_multi_component: + pass + elif from_tf or from_flax: + # TF / Flax loads never read the PyTorch formats; drop safetensors and .bin. + ignore_patterns.extend( + ( + "*.safetensors", + "*.safetensors.index.json", + "*.bin", + "*.bin.index.json", + ) + ) + elif use_safetensors is True: + # Explicit safetensors: load never reads .bin (no model_info call needed). + ignore_patterns.extend(("*.bin", "*.bin.index.json")) + elif use_safetensors is False: + # Explicit .bin: load never reads safetensors. + ignore_patterns.extend(("*.safetensors", "*.safetensors.index.json")) + else: + # Auto: skip .bin only once in-scope safetensors are confirmed (best-effort; any failure keeps both). + try: + from huggingface_hub import HfApi + + siblings = ( + HfApi() + .model_info( + model_name, + revision = revision, + token = token, + ) + .siblings + or [] + ) + # Count only in-scope model-weights safetensors (not adapters/sidecars): variant-matching if + # a variant is requested, else canonical, proving the .bin redundant. + has_safetensors = any( + _is_model_weight_safetensors(sibling.rfilename) + and _in_requested_load_scope(sibling.rfilename, subfolder) + and ( + _is_canonical_variant_model_weight_safetensors(sibling.rfilename, variant) + if variant + else _is_canonical_model_weight_safetensors(sibling.rfilename) + ) + for sibling in siblings + ) + if has_safetensors: + ignore_patterns.extend(("*.bin", "*.bin.index.json")) + except Exception: + pass + return ignore_patterns + + +def maybe_prefetch_hf_snapshot( + model_name, + token = None, + *, + revision = None, + cache_dir = None, + local_files_only = False, + fast_inference = False, + subfolder = None, + force_download = False, + use_safetensors = None, + from_tf = False, + from_flax = False, + tokenizer_only = False, + adapter_only = False, + weights_at_root = False, + variant = None, + gguf_file = None, +): + """Warm the HF cache for a remote repo before the in-process load. + + Xet can hang on a blob with no progress or exception, and a blocked native Xet thread cannot be + killed in-process. So pull the snapshot first in a killable subprocess that falls back Xet -> HTTP + on a stall (unsloth_zoo.hf_xet_fallback), making from_pretrained a cache hit. + + Returns True iff warmed (caller can clear force_download), else False (skipped: local/offline/ + local_files_only/fast_inference/old unsloth_zoo, or failed). Only a both-transports-stalled + DownloadStallError is raised; other failures are left for from_pretrained to surface. + """ + try: + from unsloth_zoo.hf_xet_fallback import ( + snapshot_download_with_xet_fallback, + DownloadStallError, + ) + except Exception: + return False # older unsloth_zoo without the helper: load normally + + if not isinstance(model_name, str) or not model_name: + return False + # Local path: nothing to download. Expand ~ first (os.path.exists does not). + model_path = os.path.expanduser(model_name) + if os.path.isdir(model_path) or os.path.exists(model_path): + return False + # Looks local but not yet on disk (e.g. an uncreated output dir): not a Hub repo id, so leave it + # for from_pretrained rather than download it. + if ( + os.path.isabs(model_path) + or model_name.startswith(("~", "./", "../", ".\\", "..\\")) + or "\\" in model_name + ): + return False + if local_files_only: # cache-only: never reach out + return False + if any( + os.environ.get(flag, "0").lower() in ("1", "true", "yes", "on") + for flag in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + ): + return False + if fast_inference: # vLLM has its own download path + return False + + # tokenizer-only / adapter-only warms allow-list exact files below, so the weight-format ignore + # list (and its auto-branch model_info call) is skipped. + ignore_patterns = ( + None + if tokenizer_only or adapter_only or gguf_file + else _prefetch_ignore_patterns( + model_name, + token = token, + revision = revision, + subfolder = subfolder, + use_safetensors = use_safetensors, + from_tf = from_tf, + from_flax = from_flax, + variant = variant, + weights_at_root = weights_at_root, + ) + ) + # Narrow the warm to what the load reads (skip extra checkpoints/precisions); every branch still warms + # root tokenizer/config/custom-code so those never fall in-process. + allow_patterns = None + if gguf_file: + # gguf_file=NAME reads exactly that GGUF, but the static ignore list drops *.gguf; so warm just + # that file (plus root aux), under / if set. + _gguf_path = ( + f"{subfolder.strip('/')}/{gguf_file}" + if isinstance(subfolder, str) and subfolder.strip("/") + else gguf_file + ) + allow_patterns = [_gguf_path, *_ROOT_AUX_PREFETCH_PATTERNS] + elif tokenizer_only: + # A distinct tokenizer repo: warm only tokenizer / config / vocab files, never its weights. + allow_patterns = list(_ROOT_AUX_PREFETCH_PATTERNS) + elif adapter_only: + # A PEFT adapter load reads only adapter_config.json + adapter_model.* (plus root aux), not any + # merged weights the repo may also publish. + allow_patterns = [*_ADAPTER_PREFETCH_PATTERNS, *_ROOT_AUX_PREFETCH_PATTERNS] + # PeftModel reads one format (safetensors when present): explicit use_safetensors wins, else + # prefer safetensors when shipped (best-effort; any failure keeps both). + if use_safetensors is False: + ignore_patterns = [ + "adapter_model*.safetensors", + "adapter_model*.safetensors.index.json", + ] + elif use_safetensors is True or _adapter_repo_has_safetensors( + model_name, token = token, revision = revision + ): + ignore_patterns = ["adapter_model*.bin", "adapter_model*.bin.index.json"] + elif isinstance(subfolder, str) and subfolder.strip("/"): + # subfolder=X: load resolves every weight under X/, so warm that subfolder (plus root aux). + allow_patterns = [f"{subfolder.strip('/')}/*", *_ROOT_AUX_PREFETCH_PATTERNS] + elif weights_at_root: + # A bare load reads only root weights: drop subdir weights (fp16/, checkpoint dirs) while keeping + # subdir configs. Diffusion leaves weights_at_root False. + ignore_patterns = [*(ignore_patterns or []), *_SUBDIR_WEIGHT_IGNORE_PATTERNS] + try: + snapshot_download_with_xet_fallback( + model_name, + token = token, + revision = revision, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + variant = variant, + ) + return True + except DownloadStallError: + # Both transports stalled: surface a clear network error, not a silent in-process hang. + raise + except Exception as exception: + logger.warning_once( + f"Unsloth: Could not pre-download {model_name} " + f"({type(exception).__name__}: {exception}); continuing with the normal load." + ) + return False + + # Ignore logging messages class HideLoggingMessage(logging.Filter): __slots__ = ("text",) diff --git a/unsloth/models/diffusion.py b/unsloth/models/diffusion.py index 12596b432e..955bf55987 100644 --- a/unsloth/models/diffusion.py +++ b/unsloth/models/diffusion.py @@ -24,7 +24,7 @@ import os import torch from transformers import AutoConfig, AutoProcessor, AutoTokenizer -from ._utils import is_bfloat16_supported +from ._utils import is_bfloat16_supported, maybe_prefetch_hf_snapshot from .llama import logger __all__ = ["FastDiffusionModel", "DIFFUSION_MODEL_TYPES", "is_diffusion_model_type"] @@ -79,7 +79,14 @@ def _resolve_diffusion_model_class(config): ) -def _load_diffusion_config(model_name, token, trust_remote_code, revision, local_files_only): +def _load_diffusion_config( + model_name, + token, + trust_remote_code, + revision, + local_files_only, + cache_dir = None, +): """Load the config, aliasing the legacy ``diffusion_gemma`` model_type to the ``diffusion_gemma4`` classes current transformers ships. AutoConfig raises on the legacy type; catch that, rewrite the type/arch names in-memory, and rebuild.""" @@ -90,6 +97,7 @@ def _load_diffusion_config(model_name, token, trust_remote_code, revision, local trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) except ValueError as e: if "diffusion_gemma" not in str(e): @@ -103,6 +111,7 @@ def _load_diffusion_config(model_name, token, trust_remote_code, revision, local token = token, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) with open(cfg_path, encoding = "utf-8") as f: cd = json.load(f) @@ -152,12 +161,16 @@ class FastDiffusionModel: os.environ.get("HF_HUB_OFFLINE", "0") == "1" or os.environ.get("TRANSFORMERS_OFFLINE", "0") == "1" ) + + cache_dir = kwargs.get("cache_dir") + config = _load_diffusion_config( model_name, token, trust_remote_code, revision, local_files_only, + cache_dir = cache_dir, ) model_type = getattr(config, "model_type", None) if not is_diffusion_model_type(model_type): @@ -168,6 +181,21 @@ class FastDiffusionModel: model_cls = _resolve_diffusion_model_class(config) + # Prefetch the whole repo root so the weight load is a cache hit. No subfolder: the pipeline + # loads every component subfolder, so narrowing would leave unet/vae/text_encoder to Xet. + maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = revision, + cache_dir = cache_dir, + local_files_only = local_files_only, + fast_inference = False, + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + # Forward variant (e.g. "fp16") so the warm keeps variant weights. + variant = kwargs.get("variant"), + ) + load_kwargs = dict( dtype = dtype, device_map = device_map, @@ -176,7 +204,14 @@ class FastDiffusionModel: attn_implementation = attn_implementation, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) + # Match the load's weight format to the warm (None/auto already matches). + if kwargs.get("use_safetensors") is not None: + load_kwargs["use_safetensors"] = kwargs["use_safetensors"] + # Forward variant to the real load so it reads the warmed variant weights. + if kwargs.get("variant") is not None: + load_kwargs["variant"] = kwargs["variant"] # Optional bitsandbytes quant. The MoE experts (3D Parameters) are not nn.Linear so bnb skips # them; only attention + dense MLP Linears quantize, lm_head/embeddings stay full precision. @@ -222,6 +257,7 @@ class FastDiffusionModel: trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) except Exception: tokenizer = AutoTokenizer.from_pretrained( @@ -230,6 +266,7 @@ class FastDiffusionModel: trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) return model, tokenizer diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index bb7289dfa8..6bec95b577 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2420,6 +2420,73 @@ class FastLlamaModel: preferred_attn_impl = resolve_attention_implementation(model_function, model_config) + # Prefetch the repo (killable child) so the weight load is a cache hit. Runs after the + # AutoConfig/model-class check so an unsupported repo fails on its small config fetch. No + # revision: the load resolves model_name (maybe a remapped prequant repo) on its default branch. + _prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + # Skip the warm only for a real vLLM load; a num_labels classification load still goes + # in-process below, so it must be warmed even under fast_inference. + fast_inference = fast_inference and num_labels is None, + subfolder = kwargs.get("subfolder"), + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + from_tf = kwargs.get("from_tf", False), + from_flax = kwargs.get("from_flax", False), + # Bare load reads only ROOT weights; skip subdir weights. Ignored when a subfolder is set. + weights_at_root = True, + variant = kwargs.get("variant"), # forward so the warm keeps the variant .bin + gguf_file = kwargs.get( + "gguf_file" + ), # forward so the warm fetches the GGUF (else ignored) + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + + # Tokenizer always loads in-process. Resolve the cache_dir the tokenizer load will actually + # use, mirroring load_correct_tokenizer: without an explicit cache_dir, Colab/Kaggle route to + # a special tokenizer cache (huggingface_tokenizers_cache / Kaggle tmp), NOT the HF-default + # cache the base snapshot warmed. So the base warm does not cover the tokenizer there. + from ..tokenizer_utils import ( + IS_COLAB_ENVIRONMENT, + IS_KAGGLE_ENVIRONMENT, + KAGGLE_TMP, + ) + + _tokenizer_repo = ( + tokenizer_name if (isinstance(tokenizer_name, str) and tokenizer_name) else model_name + ) + _tokenizer_cache_dir = kwargs.get("cache_dir") + if _tokenizer_cache_dir is None: + if IS_COLAB_ENVIRONMENT: + _tokenizer_cache_dir = "huggingface_tokenizers_cache" + elif IS_KAGGLE_ENVIRONMENT: + _tokenizer_cache_dir = os.path.join(KAGGLE_TMP, "huggingface_tokenizers_cache") + # Warm the tokenizer repo into the cache the load will use whenever the base warm did not + # cover it: a distinct tokenizer repo, fast_inference (base warm skipped), or a tokenizer + # cache_dir that differs from the base-warm cache_dir (Colab/Kaggle special cache). + _warm_tokenizer_repo = ( + isinstance(_tokenizer_repo, str) + and bool(_tokenizer_repo) + and ( + _tokenizer_repo != model_name + or fast_inference + or _tokenizer_cache_dir != kwargs.get("cache_dir") + ) + ) + if _warm_tokenizer_repo: + maybe_prefetch_hf_snapshot( + _tokenizer_repo, + token = token, + cache_dir = _tokenizer_cache_dir, + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + has_rope_scaling = False try: with open(inspect.getfile(model_function), "r", encoding = "utf-8") as file: @@ -2672,6 +2739,10 @@ class FastLlamaModel: # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name + # Route the tokenizer load to the custom cache_dir the prefetch warmed. + _tokenizer_cache_kwargs = {} + if kwargs.get("cache_dir") is not None: + _tokenizer_cache_kwargs["cache_dir"] = kwargs["cache_dir"] tokenizer = load_correct_tokenizer( tokenizer_name = tokenizer_name, model_max_length = max_position_embeddings, @@ -2679,6 +2750,7 @@ class FastLlamaModel: token = token, trust_remote_code = trust_remote_code, fix_tokenizer = fix_tokenizer, + **_tokenizer_cache_kwargs, ) model, tokenizer = patch_tokenizer(model, tokenizer) @@ -2805,6 +2877,7 @@ class FastLlamaModel: model_max_length = max_position_embeddings, padding_side = "right", token = token, + cache_dir = kwargs.get("cache_dir"), ) patch_saving_functions(tokenizer) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 562afdd645..84f808d2b5 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -106,6 +106,7 @@ from ._utils import ( _is_family_text_decoder, _apply_text_only_key_mapping, set_task_config_attr, + maybe_prefetch_hf_snapshot, ) # Single source of truth is unsloth_zoo.model_lists. Re-exported so callers @@ -865,6 +866,28 @@ class FastLanguageModel(FastLlamaModel): if is_peft: # From https://github.com/huggingface/peft/issues/184 # Now add PEFT adapters + # Warm the adapter repo: PeftModel downloads it in-process and can hang on Xet. + _prefetched = maybe_prefetch_hf_snapshot( + old_model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir"), + local_files_only = local_files_only, + # Adapter always loads in-process via PeftModel, so warm it even under fast_inference. + fast_inference = False, + force_download = kwargs.get("force_download", False), + # Leave use_safetensors auto (inheriting base format could skip a safetensors-only + # adapter). adapter_only restricts the warm to the adapter files + root aux. + adapter_only = True, + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # Forward cache_dir so the load reads the warmed adapter. No subfolder (that targets the + # base checkpoint; adapters live at the root). + peft_load_kwargs = {} + if kwargs.get("cache_dir") is not None: + peft_load_kwargs["cache_dir"] = kwargs["cache_dir"] model = PeftModel.from_pretrained( model, old_model_name, @@ -873,6 +896,7 @@ class FastLanguageModel(FastLlamaModel): local_files_only = local_files_only, is_trainable = True, trust_remote_code = trust_remote_code, + **peft_load_kwargs, ) # Patch it as well! model = dispatch_model.patch_peft_model(model, use_gradient_checkpointing) @@ -1790,6 +1814,28 @@ class FastModel(FastBaseModel): _LoraModel._create_and_replace = _patched_car + # Warm the adapter repo: PeftModel downloads it in-process and can hang on Xet. + _prefetched = maybe_prefetch_hf_snapshot( + old_model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir"), + local_files_only = local_files_only, + # Adapter always loads in-process via PeftModel, so warm it even under fast_inference. + fast_inference = False, + force_download = kwargs.get("force_download", False), + # Leave use_safetensors auto (inheriting base format could skip a safetensors-only + # adapter). adapter_only restricts the warm to the adapter files + root aux. + adapter_only = True, + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # Forward cache_dir so the load reads the warmed adapter. No subfolder (that targets the + # base checkpoint; adapters live at the root). + peft_load_kwargs = {} + if kwargs.get("cache_dir") is not None: + peft_load_kwargs["cache_dir"] = kwargs["cache_dir"] try: model = PeftModel.from_pretrained( model, @@ -1799,6 +1845,7 @@ class FastModel(FastBaseModel): local_files_only = local_files_only, is_trainable = True, trust_remote_code = trust_remote_code, + **peft_load_kwargs, ) finally: # Always restore original PEFT method, even if loading fails diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index 7e43442bfd..c1172faa94 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -19,6 +19,7 @@ from ._utils import ( SUPPORTS_BFLOAT16, resolve_model_class, resolve_encoder_attention_implementation, + maybe_prefetch_hf_snapshot, ) import inspect import json @@ -541,7 +542,12 @@ class FastSentenceTransformer(FastModel): return transformer_module @staticmethod - def _read_pooling_mode(model_name, token): + def _read_pooling_mode( + model_name, + token, + cache_dir = None, + revision = None, + ): """Read the pooling mode from modules.json, else return "mean".""" try: if os.path.exists(model_name) and os.path.exists( @@ -549,7 +555,13 @@ class FastSentenceTransformer(FastModel): ): modules_json_path = os.path.join(model_name, "modules.json") else: - modules_json_path = hf_hub_download(model_name, "modules.json", token = token) + modules_json_path = hf_hub_download( + model_name, + "modules.json", + token = token, + cache_dir = cache_dir, + revision = revision, + ) with open(modules_json_path, "r", encoding = "utf-8") as f: modules_config = json.load(f) @@ -571,6 +583,8 @@ class FastSentenceTransformer(FastModel): model_name, os.path.join(pooling_path, "config.json"), token = token, + cache_dir = cache_dir, + revision = revision, ) break @@ -950,7 +964,12 @@ class FastSentenceTransformer(FastModel): f.write(content) @staticmethod - def _module_path(model_name, token = None): + def _module_path( + model_name, + token = None, + cache_dir = None, + revision = None, + ): """Return the path to the modules.json file, or None.""" try: if os.path.exists(model_name) and os.path.isdir(model_name): @@ -958,7 +977,13 @@ class FastSentenceTransformer(FastModel): return path if os.path.exists(path) else None else: try: - return hf_hub_download(model_name, "modules.json", token = token) + return hf_hub_download( + model_name, + "modules.json", + token = token, + cache_dir = cache_dir, + revision = revision, + ) except: return None except: @@ -1135,6 +1160,8 @@ class FastSentenceTransformer(FastModel): max_seq_length, pooling_mode, trust_remote_code = False, + cache_dir = None, + revision = None, ) -> tuple[OrderedDict, bool]: """Load modules from modules.json, else fall back to hard-coded modules. @@ -1145,7 +1172,9 @@ class FastSentenceTransformer(FastModel): from sentence_transformers.models import Pooling, Normalize modules = OrderedDict() - modules_json_path = FastSentenceTransformer._module_path(model_name, token) + modules_json_path = FastSentenceTransformer._module_path( + model_name, token, cache_dir = cache_dir, revision = revision + ) if modules_json_path: with open(modules_json_path, encoding = "utf8") as f: @@ -1171,7 +1200,13 @@ class FastSentenceTransformer(FastModel): load_path = os.path.join(model_name, module_path) else: try: - load_path = load_dir_path(model_name, module_path, token = token) + load_path = load_dir_path( + model_name, + module_path, + token = token, + cache_folder = cache_dir, + revision = revision, + ) except Exception as e: print(f"Unsloth Warning: Could not download module {module_path}: {e}") continue @@ -1198,7 +1233,9 @@ class FastSentenceTransformer(FastModel): hidden_size = getattr(model.config, "hidden_size", 768) if pooling_mode == "mean": - pooling_mode = FastSentenceTransformer._read_pooling_mode(model_name, token) + pooling_mode = FastSentenceTransformer._read_pooling_mode( + model_name, token, cache_dir = cache_dir, revision = revision + ) modules["1"] = Pooling(word_embedding_dimension = hidden_size, pooling_mode = pooling_mode) modules["2"] = Normalize() @@ -1386,6 +1423,45 @@ class FastSentenceTransformer(FastModel): "Run `pip install sentence-transformers` to install it." ) + # Validate the load modes BEFORE the prefetch so a bad config fails without downloading weights. + # Guard on not for_inference: that branch below never used these flags. + if not for_inference: + # sanity check, thanks Etherl: + if full_finetuning and (load_in_4bit or load_in_8bit): + print( + "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA." + ) + load_in_4bit = False + load_in_8bit = False + load_in_fp8 = False + load_in_16bit = False + + if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: + raise RuntimeError( + "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n" + "Also, we by default set `load_in_16bit = True`.\n" + "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n" + "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`" + ) + + # Prefetch so the ST load below is a cache hit. weights_at_root stays False (ST component + # weights live in per-module subfolders). Resolve the same cache the load uses: HF cache_dir, + # else cache_folder, else SENTENCE_TRANSFORMERS_HOME, else default -- a wrong cache misses the warm. + _st_prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + local_files_only = kwargs.get("local_files_only", False), + # Forward force_download so the refresh happens in the killable child, then clear it so the + # in-process ST load reuses the warm cache instead of re-downloading over unguarded Xet. + force_download = kwargs.get("force_download", False), + ) + if _st_prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # if for_inference == True, skip Unsloth optimizations to avoid torch compile issues if for_inference: st_device = device_map @@ -1416,27 +1492,16 @@ class FastSentenceTransformer(FastModel): if k in kwargs: st_kwargs[k] = kwargs[k] + # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm + # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch). + _st_cache = kwargs.get("cache_dir") or kwargs.get("cache_folder") + if _st_cache is not None: + st_kwargs["cache_folder"] = _st_cache + st_model = SentenceTransformer(model_name, **st_kwargs) return st_model - # sanity check, thanks Etherl: - if full_finetuning and (load_in_4bit or load_in_8bit): - print( - "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA." - ) - load_in_4bit = False - load_in_8bit = False - load_in_fp8 = False - load_in_16bit = False - - if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: - raise RuntimeError( - "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n" - "Also, we by default set `load_in_16bit = True`.\n" - "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n" - "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`" - ) - + # Load-mode validation already ran before the prefetch above. if "auto_model" not in kwargs: kwargs["auto_model"] = AutoModel @@ -1533,7 +1598,8 @@ class FastSentenceTransformer(FastModel): elif is_mpnet: FastSentenceTransformer._patch_mpnet_v5() - # Load via native SentenceTransformer (bypasses Unsloth patching) + # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm + # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch). st_model = SentenceTransformer( model_name, device = st_device, @@ -1541,6 +1607,7 @@ class FastSentenceTransformer(FastModel): token = token, revision = revision, model_kwargs = model_kwargs, + cache_folder = kwargs.get("cache_dir") or kwargs.get("cache_folder"), ) # Store metadata for get_peft_model @@ -1646,7 +1713,18 @@ class FastSentenceTransformer(FastModel): # No modules.json -> force 16-bit: saving is custom for these models and # 4-bit would need dequant in save_pretrained_merged, not worth it. - has_modules_json = FastSentenceTransformer._module_path(model_name, token) is not None + # Resolve the warmed cache: hf_hub_download ignores SENTENCE_TRANSFORMERS_HOME, so pass it as cache_dir. + has_modules_json = ( + FastSentenceTransformer._module_path( + model_name, + token, + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + revision = revision, + ) + is not None + ) if not has_modules_json and load_in_4bit: print( @@ -1656,6 +1734,12 @@ class FastSentenceTransformer(FastModel): load_in_4bit = False load_in_16bit = True + # The fallback FastModel load reads HF cache_dir, not ST's cache_folder/SENTENCE_TRANSFORMERS_HOME. + # Point it at the warmed cache, but only when no explicit cache_dir was passed (which wins). + _st_cache_dir = kwargs.get("cache_folder") or os.environ.get("SENTENCE_TRANSFORMERS_HOME") + if _st_cache_dir is not None and "cache_dir" not in kwargs: + kwargs["cache_dir"] = _st_cache_dir + try: model, tokenizer = FastModel.from_pretrained( model_name = model_name, @@ -1697,6 +1781,12 @@ class FastSentenceTransformer(FastModel): max_seq_length, pooling_mode, trust_remote_code = trust_remote_code, + # Same resolved cache as above so the fallback module loads hit the warm, not Xet. + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + # Same revision as the weight load so modules hit the warm (None = default branch). + revision = revision, ) st_device = device_map diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 689e362f95..179bd0b650 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -551,6 +551,7 @@ def _construct_vlm_processor_fallback( model_type, token, trust_remote_code, + cache_dir = None, local_files_only = False, ): """Build a VLM processor manually when AutoProcessor.from_pretrained fails (some VLMs @@ -568,6 +569,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) # Load tokenizer via PreTrainedTokenizerFast (bypasses tokenizer_class check) @@ -576,6 +578,7 @@ def _construct_vlm_processor_fallback( padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) # Read tokenizer_config.json for special tokens: prefer the local file (offline @@ -601,6 +604,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, "tokenizer_config.json", token = token, + cache_dir = cache_dir, local_files_only = local_files_only, ) with open(config_path, "r", encoding = "utf-8") as f: @@ -632,6 +636,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) proc_class_name = PROCESSOR_MAPPING_NAMES.get(config.model_type) @@ -872,6 +877,9 @@ class FastBaseModel: # For debugging - we use a download counter to see if environments are not breaking or if HF is down get_statistics(kwargs.get("local_files_only", False)) + # The base + tokenizer prefetch runs AFTER the load-mode validation below, so an invalid + # load_in_* combination fails without first downloading a snapshot. + if dtype is None: dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 elif os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": @@ -968,6 +976,53 @@ class FastBaseModel: raise RuntimeError( "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!" ) + + # Prefetch the repo (killable child) so the in-process load below is a cache hit. vLLM owns the + # weight download only when actually available; if fast_inference was requested but vLLM is + # missing, the load falls through in-process, so weights must still be warmed here. + _vllm_owns_weights = fast_inference and is_vLLM_available() + _prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = kwargs.get("revision"), + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + fast_inference = _vllm_owns_weights, + subfolder = kwargs.get("subfolder"), + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + from_tf = kwargs.get("from_tf", False), + from_flax = kwargs.get("from_flax", False), + # Bare load reads only ROOT weights; skip subdir weights. Ignored when a subfolder is set. + weights_at_root = True, + variant = kwargs.get("variant"), # forward so the warm keeps the variant .bin + gguf_file = kwargs.get( + "gguf_file" + ), # forward so the warm fetches the GGUF (else ignored) + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + + # Warm a SEPARATE tokenizer repo only (model_name is covered above). Not model_name here: this + # runs before fast_inference_setup may remap the repo, so it would warm the wrong one. + _tokenizer_repo = ( + tokenizer_name if (isinstance(tokenizer_name, str) and tokenizer_name) else model_name + ) + _warm_tokenizer_repo = ( + isinstance(_tokenizer_repo, str) + and bool(_tokenizer_repo) + and _tokenizer_repo != model_name + ) + if _warm_tokenizer_repo: + maybe_prefetch_hf_snapshot( + _tokenizer_repo, + token = token, + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + _skip_modules = SKIP_QUANTIZATION_MODULES.copy() # Nemotron-H uses 'mixer' (not 'mamba') for Mamba layers. # Mamba fused kernels pass out_proj.weight directly to F.linear, @@ -1278,6 +1333,18 @@ class FastBaseModel: # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name + # On the vLLM path the tokenizer warm was deferred (fast_inference_setup may remap model_name). + # Warm the now-final tokenizer repo so the load below hits the cache (a cached/local repo is a no-op). + if _vllm_owns_weights and isinstance(tokenizer_name, str) and tokenizer_name: + maybe_prefetch_hf_snapshot( + tokenizer_name, + token = token, + revision = kwargs.get("revision"), + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + # Fix _Unsloth_Patched_ prefix in local config files from old saves (issue #4085) if os.path.isdir(tokenizer_name): import json as _json @@ -1315,6 +1382,7 @@ class FastBaseModel: language = whisper_language, task = whisper_task, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _e: @@ -1327,6 +1395,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _e: @@ -1337,6 +1406,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception: @@ -1355,6 +1425,7 @@ class FastBaseModel: model_type_arch, token, trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _fe: @@ -1440,6 +1511,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = local_files_only, ) model, _fallback_tok = patch_tokenizer(model, _fallback_tok) @@ -1469,6 +1541,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception: @@ -1478,6 +1551,7 @@ class FastBaseModel: padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 93dfa9b2ad..3a91ef188d 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -563,8 +563,11 @@ def _load_correct_tokenizer( # /tmp of Kaggle seems has a 80GB limit! # Let's utilize them cache_dir = os.path.join(KAGGLE_TMP, cache_dir) - else: + elif cache_dir == "huggingface_tokenizers_cache": + # This default name is Colab/Kaggle-only; elsewhere use the HF default cache. cache_dir = None + # else: keep a caller-supplied cache_dir so the tokenizer loads from the prefetch-warmed dir instead + # of risking an in-process Hub/Xet transfer. # Try loading the slow tokenizer. If it fails, then try Fast only # Mainly to solve Deepseek models with no tokenizer.model file @@ -1323,6 +1326,7 @@ def check_tokenizer( padding_side = "right", token = None, _reload = True, + cache_dir = None, ): # Checks tokenizer for out of bounds ids. # Mainly a fix for https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha @@ -1413,10 +1417,11 @@ def check_tokenizer( f"Fix your tokenizer since it'll perform out of bounds memory accesses." ) - if IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT: - cache_dir = "huggingface_tokenizers_cache" - else: - cache_dir = None + # Reuse a caller-supplied cache_dir (warmed cache) for the repair reload; else the + # Colab/Kaggle sentinel (HF default elsewhere), as load_correct_tokenizer does. + reload_cache_dir = cache_dir + if reload_cache_dir is None and (IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT): + reload_cache_dir = "huggingface_tokenizers_cache" # Sometimes slow tokenizer does not work like Deepseek try: @@ -1430,7 +1435,7 @@ def check_tokenizer( use_fast = False, legacy = False, from_slow = True, - cache_dir = cache_dir, + cache_dir = reload_cache_dir, ) return check_tokenizer( model = model, @@ -1440,6 +1445,7 @@ def check_tokenizer( padding_side = padding_side, token = token, _reload = False, + cache_dir = cache_dir, ) break except: From 9407d491933d82ea377687e93c672d03e0387c1b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:17:39 -0700 Subject: [PATCH 003/113] GRPO: sequence packing for the no-grad old/ref logp path (default-on) (#6738) * GRPO: optional sequence packing for the no-grad old/ref logp path Add an opt-in sequence-packing fast path to _get_per_token_logps_and_entropies, enabled with UNSLOTH_GRPO_SEQ_PACKING=1. When the batch is text-only, the padded [B, Lmax] per-chunk forward is replaced by a single varlen [1, sum L] forward (BlockDiagonalCausalMask via packed_seq_lengths with reset position_ids). Per-token logps use the same float32 chunked_hidden_states_selective_log_softmax as the padded path, so the old and reference logps are bit-for-bit identical. Safety: the packed path is self-verified once against the padded ground truth on a batch that has at least two rows with real completion tokens (self._unsloth_seq_packing_nograd_ok), so cross-sample contamination would actually manifest; a degenerate all-pad / fully tool-masked batch leaves the verdict unset and re-verifies later. If a backend silently ignores packed_seq_lengths (flat batch run under a normal causal mask, samples leaking across boundaries), the packed logps will not match and packing is disabled instead of corrupting logps. It also forces use_cache=False (a populated past_key_value disables varlen packing), skips packing when a sliding window is shorter than the packed stream, runs the same GPT-OSS offload device_synchronize the padded loop uses, and falls back on any exception (UNSLOTH_GRPO_SEQ_PACKING_DEBUG=1 prints the reason). Default off, so existing behavior is unchanged. Pairs with the matching gradient-path change in unsloth_zoo so the full GRPO logp + loss + backward can run packed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: address review feedback - Cache the packed-vs-padded verdict per unwrapped model instead of on the trainer, so a separately forwarded reference model is verified on its own forward path rather than inheriting the policy model's verdict. - Force the padded path when token_type_ids or mm_token_type_ids are present, matching the extra vision kwargs the padded loop forwards. - Require the xformers varlen backend before packing. Without it the packed mask falls back to a dense O(T^2) SDPA mask that can OOM on the flattened batch, so we keep the padded loop in that case. - On any packed-forward failure (missing backend, OOM, unsupported forward) empty the cache on OOM, disable packing for that model, and fall back to the chunked padded loop instead of retrying every step. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: default-on, verify against per-row reference Redesign of the optional sequence-packing fast path for the no-grad old/ref logprob recompute, after establishing that the packed forward is the exact per-row computation and the padded batch forward is the side that mis-positions left-padded rows on long completions. - Default the packing on (UNSLOTH_GRPO_SEQ_PACKING, disable with 0). - Verify the packed logprobs against the per-row clean forward (each row's real tokens alone, reset 0-based positions, no padding), not the padded batch which is itself wrong for left-padding. Cross-sample contamination (a backend ignoring packed_seq_lengths) shows up as a large mismatch and falls back to the padded loop. - Make the trust decision shape and RoPE aware: re-verify whenever the packed total length or the longest segment grows past what was verified, so a later batch crossing a LongRoPE short/long cache boundary is re-checked instead of trusted blindly. - Run lm_head only on completion-prediction positions instead of every packed prompt token, so long-prompt/short-completion batches do not pay for projecting the whole packed prompt. - Drop the hard xformers import so the path also runs in FlashAttention-only environments; the per-row verification guards correctness regardless of backend. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: disable entirely on cross-sample mismatch When the per-row verification fails, distinguish the two failure modes by magnitude instead of by sequence length: - A large mismatch (>= 1.5) is the cross-sample contamination signature: the model's attention does not honor the block-diagonal packed mask (seen on some MoE / custom-attention models, e.g. qwen2_moe). Disable packing entirely for the model so later batches do not pay the verification cost again. - A moderate mismatch is more likely a length-boundary effect (a LongRoPE short/long cache switch): keep marking just that length region unsafe so packing still runs for smaller shapes. Validated: Qwen1.5-MoE falls back after a single verification (grad and no-grad ok flags go False, no re-verify on later steps); dense Llama-3.2 and Qwen3 still verify and engage packing. * GRPO no-grad packing: trim comments to be concise * GRPO no-grad packing: fix per-row completion boundary for left-padded rows The completion-target selection used a single global boundary (col >= L - logits_to_keep). After left-packing, each row's completion starts at (L - logits_to_keep) - left_pad[row], so for left-padded rows the first left_pad completion tokens fall below the global boundary and were dropped, leaving 0 logprobs at real completion positions that the loss mask keeps. Use the per-row boundary so packed coverage matches create_completion_attention_mask exactly, and widen the self-verify mask to the full per-row completion region so it can catch coverage gaps. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: gate verification on real completion rows Count active rows via create_completion_attention_mask (the same mask the loss uses) instead of any non-pad token in the packed window. Prompt-only rows carry prompt-overflow tokens in the window and could otherwise satisfy the >= 2 verification guard, letting a batch with a single real completion row cache a trust decision. This matches the gradient path, which already gates on the completion mask. The same mask is reused for the self-verify comparison. * GRPO no-grad packing: gate debug logging on UNSLOTH_ENABLE_LOGGING Use the shared UNSLOTH_ENABLE_LOGGING global (import_fixes, re-exported by _utils) instead of a bespoke UNSLOTH_GRPO_SEQ_PACKING_DEBUG env var for the packing debug prints, matching the rest of the codebase. * GRPO packing: import UNSLOTH_ENABLE_LOGGING inside the injected logp function _get_per_token_logps_and_entropies is copied verbatim into the generated GRPO trainer via inspect.getsource, and that module never imported UNSLOTH_ENABLE_LOGGING, so the default-on packing verify path raised NameError (and the except handler re-raised it). Import the flag locally, before the try, so the name is defined in the generated module too. Drop it from the now-unused module-level import. * GRPO no-grad packing: harden unsafe-length skip, verify guard, fallback cleanup Three fixes to the no-grad logp packing path, mirroring the grad path: - skip the packed forward for known-unsafe lengths by reading unsafe_T and gating on it before the forward, instead of running the full packed pass and the result build only to discard them (wastes a pass, can OOM at large T) - only widen the verified T/seg envelope when >= 2 completion rows actually exercised cross-sample packing; a < 2 row batch cannot expose leakage, so it must not extend the trusted shape that later multi-row batches skip verify for - drop the packed intermediates (hidden/sel/result/ref) before the padded fallback loop so it does not run with the flattened hidden state still resident * GRPO no-grad packing: cap the flattened forward at one mini-batch budget The packed path built a single [1, sum L] forward over every row before any size check, so a large batch could exceed the memory the padded path bounds per mini-batch. Gate packing on _pk_T <= _pk_cap (B * seq_len, one padded mini-batch's token budget); larger batches fall back to the chunked padded loop. * GRPO no-grad packing: disable unless unsloth_zoo has the masked-column guard The packed path leaves masked prompt/pad logprob columns at 0, which only stays finite if unsloth_zoo grpo_compute_loss zeroes them before exp() (zoo#840). An older unsloth_zoo without that guard would NaN. Detect the guard once (cached on the model) via inspect.getsource and gate packing on it, so #6738 is safe with any unsloth_zoo version and re-enables packing automatically once a guarded zoo is installed, independent of the pinned lower bound. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO packing: hoist env gates and zoo-guard detection to one-time module checks Read UNSLOTH_GRPO_SEQ_PACKING and detect the unsloth_zoo masked-column guard once at import time (module constants plus RL_PRE_ITEMS for the generated trainer cache) instead of per call, and drop the in-function UNSLOTH_ENABLE_LOGGING import for a module-top one. The UNSLOTH_GRPO_SEQ_PACKING_VERIFY force-verify debug knob is commented out, kept in place for hand re-enable; the first-use and envelope-growth self-verify stays active. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO packing: cap the flattened forward by the padded chunk rows B counts chunks at this point, so B * seq_len understated (small runs) or overstated (large runs) the padded mini-batch token budget; use batch_size * seq_len, the rows the padded loop actually forwards per chunk. * GRPO sequence packing: tighten comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- unsloth/models/rl_replacements.py | 245 +++++++++++++++++++++++++++++- 1 file changed, 243 insertions(+), 2 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 3be614cf4a..0573fd5fd8 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -29,6 +29,7 @@ from collections import defaultdict from unsloth_zoo.rl_replacements import ( RL_REPLACEMENTS, left_pack_padding, + create_completion_attention_mask, chunked_selective_log_softmax, _unsloth_get_mm_token_id, _unsloth_fix_mm_token_type_ids, @@ -48,7 +49,22 @@ from ..device_type import ( ALLOW_PREQUANTIZED_MODELS, ) import textwrap -from ._utils import _get_inference_mode_context_manager +from ._utils import _get_inference_mode_context_manager, UNSLOTH_ENABLE_LOGGING + +# One-time GRPO sequence-packing gates; mirrored into the generated trainer cache via RL_PRE_ITEMS. +UNSLOTH_GRPO_SEQ_PACKING_ON = os.environ.get("UNSLOTH_GRPO_SEQ_PACKING", "1").lower() not in ( + "0", + "false", + "no", + "off", +) +# Packing needs zoo#840's masked-column guard in grpo_compute_loss (installed zoo is fixed per-process). +try: + UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = "torch.where(_keep, new" in inspect.getsource( + RL_REPLACEMENTS["grpo_compute_loss"] + ) +except Exception: + UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False RL_EXTRA_ARGS = defaultdict(list) RL_FUNCTIONS = defaultdict(list) @@ -1359,6 +1375,212 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): ) os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ---- + # One varlen [1, sum L] forward replaces the padded [B, Lmax] loop (also fixes the + # left-pad RoPE error). Self-verified against the per-row forward, re-checked as T + # grows; falls back if a backend ignores packed_seq_lengths. + logprobs = None + _pk_result = None + _pk_use = False + _pk_enabled = UNSLOTH_GRPO_SEQ_PACKING_ON + # Without zoo#840's masked-column guard, zeroed prompt/pad columns turn NaN in exp(). + _pk_enabled = _pk_enabled and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD + _pk_ok = getattr(unwrapped_model, "_unsloth_seq_packing_nograd_ok", None) + if ( + _pk_enabled + and pixel_values is None + and token_type_ids is None + and mm_token_type_ids is None + and _pk_ok is not False + ): + try: + _pk_pad = self.processing_class.pad_token_id + _pk_keep = input_ids != _pk_pad + _pk_len = _pk_keep.sum(dim = 1) + _pk_len_cpu = _pk_len.tolist() # single GPU->CPU sync, reused below + _pk_nz_cpu = [_n for _n in _pk_len_cpu if _n > 0] + _pk_flat = input_ids[_pk_keep].unsqueeze(0) + _pk_T = _pk_flat.shape[1] + _pk_L = input_ids.shape[1] + _pk_W = logits_to_keep + max_left_pad + _pk_maxseg = max(_pk_nz_cpu) if _pk_nz_cpu else 0 + # sliding-window models lose the per-sequence local window in a packed stream + _pk_sw = getattr( + getattr(unwrapped_model, "config", None), "sliding_window", None + ) + _pk_sw_ok = not (isinstance(_pk_sw, int) and _pk_sw > 0 and _pk_maxseg > _pk_sw) + # per-row completion mask (same as the loss); prompt-only rows count as inactive + _pk_cmask = create_completion_attention_mask( + input_ids[:, -_pk_W:], left_pad_tokens_per_prompt, max_left_pad, _pk_pad + ) + _pk_active = int(_pk_cmask.any(dim = 1).sum()) + # skip the packed forward entirely at known-unsafe lengths (avoids a wasted pass / OOM) + _pk_unsafe = getattr( + unwrapped_model, "_unsloth_seq_packing_nograd_unsafe_T", None + ) + # cap the flattened forward at one padded [batch_size, seq_len] mini-batch's + # token budget; anything larger uses the chunked padded loop + _pk_cap = batch_size * seq_len + if ( + _pk_T >= 2 + and _pk_T <= _pk_cap + and len(_pk_nz_cpu) > 0 + and _pk_sw_ok + and not (_pk_unsafe is not None and _pk_T >= _pk_unsafe) + and (_pk_ok is True or _pk_active >= 2) + ): + # reset 0-based position_ids per segment + _pk_pos = (_pk_keep.cumsum(dim = 1) - 1)[_pk_keep].unsqueeze(0) + _pk_chunks = max(1, total_rows * multiplier) + _pk_nz_idx = _pk_keep.nonzero( + as_tuple = False + ) # [T, 2] = (row, col), row-major + _pk_within = _pk_nz_idx[1:, 0] == _pk_nz_idx[:-1, 0] # [T-1] + # per-row completion start after left-packing (matches create_completion_attention_mask) + _pk_cstart = (_pk_L - logits_to_keep) - left_pad_tokens_per_prompt # [rows] + _pk_ctgt = (_pk_nz_idx[1:, 1] >= _pk_cstart[_pk_nz_idx[1:, 0]]) & _pk_within + with _get_inference_mode_context_manager(model): + with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): + # use_cache=False: a KV cache silently disables varlen packing + _pk_hidden = unwrapped_model( + input_ids = _pk_flat, + position_ids = _pk_pos, + packed_seq_lengths = torch.tensor( + _pk_nz_cpu, dtype = torch.int32, device = input_ids.device + ), + use_cache = False, + ).logits + _pk_sel = chunked_hidden_states_selective_log_softmax( + _pk_hidden[0, :-1, :][_pk_ctgt].unsqueeze(0), + lm_head, + _pk_flat[0, 1:][_pk_ctgt].unsqueeze(0), + _pk_chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + )[0] + # GPT-OSS offload race guard (matches the padded loop) + device_synchronize() + # scatter each completion logprob back to its (row, col) so [:, -_pk_W:] matches padded + _pk_tgt = (_pk_nz_idx[1:, 0] * _pk_L + _pk_nz_idx[1:, 1])[_pk_ctgt] + _pk_result = ( + torch.zeros( + total_rows * _pk_L, + dtype = torch.float32, + device = input_ids.device, + ) + .index_put((_pk_tgt,), _pk_sel.to(torch.float32)) + .view(total_rows, _pk_L)[:, -_pk_W:] + ) + # re-verify when T or the longest segment grows past what was verified + # (a LongRoPE cache switch can change the result) + _pk_vT = int( + getattr(unwrapped_model, "_unsloth_seq_packing_nograd_verified_T", 0) + ) + _pk_vS = int( + getattr(unwrapped_model, "_unsloth_seq_packing_nograd_verified_seg", 0) + ) + # debug: hand-edit this condition to force re-verify every step + if _pk_ok is True and _pk_T <= _pk_vT and _pk_maxseg <= _pk_vS: + _pk_use = True # already verified for this shape + else: + # verify against the per-row forward (ground truth) + _pk_ref = torch.zeros_like(_pk_result) + with _get_inference_mode_context_manager(model): + with torch.amp.autocast( + device_type = "cuda", dtype = self._autocast_dtype + ): + for _pk_i in range(total_rows): + _pk_ni = _pk_len_cpu[_pk_i] + if _pk_ni < 2: + continue + _pk_rmask = _pk_keep[_pk_i] + _pk_real = input_ids[_pk_i][_pk_rmask].unsqueeze(0) + _pk_rpos = torch.arange( + _pk_ni, device = input_ids.device + ).unsqueeze(0) + _pk_rh = unwrapped_model( + input_ids = _pk_real, + position_ids = _pk_rpos, + use_cache = False, + ).logits + _pk_rsel = chunked_hidden_states_selective_log_softmax( + _pk_rh[:, :-1, :], + lm_head, + _pk_real[:, 1:], + 1, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + )[0] + _pk_rcols = _pk_rmask.nonzero(as_tuple = False).squeeze(1)[ + 1: + ] - (_pk_L - _pk_W) + _pk_rkeep = _pk_rcols >= 0 + _pk_ref[_pk_i, _pk_rcols[_pk_rkeep]] = _pk_rsel[ + _pk_rkeep + ].to(torch.float32) + device_synchronize() + # compare over the loss-mask region only + _pk_cm = _pk_cmask.float() + _pk_diff = float(((_pk_result - _pk_ref).abs() * _pk_cm).max()) + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO seq-packing (no-grad) verify: T={_pk_T} maxseg={_pk_maxseg} packed-vs-perrow max|d|={_pk_diff:.4f}", + flush = True, + ) + # kernel-noise floor ~0.25; cross-sample contamination is >= 2.4 + if _pk_diff < 7e-1: + unwrapped_model._unsloth_seq_packing_nograd_ok = True + # widen the trusted shape only when >= 2 completion rows exercised + # cross-sample packing; single-row passes prove nothing + if _pk_active >= 2: + unwrapped_model._unsloth_seq_packing_nograd_verified_T = max( + _pk_vT, _pk_T + ) + unwrapped_model._unsloth_seq_packing_nograd_verified_seg = max( + _pk_vS, _pk_maxseg + ) + _pk_ok = True + _pk_use = True + else: + _pk_use = False + if _pk_diff >= 1.5: + # contamination (attention ignores the packed mask): disable packing + unwrapped_model._unsloth_seq_packing_nograd_ok = False + else: + # likely a length boundary (LongRoPE): mark unsafe, keep smaller shapes + unwrapped_model._unsloth_seq_packing_nograd_unsafe_T = ( + _pk_T if _pk_unsafe is None else min(_pk_unsafe, _pk_T) + ) + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO seq-packing (no-grad) fell back at T={_pk_T} (diff={_pk_diff:.3f})", + flush = True, + ) + except Exception as _pk_err: + # any failure: drop intermediates, use the padded loop, do not retry + _pk_hidden = None + _pk_sel = None + _pk_result = None + _pk_use = False + if isinstance(_pk_err, torch.cuda.OutOfMemoryError): + torch.cuda.empty_cache() + unwrapped_model._unsloth_seq_packing_nograd_ok = False + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO sequence-packing (no-grad) disabled (fell back to padded): {_pk_err!r}", + flush = True, + ) + if _pk_use and _pk_result is not None: + logprobs = _pk_result # verified -> skip the loop + zipped_inputs = [] + else: + # free packed intermediates before running the padded loop + _pk_hidden = _pk_sel = _pk_result = _pk_ref = None + with _get_inference_mode_context_manager(model): for ( input_ids_chunk, @@ -1443,7 +1665,8 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): # However, it seems that this line does not slow down or disrupt models. device_synchronize() all_logprobs_list.append(logprobs_chunk) - logprobs = torch.cat(all_logprobs_list, dim = 0) + if logprobs is None: # padded fallback when packing was not used + logprobs = torch.cat(all_logprobs_list, dim = 0) entropies = None os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" @@ -1523,6 +1746,24 @@ RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_accumulated_loss)) RL_PRE_ITEMS["grpo_trainer"].append(grpo_compute_loss_slow) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(grpo_update_SamplingParams)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_get_inference_mode_context_manager)) +# inspect.getsource inlines function bodies but not module imports, so constants the inlined +# grpo functions reference (e.g. UNSLOTH_ENABLE_LOGGING) must be redefined in the generated cache. +RL_PRE_ITEMS["grpo_trainer"].append( + "import os as _unsloth_os\n" + "UNSLOTH_ENABLE_LOGGING = _unsloth_os.environ.get('UNSLOTH_ENABLE_LOGGING', '0') in ('1', 'True', 'true')\n" +) +# One-time sequence-packing gates, same values as the module-top constants above. +RL_PRE_ITEMS["grpo_trainer"].append( + "UNSLOTH_GRPO_SEQ_PACKING_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_SEQ_PACKING', '1').lower() not in ('0', 'false', 'no', 'off')\n" +) +RL_PRE_ITEMS["grpo_trainer"].append( + "try:\n" + " import inspect as _unsloth_inspect\n" + " from unsloth_zoo.rl_replacements import RL_REPLACEMENTS as _unsloth_zoo_RL\n" + " UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = 'torch.where(_keep, new' in _unsloth_inspect.getsource(_unsloth_zoo_RL['grpo_compute_loss'])\n" + "except Exception:\n" + " UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False\n" +) # Edit _get_per_token_logps to handle mixed precision From 08e133cd6b035adc75474985cb8c097c5b1a3f95 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:26:24 -0700 Subject: [PATCH 004/113] Add PrefixGrouper for GRPO: dedup the shared prompt across a group's completions (#6871) * GRPO: optional sequence packing for the no-grad old/ref logp path Add an opt-in sequence-packing fast path to _get_per_token_logps_and_entropies, enabled with UNSLOTH_GRPO_SEQ_PACKING=1. When the batch is text-only, the padded [B, Lmax] per-chunk forward is replaced by a single varlen [1, sum L] forward (BlockDiagonalCausalMask via packed_seq_lengths with reset position_ids). Per-token logps use the same float32 chunked_hidden_states_selective_log_softmax as the padded path, so the old and reference logps are bit-for-bit identical. Safety: the packed path is self-verified once against the padded ground truth on a batch that has at least two rows with real completion tokens (self._unsloth_seq_packing_nograd_ok), so cross-sample contamination would actually manifest; a degenerate all-pad / fully tool-masked batch leaves the verdict unset and re-verifies later. If a backend silently ignores packed_seq_lengths (flat batch run under a normal causal mask, samples leaking across boundaries), the packed logps will not match and packing is disabled instead of corrupting logps. It also forces use_cache=False (a populated past_key_value disables varlen packing), skips packing when a sliding window is shorter than the packed stream, runs the same GPT-OSS offload device_synchronize the padded loop uses, and falls back on any exception (UNSLOTH_GRPO_SEQ_PACKING_DEBUG=1 prints the reason). Default off, so existing behavior is unchanged. Pairs with the matching gradient-path change in unsloth_zoo so the full GRPO logp + loss + backward can run packed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: address review feedback - Cache the packed-vs-padded verdict per unwrapped model instead of on the trainer, so a separately forwarded reference model is verified on its own forward path rather than inheriting the policy model's verdict. - Force the padded path when token_type_ids or mm_token_type_ids are present, matching the extra vision kwargs the padded loop forwards. - Require the xformers varlen backend before packing. Without it the packed mask falls back to a dense O(T^2) SDPA mask that can OOM on the flattened batch, so we keep the padded loop in that case. - On any packed-forward failure (missing backend, OOM, unsupported forward) empty the cache on OOM, disable packing for that model, and fall back to the chunked padded loop instead of retrying every step. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: default-on, verify against per-row reference Redesign of the optional sequence-packing fast path for the no-grad old/ref logprob recompute, after establishing that the packed forward is the exact per-row computation and the padded batch forward is the side that mis-positions left-padded rows on long completions. - Default the packing on (UNSLOTH_GRPO_SEQ_PACKING, disable with 0). - Verify the packed logprobs against the per-row clean forward (each row's real tokens alone, reset 0-based positions, no padding), not the padded batch which is itself wrong for left-padding. Cross-sample contamination (a backend ignoring packed_seq_lengths) shows up as a large mismatch and falls back to the padded loop. - Make the trust decision shape and RoPE aware: re-verify whenever the packed total length or the longest segment grows past what was verified, so a later batch crossing a LongRoPE short/long cache boundary is re-checked instead of trusted blindly. - Run lm_head only on completion-prediction positions instead of every packed prompt token, so long-prompt/short-completion batches do not pay for projecting the whole packed prompt. - Drop the hard xformers import so the path also runs in FlashAttention-only environments; the per-row verification guards correctness regardless of backend. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: disable entirely on cross-sample mismatch When the per-row verification fails, distinguish the two failure modes by magnitude instead of by sequence length: - A large mismatch (>= 1.5) is the cross-sample contamination signature: the model's attention does not honor the block-diagonal packed mask (seen on some MoE / custom-attention models, e.g. qwen2_moe). Disable packing entirely for the model so later batches do not pay the verification cost again. - A moderate mismatch is more likely a length-boundary effect (a LongRoPE short/long cache switch): keep marking just that length region unsafe so packing still runs for smaller shapes. Validated: Qwen1.5-MoE falls back after a single verification (grad and no-grad ok flags go False, no re-verify on later steps); dense Llama-3.2 and Qwen3 still verify and engage packing. * GRPO no-grad packing: trim comments to be concise * GRPO no-grad packing: fix per-row completion boundary for left-padded rows The completion-target selection used a single global boundary (col >= L - logits_to_keep). After left-packing, each row's completion starts at (L - logits_to_keep) - left_pad[row], so for left-padded rows the first left_pad completion tokens fall below the global boundary and were dropped, leaving 0 logprobs at real completion positions that the loss mask keeps. Use the per-row boundary so packed coverage matches create_completion_attention_mask exactly, and widen the self-verify mask to the full per-row completion region so it can catch coverage gaps. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO no-grad packing: gate verification on real completion rows Count active rows via create_completion_attention_mask (the same mask the loss uses) instead of any non-pad token in the packed window. Prompt-only rows carry prompt-overflow tokens in the window and could otherwise satisfy the >= 2 verification guard, letting a batch with a single real completion row cache a trust decision. This matches the gradient path, which already gates on the completion mask. The same mask is reused for the self-verify comparison. * GRPO no-grad packing: gate debug logging on UNSLOTH_ENABLE_LOGGING Use the shared UNSLOTH_ENABLE_LOGGING global (import_fixes, re-exported by _utils) instead of a bespoke UNSLOTH_GRPO_SEQ_PACKING_DEBUG env var for the packing debug prints, matching the rest of the codebase. * GRPO packing: import UNSLOTH_ENABLE_LOGGING inside the injected logp function _get_per_token_logps_and_entropies is copied verbatim into the generated GRPO trainer via inspect.getsource, and that module never imported UNSLOTH_ENABLE_LOGGING, so the default-on packing verify path raised NameError (and the except handler re-raised it). Import the flag locally, before the try, so the name is defined in the generated module too. Drop it from the now-unused module-level import. * GRPO no-grad packing: harden unsafe-length skip, verify guard, fallback cleanup Three fixes to the no-grad logp packing path, mirroring the grad path: - skip the packed forward for known-unsafe lengths by reading unsafe_T and gating on it before the forward, instead of running the full packed pass and the result build only to discard them (wastes a pass, can OOM at large T) - only widen the verified T/seg envelope when >= 2 completion rows actually exercised cross-sample packing; a < 2 row batch cannot expose leakage, so it must not extend the trusted shape that later multi-row batches skip verify for - drop the packed intermediates (hidden/sel/result/ref) before the padded fallback loop so it does not run with the flattened hidden state still resident * GRPO no-grad packing: cap the flattened forward at one mini-batch budget The packed path built a single [1, sum L] forward over every row before any size check, so a large batch could exceed the memory the padded path bounds per mini-batch. Gate packing on _pk_T <= _pk_cap (B * seq_len, one padded mini-batch's token budget); larger batches fall back to the chunked padded loop. * GRPO no-grad packing: disable unless unsloth_zoo has the masked-column guard The packed path leaves masked prompt/pad logprob columns at 0, which only stays finite if unsloth_zoo grpo_compute_loss zeroes them before exp() (zoo#840). An older unsloth_zoo without that guard would NaN. Detect the guard once (cached on the model) via inspect.getsource and gate packing on it, so #6738 is safe with any unsloth_zoo version and re-enables packing automatically once a guarded zoo is installed, independent of the pinned lower bound. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO packing: hoist env gates and zoo-guard detection to one-time module checks Read UNSLOTH_GRPO_SEQ_PACKING and detect the unsloth_zoo masked-column guard once at import time (module constants plus RL_PRE_ITEMS for the generated trainer cache) instead of per call, and drop the in-function UNSLOTH_ENABLE_LOGGING import for a module-top one. The UNSLOTH_GRPO_SEQ_PACKING_VERIFY force-verify debug knob is commented out, kept in place for hand re-enable; the first-use and envelope-growth self-verify stays active. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO packing: cap the flattened forward by the padded chunk rows B counts chunks at this point, so B * seq_len understated (small runs) or overstated (large runs) the padded mini-batch token budget; use batch_size * seq_len, the rows the padded loop actually forwards per chunk. * Add PrefixGrouper for GRPO: dedup the shared prompt across a group's completions In GRPO every prompt spawns G=num_generations completions that share the prompt prefix, so the trunk logprob forward re-encodes that prefix G times. PrefixGrouper stores the prefix once and concatenates only the G suffixes behind a FlexAttention shared-prefix mask, cutting the forward from G*(P+R) to P+G*R tokens across both the no-grad old/ref forwards and the grad logp forward. Default off behind the UNSLOTH_GRPO_PREFIX_GROUPER env gate, so the gate-unset path is byte-identical to today. A tok_r auto-gate and a first-use self-verify (fall back and mark the shape unsafe on mismatch) keep it from ever shipping wrong logprobs silently. Wired for llama, mistral, qwen3, gemma2, cohere, granite and falcon_h1, plus qwen2 and gemma through the shared LlamaAttention_fast_forward. Stacked on the GRPO sequence-packing PR (#6738); the grad path lands in a companion unsloth-zoo PR. Also fixes a latent UNSLOTH_ENABLE_LOGGING NameError in the seq-packing no-grad verify path by defining the name as a generated-cache pre-item. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: enforce the sliding-window cap, gate softcap models, bound the mask cache Add a max_segment_cap kwarg to build_group_layout so it falls back when a group's span (prefix + longest suffix) exceeds the model's local window, and pass the config sliding_window into the no-grad engage gate the same way the packed _pk guard derives it. Skip PrefixGrouper entirely for attn_logit_softcapping models, since the FlexAttention kernel never applies logit softcapping. Bound _BLOCK_MASK_CACHE to a FIFO of 8 so per-step lengths cannot pin BlockMasks forever, release the PG hidden before the verify forward, and align the UNSLOTH_ENABLE_LOGGING pre-item truthiness with the canonical form. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: vectorize the real-column scan in build_group_layout Replace the per-row O(B*L) Python scan of the keep mask with a GPU-derived contiguous-run fast path (first real column + count per row), keeping the general scan only as a fallback for non-contiguous rows. Works for both call sites: the no-grad layout (left-padded prompt + right-padded completion, run does not start at column 0) and the grad layout (left-packed). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: hoist the gate and kernel imports to one-time module checks, AGPLv3 headers Read UNSLOTH_GRPO_PREFIX_GROUPER and resolve the prefix_grouper imports once at module level (source constants plus an RL_PRE_ITEMS entry for the generated trainer cache) instead of per call, matching the sequence-packing gates. The prefix_grouper env helpers become one-time module reads with unchanged signatures, and attention_dispatch resolves the FlexAttention kernel once behind the same gate (lazy fallback kept). The two new prefix_grouper files move to AGPLv3 headers. * PrefixGrouper: length-envelope trust and hybrid SSM exclusion Verified signatures now record (max T, max segment) and re-verify when either grows, matching the packed path's envelope. Hybrid SSM models (FalconH1 etc.) are excluded at the gate since only attention gets the shared-prefix isolation, and the FalconH1 wiring is removed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: defer the unverified no-grad forward until the packed reference exists Unverified shapes no longer run the whole-batch shared-prefix forward up front; it now runs at the verify site, only when the packed path produced a reference. A declined packed path (budget, window) therefore costs no wasted PG forward per step. Trusted shapes still run it first to skip the full-row forward, with the same fallback. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * PrefixGrouper: disable under vLLM (fast_inference=True) With colocated vLLM generation the rollout dominates the GRPO step, so the shared-prefix training forward saves little end-to-end and its first-use self-verify (which also runs the full-row path) is net overhead. Gate PG on not use_vllm so it only engages on the raw transformers path, where the training forward is on the critical path. Packing is unaffected. * PrefixGrouper: compile the FlexAttention kernel with dynamic shapes GRPO changes the packed length T almost every batch. With dynamic=False the flex forward+backward kernel recompiled on every new T (~14s each on a 4B trunk), which dominated the step and made PG a net loss. dynamic=True compiles once, then reuses the kernel across all lengths recompile-free (a new shape drops from ~14s to ~1.4ms after a two-graph warmup). T is still padded to a multiple of 128 for the backward block assertion. * PrefixGrouper: default on Enable PrefixGrouper by default (UNSLOTH_GRPO_PREFIX_GROUPER defaults to 1; set 0 to disable). Still auto-disabled under vLLM (fast_inference=True) and by the arch/softcap/ SSM/tok_r gates, and the first-use self-verify falls back on any mismatch, so this is a memory-first default on the raw-transformers path with no correctness risk. * GRPO PrefixGrouper: gate on zoo masked-column guard and exclude MoE - Require the zoo masked-column guard (zoo#840) before PrefixGrouper can engage. PG rides the sequence-packing path, so when the first-step self-verify is off the fast path trusts PG output directly; without the guard those masked columns feed NaN into the packed loss. Gate PG on the same UNSLOTH_ZOO_HAS_MASKED_COL_GUARD the packing path already checks. - Exclude MoE configs (num_experts, num_local_experts, n_routed_experts, moe_intermediate_size) alongside the hybrid-SSM markers. Only the threaded attention forwards carry the shared-prefix isolation, so a MoE decoder that does not forward prefix_seg_info would let suffixes leak across completions. - Refresh the stale default-off comments now that UNSLOTH_GRPO_PREFIX_GROUPER is on by default. * GRPO PrefixGrouper: import chunked_hidden_states_selective_log_softmax The shared-prefix forward passes chunked_hidden_states_selective_log_softmax into extract_logps, but the name was only ever provided by the generated trainer cache (rl.py injects grpo_selective_log_softmax_code), never bound in this module. Import it from unsloth_zoo.rl_replacements next to its sibling chunked_selective_log_softmax so the source resolves the name in every scope (the new _pg_run_forward closure included). No runtime change: the cache still defines the function via template injection. * GRPO PrefixGrouper: dropout gate, device-safe layout, Mistral mask skip Addresses three review findings on the shared-prefix path: - Skip PrefixGrouper when the model sets a nonzero attention_dropout. The normal backends apply config.attention_dropout while training (e.g. Granite dense flash/sdpa/xformers), but the FlexAttention shared-prefix path is deterministic, so gate PG off for those configs rather than train on mismatched activations. - Move the shared-prefix mask labels to the consumer (Q) device in get_block_mask and the target index maps to hidden.device in extract_logps, mirroring the packed path moving its metadata to the consumer device. Prevents cross-device indexing when the model is sharded across GPUs. - Do not synthesize a causal attention_mask in the Mistral forward when prefix_seg_info is present. On the no-xFormers path that synthetic mask tripped resolve_prefix_seg_info and forced PG to always fall back to the packed forward. * GRPO sequence packing: tighten comments * GRPO PrefixGrouper: tighten comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO PrefixGrouper: persistent disable on runtime failure; build block-mask labels with inference mode disabled - rl_replacements: on a PG forward exception (FlexAttention/Triton compile failure or OOM), set a model-level _unsloth_prefix_grouper_nograd_disabled flag and consult it in the engage gate, mirroring the seq-packing handler, so a GPU-wide failure is not retried and re-paid every step. - prefix_grouper_kernel: move the .to(device) label copies inside the inference_mode(False) block so a cross-device (model-parallel shard) first build does not capture inference tensors, which otherwise cannot be saved for backward when the grad training forward reuses the cached BlockMask. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- unsloth/models/cohere.py | 5 + unsloth/models/gemma2.py | 7 + unsloth/models/granite.py | 5 + unsloth/models/llama.py | 12 +- unsloth/models/mistral.py | 13 +- unsloth/models/qwen3.py | 5 + unsloth/models/rl_replacements.py | 263 ++++++++++++++- unsloth/utils/attention_dispatch.py | 65 ++++ unsloth/utils/prefix_grouper.py | 351 ++++++++++++++++++++ unsloth/utils/prefix_grouper_kernel.py | 436 +++++++++++++++++++++++++ 10 files changed, 1156 insertions(+), 6 deletions(-) create mode 100644 unsloth/utils/prefix_grouper.py create mode 100644 unsloth/utils/prefix_grouper_kernel.py diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index 0b7f3ab973..cb367d451e 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -22,6 +22,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, ) try: @@ -151,6 +152,9 @@ def CohereAttention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -161,6 +165,7 @@ def CohereAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 68b9ebe22f..4a0531db78 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -22,6 +22,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, SDPA, ) from .gemma import ( @@ -168,6 +169,11 @@ def Gemma2Attention_fast_forward( }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. gemma2 is + # sliding-window and softcapped: the engage gate caps spans at the window and + # excludes softcap models entirely, so PG never engages here. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -179,6 +185,7 @@ def Gemma2Attention_fast_forward( attention_mask = attention_mask, causal_mask = causal_mask, sliding_window = sliding_window, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index f5b0f57aa6..4dedf642eb 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -23,6 +23,7 @@ from ..utils.attention_dispatch import ( AttentionContext, run_attention, select_attention_backend, + resolve_prefix_seg_info, SDPA, ) from .llama import ( @@ -159,6 +160,9 @@ def GraniteAttention_fast_forward( }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -169,6 +173,7 @@ def GraniteAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 6bec95b577..417a88f480 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -39,6 +39,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from torch.nn.functional import scaled_dot_product_attention from transformers import __version__ as transformers_version @@ -738,6 +739,10 @@ def LlamaAttention_fast_forward( flash_dense_kwargs = {"causal": True}, flash_varlen_kwargs = {"dropout_p": 0.0, "causal": True}, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward (same route + # as packed_seq_lengths); misuse (KV cache / padding mask) raises. None => byte-identical + # default. Reuse of this forward also carries the branch to qwen2 & gemma. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -748,6 +753,7 @@ def LlamaAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = config, context = context, Q = Q, K = K, V = V) @@ -895,8 +901,10 @@ def LlamaModel_fast_forward( seq_length_with_past = seq_length # Fix out of bounds tokenization unless we were given packed metadata - allow_overlength = getattr(self, "_unsloth_allow_packed_overlength", False) or ( - "packed_seq_lengths" in kwargs + allow_overlength = ( + getattr(self, "_unsloth_allow_packed_overlength", False) + or ("packed_seq_lengths" in kwargs) + or ("prefix_seg_info" in kwargs and kwargs["prefix_seg_info"] is not None) ) if hasattr(self, "max_seq_length") and not allow_overlength: if seq_length > self.max_seq_length: diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index df2a4de5bd..4350565fe2 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -27,6 +27,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from .llama import ( LlamaRotaryEmbedding, @@ -124,6 +125,9 @@ def MistralAttention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -134,6 +138,7 @@ def MistralAttention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) @@ -161,7 +166,13 @@ def MistralForCausalLM_fast_forward( *args, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: - if causal_mask is None and past_key_values is None: + # PrefixGrouper brings its own mask: a synthesized causal attention_mask would trip + # resolve_prefix_seg_info on the no-xFormers path and force a fallback. + if ( + causal_mask is None + and past_key_values is None + and kwargs.get("prefix_seg_info", None) is None + ): bsz, q_len = input_ids.shape sliding_window = getattr(self.config, "sliding_window", None) diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index e28e72d3ea..0d05a2d538 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -23,6 +23,7 @@ from ..utils.attention_dispatch import ( run_attention, SDPA, select_attention_backend, + resolve_prefix_seg_info, ) from .llama import ( LlamaRotaryEmbedding, @@ -146,6 +147,9 @@ def Qwen3Attention_fast_forward( "softmax_scale": getattr(self, "softmax_scale", None), }, ) + # PrefixGrouper seg table rides in **kwargs from the GRPO logprob forward; misuse + # (KV cache / padding mask) raises. None => byte-identical default. + _pg_seg = resolve_prefix_seg_info(kwargs, past_key_value, attention_mask) context = AttentionContext( bsz = bsz, q_len = q_len, @@ -156,6 +160,7 @@ def Qwen3Attention_fast_forward( seq_info = seq_info, attention_mask = attention_mask, causal_mask = causal_mask, + prefix_seg_info = _pg_seg, ) A = run_attention(config = attention_config, context = context, Q = Q, K = K, V = V) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 0573fd5fd8..098950de08 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -31,6 +31,7 @@ from unsloth_zoo.rl_replacements import ( left_pack_padding, create_completion_attention_mask, chunked_selective_log_softmax, + chunked_hidden_states_selective_log_softmax, _unsloth_get_mm_token_id, _unsloth_fix_mm_token_type_ids, ) @@ -65,6 +66,25 @@ try: ) except Exception: UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False +# One-time PrefixGrouper gate; any import failure degrades to "PrefixGrouper off". +_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None +UNSLOTH_GRPO_PREFIX_GROUPER_ON = os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ( + "0", + "false", + "no", + "off", +) +if UNSLOTH_GRPO_PREFIX_GROUPER_ON: + try: + from ..utils.prefix_grouper import ( + build_group_layout as _pg_build_layout, + prefix_grouper_enabled as _pg_enabled_fn, + verify_on as _pg_verify_on, + tol_ok as _pg_tol_ok, + TOL_KILL as _PG_TOL_KILL, + ) + except Exception: + UNSLOTH_GRPO_PREFIX_GROUPER_ON = False RL_EXTRA_ARGS = defaultdict(list) RL_FUNCTIONS = defaultdict(list) @@ -1380,6 +1400,166 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): # left-pad RoPE error). Self-verified against the per-row forward, re-checked as T # grows; falls back if a backend ignores packed_seq_lengths. logprobs = None + + # ---- PrefixGrouper (GRPO shared-prompt dedup; default ON, exact + self-verified) ---- + # G completions per prompt share the prefix; the packed path forwards it G times, + # PrefixGrouper stores it once (FlexAttention shared-prefix mask), cutting the trunk + # forward from G*(P+R) to P+G*R tokens. Gated by UNSLOTH_GRPO_PREFIX_GROUPER (needs + # seq-packing), tok_r auto-gate, and first-use self-verify vs the packed path + # (mismatch => fall back + mark unsafe), so a mask/isolation regression cannot ship + # silently. When off / ungrouped / unverified, the packed path below runs as before. + _pg_result = None + _pg_use = False + _pg_skip_pk = False # once a shape is PG-verified, skip the full-row forward + _pg_forward_fn = None # deferred PG forward (runs at the verify site below) + _pg_num_gen = getattr(self, "num_generations", None) + # Env gate hoisted to module level (mirrored via RL_PRE_ITEMS). Skip PG under vLLM + # (fast_inference=True): the rollout dominates the step, so PG saves little and its + # first-use self-verify is net overhead. + _pg_engage = ( + UNSLOTH_GRPO_PREFIX_GROUPER_ON + and not getattr(self, "use_vllm", False) + and not getattr(unwrapped_model, "_unsloth_prefix_grouper_nograd_disabled", False) + ) + if _pg_engage: + try: + # Skip softcap models (the flex kernel never applies attn_logit_softcapping) + # and hybrid SSM / MoE models: only the threaded attention forwards get the + # shared-prefix isolation, so a Mamba or MoE decoder that does not forward + # prefix_seg_info would leak suffixes across completions. PG also rides on + # sequence packing, so it needs the same zoo masked-column guard. + _pg_cfg = getattr(unwrapped_model, "config", None) + _pg_engage = ( + _pg_enabled_fn() + and UNSLOTH_ZOO_HAS_MASKED_COL_GUARD + and pixel_values is None + and token_type_ids is None + and mm_token_type_ids is None + and _pg_num_gen is not None + and _pg_num_gen >= 2 + and not getattr(_pg_cfg, "attn_logit_softcapping", None) + # normal backends apply config.attention_dropout in training; the flex + # path is deterministic, so skip PG when it is set. + and not getattr(_pg_cfg, "attention_dropout", 0) + and not any( + getattr(_pg_cfg, _pg_a, None) is not None + for _pg_a in ( + "mamba_d_ssm", + "mamba_d_state", + "mamba_expand", + "num_experts", + "num_local_experts", + "n_routed_experts", + "moe_intermediate_size", + ) + ) + ) + except Exception: + _pg_engage = False + if _pg_engage: + try: + _pg_pad = self.processing_class.pad_token_id + # cap the PG span (P+max(R)) at the sliding window, like the packed _pk_sw guard. + _pg_sw = getattr( + getattr(unwrapped_model, "config", None), "sliding_window", None + ) + if not (isinstance(_pg_sw, int) and _pg_sw > 0): + _pg_sw = None + _pg_layout = _pg_build_layout( + input_ids, + logits_to_keep, + _pg_pad, + _pg_num_gen, + left_pad_tokens_per_prompt, + max_segment_cap = _pg_sw, + ) + _pg_unsafe = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None + ) + if _pg_unsafe is None: + _pg_unsafe = set() + if _pg_layout is not None and _pg_layout.signature not in _pg_unsafe: + _pg_sig = _pg_layout.signature + _pg_verified = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None + ) + if _pg_verified is None: + _pg_verified = set() + _pg_chunks = max(1, total_rows * multiplier) + + def _pg_run_forward(_pg_layout = _pg_layout, _pg_chunks = _pg_chunks): + with _get_inference_mode_context_manager(model): + with torch.amp.autocast( + device_type = "cuda", dtype = self._autocast_dtype + ): + _pg_hidden = unwrapped_model( + input_ids = _pg_layout.flat_ids, + position_ids = _pg_layout.position_ids, + prefix_seg_info = _pg_layout.prefix_seg_info, + use_cache = False, + ).logits + _pg_r = _pg_layout.extract_logps( + _pg_hidden, + lm_head, + chunked_hidden_states_selective_log_softmax, + _pg_chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + ) + _pg_hidden = None # release before any verify forward + device_synchronize() + # clip to the loss window [B, logits_to_keep+max_left_pad] + _pg_w = logits_to_keep + max_left_pad + if _pg_r.shape[1] > _pg_w: + _pg_r = _pg_r[:, -_pg_w:] + return _pg_r + + # trust only within the verified envelope: re-verify when T or the + # longest segment grows, like the packed path + _pg_T = int(_pg_layout.flat_ids.shape[1]) + _pg_maxseg = int(_pg_layout.position_ids.max()) + 1 + _pg_env = ( + _pg_verified.get(_pg_sig) if isinstance(_pg_verified, dict) else None + ) + if (not _pg_verify_on()) or ( + _pg_env is not None and _pg_T <= _pg_env[0] and _pg_maxseg <= _pg_env[1] + ): + # trusted shape: run PG now and skip the full-row forward below + _pg_result = _pg_run_forward() + _pg_use = True + _pg_skip_pk = True + else: + # unverified shape: defer the forward until the packed reference + # exists (verify site below), so a declined packed path never wastes + # a whole-batch PG forward + _pg_forward_fn = _pg_run_forward + except Exception as _pg_err: + _pg_result = None + _pg_use = False + _pg_skip_pk = False + _pg_forward_fn = None + # A FlexAttention/Triton compile failure or OOM here is GPU-wide, not + # layout-specific, so retrying the same PG forward every step just re-pays + # the failure. Persistently disable PG (mirrors the seq-packing handler + # setting _unsloth_seq_packing_nograd_ok = False); the packed/padded path + # below still produces the exact result. + unwrapped_model._unsloth_prefix_grouper_nograd_disabled = True + if isinstance(_pg_err, torch.cuda.OutOfMemoryError): + torch.cuda.empty_cache() + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) disabled (fell back to packed): {_pg_err!r}", + flush = True, + ) + + # ---- Sequence packing (default-on; disable with UNSLOTH_GRPO_SEQ_PACKING=0) ---- + # One varlen [1, sum L] block-diagonal forward replaces the padded [B, Lmax] loop + # (exact per-row result; also fixes the padded path's left-pad RoPE error). + # Self-verified vs the per-row forward, re-checked as T grows; falls back if a + # backend ignores packed_seq_lengths. lm_head runs on completion positions only. _pk_result = None _pk_use = False _pk_enabled = UNSLOTH_GRPO_SEQ_PACKING_ON @@ -1388,6 +1568,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): _pk_ok = getattr(unwrapped_model, "_unsloth_seq_packing_nograd_ok", None) if ( _pk_enabled + and not _pg_skip_pk and pixel_values is None and token_type_ids is None and mm_token_type_ids is None @@ -1462,7 +1643,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): )[0] # GPT-OSS offload race guard (matches the padded loop) device_synchronize() - # scatter each completion logprob back to its (row, col) so [:, -_pk_W:] matches padded + # scatter each logprob back to its (row, col) so [:, -_pk_W:] matches padded _pk_tgt = (_pk_nz_idx[1:, 0] * _pk_L + _pk_nz_idx[1:, 1])[_pk_ctgt] _pk_result = ( torch.zeros( @@ -1574,7 +1755,73 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): f"[Unsloth] GRPO sequence-packing (no-grad) disabled (fell back to padded): {_pk_err!r}", flush = True, ) - if _pk_use and _pk_result is not None: + # ---- PrefixGrouper first-use self-verify (no-grad) ---- + # Compare the untrusted PG result to the full-row packed result (itself verified vs + # per-row) over the completion mask: < tol_ok -> trust the structure; >= TOL_KILL -> + # unsafe forever; borderline -> fall back this shape. + if _pg_forward_fn is not None and not _pg_use: + if _pk_use and _pk_result is not None: + try: + # deferred PG forward, run only now that the packed reference exists + _pg_result = _pg_forward_fn() + _pg_W2 = logits_to_keep + max_left_pad + _pg_cm = create_completion_attention_mask( + input_ids[:, -_pg_W2:], + left_pad_tokens_per_prompt, + max_left_pad, + self.processing_class.pad_token_id, + ).float() + _pg_a = _pg_result[:, -_pg_W2:].float() + _pg_b = _pk_result[:, -_pg_W2:].float() + _pg_diff = float(((_pg_a - _pg_b).abs() * _pg_cm).max()) + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) verify: sig={_pg_layout.signature} " + f"shared-prefix vs full-row-packed max|d|={_pg_diff:.4f}", + flush = True, + ) + if _pg_diff < _pg_tol_ok(): + _pg_v = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_verified", None + ) + if not isinstance(_pg_v, dict): + _pg_v = {} + _pg_vT = int(_pg_layout.flat_ids.shape[1]) + _pg_vS = int(_pg_layout.position_ids.max()) + 1 + _pg_old = _pg_v.get(_pg_layout.signature, (0, 0)) + _pg_v[_pg_layout.signature] = ( + max(_pg_vT, _pg_old[0]), + max(_pg_vS, _pg_old[1]), + ) + unwrapped_model._unsloth_prefix_grouper_nograd_verified = _pg_v + _pg_use = True + else: + _pg_u = getattr( + unwrapped_model, "_unsloth_prefix_grouper_nograd_unsafe", None + ) + if _pg_u is None: + _pg_u = set() + if _pg_diff >= _PG_TOL_KILL: + _pg_u.add(_pg_layout.signature) + unwrapped_model._unsloth_prefix_grouper_nograd_unsafe = _pg_u + _pg_use = False + except Exception as _pg_err3: + _pg_result = None + _pg_use = False + if isinstance(_pg_err3, torch.cuda.OutOfMemoryError): + torch.cuda.empty_cache() + os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1" + if UNSLOTH_ENABLE_LOGGING: + print( + f"[Unsloth] GRPO PrefixGrouper (no-grad) verify failed (fell back to packed): {_pg_err3!r}", + flush = True, + ) + # else: no packed reference (packing off/failed) -> cannot verify; fall back. + + if _pg_use and _pg_result is not None: + logprobs = _pg_result # PrefixGrouper verified/trusted -> skip the loop + zipped_inputs = [] + elif _pk_use and _pk_result is not None: logprobs = _pk_result # verified -> skip the loop zipped_inputs = [] else: @@ -1752,7 +1999,7 @@ RL_PRE_ITEMS["grpo_trainer"].append( "import os as _unsloth_os\n" "UNSLOTH_ENABLE_LOGGING = _unsloth_os.environ.get('UNSLOTH_ENABLE_LOGGING', '0') in ('1', 'True', 'true')\n" ) -# One-time sequence-packing gates, same values as the module-top constants above. +# Sequence-packing gates, same values as the module-top constants. RL_PRE_ITEMS["grpo_trainer"].append( "UNSLOTH_GRPO_SEQ_PACKING_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_SEQ_PACKING', '1').lower() not in ('0', 'false', 'no', 'off')\n" ) @@ -1764,6 +2011,16 @@ RL_PRE_ITEMS["grpo_trainer"].append( "except Exception:\n" " UNSLOTH_ZOO_HAS_MASKED_COL_GUARD = False\n" ) +# PrefixGrouper gate, same shape as the module-top constants. +RL_PRE_ITEMS["grpo_trainer"].append( + "_pg_build_layout = _pg_enabled_fn = _pg_verify_on = _pg_tol_ok = _PG_TOL_KILL = None\n" + "UNSLOTH_GRPO_PREFIX_GROUPER_ON = _unsloth_os.environ.get('UNSLOTH_GRPO_PREFIX_GROUPER', '1').lower() not in ('0', 'false', 'no', 'off')\n" + "if UNSLOTH_GRPO_PREFIX_GROUPER_ON:\n" + " try:\n" + " from unsloth.utils.prefix_grouper import build_group_layout as _pg_build_layout, prefix_grouper_enabled as _pg_enabled_fn, verify_on as _pg_verify_on, tol_ok as _pg_tol_ok, TOL_KILL as _PG_TOL_KILL\n" + " except Exception:\n" + " UNSLOTH_GRPO_PREFIX_GROUPER_ON = False\n" +) # Edit _get_per_token_logps to handle mixed precision diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 2e984bad0a..68fb33dad9 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -17,6 +17,7 @@ from __future__ import annotations +import os from dataclasses import dataclass from typing import Any, Optional, Tuple @@ -42,6 +43,17 @@ if HAS_XFORMERS and torch.cuda.is_available(): HAS_XFORMERS = False SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "") +# PrefixGrouper kernel, resolved once when the env gate is on so PG-off users never load +# torch flex_attention. +_flex_shared_prefix_attention = None +if os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER", "1").lower() not in ("0", "false", "no", "off"): + try: + from .prefix_grouper_kernel import ( + flex_shared_prefix_attention as _flex_shared_prefix_attention, + ) + except Exception: + _flex_shared_prefix_attention = None + FLASH_VARLEN = "flash_varlen" FLASH_DENSE = "flash_dense" XFORMERS = "xformers" @@ -84,6 +96,9 @@ class AttentionContext: attention_mask: Optional[Tensor] causal_mask: Optional[Any] sliding_window: Optional[int] = None + # PrefixGrouper: non-None routes Q/K/V through the FlexAttention shared-prefix kernel; + # None leaves every existing construction/behavior unchanged. + prefix_seg_info: Optional[Any] = None def select_attention_backend(use_varlen: bool = False) -> str: @@ -99,6 +114,33 @@ def select_attention_backend(use_varlen: bool = False) -> str: return SDPA +def resolve_prefix_seg_info(kwargs, past_key_value, attention_mask): + """PrefixGrouper shared-prefix segment table resolver for the arch attention forwards. + + The GRPO PrefixGrouper packed path rides a ``PrefixSegInfo`` in through ``**kwargs`` + (same route as ``packed_seq_lengths``). When present, the forward must route Q/K/V + through the FlexAttention shared-prefix kernel via ``AttentionContext.prefix_seg_info``. + + Returns the seg table (or ``None`` when PrefixGrouper did not group this batch -- the + unchanged path). Hardened: the shared-prefix stream is NOT a plain causal sequence, so running + it under a KV cache or an explicit padding mask would silently produce wrong logprobs. + That combination can only arise from misuse (PrefixGrouper only rides in via the GRPO + logprob forward, which is mask-free prefill), so we RAISE loudly instead of degrading + to a wrong result. + + Factored here so every arch (llama/mistral/qwen3/gemma2/cohere/granite/falcon_h1) + shares one implementation and cannot drift. + """ + seg = kwargs.get("prefix_seg_info", None) + if seg is not None and (past_key_value is not None or attention_mask is not None): + raise RuntimeError( + "PrefixGrouper: prefix_seg_info requires prefill with no KV cache and no " + f"attention_mask (got past_key_value={past_key_value is not None}, " + f"attention_mask={attention_mask is not None})." + ) + return seg + + def run_attention( *, config: AttentionConfig, context: AttentionContext, Q: Tensor, K: Tensor, V: Tensor ) -> Tensor: @@ -111,6 +153,28 @@ def run_attention( and SDPA handle packing via a block-diagonal mask. """ + # PrefixGrouper shared-prefix attention (GRPO dedup). Q/K/V here are [bsz, H, T, D]; + # the kernel takes/returns [1, T, H, D], matching the other backends. The field is + # only set when the env gate is on and grouping succeeded; None keeps every backend + # byte-identical. + if context.prefix_seg_info is not None: + flex_shared_prefix_attention = _flex_shared_prefix_attention + if flex_shared_prefix_attention is None: + # gate flipped on after import (or one-time load failed): resolve lazily. + from ..utils.prefix_grouper_kernel import flex_shared_prefix_attention + + scale = None + if config.flash_varlen_kwargs: + scale = config.flash_varlen_kwargs.get("softmax_scale") + A = flex_shared_prefix_attention( + Q.transpose(1, 2), + K.transpose(1, 2), + V.transpose(1, 2), + context.prefix_seg_info, + scale = scale, + ) + return A # [1, T, n_heads, head_dim] + backend = config.backend if backend == FLASH_VARLEN and context.seq_info is None: backend = FLASH_DENSE if HAS_FLASH_ATTENTION else SDPA @@ -337,5 +401,6 @@ __all__ = [ "AttentionConfig", "AttentionContext", "select_attention_backend", + "resolve_prefix_seg_info", "run_attention", ] diff --git a/unsloth/utils/prefix_grouper.py b/unsloth/utils/prefix_grouper.py new file mode 100644 index 0000000000..4e6ff9672c --- /dev/null +++ b/unsloth/utils/prefix_grouper.py @@ -0,0 +1,351 @@ +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""PrefixGrouper layout builder + completion-logprob extraction for the Unsloth GRPO +packed path (all archs that route through the varlen attention dispatch). + +Given the de-padded, LEFT-PACKED input_ids the packed GRPO path already works with, this +module: + + 1. Detects consecutive ``num_generations`` rows that share a prompt prefix (byte- + identical prompt precondition; falls back / returns None otherwise). + 2. Builds ONE flat shared-prefix stream across all groups + ``[ prefix_g0, suf_g0_0 .. suf_g0_{G-1}, prefix_g1, ... ]`` with position_ids that + continue each prefix positionally, plus a ``PrefixSegInfo`` segment table for the + FlexAttention shared-prefix kernel. + 3. Extracts completion logprobs via the index map (completion pos ``j==0`` predicted + from the shared prefix's last token; ``j>=1`` from the preceding suffix token) and + scatters them back into ``[total_rows, W]`` EXACTLY where the full-row packed path + puts them (dest = ``orig_row*L + orig_col``), so grpo_compute_loss / completion_mask + / TIS / metrics are byte-untouched. + +The flat stream is built by GATHERING original (row, col) coordinates out of input_ids, +so the grad path's autograd flows to the same embedding rows as today (the shared prefix +now contributes grad once = the sum of the G repeats, which is mathematically identical). + +``chunked_hidden_states_selective_log_softmax`` (from unsloth_zoo, passed in) is reused +verbatim over the gathered predicting-position hidden states, so fp32 accumulation, +logit_scale/softcapping/temperature are all preserved. + +Env: + UNSLOTH_GRPO_PREFIX_GROUPER=1 engage (default ON; set 0 to disable). Auto-off under vLLM. + UNSLOTH_GRPO_PREFIX_GROUPER_TOKR=1.3 tok_r auto-gate threshold (env-overridable) + UNSLOTH_GRPO_PREFIX_GROUPER_VERIFY=1 first-step self-verify (default ON) + UNSLOTH_GRPO_PREFIX_GROUPER_TOL=0.7 self-verify PASS band (nats) +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch + +from .prefix_grouper_kernel import build_seg_info_multigroup, PrefixSegInfo + + +# --------------------------------------------------------------------------- +# Env helpers +# --------------------------------------------------------------------------- +def env_on(name: str, default: str = "0") -> bool: + return os.environ.get(name, default).lower() not in ("0", "false", "no", "off") + + +# One-time env reads; the helpers stay callable since unsloth_zoo imports and calls them. +_ENABLED = env_on("UNSLOTH_GRPO_SEQ_PACKING", "1") and env_on("UNSLOTH_GRPO_PREFIX_GROUPER", "1") +_VERIFY_ON = env_on("UNSLOTH_GRPO_PREFIX_GROUPER_VERIFY", "1") +_TOKR_THRESHOLD = float(os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER_TOKR", "1.3")) +_TOL_OK = float(os.environ.get("UNSLOTH_GRPO_PREFIX_GROUPER_TOL", "0.7")) + + +def prefix_grouper_enabled() -> bool: + """PrefixGrouper requires seq-packing on (it reuses its de-pad + scatter machinery).""" + return _ENABLED + + +def verify_on() -> bool: + return _VERIFY_ON + + +def tokr_threshold() -> float: + return _TOKR_THRESHOLD + + +def tol_ok() -> float: + return _TOL_OK + + +# diff >= TOL_KILL = broken mask/isolation -> structure permanently unsafe; between +# tol_ok and TOL_KILL -> fall back for this shape but keep trying others. +TOL_KILL = 1.5 + + +@dataclass +class GroupLayout: + """Everything the GRPO forward needs to run + extract the shared-prefix path.""" + + flat_ids: torch.Tensor # [1, T] (T == seg.T) + position_ids: torch.Tensor # [1, T] + prefix_seg_info: PrefixSegInfo + # per completion target token, aligned 1:1: + tgt_rows: torch.Tensor # [N] original row index + tgt_cols: torch.Tensor # [N] original padded column in that row + tgt_pred: torch.Tensor # [N] flat predicting index (into the T stream) + tgt_flat: torch.Tensor # [N] flat index of the target token itself (into T) + total_rows: int + L: int # original padded seq length (input_ids.shape[1]) + W: int # logits_to_keep + max_left_pad (scatter width) + tok_r: float + signature: Tuple + + def extract_logps( + self, + hidden, + lm_head, + chunked_fn, + chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + ) -> torch.Tensor: + """hidden: [1, T, Hdim] (pre-lm_head hidden states, UNSLOTH_RETURN_HIDDEN_STATES=1). + Returns [total_rows, W] float32, byte-compatible with the packed path result.""" + # In a sharded model hidden may live on the lm-head device; move the small index + # maps to hidden.device before indexing. + device = hidden.device + pred_h = hidden[0, self.tgt_pred.to(device), :].unsqueeze(0) # [1, N, Hdim] + tgt_ids = self.flat_ids[0, self.tgt_flat].to(device).unsqueeze(0) # [1, N] + sel = chunked_fn( + pred_h, + lm_head, + tgt_ids, + chunks, + logit_scale_multiply, + logit_scale_divide, + logit_softcapping, + temperature, + )[0] # [N] logprobs + dest = self.tgt_rows.to(device) * self.L + self.tgt_cols.to(device) + result = ( + torch.zeros(self.total_rows * self.L, dtype = torch.float32, device = device) + .index_put((dest,), sel.to(torch.float32)) + .view(self.total_rows, self.L)[:, -self.W :] + ) + return result + + +def _build_groups(ids_cpu, real_cols_cpu, cstart_cpu, num_generations, total_rows): + """CPU-side grouping. Returns group dicts or None. Mirrors the packed _pk_* partition. + + A row's REAL tokens are the columns where input != pad. Its completion region (what + the packed path scatters, then completion_mask masks) is the real columns with + original col >= cstart_r, where cstart_r = (L - logits_to_keep) - left_pad_r. The + prompt is the real columns < cstart_r. Within a GRPO group all G rows share the same + prompt => same left_pad => same cstart => the prompt real columns are BYTE-IDENTICAL + across the group (the shared prefix). We require that byte-identity (falls back + otherwise). No prompt-tail special-casing: every suffix token is scattered exactly + like the packed path; completion_mask masks the leading prompt-tail positions. + """ + G = num_generations + if G is None or G < 2 or total_rows % G != 0: + return None + groups = [] + for g0 in range(0, total_rows, G): + rows = list(range(g0, g0 + G)) + prompt_cols_per_row = [] # real cols < cstart + prompt_toks_per_row = [] + comp_cols_per_row = [] # real cols >= cstart (the completion region packed scatters) + for r in rows: + cs = cstart_cpu[r] + rc = real_cols_cpu[r] + p_cols = [c for c in rc if c < cs] + c_cols = [c for c in rc if c >= cs] + prompt_cols_per_row.append(p_cols) + prompt_toks_per_row.append([ids_cpu[r][c] for c in p_cols]) + comp_cols_per_row.append(c_cols) + if any(len(p) == 0 for p in prompt_toks_per_row): + return None + # require BYTE-IDENTICAL prompts across the group (shared-prefix precondition). + P = len(prompt_toks_per_row[0]) + if any(len(prompt_toks_per_row[k]) != P for k in range(1, G)): + return None + p0 = prompt_toks_per_row[0] + if any(prompt_toks_per_row[k] != p0 for k in range(1, G)): + return None + if P == 0: + return None + R_list = [len(c) for c in comp_cols_per_row] + if sum(R_list) == 0: + return None + groups.append( + dict( + rows = rows, + P = P, + prefix_cols = prompt_cols_per_row[0], # shared prompt real columns (row0) + prefix_row = rows[0], + R_list = R_list, + suf_cols = comp_cols_per_row, # per-row completion-region real columns + ) + ) + return groups + + +def _tok_r(groups) -> float: + tok_full = 0 + tok_sp = 0 + for gm in groups: + P = gm["P"] + Rs = gm["R_list"] + tok_full += sum(P + r for r in Rs) # G*P + sumR + tok_sp += P + sum(Rs) # P + sumR + return (tok_full / tok_sp) if tok_sp else 1.0 + + +def build_group_layout( + input_ids, + logits_to_keep, + pad_id, + num_generations, + left_pad_tokens_per_prompt, + *, + apply_tokr_gate = True, + max_segment_cap = None, +): + """Build the shared-prefix GroupLayout, or return None to fall back to the packed path. + + input_ids : [B, L]. GRPO's layout is left-padded in the prompt and right-padded in + the completion. Real tokens of a row are a contiguous run not necessarily + starting at column 0. + logits_to_keep : int + left_pad_tokens_per_prompt : [B] long tensor (per-row left-pad count in the prompt). + """ + device = input_ids.device + total_rows, L = input_ids.shape + keep = input_ids != pad_id + # completion start column per row (matches create_completion_attention_mask / _pk_cstart). + cstart = ((L - logits_to_keep) - left_pad_tokens_per_prompt).to(torch.long) + cstart_cpu = cstart.tolist() + ids_cpu = input_ids.tolist() + # per-row real (non-pad) columns. GRPO rows are one contiguous real run, so derive + # [first, first+n) on GPU; the O(B*L) scan is only a non-contiguous fallback. + n_real = keep.sum(dim = 1) + first = torch.argmax(keep.to(torch.int8), dim = 1) + ar = torch.arange(L, device = device) + contiguous = bool( + (keep == ((ar >= first.unsqueeze(1)) & (ar < (first + n_real).unsqueeze(1)))).all() + ) + if contiguous: + real_cols_cpu = [list(range(f, f + n)) for f, n in zip(first.tolist(), n_real.tolist())] + else: + keep_cpu = keep.tolist() + real_cols_cpu = [[c for c in range(L) if keep_cpu[r][c]] for r in range(total_rows)] + + groups = _build_groups(ids_cpu, real_cols_cpu, cstart_cpu, num_generations, total_rows) + if groups is None: + return None + + # sliding-window guard: a group's PG span is P + max(R); fall back if it exceeds the window. + if max_segment_cap is not None: + for gm in groups: + if gm["P"] + max(gm["R_list"]) > max_segment_cap: + return None + + tok_r = _tok_r(groups) + if apply_tokr_gate and tok_r < tokr_threshold(): + return None # low reuse -> not worth it; use the full-row packed path + + # Build flat stream by gathering original (row, col) coordinates. + group_specs = [(gm["P"], gm["R_list"]) for gm in groups] + seg, group_meta = build_seg_info_multigroup(group_specs, device) + + flat_src_rows: List[int] = [] + flat_src_cols: List[int] = [] + pos_list: List[int] = [] + tgt_rows: List[int] = [] + tgt_cols: List[int] = [] + tgt_pred: List[int] = [] + tgt_flat: List[int] = [] + + for gm, meta in zip(groups, group_meta): + rows = gm["rows"] + P = gm["P"] + r0 = gm["prefix_row"] + prefix_cols = gm["prefix_cols"] # ORIGINAL real prompt columns (len P) of row0 + plast = meta["prefix_last_index"] # base + P - 1 + # gather the shared prefix once, from row0. + flat_src_rows.extend([r0] * P) + flat_src_cols.extend(prefix_cols) + pos_list.extend(range(P)) + # suffixes: every suffix token is a completion-region target (scattered like the + # packed path; completion_mask hides prompt-tail positions). + for i, r in enumerate(rows): + cols = gm["suf_cols"][i] + r_i = len(cols) + s, e = meta["suffix_slices"][i] # flat offsets [s, e) + flat_src_rows.extend([r] * r_i) + flat_src_cols.extend(cols) + pos_list.extend(range(P, P + r_i)) + for j in range(r_i): + # pos 0 is predicted from the prefix's last token; j>=1 from the previous suffix token. + pred = plast if j == 0 else (s + j - 1) + tgt_rows.append(r) + tgt_cols.append(cols[j]) # ORIGINAL padded column in row r + tgt_pred.append(pred) + tgt_flat.append(s + j) # flat index of the target token itself + + T = len(flat_src_rows) + assert T == seg.T, f"flat stream len {T} != seg.T {seg.T}" + fr = torch.tensor(flat_src_rows, device = device, dtype = torch.long) + fc = torch.tensor(flat_src_cols, device = device, dtype = torch.long) + flat_ids = input_ids[fr, fc].unsqueeze(0) # [1, T] (grad-safe gather) + position_ids = torch.tensor(pos_list, device = device, dtype = torch.long).unsqueeze(0) + + max_left_pad = int(left_pad_tokens_per_prompt.max().item()) if total_rows else 0 + W = logits_to_keep + max_left_pad + + # self-verify cache key: the mask/index-map/scatter logic is structural, so key on + # (num_groups, group_sizes), not exact lengths -- GRPO lengths change every step and + # keying on T would re-verify forever ("verify once, then trust", like the packed path). + grp_sizes = tuple(sorted(len(gm["R_list"]) for gm in groups)) + sig = (len(groups), grp_sizes) + + return GroupLayout( + flat_ids = flat_ids, + position_ids = position_ids, + prefix_seg_info = seg, + tgt_rows = torch.tensor(tgt_rows, device = device, dtype = torch.long), + tgt_cols = torch.tensor(tgt_cols, device = device, dtype = torch.long), + tgt_pred = torch.tensor(tgt_pred, device = device, dtype = torch.long), + tgt_flat = torch.tensor(tgt_flat, device = device, dtype = torch.long), + total_rows = total_rows, + L = L, + W = W, + tok_r = tok_r, + signature = sig, + ) + + +__all__ = [ + "GroupLayout", + "build_group_layout", + "prefix_grouper_enabled", + "verify_on", + "tokr_threshold", + "tol_ok", + "TOL_KILL", + "env_on", +] diff --git a/unsloth/utils/prefix_grouper_kernel.py b/unsloth/utils/prefix_grouper_kernel.py new file mode 100644 index 0000000000..9a9719b015 --- /dev/null +++ b/unsloth/utils/prefix_grouper_kernel.py @@ -0,0 +1,436 @@ +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""FlexAttention shared-prefix kernel for PrefixGrouper (GRPO shared-prompt dedup). + +In GRPO every prompt spawns ``G = num_generations`` completions that share the same +prompt prefix. The full-row packed path forwards the identical prefix ``G`` times. +PrefixGrouper stores the prefix ONCE and concatenates only the ``G`` suffixes, with an +attention layout where each suffix token attends to ``[the single shared prefix] + +[causal within its own suffix]``. This kernel expresses that one-prefix -> many-suffix +fan-out via a ``torch.nn.attention.flex_attention`` block mask, so the masked-out +cross-suffix / cross-group blocks are never computed and the ``P + G*R`` FLOP saving is +realised (not merely a masked dense ``O(T^2)``). + +Mask semantics (identical to the certified SDPA oracle): + + keep(q_idx, kv_idx) = same_group(q, kv) AND + ( is_prefix[kv_idx] # full prefix visibility + OR ( suffix_of_kv[kv_idx] == suffix_of_kv[q_idx] # same suffix ... + AND kv_idx <= q_idx ) ) # ... causal within it + +This module is self-contained (no dependency on any temp/ scratch dir) so PrefixGrouper +works from the installed source after a fresh compile. It is only imported lazily from +``attention_dispatch.run_attention`` when ``prefix_seg_info`` is present, which itself is +only ever set when ``UNSLOTH_GRPO_PREFIX_GROUPER`` is on and grouping succeeded, so the +default (off) path never touches this file. + +Provided entry points: + * ``PrefixSegInfo`` : per-flat-token segment metadata + cache signature. + * ``build_seg_info_multigroup``: build PrefixSegInfo for many groups packed flat. + * ``build_seg_info_from_layout``: build PrefixSegInfo for ONE group (test helper). + * ``get_block_mask`` : cached create_block_mask keyed on the signature. + * ``flex_shared_prefix_attention(Q, K, V, prefix_seg_info)`` + Q/K/V of shape [1, T, n_heads, head_dim]; returns [1, T, n_heads, head_dim], + IDENTICAL semantics to the SDPA oracle. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import torch +from torch.nn.attention.flex_attention import ( + BlockMask, + create_block_mask, + flex_attention, +) + +# GRPO feeds many distinct segment lengths; at dynamo's default recompile_limit (8) the +# compiled kernel silently reuses a mismatched specialisation (wrong results). Raise it. +torch._dynamo.config.recompile_limit = max(getattr(torch._dynamo.config, "recompile_limit", 8), 256) +torch._dynamo.config.accumulated_recompile_limit = max( + getattr(torch._dynamo.config, "accumulated_recompile_limit", 256), 2048 +) + + +# Compiled kernels: torch.compile fuses the sparse mask into one kernel. dynamic=True is +# required: T changes almost every GRPO batch and dynamic=False recompiles per T (~14s +# each). T is still padded to a multiple of 128 (_pad_len) for the backward kernel. +_flex_attention_compiled = torch.compile(flex_attention, dynamic = True) +_create_block_mask_compiled = torch.compile(create_block_mask, dynamic = True) + +# Flash block sizes by Q dtype (env-overridable). The two disjoint key runs (prefix + +# own-suffix) stress online-softmax accumulation: fp32 needs 32/32 for a ~1e-6 floor; +# bf16 passes parity at 128/64 and is ~5x faster (128/128 OOMs Triton on B200). +_FP32_BLOCK_M = int(os.environ.get("PG_FLEX_BLOCK_M", "32")) +_FP32_BLOCK_N = int(os.environ.get("PG_FLEX_BLOCK_N", "32")) +_BF16_BLOCK_M = int(os.environ.get("PG_FLEX_BF16_BLOCK_M", "128")) +_BF16_BLOCK_N = int(os.environ.get("PG_FLEX_BF16_BLOCK_N", "64")) + + +def _kernel_options_for_dtype(dtype): + """Pick the numerically-safe flash block sizes for the Q dtype.""" + if dtype == torch.bfloat16 or dtype == torch.float16: + return {"BLOCK_M": _BF16_BLOCK_M, "BLOCK_N": _BF16_BLOCK_N} + return {"BLOCK_M": _FP32_BLOCK_M, "BLOCK_N": _FP32_BLOCK_N} + + +# Backward-compat constant (fp32 default). +_FLEX_KERNEL_OPTIONS = {"BLOCK_M": _FP32_BLOCK_M, "BLOCK_N": _FP32_BLOCK_N} + +# The compiled backward trips an Inductor assertion when T is not a multiple of 128, so +# pad the flat sequence. Pad tokens form a group that attends to / is attended by nothing +# (all-masked rows return 0, not NaN) and are sliced off the output. +_PAD_MULTIPLE = 128 +_PAD_GROUP = -99 # sentinel group id / suffix id for pad tokens + + +def _pad_len(T: int) -> int: + return ((T + _PAD_MULTIPLE - 1) // _PAD_MULTIPLE) * _PAD_MULTIPLE + + +# --------------------------------------------------------------------------- +# Segment metadata +# --------------------------------------------------------------------------- + + +@dataclass +class PrefixSegInfo: + """Per-flat-token segment metadata driving the shared-prefix block mask. + + The label tensors are 1-D of length ``T_pad`` (>= real ``T``, padded up to a multiple + of 128 so the backward kernel compiles). Positions ``[T:T_pad)`` are pad tokens + (group/suffix == _PAD_GROUP) that attend to nothing. + + Attributes + ---------- + group_of_kv : LongTensor [T_pad] + Group id per flat token (0..num_groups-1); _PAD_GROUP for pad tokens. + is_prefix : BoolTensor [T_pad] + True iff the token is a prefix token of its group (False for pad). + suffix_of_kv : LongTensor [T_pad] + Suffix id per flat token; -1 for prefix, _PAD_GROUP for pad. Suffix ids are + globally unique across groups. + signature : hashable + Cache key for the block mask (depends only on the labels + T_pad). + T : int + Real flat sequence length (Q/K/V of this length are padded internally). + T_pad : int + Padded length (multiple of 128) at which the block mask is built. + """ + + group_of_kv: torch.Tensor + is_prefix: torch.Tensor + suffix_of_kv: torch.Tensor + signature: Tuple + T: int + T_pad: int + + +def _pad_labels(group_of_kv, is_prefix, suffix_of_kv, device): + """Pad the label tensors up to a multiple of 128 with pad-token sentinels.""" + T = int(group_of_kv.numel()) + T_pad = _pad_len(T) + if T_pad == T: + return group_of_kv, is_prefix, suffix_of_kv, T, T_pad + pad = T_pad - T + group_of_kv = torch.cat( + [group_of_kv, torch.full((pad,), _PAD_GROUP, dtype = torch.long, device = device)] + ) + is_prefix = torch.cat([is_prefix, torch.zeros(pad, dtype = torch.bool, device = device)]) + suffix_of_kv = torch.cat( + [suffix_of_kv, torch.full((pad,), _PAD_GROUP, dtype = torch.long, device = device)] + ) + return group_of_kv, is_prefix, suffix_of_kv, T, T_pad + + +def build_seg_info_from_layout(layout, device: Optional[torch.device] = None) -> PrefixSegInfo: + """Build PrefixSegInfo for ONE group from an object with ``.flat_ids``, ``.P`` and + ``.suffix_slices`` (used by the parity test / oracle helpers).""" + if device is None: + device = layout.flat_ids.device + T = int(layout.flat_ids.shape[1]) + P = int(layout.P) + + group_of_kv = torch.zeros(T, dtype = torch.long, device = device) # single group -> 0 + is_prefix = torch.zeros(T, dtype = torch.bool, device = device) + is_prefix[:P] = True + suffix_of_kv = torch.full((T,), -1, dtype = torch.long, device = device) + for i, (s, e) in enumerate(layout.suffix_slices): + suffix_of_kv[s:e] = i + + group_of_kv, is_prefix, suffix_of_kv, T, T_pad = _pad_labels( + group_of_kv, is_prefix, suffix_of_kv, device + ) + sig = ("single", T_pad, P, tuple((s, e) for (s, e) in layout.suffix_slices)) + return PrefixSegInfo( + group_of_kv = group_of_kv, + is_prefix = is_prefix, + suffix_of_kv = suffix_of_kv, + signature = sig, + T = T, + T_pad = T_pad, + ) + + +def build_seg_info_multigroup( + group_specs: List[Tuple[int, List[int]]], device: torch.device +) -> Tuple[PrefixSegInfo, List[dict]]: + """Build PrefixSegInfo for several shared-prefix groups packed block-diagonally. + + Parameters + ---------- + group_specs : list of (P_g, [R_{g,0}, R_{g,1}, ...]) + For each group: prefix length and the list of suffix lengths. + + Returns + ------- + seg : PrefixSegInfo + group_meta : list of dicts with 'base', 'P', 'prefix_last_index', 'suffix_slices' + (flat offsets), enough to build the completion index map. + """ + group_of_list = [] + is_prefix_list = [] + suffix_of_list = [] + group_meta = [] + + base = 0 + suffix_counter = 0 + sig_parts = [] + for gid, (P, R_list) in enumerate(group_specs): + # prefix + group_of_list.append(torch.full((P,), gid, dtype = torch.long, device = device)) + is_prefix_list.append(torch.ones(P, dtype = torch.bool, device = device)) + suffix_of_list.append(torch.full((P,), -1, dtype = torch.long, device = device)) + prefix_last_index = base + P - 1 + suffix_slices = [] + cursor = base + P + for r in R_list: + group_of_list.append(torch.full((r,), gid, dtype = torch.long, device = device)) + is_prefix_list.append(torch.zeros(r, dtype = torch.bool, device = device)) + suffix_of_list.append(torch.full((r,), suffix_counter, dtype = torch.long, device = device)) + suffix_slices.append((cursor, cursor + r)) + cursor += r + suffix_counter += 1 + group_meta.append( + { + "base": base, + "P": P, + "prefix_last_index": prefix_last_index, + "suffix_slices": suffix_slices, + } + ) + sig_parts.append((P, tuple(R_list))) + base = cursor + + group_of_kv = torch.cat(group_of_list) + is_prefix = torch.cat(is_prefix_list) + suffix_of_kv = torch.cat(suffix_of_list) + group_of_kv, is_prefix, suffix_of_kv, T, T_pad = _pad_labels( + group_of_kv, is_prefix, suffix_of_kv, device + ) + sig = ("multi", T_pad, tuple(sig_parts)) + seg = PrefixSegInfo( + group_of_kv = group_of_kv, + is_prefix = is_prefix, + suffix_of_kv = suffix_of_kv, + signature = sig, + T = T, + T_pad = T_pad, + ) + return seg, group_meta + + +# --------------------------------------------------------------------------- +# Block-mask builder + cache, keyed on (signature, device): the mask depends only on the +# per-token labels and T, so it is reused across layers and steps. + +_BLOCK_MASK_CACHE: Dict[Tuple, BlockMask] = {} + + +def _make_mask_mod(group_of_kv, is_prefix, suffix_of_kv): + """Return a mask_mod closure over the (device) label tensors. + + keep(q, kv) = same_group AND + ( is_prefix[kv] AND kv <= q # causal within/ into prefix + OR ( suffix_of_kv[kv] == suffix_of_kv[q] # same suffix ... + AND (not is_prefix[q]) # q is a suffix token ... + AND kv <= q ) ) # ... causal within it + + The single ``kv <= q`` guard on the is_prefix branch gives BOTH prefix-causal + behaviour (a prefix q sees only earlier prefix tokens) AND full-prefix-visibility for + suffixes (every prefix index < every suffix index in a group, so kv <= q always holds + for a suffix q vs a prefix kv of its group), matching the SDPA oracle exactly. + """ + + def mask_mod(b, h, q_idx, kv_idx): + same_group = group_of_kv[q_idx] == group_of_kv[kv_idx] + kv_is_prefix = is_prefix[kv_idx] + causal = kv_idx <= q_idx + same_suffix = (suffix_of_kv[kv_idx] == suffix_of_kv[q_idx]) & (~is_prefix[q_idx]) + keep = same_group & ((kv_is_prefix & causal) | (same_suffix & causal)) + return keep + + return mask_mod + + +def get_block_mask( + seg: PrefixSegInfo, + device: torch.device, + compile_mask: bool = True, +) -> BlockMask: + """Return a cached BlockMask for the segment signature (built once, reused). + + CRITICAL: the block mask is cached and shared across BOTH the no-grad old/ref logprob + forward (which runs under torch.inference_mode) and the grad training forward. If the + mask were first built under inference_mode, its tensors would be INFERENCE tensors that + "cannot be saved for backward" when reused in the grad forward. We therefore build the + mask with inference mode explicitly DISABLED, so the same cached BlockMask is a normal + tensor usable by autograd. (The mask depends only on integer labels; it needs no grad.) + """ + key = (seg.signature, str(device)) + bm = _BLOCK_MASK_CACHE.get(key) + if bm is not None: + return bm + + # Move labels to the consumer (Q) device: with a sharded model the seg tensors live on + # input_ids.device and would index cross-device. Copies once per (signature, device). + # These copies must also run with inference mode DISABLED (same reason as the mask build): + # when this entry is first built under the no-grad old/ref forward's inference_mode and + # device != seg.device, a .to(device) copy would be an inference tensor that mask_mod + # captures, which then cannot be saved for backward when the grad training forward reuses + # the cached mask. + builder = _create_block_mask_compiled if compile_mask else create_block_mask + with torch.inference_mode(False): + mask_mod = _make_mask_mod( + seg.group_of_kv.to(device), seg.is_prefix.to(device), seg.suffix_of_kv.to(device) + ) + bm = builder( + mask_mod, + B = 1, + H = None, + Q_LEN = seg.T_pad, + KV_LEN = seg.T_pad, + device = device, + ) + # FIFO bound: GRPO lengths change nearly every step, so evict the oldest to cap GPU pins. + if len(_BLOCK_MASK_CACHE) >= 8: + _BLOCK_MASK_CACHE.pop(next(iter(_BLOCK_MASK_CACHE))) + _BLOCK_MASK_CACHE[key] = bm + return bm + + +def clear_block_mask_cache(): + _BLOCK_MASK_CACHE.clear() + + +def _pad_qkv_seq(x: torch.Tensor, T_pad: int) -> torch.Tensor: + """Zero-pad a [B, H, T, D] tensor along the sequence dim up to T_pad.""" + T = x.shape[2] + if T_pad == T: + return x + pad = torch.zeros(x.shape[0], x.shape[1], T_pad - T, x.shape[3], device = x.device, dtype = x.dtype) + return torch.cat([x, pad], dim = 2) + + +def _run_flex(q, k, v, block_mask, enable_gqa, scale, compiled, T, T_pad): + """Pad q/k/v to T_pad, run flex, slice the output back to T. q/k/v: [B,H,T,D].""" + qp = _pad_qkv_seq(q, T_pad) + kp = _pad_qkv_seq(k, T_pad) + vp = _pad_qkv_seq(v, T_pad) + if compiled: + out = _flex_attention_compiled( + qp, + kp, + vp, + block_mask = block_mask, + enable_gqa = enable_gqa, + scale = scale, + kernel_options = _kernel_options_for_dtype(qp.dtype), + ) + else: + # eager path (fp64 parity): dense scores, no kernel_options. + out = flex_attention( + qp, + kp, + vp, + block_mask = block_mask, + enable_gqa = enable_gqa, + scale = scale, + ) + return out[:, :, :T, :] + + +# --------------------------------------------------------------------------- +# The kernel entry point +# --------------------------------------------------------------------------- + + +def flex_shared_prefix_attention( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + prefix_seg_info: PrefixSegInfo, + scale: Optional[float] = None, + block_mask: Optional[BlockMask] = None, + compiled: bool = True, +) -> torch.Tensor: + """Shared-prefix attention via FlexAttention. + + Parameters + ---------- + Q, K, V : Tensor [1, T, n_heads, head_dim] + (Q has n_heads, K/V have n_kv_heads for GQA). + prefix_seg_info : PrefixSegInfo + scale : optional float, softmax scale (defaults to 1/sqrt(head_dim)). + block_mask : optional precomputed BlockMask (else built/cached from seg info). + + Returns + ------- + Tensor [1, T, n_heads, head_dim], identical semantics to the SDPA oracle branch. + """ + assert Q.dim() == 4 and Q.shape[0] == 1, f"expected [1,T,H,D], got {tuple(Q.shape)}" + device = Q.device + # FlexAttention wants [B, H, T, D]. + q = Q.transpose(1, 2) # [1, n_heads, T, D] + k = K.transpose(1, 2) # [1, n_kv_heads, T, D] + v = V.transpose(1, 2) + + n_heads = q.shape[1] + n_kv = k.shape[1] + enable_gqa = n_heads != n_kv + T = q.shape[2] + T_pad = prefix_seg_info.T_pad + assert T == prefix_seg_info.T, f"Q length {T} != seg.T {prefix_seg_info.T}" + + if block_mask is None: + block_mask = get_block_mask(prefix_seg_info, device, compile_mask = compiled) + + out = _run_flex(q, k, v, block_mask, enable_gqa, scale, compiled, T, T_pad) + # back to [1, T, n_heads, D] + return out.transpose(1, 2).contiguous() + + +__all__ = [ + "PrefixSegInfo", + "build_seg_info_multigroup", + "build_seg_info_from_layout", + "get_block_mask", + "clear_block_mask_cache", + "flex_shared_prefix_attention", +] From 22bd86ecb7b80856c4e639001b2e15394bdbe82b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:44:55 -0700 Subject: [PATCH 005/113] Handle odd shapes and non-float scales in FP8BlockQuantLinear (#6848) * Handle odd shapes and non-float scales in FP8BlockQuantLinear Small fp8 checkpoints (e.g. tiny test models) break the block-quantized linear in three ways: weight scales stored in a float8 dtype such as float8_e8m0fnu have no triton dtype mapping; activations whose hidden dim is not a multiple of the activation quant block fail act_quant's divisibility assert; and weights whose dims are not multiples of the weight block cannot be tiled by the triton dequant kernel. Cast non-float scales to float32 on entry, and when the hidden dim does not divide into the activation block, dequantize the weight and run a plain matmul instead of the fp8 block matmul. The dequant goes through a new shape-safe helper that falls back to a torch-native scale expansion when the weight does not tile evenly; backward uses the same helper so the gradient path works for every shape the forward accepts. Full-size checkpoints are unaffected. * Add tiny / e8m0 fp8 block-quant regression test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix FP8 block-quant fallback: real block size in dequant and scalar-scale fast path * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Route rectangular fp8 blocks through torch dequant and keep block_size across e8m0 upcast The triton weight_dequant kernel uses one BLOCK_SIZE for both axes, so rectangular blocks (block_size[0] != block_size[1]) mis-index the column scale and corrupt grad_X. Route those through the torch scale expansion, which handles each dimension independently, and keep the triton path for square blocks only. Also preserve a block_size attribute carried on the scale tensor across the e8m0 -> float32 upcast so the later lookup no longer falls back to [128, 128]. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_fp8_tiny_e8m0.py | 123 ++++++++++++++++++++++++++++++++++++ unsloth/kernels/fp8.py | 48 +++++++++++++- 2 files changed, 170 insertions(+), 1 deletion(-) create mode 100644 tests/test_fp8_tiny_e8m0.py diff --git a/tests/test_fp8_tiny_e8m0.py b/tests/test_fp8_tiny_e8m0.py new file mode 100644 index 0000000000..cf49c8c92f --- /dev/null +++ b/tests/test_fp8_tiny_e8m0.py @@ -0,0 +1,123 @@ +"""FP8 block-quant linear must handle tiny / non-tileable weights and e8m0 scales. + +Two things break the triton block path: + * a hidden dim not divisible by the activation block size (tiny test models), + * float8_e8m0fnu weight scales, which have no triton dtype mapping. +The forward falls back to a torch-native blockwise dequant + bf16 matmul; this +test checks that fallback runs finite forward + backward and matches a plain +dequant reference. +""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason = "needs CUDA") + + +def _reference(X, weight, scale, block): + # Expand the per-block scale to full weight shape and dequantize. + m, n = weight.shape + s = scale.to(torch.float32) + s = s.repeat_interleave(block[0], 0)[:m].repeat_interleave(block[1], 1)[:, :n] + W = (weight.to(torch.float32) * s).to(X.dtype) + return X @ W.T + + +def test_tiny_non_tileable_forward_backward_matches_reference(): + from unsloth.kernels.fp8 import FP8BlockQuantLinear + + torch.manual_seed(0) + dev = "cuda" + block = [128, 128] + m, n = 8, 8 # non-tileable, in-dim % 128 != 0 + weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) # (out=m, in=n) + scale = torch.rand(1, 1, device = dev, dtype = torch.float32) + 0.5 + X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True) + + out = FP8BlockQuantLinear.apply(X, weight, scale) + assert torch.isfinite(out).all(), "forward produced non-finite values" + + ref = _reference(X.detach(), weight, scale, block) + torch.testing.assert_close(out, ref, atol = 5e-2, rtol = 5e-2) + + out.sum().backward() + assert X.grad is not None and torch.isfinite(X.grad).all(), "backward non-finite" + + +def test_e8m0_scale_is_upcast_and_runs(): + from unsloth.kernels.fp8 import FP8BlockQuantLinear + + if not hasattr(torch, "float8_e8m0fnu"): + pytest.skip("torch build lacks float8_e8m0fnu") + + dev = "cuda" + m, n = 8, 8 + weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) + scale = (torch.rand(1, 1, device = dev) + 1.0).to(torch.float8_e8m0fnu) + X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True) + + out = FP8BlockQuantLinear.apply(X, weight, scale) + assert torch.isfinite(out).all() + out.sum().backward() + assert torch.isfinite(X.grad).all() + + +def test_rectangular_block_dequant_matches_reference(): + # Rectangular blocks (block_size[0] != block_size[1]) that tile evenly used to + # route through the triton weight_dequant kernel, which uses a single BLOCK_SIZE + # for both axes and mis-indexes the column scale. Verify the torch expansion path + # now matches the reference for a 64x256 weight with block [64, 128] (scale 1x2). + from unsloth.kernels.fp8 import _blockwise_weight_dequant_any_shape + + torch.manual_seed(0) + dev = "cuda" + block = [64, 128] + m, n = 64, 256 # evenly tiled: 64 % 64 == 0, 256 % 128 == 0 + weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) + # Distinct per-block column scales expose column mis-indexing. + scale = torch.tensor([[0.5, 3.0]], device = dev, dtype = torch.float32) + + W_deq = _blockwise_weight_dequant_any_shape(weight, scale, block, torch.bfloat16) + + s = scale.repeat_interleave(block[0], 0)[:m].repeat_interleave(block[1], 1)[:, :n] + ref = (weight.to(torch.float32) * s).to(torch.bfloat16) + torch.testing.assert_close(W_deq, ref, atol = 5e-3, rtol = 5e-3) + + +def test_e8m0_scale_preserves_non_default_block_size_attr(): + # An e8m0 scale carrying a non-default block_size attribute must keep it across + # the float32 upcast in forward; otherwise the lookup falls back to [128, 128] + # and a compatible layout is wrongly rejected as incompatible. + from unsloth.kernels.fp8 import FP8BlockQuantLinear + + if not hasattr(torch, "float8_e8m0fnu"): + pytest.skip("torch build lacks float8_e8m0fnu") + + torch.manual_seed(0) + dev = "cuda" + block = [64, 64] + # in-dim 96 is not divisible by block[1]=64 -> forward takes the torch dequant + # fallback (no fp8 matmul kernel). Scale shape (2, 2) validates for [64, 64] but + # not [128, 128] (which expects (1, 1)). + m, n = 128, 96 + weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16) # no block_size attr + scale_f = torch.rand(2, 2, device = dev) + 1.0 + scale = scale_f.to(torch.float8_e8m0fnu) + scale.block_size = block # attribute lives on the scale, not the weight + X = torch.randn(4, n, device = dev, dtype = torch.bfloat16, requires_grad = True) + + # With [128, 128] this raises "not compatible with block size"; success proves + # the [64, 64] attribute survived the e8m0 -> float32 upcast. + out = FP8BlockQuantLinear.apply(X, weight, scale) + assert torch.isfinite(out).all() + + ref = _reference(X.detach(), weight, scale.to(torch.float32), block) + torch.testing.assert_close(out, ref, atol = 5e-2, rtol = 5e-2) + + out.sum().backward() + assert X.grad is not None and torch.isfinite(X.grad).all() + + +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, "-q"])) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index ca608fa01b..80db2f466b 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -327,11 +327,42 @@ fp8_block_matmul = ( ) +def _blockwise_weight_dequant_any_shape(weight, weight_scale, block_size, out_dtype): + """Blockwise fp8 weight dequant for any shape: triton when the weight tiles + evenly into block_size, else a torch-native per-block scale expansion.""" + m, n = weight.shape + if weight_scale.dtype not in (torch.float32, torch.float16, torch.bfloat16): + weight_scale = weight_scale.to(torch.float32) # e.g. float8_e8m0fnu scales break triton + if weight_scale.numel() == 1: + # Per-tensor scale: the normal forward stashes the un-expanded scalar, + # which repeat_interleave cannot grow to (m, n). Scale directly. + return (weight.to(torch.float32) * weight_scale.float()).to(out_dtype) + if m % block_size[0] != 0 or n % block_size[1] != 0 or block_size[0] != block_size[1]: + # Uneven tiling, or rectangular blocks. The triton kernel uses a single + # BLOCK_SIZE for both axes and derives the column scale stride from it, so + # it mis-indexes the scale when block_size[0] != block_size[1]. Expand the + # per-block scales in torch, which handles both dimensions independently. + s_full = weight_scale.repeat_interleave(block_size[0], 0)[:m] + s_full = s_full.repeat_interleave(block_size[1], 1)[:, :n] + return (weight.to(torch.float32) * s_full).to(out_dtype) + # Even tiling with square blocks: block-quant dequant with the real block size + # (weight_dequant would silently default to 128 and dequantize wrongly). + return weight_dequant_block(weight, weight_scale, block_size = block_size[0], dtype = out_dtype) + + class FP8BlockQuantLinear(torch.autograd.Function): @staticmethod def forward(ctx, X, weight, weight_scale): m, n = weight.shape + if weight_scale.dtype not in (torch.float32, torch.float16, torch.bfloat16): + # Upcast (e.g. e8m0) returns a fresh tensor and drops any Python + # attribute, so carry block_size across the cast for the lookup below. + _scale_block_size = getattr(weight_scale, "block_size", None) + weight_scale = weight_scale.to(torch.float32) # e8m0 scales break triton dtype mapping + if _scale_block_size is not None: + weight_scale.block_size = _scale_block_size + # Original scale, saved for backward before any transformation original_weight_scale = weight_scale @@ -360,6 +391,18 @@ class FP8BlockQuantLinear(torch.autograd.Function): if not weight.is_contiguous(): weight = weight.contiguous() + if X.shape[-1] % block_size[1] != 0: + # Hidden dim not divisible by the activation block: dequant + plain matmul. + # Use the original (un-expanded) scale so a scalar per-tensor scale keeps + # the fast scalar path in both forward and backward. + W_deq = _blockwise_weight_dequant_any_shape( + weight, original_weight_scale, block_size, X.dtype + ) + ctx.weight = weight + ctx.weight_scale = original_weight_scale + ctx.block_size = block_size + return torch_matmul(X, W_deq.T).to(X.dtype) + qinput, scale = act_quant(X, block_size[1]) output = fp8_block_matmul( qinput, @@ -371,11 +414,14 @@ class FP8BlockQuantLinear(torch.autograd.Function): ) ctx.weight = weight ctx.weight_scale = original_weight_scale # Save original for backward + ctx.block_size = block_size return output.to(X.dtype) @staticmethod def backward(ctx, grad_output): - W_deq = weight_dequant(ctx.weight, ctx.weight_scale) + W_deq = _blockwise_weight_dequant_any_shape( + ctx.weight, ctx.weight_scale, ctx.block_size, grad_output.dtype + ) grad_X = torch_matmul(grad_output, W_deq) del W_deq return grad_X, None, None From 7cc1752a646a1371a4d393ea7cd877417361b1d4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:45:06 -0700 Subject: [PATCH 006/113] Scope MoE expert LoRA detection to actual MLP projection targets (#6849) * Scope MoE expert LoRA detection to actual MLP projection targets _moe_target_set_from_string treated any regex containing the substring mlp or ffn as targeting the expert MLP projections. Unsloth's auto-generated attention-only regex lists mlp, ffn and feed_forward as allowed intermediate path segments while its final group matches only q_proj/k_proj/v_proj/o_proj, so attention-only finetuning on MoE models silently enabled expert LoRA as well: the experts were trained and every MoE layer paid the extra expert LoRA grouped matmuls. Detect expert intent from the projection names themselves (gate_proj/up_proj/down_proj/gate_up_proj) instead of the mlp substring. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments * Detect MoE expert LoRA via mlp path segment, not proj names The auto-generated target regex always lists every projection leaf (q/k/v/o and gate/up/down), so keying detection on a proj name mis-fired: it enabled expert LoRA for attention-only regexes and dropped the mlp/ffn path regexes. Key on the mlp/ffn/feed_forward/experts path segment instead, which is present only when the MLP/experts are actually targeted. Add a regression test for the attention-only case. * Scope expert LoRA targets to the leaves a regex names An mlp path alternative with attention-only leaves, for example (mlp|self_attn).(q_proj|o_proj), no longer enables expert LoRA, and a regex naming a single expert leaf such as .*experts.*down_proj now targets only that projection instead of the whole broad set. Generic mlp projections (.*mlp.*proj) and the auto regex mlp tag block keep the broad set for fused-expert models whose leaves are plain Parameters. * Route explicit leaf list into MoE expert detection An attention-only explicit target_modules list routed through get_peft_regex for family scoping (e.g. FastVisionModel with vision layers off) yields a regex carrying the full mlp|feed_forward|ffn|dense component block even though its leaf group only names q/k/v/o_proj. Keying expert detection on that regex trained the experts for a language-only/attention-only request. Use the caller's original leaf list for detection; only the auto path uses the regex, where the mlp block is the sole MLP-intent signal on fused-expert models. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Respect finetune_mlp_modules and finetune_language_layers scope for MoE expert detection When an explicit leaf list that names MLP projections (gate_proj/up_proj/down_proj) is routed through get_peft_regex under finetune_mlp_modules=False, the scoped regex correctly drops the MLP leaves, but MoE expert detection was still keyed on the original list and re-added mlp.experts.* via target_parameters, training the experts the caller had frozen. Same gap for finetune_language_layers=False on vision-only runs. Prefer the original list only when MLP and language families are both in scope (preserving the attention-only fix); otherwise honor the scoped result so the frozen family is respected. Factored the choice into _select_moe_detection_targets with unit tests over the full selection matrix. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_moe_lora_targets.py | 187 +++++++++++++++++++++++++++++++++ unsloth/models/_utils.py | 47 ++++++++- unsloth/models/vision.py | 27 ++++- 3 files changed, 257 insertions(+), 4 deletions(-) diff --git a/tests/test_moe_lora_targets.py b/tests/test_moe_lora_targets.py index 994d39f261..7f9b9a0485 100644 --- a/tests/test_moe_lora_targets.py +++ b/tests/test_moe_lora_targets.py @@ -49,3 +49,190 @@ def test_explicit_dotted_module_target_does_not_discover_moe_parameters(): ) is None ) + + +@pytest.mark.parametrize( + "target_modules", + [ + # Attention-only auto-regex lists every projection leaf (incl. gate/up/down) + # but its path segment is attention-only, so experts must NOT be targeted. + r"(?:\bmodel\.layers\.[\d]{1,}\.(?:self_attn|attention|attn|mixer)\.(?:q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj))", + ".*self_attn.*proj", + # An mlp path alternative with attention-only leaves is still attention-only. + r"model\.layers\.\d+\.(?:mlp|self_attn)\.(?:q_proj|k_proj|v_proj|o_proj)", + ], +) +def test_attention_only_regex_does_not_discover_moe_parameters(target_modules): + from unsloth.models._utils import get_moe_target_parameters + assert get_moe_target_parameters(_FakeMoeModel(), target_modules) is None + + +def test_single_leaf_regex_targets_only_that_projection(): + from unsloth.models._utils import get_moe_target_parameters + assert get_moe_target_parameters(_FakeMoeModel(), ".*experts.*down_proj") == [ + "mlp.experts.down_proj", + ] + assert get_moe_target_parameters(_FakeMoeModel(), ".*mlp.*gate_proj") == [ + "mlp.experts.gate_up_proj", + ] + + +def test_auto_regex_mlp_tag_block_discovers_moe_on_fused_models(): + # get_peft_regex on a fused-expert model lists only attention Linears as + # leaves; the mlp tag block is the remaining signal of MLP finetune intent. + from unsloth.models._utils import get_moe_target_parameters + both_auto = ( + r"(?:\bmodel\.layers\.[\d]{1,}\." + r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer)\." + r"(?:(?:q_proj|k_proj|v_proj|o_proj)))" + ) + assert get_moe_target_parameters(_FakeMoeModel(), both_auto) == [ + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] + + +def test_explicit_attention_only_list_does_not_discover_moe_parameters(): + # An explicit attention-only leaf list names no MLP projection, so experts + # must never be targeted. get_peft_model routes this ORIGINAL list (not the + # scoped regex) into detection precisely because family scoping makes + # get_peft_regex emit its full "mlp|feed_forward|ffn|dense" component block + # even for an attention-only request (see the regex below), which the + # string fallback cannot distinguish from the fused-expert auto regex. + from unsloth.models._utils import get_moe_target_parameters + + attn_only_list = ["q_proj", "k_proj", "v_proj", "o_proj"] + assert get_moe_target_parameters(_FakeMoeModel(), attn_only_list) is None + assert get_moe_target_parameters(_FakeMoeModel(), tuple(attn_only_list)) is None + + # The regex get_peft_regex emits for that same attention-only list under a + # vision-off family scope carries the mlp component block, so the string + # path would wrongly enable experts -- hence detection must use the list. + scoped_regex = ( + r"(?:.*?(?:language|text).*?" + r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense|mixer).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + assert get_moe_target_parameters(_FakeMoeModel(), scoped_regex) == [ + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] + + +def test_frozen_mlp_full_list_does_not_discover_moe_parameters(): + # Regression: an explicit list that names MLP leaves together with + # finetune_mlp_modules=False must NOT train experts. get_peft_regex scopes + # the MLP leaves out (its emitted regex carries no mlp tag block), so + # detection has to key on that SCOPED regex -- keying on the original list + # would let its gate/up/down leaves silently re-enable the frozen experts. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + # Representative of what get_peft_regex emits for that list under + # finetune_mlp_modules=False: attention-only path, no mlp component block. + scoped_regex = ( + r"(?:.*?(?:language|text).*?" + r"(?:self_attn|attention|attn|mixer).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = False, + finetune_language_layers = True, + ) + assert selected is scoped_regex + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None + + +def test_frozen_language_full_list_does_not_discover_moe_parameters(): + # Vision-only request (finetune_language_layers=False) with a full leaf list + # must not reach the language-model experts either. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = ["q_proj", "gate_proj", "up_proj", "down_proj"] + scoped_regex = ( + r"(?:.*?(?:vision|visual|image).*?" + r"(?:self_attn|attention|attn|mixer).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = False, + ) + assert selected is scoped_regex + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None + + +def test_in_scope_mlp_full_list_still_discovers_moe_parameters(): + # With MLP and language both in scope, an explicit list that names MLP + # leaves SHOULD enable the experts (unchanged behavior): the original list + # is preferred and carries the gate/up/down intent. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + original_list = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + scoped_regex = r".*self_attn.*proj" # unused: original list is preferred + selected = _select_moe_detection_targets( + original_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = True, + ) + assert selected is original_list + assert get_moe_target_parameters(_FakeMoeModel(), selected) == [ + "mlp.experts.gate_up_proj", + "mlp.experts.down_proj", + ] + + +def test_attention_only_list_prefers_original_when_in_scope(): + # The case the PR originally fixed: an attention-only list routed through + # get_peft_regex under a family scope (e.g. vision-off) still keeps experts + # off, because with MLP+language in scope detection uses the original + # attention-only list rather than the regex's spurious mlp component block. + from unsloth.models._utils import ( + _select_moe_detection_targets, + get_moe_target_parameters, + ) + + attn_only_list = ["q_proj", "k_proj", "v_proj", "o_proj"] + scoped_regex = ( # carries the spurious mlp block get_peft_regex always adds + r"(?:.*?(?:language|text).*?" + r"(?:self_attn|attention|attn|mixer|mlp|feed_forward|ffn|dense).*?" + r"(?:q_proj|k_proj|v_proj|o_proj))" + ) + selected = _select_moe_detection_targets( + attn_only_list, + scoped_regex, + finetune_mlp_modules = True, + finetune_language_layers = True, + ) + assert selected is attn_only_list + assert get_moe_target_parameters(_FakeMoeModel(), selected) is None diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 260fe36652..b504a19f74 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -86,6 +86,7 @@ __all__ = [ "maybe_prefetch_hf_snapshot", "is_moe_model", "get_moe_target_parameters", + "_select_moe_detection_targets", "make_fast_generate_wrapper", "_mark_unsloth_disable_data_parallel", "_patch_transformers_trainer_data_parallel", @@ -3913,8 +3914,25 @@ def _moe_target_set_from_string(target_modules: str) -> set[str]: return {target_modules} is_regex = re.search(r"[*+?()[\]{}|\\^$]", target_modules) is not None - targets_mlp = "mlp" in target_modules or "ffn" in target_modules - if is_regex and "proj" in target_modules and targets_mlp: + # Key detection on the mlp/ffn/experts path segment (absent from an + # attention-only regex), never on q/k/v/o leaves alone. + targets_mlp_path = any( + tag in target_modules for tag in ("mlp", "ffn", "feed_forward", "experts") + ) + if not is_regex or not targets_mlp_path: + return set() + # Explicit expert leaves scope the target set to exactly those leaves. + named = {name for name in _MOE_BROAD_MLP_TARGETS if name in target_modules} + if named: + return named + # A generic projection under an mlp path (e.g. ".*mlp.*proj"): any proj + # occurrence that is not an attention leaf name. + if re.search(r"(? Optional[List[str return None +def _select_moe_detection_targets( + original_target_modules, + scoped_target_modules, + finetune_mlp_modules = True, + finetune_language_layers = True, +): + """Pick what get_moe_target_parameters keys expert detection on. + + Prefer the caller's ORIGINAL explicit leaf list over the scoped regex so an + attention-only request is not pushed into the experts by get_peft_regex's + ``mlp|feed_forward|ffn|dense`` component block (which the string fallback + cannot tell apart from a fused-expert auto regex). + + But only when the MLP and language families are BOTH still in scope. If the + caller scoped MLP or language OFF (``finetune_mlp_modules=False`` or + ``finetune_language_layers=False``) the scoped regex already drops the MoE + experts, and reusing the original list -- which may still name gate/up/down + leaves -- would wrongly re-introduce them. In that case honor the scoped + result so the frozen-MLP / vision-only request is respected. + """ + if original_target_modules is not None and finetune_mlp_modules and finetune_language_layers: + return original_target_modules + return scoped_target_modules + + def make_fast_generate_wrapper(original_generate): """ Creates a wrapper around model.generate that checks for incorrect diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 179bd0b650..e80ed8917b 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -37,6 +37,7 @@ from ._utils import ( _get_text_only_config, _is_family_text_decoder, _apply_text_only_key_mapping, + _select_moe_detection_targets, set_task_config_attr, ) from ._utils import * @@ -1703,6 +1704,16 @@ class FastBaseModel: ) else: _audio_kwargs = {} + # Remember the caller's ORIGINAL explicit leaf list for MoE expert + # detection. When an explicit list is routed through get_peft_regex for + # family scoping below, the generated regex carries get_peft_regex's full + # "mlp|feed_forward|ffn|dense" component block even when the caller named + # only attention leaves (q/k/v/o_proj). Keying expert detection on that + # regex would train the experts for an attention-only request. The + # original list carries the true leaf intent, so use it for MoE detection; + # only the auto (None / "all-linear") path relies on the regex, whose mlp + # block is the sole remaining MLP-intent signal on fused-expert models. + _moe_detect_target = target_modules if type(target_modules) in (list, tuple) else None if target_modules is None or target_modules == "all-linear": target_modules = get_peft_regex( model, @@ -1780,9 +1791,21 @@ class FastBaseModel: loftq_config, lora_dropout, bias, init_lora_weights, model ) - # Auto-detect MoE models and populate target_parameters for expert layers + # Auto-detect MoE models and populate target_parameters for expert layers. + # Prefer the caller's ORIGINAL explicit leaf list over the scoped regex so an + # attention-only request does not train experts via get_peft_regex's mlp block, + # but only when MLP and language families are both still in scope. If the caller + # scoped MLP or language OFF (finetune_mlp_modules / finetune_language_layers + # False), the scoped regex already dropped the experts, so honor it instead of + # re-introducing the original list's gate/up/down leaves. if target_parameters is None: - target_parameters = get_moe_target_parameters(model, target_modules) + _moe_targets = _select_moe_detection_targets( + _moe_detect_target, + target_modules, + finetune_mlp_modules = finetune_mlp_modules, + finetune_language_layers = finetune_language_layers, + ) + target_parameters = get_moe_target_parameters(model, _moe_targets) if finetune_last_n_layers is not None and layers_to_transform is None: _total_layers = _get_total_transformer_layers(model) From c520662c12dd01fb3f83eb871115b3611e39c172 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:45:45 -0700 Subject: [PATCH 007/113] Honor an explicit sdpa or flex_attention request when flash is disabled (#6847) * Honor an explicit sdpa or flex_attention request when flash is disabled When flash attention is disabled for a model, the fallback selection could downgrade a caller who explicitly passed attn_implementation='sdpa' or 'flex_attention' to a different backend, because the disable reason is flash-specific. Keep an explicit non-flash request as-is; flash requests still fall back as before. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments * Gate honor-explicit attention on provenance and flex support Only honor an explicit non-flash attention request when it comes from the caller argument, not from a config value the loaders synthesize (the language path seeds attn_implementation=sdpa). Honor explicit flex_attention only when supports_flex_attention is True so excluded/broken configs (e.g. gpt_oss) fall back instead of selecting a known-broken backend. Explicit sdpa stays honored. * Honor explicit sdpa through the resolver guard * Keep SDPA exclusions when honoring an explicit sdpa request An explicit attn_implementation="sdpa" was re-enabling sdpa for models in _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) where sdpa is known-broken: the helper honored the request and the resolver's final not-supports_sdpa guard skipped the eager downgrade for any explicit request. Honor an explicit sdpa only when the model is not sdpa-excluded, mirroring the flex guard that already falls back for _FLEX_EXCLUDED_MODELS via supports_flex_attention. Conservative supports_sdpa=False (large head dim / attention-sink models) still honors an explicit sdpa; a synthesized/default sdpa still downgrades to eager. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honor DISABLE_SDPA_MODEL_NAMES when honoring explicit sdpa The honor-explicit-sdpa guard only skipped the sdpa->eager downgrade for models in _SDPA_EXCLUDED_MODELS (gpt_oss). Gemma3/Gemma3Text disable SDPA through the loader's DISABLE_SDPA_MODEL_NAMES (their bundled SDPA modules are wrong), so an explicit sdpa request bypassed the downgrade and re-enabled a known-wrong path. Extend _is_sdpa_excluded to also treat DISABLE_SDPA_MODEL_NAMES membership as excluded, replicating the loader's trailing-comma substring match so gemma3 and gemma3_text match but gemma3n does not. Move the constant into _utils.py (single source of truth, re-exported from loader.py) to avoid a loader -> _utils cycle. Conservative supports_sdpa=False models not in either list still honor explicit sdpa. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/test_attn_impl_honor_explicit.py | 190 +++++++++++++++++++++++++ unsloth/models/_utils.py | 64 ++++++++- unsloth/models/loader.py | 12 +- 3 files changed, 255 insertions(+), 11 deletions(-) create mode 100644 tests/test_attn_impl_honor_explicit.py diff --git a/tests/test_attn_impl_honor_explicit.py b/tests/test_attn_impl_honor_explicit.py new file mode 100644 index 0000000000..3fb7a2208f --- /dev/null +++ b/tests/test_attn_impl_honor_explicit.py @@ -0,0 +1,190 @@ +"""An explicit non-flash attention request must survive the flash disable path. + +When flash attention is disabled for a model, a caller who explicitly asked for +"sdpa" or "flex_attention" should keep that choice instead of being downgraded +to whatever the conservative supports_* fallback would pick. +""" + +import pytest + +from unsloth.models._utils import ( + _disable_flash_attention_if_needed, + resolve_attention_implementation, +) + + +def test_explicit_sdpa_is_honored_even_when_not_marked_supported(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "sdpa", + supports_sdpa = False, # conservative flag would have skipped sdpa + supports_flex_attention = False, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_explicit_flex_is_honored_when_supported(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flex_attention", + supports_sdpa = True, + supports_flex_attention = True, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "flex_attention" + assert config.get("_attn_implementation") == "flex_attention" + + +def test_explicit_flex_falls_back_when_not_supported(): + # flex_attention is False for known-broken/excluded configs (e.g. gpt_oss), + # so an explicit flex request must not select that backend - it falls back. + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flex_attention", + supports_sdpa = True, + supports_flex_attention = False, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + + +def test_synthesized_config_sdpa_is_not_treated_as_explicit(): + # The language loader seeds the config with attn_implementation="sdpa"; when the + # caller passes nothing, that synthesized value must not override the flex fallback + # for a model that supports flex but not sdpa. + config = {"attn_implementation": "sdpa"} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = None, + supports_sdpa = False, + supports_flex_attention = True, + would_use_flash_attention = False, + disable_reason = "unit test forces flash disabled", + ) + assert result == "flex_attention" + + +def test_no_disable_reason_returns_request_untouched(): + result = _disable_flash_attention_if_needed( + {}, + attn_implementation = "flash_attention_2", + disable_reason = None, + ) + assert result == "flash_attention_2" + + +def test_flash_request_still_falls_back_when_disabled(): + config = {} + result = _disable_flash_attention_if_needed( + config, + attn_implementation = "flash_attention_2", + supports_sdpa = True, + would_use_flash_attention = True, + disable_reason = "unit test forces flash disabled", + ) + assert result == "sdpa" + + +def test_resolver_honors_explicit_sdpa_when_not_supported_and_flash_disabled(): + # End-to-end through the public resolver: an explicit sdpa request with a + # flash-disabled config (oversized head dim) and supports_sdpa=False must not be + # rewritten to eager by the resolver's own not-supports_sdpa guard. + config = {"model_type": "test", "head_dim": 512} # head_dim > 256 disables flash + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_resolver_downgrades_non_explicit_sdpa_when_not_supported(): + # No explicit request: the model resolution seeds sdpa/eager and the guard must + # still downgrade a synthesized sdpa to eager for a model that cannot run it. + config = {"model_type": "test", "attn_implementation": "sdpa"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = None, + supports_sdpa = False, + ) + assert result == "eager" + + +def test_resolver_downgrades_explicit_sdpa_for_sdpa_excluded_model(): + # gpt_oss is in _SDPA_EXCLUDED_MODELS (sdpa is known-broken) and _FLASH_EXCLUDED_MODELS + # (flash disabled). Honoring an explicit sdpa request must not re-enable that broken + # backend: it downgrades to eager, mirroring how an explicit flex request falls back + # for _FLEX_EXCLUDED_MODELS. supports_sdpa=True proves the exclusion overrides even a + # model that otherwise advertises SDPA support. + config = {"model_type": "gpt_oss"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = True, + ) + assert result == "eager" + assert config.get("_attn_implementation") == "eager" + + +@pytest.mark.parametrize("model_type", ["gemma3", "gemma3_text"]) +def test_resolver_downgrades_explicit_sdpa_for_disable_sdpa_model(model_type): + # gemma3 / gemma3_text are in DISABLE_SDPA_MODEL_NAMES: the loader forces + # supports_sdpa=False because their bundled SDPA modules are wrong. An explicit + # sdpa request with flash disabled must NOT re-enable that known-wrong path - it + # downgrades to eager, exactly like _SDPA_EXCLUDED_MODELS (gpt_oss). head_dim>256 + # disables flash to mirror the real flash-disabled scenario. + config = {"model_type": model_type, "head_dim": 512} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "eager" + assert config.get("_attn_implementation") == "eager" + + +def test_resolver_does_not_overmatch_gemma3n_for_explicit_sdpa(): + # The "gemma3," trailing-comma guard must not match gemma3n: gemma3n is not in + # DISABLE_SDPA_MODEL_NAMES, so it stays a conservative (not known-wrong) model and an + # explicit sdpa request is still honored. Proves the substring match neither over- nor + # under-matches. + config = {"model_type": "gemma3n", "head_dim": 512} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = "sdpa", + supports_sdpa = False, + ) + assert result == "sdpa" + assert config.get("_attn_implementation") == "sdpa" + + +def test_resolver_downgrades_synthesized_sdpa_for_disable_sdpa_model(): + # A synthesized/default sdpa (requested is None; the value came from config) on a + # DISABLE_SDPA_MODEL_NAMES model must still downgrade to eager. + config = {"model_type": "gemma3", "attn_implementation": "sdpa"} + result = resolve_attention_implementation( + model_class = None, + config = config, + requested_attn_implementation = None, + supports_sdpa = False, + ) + assert result == "eager" + + +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, "-q"])) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b504a19f74..169b610988 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -423,6 +423,18 @@ def apply_unsloth_gradient_checkpointing(use_gradient_checkpointing, max_seq_len _FLEX_EXCLUDED_MODELS = ("gpt_oss", "mllama", "nemotron_h", "modernbert") _FLEX_PREFERRED_MODELS = ("gemma3", "gemma3_text", "shieldgemma2") _SDPA_EXCLUDED_MODELS = ("gpt_oss",) +# The loader (loader.py) forces supports_sdpa=False for these because their bundled +# SDPA modules are wrong. Kept here, not in loader.py, so _is_sdpa_excluded can honor +# them without a loader -> _utils import cycle (loader.py already imports from _utils +# and re-exports this name for callers like sentence_transformer.py). Entries are matched +# as substrings against a comma-joined model_types string ending in a comma, so "gemma3," +# matches a distinct "gemma3" entry but not "gemma3n", and "gemma3_text" matches the +# EmbeddingGemma text model. +DISABLE_SDPA_MODEL_NAMES = [ + "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore + "gpt_oss", +] _FLASH_EXCLUDED_MODELS = ("gpt_oss",) _EAGER_ONLY_PREFIXES = ("gemma3n",) _FLASH_ATTENTION_MAX_HEAD_DIM = 256 @@ -433,8 +445,23 @@ def _is_flex_excluded(model_type): return model_type in _FLEX_EXCLUDED_MODELS +def _is_sdpa_disabled_by_name(model_type): + # Mirror the loader's DISABLE_SDPA_MODEL_NAMES check: loader.py builds + # model_types_all = ",".join(model_types) + "," and tests `name in model_types_all`. + # Rebuild the same trailing-comma form for a single model_type so the match is + # identical (e.g. "gemma3," matches "gemma3" but not "gemma3n", and "gemma3_text" + # still matches "gemma3_text"). + model_types_all = model_type.lower() + "," + return any(name.lower() in model_types_all for name in DISABLE_SDPA_MODEL_NAMES) + + def _is_sdpa_excluded(model_type): - return model_type in _SDPA_EXCLUDED_MODELS + # SDPA is known-broken for these models, so an explicit sdpa request must not + # re-enable it. Two sources: _SDPA_EXCLUDED_MODELS (resolver-level, e.g. gpt_oss) + # and DISABLE_SDPA_MODEL_NAMES (loader-level, e.g. gemma3 / gemma3_text, which the + # loader also forces to supports_sdpa=False). + lowered = model_type.lower() + return lowered in _SDPA_EXCLUDED_MODELS or _is_sdpa_disabled_by_name(lowered) def _is_flash_excluded(model_type): @@ -610,6 +637,12 @@ def _disable_flash_attention_if_needed( if disable_reason is None: return attn_implementation + # Only an implementation passed by the caller counts as an explicit request. + # Values read from the config are synthesized by the loaders (the language path + # seeds the config with attn_implementation="sdpa") or come from Transformers + # defaults, so they must not be treated as a deliberate user choice. + explicit_request = attn_implementation + requested_attn_implementation = attn_implementation if requested_attn_implementation is None: requested_attn_implementation = _config_get(config, "_attn_implementation", None) @@ -619,6 +652,20 @@ def _disable_flash_attention_if_needed( if requested_attn_implementation == "eager": return _set_attn_impl(config, "eager") + model_type = _config_get(config, "model_type", "") + + # The disable reason is flash-specific: honor an explicit non-flash request from + # the caller instead of downgrading it. SDPA is honored unless the model's SDPA is + # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES + # (e.g. gemma3 / gemma3_text); flex_attention + # is honored only when it is actually usable, since supports_flex_attention already + # rejects the excluded/broken/unavailable configs. This keeps an explicit request + # from selecting a backend the repo marks as wrong. + if explicit_request == "sdpa" and not _is_sdpa_excluded(model_type.lower()): + return _set_attn_impl(config, "sdpa") + if explicit_request == "flex_attention" and supports_flex_attention: + return _set_attn_impl(config, "flex_attention") + if supports_sdpa: fallback_attn_implementation = "sdpa" elif supports_flex_attention: @@ -631,7 +678,6 @@ def _disable_flash_attention_if_needed( if _is_flash_attention_requested(requested_attn_implementation) else "flash_attention_2" ) - model_type = _config_get(config, "model_type", "") warning_key = ( model_type, logged_attn_implementation, @@ -845,7 +891,19 @@ def resolve_attention_implementation( final_attn_impl = requested_attn_implementation _set_attn_impl(config, final_attn_impl) - if not supports_sdpa and final_attn_impl == "sdpa": + # A caller who explicitly passes requested_attn_implementation="sdpa" keeps it even + # on a conservatively unsupported model, mirroring _disable_flash_attention_if_needed + # which honors an explicit sdpa request. The exception is a model whose SDPA is + # known-broken - _SDPA_EXCLUDED_MODELS (e.g. gpt_oss) or DISABLE_SDPA_MODEL_NAMES + # (e.g. gemma3 / gemma3_text, which the loader also forces to supports_sdpa=False): + # an explicit request must not re-enable it, so it still downgrades to eager, just + # like flex falls back for _FLEX_EXCLUDED_MODELS. A synthesized/default sdpa + # (requested is None, so the value came from the model resolution above or the + # config) also downgrades. + honor_explicit_sdpa = requested_attn_implementation == "sdpa" and not _is_sdpa_excluded( + model_type + ) + if not supports_sdpa and final_attn_impl == "sdpa" and not honor_explicit_sdpa: print( f"Unsloth: {(model_type_name or 'model').title()} does not support SDPA - switching to fast eager." ) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 84f808d2b5..22cb65dc4a 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -21,6 +21,10 @@ from ._utils import ( USE_MODELSCOPE, get_transformers_model_type, hf_login, + # Single source of truth is _utils.py; re-exported here so callers doing + # `from unsloth.models.loader import DISABLE_SDPA_MODEL_NAMES` keep working and so + # _is_sdpa_excluded (in _utils) can honor it without a loader -> _utils cycle. + DISABLE_SDPA_MODEL_NAMES, ) from .granite import FastGraniteModel from .llama import FastLlamaModel, logger @@ -196,14 +200,6 @@ DISABLE_COMPILE_MODEL_NAMES = [ "granite,llava_next", # Granite-vision 3 ] -global DISABLE_SDPA_MODEL_NAMES -# Disables some SDPA modules since it's wrong -DISABLE_SDPA_MODEL_NAMES = [ - "gemma3,", # Add comma bc gemma3 will match gemma3n - "gemma3_text", # Gemma3TextModel (EmbeddingGemma) - substring match, keep underscore - "gpt_oss", -] - def _fix_rope_inv_freq(model): """Fix inv_freq corruption caused by transformers v5 meta-device loading. From c7b8666ce43d223fb3ffb1ccce89d73e26291e4a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:46:30 -0700 Subject: [PATCH 008/113] Auto-enable grouped MoE on loaded / PEFT'd models via loader hook (#6727) * Auto-enable grouped MoE on loaded / PEFT'd models via loader hook Wraps the FastLlamaModel and FastBaseModel from_pretrained / get_peft_model leaves with wrap_loader_for_grouped_moe so the grouped-GEMM MoE forward is installed on the live instance after the model and its compiled module are built. Gated by UNSLOTH_MOE_GROUPED and wrapped in try/except, so it is a no-op when the unsloth_zoo module is absent or no eligible MoE block exists. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Install grouped-MoE loader wrappers before PatchFastRL * Re-evaluate grouped MoE after loading a PEFT adapter When loading an existing adapter through FastLanguageModel.from_pretrained, the base model is evaluated for grouped MoE when the wrapped from_pretrained leaf returns, but the adapter is attached afterwards via PeftModel and patch_peft_model. Re-run auto_enable_grouped_moe on the final model so blocks whose experts gained LoRA are restored to the original loop, attention-only adapters keep the grouped path on their frozen experts, and recompute is re-derived from the final gradient-checkpointing state. Guarded so it never blocks adapter loading. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim comments in the grouped MoE loader hooks Shorten the loader re-eval and llama.py wrapper comments; code is unchanged (verified comment-only). * Re-evaluate grouped MoE after loading a PEFT adapter on the vision path --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/llama.py | 13 +++++++++++++ unsloth/models/loader.py | 18 ++++++++++++++++++ unsloth/models/vision.py | 13 +++++++++++++ 3 files changed, 44 insertions(+) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 417a88f480..564be09578 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3824,4 +3824,17 @@ class FastLlamaModel: from .rl import PatchFastRL +# Auto-enable grouped-GEMM MoE (tf<5 ModuleList experts) on built / PEFT'd models. Wrap the +# loader leaves before PatchFastRL so downstream patchers see the wrapped versions. Guarded. +try: + from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe + FastLlamaModel.from_pretrained = staticmethod( + wrap_loader_for_grouped_moe(FastLlamaModel.from_pretrained) + ) + FastLlamaModel.get_peft_model = staticmethod( + wrap_loader_for_grouped_moe(FastLlamaModel.get_peft_model) + ) +except Exception: + pass + PatchFastRL(FastLanguageModel = FastLlamaModel) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 22cb65dc4a..2818b6ee80 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -896,6 +896,15 @@ class FastLanguageModel(FastLlamaModel): ) # Patch it as well! model = dispatch_model.patch_peft_model(model, use_gradient_checkpointing) + # Re-evaluate grouped MoE now the adapter is attached: an expert-LoRA block falls back + # to the original loop, an attention-only adapter keeps the grouped path. Guarded. + try: + from unsloth_zoo.temporary_patches.moe_grouped_modulelist import ( + auto_enable_grouped_moe, + ) + auto_enable_grouped_moe(model) + except Exception: + pass # optional speedup; never block model loading # Patch Tiled MLP # to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}"" @@ -1852,6 +1861,15 @@ class FastModel(FastBaseModel): model = FastBaseModel.post_patch_model( model, use_gradient_checkpointing, trust_remote_code = trust_remote_code ) + # Re-evaluate grouped MoE now the adapter is attached: an expert-LoRA block falls back + # to the original loop, an attention-only adapter keeps the grouped path. Guarded. + try: + from unsloth_zoo.temporary_patches.moe_grouped_modulelist import ( + auto_enable_grouped_moe, + ) + auto_enable_grouped_moe(model) + except Exception: + pass # optional speedup; never block model loading # Apply QAT if specified if qat_scheme is not None: diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index e80ed8917b..0a68a49fee 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -2304,3 +2304,16 @@ def check_dataset_for_missing_videos( warnings.warn(error_msg, stacklevel = 2) return missing + + +# Auto-enable grouped-GEMM MoE (transformers<5 ModuleList experts); see llama.py. +try: + from unsloth_zoo.temporary_patches.moe_grouped_modulelist import wrap_loader_for_grouped_moe + FastBaseModel.from_pretrained = staticmethod( + wrap_loader_for_grouped_moe(FastBaseModel.from_pretrained) + ) + FastBaseModel.get_peft_model = staticmethod( + wrap_loader_for_grouped_moe(FastBaseModel.get_peft_model) + ) +except Exception: + pass From 95a73f0f515906d6d691c2703a6a04902f0efdb0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:47:26 -0700 Subject: [PATCH 009/113] Honor explicit load_in_16bit for local -bf16 directories (#6726) A model path ending in -bf16 unconditionally forced 16-bit loading, so a LOCAL checkpoint directory whose name happens to end in -bf16 could never be loaded in 4-bit, 8-bit or fp8: the suffix rule silently overrode the caller's quantization flags. Hub repo ids keep the existing behavior (the suffix is a publishing convention there), but for a local directory (expanduser-aware, so tilde paths are detected too) the requested quantization is preserved unless the caller explicitly passes load_in_16bit=True. --- unsloth/models/loader.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 2818b6ee80..9ce74c4d02 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -466,8 +466,10 @@ class FastLanguageModel(FastLlamaModel): ("-unsloth-bnb-4bit", "-bnb-4bit") ): model_name = _strip_unsloth_bnb_4bit_suffix(model_name) - # Change -BF16 to all False for 4bit, 8bit etc - if model_name.lower().endswith("-bf16"): + # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set + if model_name.lower().endswith("-bf16") and ( + load_in_16bit or not os.path.isdir(os.path.expanduser(model_name)) + ): load_in_4bit = False load_in_8bit = False load_in_fp8 = False @@ -625,8 +627,10 @@ class FastLanguageModel(FastLlamaModel): ("-unsloth-bnb-4bit", "-bnb-4bit") ): model_name = _strip_unsloth_bnb_4bit_suffix(model_name) - # Change -BF16 to all False for 4bit, 8bit etc - if model_name.lower().endswith("-bf16"): + # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set + if model_name.lower().endswith("-bf16") and ( + load_in_16bit or not os.path.isdir(os.path.expanduser(model_name)) + ): load_in_4bit = False load_in_8bit = False load_in_fp8 = False @@ -1145,8 +1149,10 @@ class FastModel(FastBaseModel): ("-unsloth-bnb-4bit", "-bnb-4bit") ): model_name = _strip_unsloth_bnb_4bit_suffix(model_name) - # Change -BF16 to all False for 4bit, 8bit etc - if model_name.lower().endswith("-bf16"): + # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set + if model_name.lower().endswith("-bf16") and ( + load_in_16bit or not os.path.isdir(os.path.expanduser(model_name)) + ): load_in_4bit = False load_in_8bit = False load_in_fp8 = False @@ -1503,8 +1509,10 @@ class FastModel(FastBaseModel): ("-unsloth-bnb-4bit", "-bnb-4bit") ): model_name = _strip_unsloth_bnb_4bit_suffix(model_name) - # Change -BF16 to all False for 4bit, 8bit etc - if model_name.lower().endswith("-bf16"): + # '-bf16' hub repos load bf16; a local dir keeps the requested quant unless 16bit is set + if model_name.lower().endswith("-bf16") and ( + load_in_16bit or not os.path.isdir(os.path.expanduser(model_name)) + ): load_in_4bit = False load_in_8bit = False load_in_fp8 = False From efcaffb17b5ae085f98bc35bfa8570c5192130d1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:48:10 -0700 Subject: [PATCH 010/113] Sync FORCE_FLOAT32 fallback with unsloth-zoo (gemma4, glm4_moe, qwen3_moe) (#6865) * Add gemma4, glm4_moe and qwen3_moe to the FORCE_FLOAT32 fallback list Keeps the fallback list (used only if the unsloth_zoo import fails) in sync with unsloth_zoo/model_lists.py, which now force-float32s these MoE archs so a float16 request loads bf16 and trains finite instead of NaNing the grad_norm. * Union FORCE_FLOAT32 fallback so new archs force float32 with older unsloth_zoo --- unsloth/models/loader.py | 34 ++++++++++++++++++++-------------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 9ce74c4d02..5ba9b54ce1 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -113,22 +113,28 @@ from ._utils import ( maybe_prefetch_hf_snapshot, ) -# Single source of truth is unsloth_zoo.model_lists. Re-exported so callers -# doing `from unsloth.models.loader import FORCE_FLOAT32` keep working. -# Fallback list mirrors zoo for users who upgrade unsloth without upgrading -# unsloth_zoo (so this module never fails at import). +# Source of truth is unsloth_zoo.model_lists. Re-exported so callers doing +# `from unsloth.models.loader import FORCE_FLOAT32` keep working. The fallback +# list is also unioned in so a newer unsloth still forces float32 for these +# archs when paired with an older unsloth_zoo that predates them (upgrade skew). +_FORCE_FLOAT32_FALLBACK = [ + "gemma3,", # Add comma bc gemma3 will match gemma3n + "gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3) + "gemma3n", + "gemma4", # Gemma4 (gemma4 / gemma4_text): float16 NaNs grad norms in the backward + "glm4_moe", # GLM-4.x MoE (glm4_moe / glm4_moe_lite): float16 NaNs grad norms + "gpt_oss", + "qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training + "qwen3_moe", # Qwen3-MoE (Qwen3-30B-A3B): float16 NaNs grad norms in the backward +] try: - from unsloth_zoo import FORCE_FLOAT32 # noqa: F401 + from unsloth_zoo import FORCE_FLOAT32 as _ZOO_FORCE_FLOAT32 + FORCE_FLOAT32 = list(_ZOO_FORCE_FLOAT32) except ImportError: - global FORCE_FLOAT32 - # Forces float32 precision since float16 goes to infinity - FORCE_FLOAT32 = [ - "gemma3,", # Add comma bc gemma3 will match gemma3n - "gemma3text", # Gemma3TextModel (EmbeddingGemma, standalone text-only Gemma3) - "gemma3n", - "gpt_oss", - "qwen3_5", # Qwen3.5 GDN layers produce NaN grad norms in float16 training - ] + FORCE_FLOAT32 = [] +for _mt in _FORCE_FLOAT32_FALLBACK: + if not any(_mt in _entry for _entry in FORCE_FLOAT32): + FORCE_FLOAT32.append(_mt) global DISABLE_COMPILE_MODEL_NAMES # Must be alphabetically sorted for each entry From cf4906dbe60d4cb7f21bafa75e147a6fa907b557 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 05:50:15 -0700 Subject: [PATCH 011/113] Note bundled flash-linear-attention kernels for gated-deltanet models (#6850) * Note the bundled flash-linear-attention kernels for gated-deltanet models Unsloth Zoo now bundles the flash-linear-attention (fla) gated-delta Triton kernels and injects them automatically, so gated-deltanet models (Qwen3-Next, Qwen3.5, Kimi-Linear) get the fast path with no pip install. Replace the old install advisory with a one-time note that fires only when the bundled kernels could not be enabled on the current setup (no CUDA, or torch < 2.7 / triton < 3.3), i.e. exactly when transformers falls back to the slow pure PyTorch path. * Tighten comments * Normalize model_types in fla install advisory for None and single string * Cover olmo_hybrid in the gated-deltanet fla advisory --- unsloth/models/loader.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 5ba9b54ce1..ba23197861 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -206,6 +206,44 @@ DISABLE_COMPILE_MODEL_NAMES = [ "granite,llava_next", # Granite-vision 3 ] +# Architectures with gated-deltanet (linear attention) layers. Unsloth bundles the +# flash-linear-attention Triton kernels (unsloth_zoo/_vendored/fla), so no install is +# needed; transformers uses the much slower pure PyTorch path only when they can't be enabled. +FLA_MODEL_TYPE_PREFIXES = ("qwen3_next", "qwen3_5", "kimi_linear", "olmo_hybrid") +_fla_advised = False + + +def _maybe_advise_fla_install(model_types): + """One-time note when a gated-deltanet model loads without the fast kernels. + + The kernels ship with Unsloth (no install needed); this fires only when they + could not be enabled on this platform (e.g. no CUDA, torch < 2.7 or + triton < 3.3), i.e. exactly when transformers uses the slow pure PyTorch path. + """ + global _fla_advised + if _fla_advised: + return + if model_types is None: + return + if isinstance(model_types, str): + model_types = [model_types] # a lone string would otherwise iterate chars + try: + if not any( + isinstance(t, str) and t.startswith(FLA_MODEL_TYPE_PREFIXES) for t in model_types + ): + return + from transformers.utils.import_utils import is_flash_linear_attention_available + if is_flash_linear_attention_available(): + return # bundled (or user-installed) fast kernels are active + except Exception: + return + _fla_advised = True + print( + "Unsloth: This model uses gated-deltanet linear attention layers. Unsloth\n" + "bundles the flash-linear-attention kernels, but they could not be enabled\n" + "on this setup (they need CUDA with torch >= 2.7 and triton >= 3.3), so\n" + "transformers will use a slower pure PyTorch path." + ) def _fix_rope_inv_freq(model): """Fix inv_freq corruption caused by transformers v5 meta-device loading. @@ -1304,6 +1342,7 @@ class FastModel(FastBaseModel): trust_remote_code = trust_remote_code, ) model_types_all = ",".join(model_types) + "," + _maybe_advise_fla_install(model_types) # ---- Text-diffusion models (e.g. DiffusionGemma) take a transformers-only slow path. ---- # These use a custom block-diffusion `generate` and a novel backbone, so we skip Unsloth's From 487b420948ef6fea0a77dac7b7970aa3e6b2f5c0 Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:41:33 +0530 Subject: [PATCH 012/113] CI: pin lockfile-audit actions to commit SHAs (#6902) --- .github/workflows/lockfile-audit.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/lockfile-audit.yml b/.github/workflows/lockfile-audit.yml index 9c28e21672..aaf258d615 100644 --- a/.github/workflows/lockfile-audit.yml +++ b/.github/workflows/lockfile-audit.yml @@ -60,11 +60,11 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 with: persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 with: python-version: '3.12' From f4d1dc541fbbca17e7ed59daea07465696877255 Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:41:41 +0530 Subject: [PATCH 013/113] fix(fp8): use int64 offsets in weight_dequant_kernel (#6884) --- unsloth/kernels/fp8.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 80db2f466b..935ffbb447 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -68,7 +68,9 @@ def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): n = tl.cdiv(N, BLOCK_SIZE) offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - offs = offs_m[:, None] * N + offs_n[None, :] + # tl.arange is int32, so offs_m * N overflows for tensors with more than + # 2**31 elements (e.g. flattened MoE expert stacks); index in int64. + offs = offs_m[:, None].to(tl.int64) * N + offs_n[None, :].to(tl.int64) mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) x = tl.load(x_ptr + offs, mask = mask).to(tl.float32) s = tl.load(s_ptr + pid_m * n + pid_n) From c44d94f1ae8971b644e2b4807285ab65dc09713b Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:41:49 +0530 Subject: [PATCH 014/113] fix: map None quant method to q8_0 before lowercasing in GGUF export (#6889) --- .../test_quant_method_none_normalization.py | 81 +++++++++++++++++++ unsloth/save.py | 14 ++-- 2 files changed, 89 insertions(+), 6 deletions(-) create mode 100644 tests/saving/test_quant_method_none_normalization.py diff --git a/tests/saving/test_quant_method_none_normalization.py b/tests/saving/test_quant_method_none_normalization.py new file mode 100644 index 0000000000..c1c5fd3686 --- /dev/null +++ b/tests/saving/test_quant_method_none_normalization.py @@ -0,0 +1,81 @@ +"""CPU-only regression for the quant-method normalization loops in save.py. + +`unsloth_save_pretrained_gguf` and `save_to_gguf_generic` each normalize the +`quantization_method` list, mapping a ``None`` element to ``"q8_0"``. The mapping +used to call ``quant_method.lower()`` as the first statement of the loop, so a +``None`` element (e.g. ``quantization_method=[None]`` or ``["q4_k_m", None]``) +raised ``AttributeError: 'NoneType' object has no attribute 'lower'`` and the +``elif quant_method is None`` branch was unreachable dead code. + +The loop is inline inside two heavy functions (importing unsloth needs +unsloth_zoo / a GPU), so - like test_is_gpt_oss_detection.py - we extract just the +loop source via ``ast`` and exec it against sample inputs. That exercises the real +source: it fails on the old ordering and passes once ``None`` is handled first. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" +SAVE_SRC = SAVE_PY.read_text(encoding = "utf-8") +SAVE_TREE = ast.parse(SAVE_SRC, filename = str(SAVE_PY)) + +# The target functions and the list variable each one appends the normalized method to. +TARGETS = ( + ("unsloth_save_pretrained_gguf", "quantization_methods"), + ("save_to_gguf_generic", "new_quantization_methods"), +) + + +def _func(tree, name): + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"function {name!r} not found in {SAVE_PY.name}") + + +def _quant_loop(func_name): + # The quant-normalization `for` loop iterates `quantization_method`; grab its source. + func = _func(SAVE_TREE, func_name) + for node in ast.walk(func): + if ( + isinstance(node, ast.For) + and isinstance(node.iter, ast.Call) + and isinstance(node.iter.func, ast.Name) + and node.iter.func.id == "enumerate" + and isinstance(node.iter.args[0], ast.Name) + and node.iter.args[0].id == "quantization_method" + ): + return node + raise AssertionError(f"quant-normalization loop not found in {func_name}") + + +def _run_loop(func_name, out_var, quantization_method): + # exec just the extracted loop against a given input, returning the appended methods. + loop_src = ast.get_source_segment(SAVE_SRC, _quant_loop(func_name)) + namespace = {out_var: [], "quantization_method": quantization_method} + exec(loop_src, {"__builtins__": __builtins__}, namespace) + return namespace[out_var] + + +@pytest.mark.parametrize("func_name, out_var", TARGETS) +def test_none_element_maps_to_q8_0(func_name, out_var): + # A bare None inside the list must map to q8_0, not raise AttributeError. + assert _run_loop(func_name, out_var, [None]) == ["q8_0"] + + +@pytest.mark.parametrize("func_name, out_var", TARGETS) +def test_none_mixed_with_strings(func_name, out_var): + # None resolves to q8_0 while sibling string methods are still normalized (lowercased). + assert _run_loop(func_name, out_var, ["Q4_K_M", None]) == ["q4_k_m", "q8_0"] + + +@pytest.mark.parametrize("func_name, out_var", TARGETS) +def test_string_methods_unchanged(func_name, out_var): + # The fix must not alter behavior for the ordinary string inputs. + methods = ["not_quantized", "fast_quantized", "quantized", "Q8_0"] + assert _run_loop(func_name, out_var, methods) == ["f16", "q8_0", "q4_k_m", "q8_0"] diff --git a/unsloth/save.py b/unsloth/save.py index a6697e98a1..020c63a9e2 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -2926,15 +2926,16 @@ def unsloth_save_pretrained_gguf( "Unsloth: quantization_method can only be a string or a list of strings" ) for i, quant_method in enumerate(quantization_method): - quant_method = quant_method.lower() + if quant_method is None: + quant_method = "q8_0" + else: + quant_method = quant_method.lower() if quant_method == "not_quantized": quant_method = "f16" elif quant_method == "fast_quantized": quant_method = "q8_0" elif quant_method == "quantized": quant_method = "q4_k_m" - elif quant_method is None: - quant_method = "q8_0" quantization_methods.append(quant_method.lower()) try: @@ -3727,15 +3728,16 @@ def save_to_gguf_generic( "Unsloth: quantization_method can only be a string or a list of strings" ) for i, quant_method in enumerate(quantization_method): - quant_method = quant_method.lower() + if quant_method is None: + quant_method = "q8_0" + else: + quant_method = quant_method.lower() if quant_method == "not_quantized": quant_method = "f16" elif quant_method == "fast_quantized": quant_method = "q8_0" elif quant_method == "quantized": quant_method = "q4_k_m" - elif quant_method is None: - quant_method = "q8_0" new_quantization_methods.append(quant_method.lower()) else: new_quantization_methods.append(quantization_type.lower()) From cc99aab607cd5310ae0dd9468475e1d084fa9eda Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Mon, 6 Jul 2026 19:41:57 +0530 Subject: [PATCH 015/113] fix: correct class name in SyntheticDataKit.chunk_data guard message (#6901) --- tests/test_synthetic_chunk_data.py | 26 ++++++++++++++++++++++++++ unsloth/dataprep/synthetic.py | 2 +- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/test_synthetic_chunk_data.py b/tests/test_synthetic_chunk_data.py index b9167d214f..abc2c01443 100644 --- a/tests/test_synthetic_chunk_data.py +++ b/tests/test_synthetic_chunk_data.py @@ -104,10 +104,36 @@ def test_chunk_data_rejects_overlap_not_smaller_than_chunk(): os.unlink(path) +def test_chunk_data_uninitialized_error_names_real_class(): + # Without max_seq_length the guard tells the user which method to call first. + # The message must name the real class (SyntheticDataKit) so copying it works; + # a misspelling would raise NameError when the user follows it verbatim. + kit = SyntheticDataKit.__new__(SyntheticDataKit) + kit.tokenizer = _MockTokenizer() # max_seq_length intentionally unset + with tempfile.NamedTemporaryFile("w", suffix = ".txt", delete = False) as f: + f.write("word " * 50) + path = f.name + try: + try: + kit.chunk_data(filename = path) + raise AssertionError("expected RuntimeError when max_seq_length is unset") + except RuntimeError as e: + msg = str(e) + assert ( + "SyntheticDataKit.from_pretrained" in msg + ), f"error must name SyntheticDataKit.from_pretrained, got: {msg}" + assert ( + "SynthetidDataKit" not in msg + ), f"error must not misspell the class name, got: {msg}" + finally: + os.unlink(path) + + if __name__ == "__main__": test_chunk_data_keeps_single_chunk_document() test_chunk_data_still_splits_long_document() test_chunk_data_empty_document_yields_no_chunks() test_chunk_data_short_document_is_not_split_into_fragments() test_chunk_data_rejects_overlap_not_smaller_than_chunk() + test_chunk_data_uninitialized_error_names_real_class() print("OK") diff --git a/unsloth/dataprep/synthetic.py b/unsloth/dataprep/synthetic.py index 10690810fe..6f025343f5 100644 --- a/unsloth/dataprep/synthetic.py +++ b/unsloth/dataprep/synthetic.py @@ -391,7 +391,7 @@ class SyntheticDataKit: assert os.path.exists(filename) assert hasattr(self, "tokenizer") if not hasattr(self, "max_seq_length"): - raise RuntimeError("Please use SynthetidDataKit.from_pretrained(...) first!") + raise RuntimeError("Please use SyntheticDataKit.from_pretrained(...) first!") if not hasattr(self, "overlap") or not hasattr(self, "max_generation_tokens"): raise RuntimeError("Please use prepare_qa_generation first!") From 46e2cf5deea00419539a870c9113e3db68ab240b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 08:27:26 -0700 Subject: [PATCH 016/113] studio: label RAM and VRAM readouts as GiB not GB (#6895) The live resource monitor and GPU readouts derive memory from binary byte counts (bytes / 1024**3 for torch and psutil, MiB / 1024 for the nvidia-smi path), which is GiB, but the UI labeled the values "GB". On a B200 this showed "178.35 GB" for a card whose nvidia-smi total is 183359 MiB (179 GiB), so it looked like memory was missing. Relabel the measured RAM and VRAM readouts to GiB across the floating monitor, the resources tab, the studio live GPU panel, the hub header, the about tab and the onboarding summary. The numeric values are unchanged, so the training GPU selection and memory-fit logic that read the same fields are unaffected. Disk stays labeled GB because the backend reports it in decimal GB (bytes / 1e9), and model file sizes and download progress keep their decimal GB labels to match Hugging Face. --- .../src/components/floating-monitor.tsx | 10 ++++---- studio/frontend/src/features/hub/hub-page.tsx | 4 ++-- .../components/steps/summary-step.tsx | 2 +- .../src/features/settings/tabs/about-tab.tsx | 2 +- .../features/settings/tabs/resources-tab.tsx | 23 +++++++++++++------ .../studio/sections/progress-section.tsx | 4 ++-- 6 files changed, 28 insertions(+), 17 deletions(-) diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx index f02da6612e..bce4bf2831 100644 --- a/studio/frontend/src/components/floating-monitor.tsx +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -27,9 +27,11 @@ function usageTextClass(percent: number): string { return "text-primary"; } -function formatGb(value: number): string { +function formatGiB(value: number): string { + // RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching + // nvidia-smi and PyTorch, so label the readout GiB rather than GB. const digits = value >= 10 ? 1 : 2; - return `${value.toFixed(digits)} GB`; + return `${value.toFixed(digits)} GiB`; } export function FloatingMonitor() { @@ -116,7 +118,7 @@ export function FloatingMonitor() {
- {formatGb(ramUsed)} / {formatGb(ramTotal)} + {formatGiB(ramUsed)} / {formatGiB(ramTotal)}
- {formatGb(vramUsed)} / {formatGb(vramTotal)} + {formatGiB(vramUsed)} / {formatGiB(vramTotal)}
0 - ? `${Math.round(gpu.systemRamTotalGb)} GB` + ? `${Math.round(gpu.systemRamTotalGb)} GiB` : "Unavailable"; const coreLabel = gpu.cpuCore > 0 && gpu.cpuThread > 0 diff --git a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx index 6b6b11bf1f..eb50d398b6 100644 --- a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx @@ -125,7 +125,7 @@ export function SummaryStep() { GPU
{hw.gpuName ?? "---"} - {hw.vramTotalGb != null ? `${hw.vramTotalGb} GB` : "---"} + {hw.vramTotalGb != null ? `${hw.vramTotalGb} GiB` : "---"}
diff --git a/studio/frontend/src/features/settings/tabs/about-tab.tsx b/studio/frontend/src/features/settings/tabs/about-tab.tsx index ff751e3cd6..1f323a9e8d 100644 --- a/studio/frontend/src/features/settings/tabs/about-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/about-tab.tsx @@ -158,7 +158,7 @@ export function AboutTab() { {gpu.name ?? "—"} {gpu.vramTotalGb != null - ? ` · ${Math.round(gpu.vramTotalGb)} GB` + ? ` · ${Math.round(gpu.vramTotalGb)} GiB` : ""} diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx index d5e19cc51c..6c30858c63 100644 --- a/studio/frontend/src/features/settings/tabs/resources-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx @@ -47,6 +47,15 @@ function formatGb(value: number | null | undefined): string { return `${safe.toFixed(digits)} GB`; } +// RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching +// nvidia-smi and PyTorch, so label those readouts GiB. Disk stays on formatGb +// because the backend reports disk in decimal GB (bytes / 1e9). +function formatGiB(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + const digits = safe >= 10 ? 1 : 2; + return `${safe.toFixed(digits)} GiB`; +} + function formatMb(value: number | null | undefined): string { const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; return `${Math.round(safe).toLocaleString()} MB`; @@ -300,9 +309,9 @@ export function ResourcesTab() { /> @@ -318,13 +327,13 @@ export function ResourcesTab() { label={t("settings.resources.liveMonitor.vram")} value={ hasGpu - ? `${formatGb(metrics.vramUsed)} / ${formatGb(metrics.vramTotal)}` + ? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}` : t("settings.resources.liveMonitor.noGpu") } detail={ hasGpu ? t("settings.resources.liveMonitor.free", { - value: formatGb(metrics.vramFree), + value: formatGiB(metrics.vramFree), }) : backendLabel } @@ -373,17 +382,17 @@ export function ResourcesTab() {
{t("settings.resources.gpu.used", { - value: formatGb(used), + value: formatGiB(used), })} {t("settings.resources.gpu.free", { - value: formatGb(free), + value: formatGiB(free), })} {t("settings.resources.gpu.total", { - value: formatGb(total), + value: formatGiB(total), })}
diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index abab35db93..9c9398688f 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -411,7 +411,7 @@ function LiveGpuPanel({ value={index} className="bg-popover text-popover-foreground dark:bg-zinc-900 dark:text-zinc-100" > - GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GB` : "N/A"}) + GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GiB` : "N/A"}) ))} @@ -446,7 +446,7 @@ function LiveGpuPanel({ icon={} value={ currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null - ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB` + ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GiB` : "--" } pct={currentGpu.vram_utilization_pct ?? 0} From 2fada48ef5fb1f827cbb78684e42c9a06bd29020 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 09:13:14 -0700 Subject: [PATCH 017/113] Fix llama3 RoPE scaling dropped on transformers v5 (#6907) * Fix llama3 RoPE scaling dropped on transformers v5 transformers v5 loads on meta then blanks non-persistent buffers, so _fix_rope_inv_freq rebuilds inv_freq after load. It recomputed a vanilla inv_freq and applied _apply_inv_freq_scaling, a no-op on the base LlamaRotaryEmbedding used by the config/llama3 path, so inv_freq ended up divided by 1 instead of the config factor (8 for Llama 3.1, 32 for Llama 3.2). This corrupts long-range positions and inflates long-context loss about 3-5x. transformers 4.x was unaffected. Route __init__ and the v5 repair through one _unsloth_recompute_inv_freq so they cannot diverge, and stash the config on the rotary module so the repair can rebuild the same scaled value. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add test for llama3 RoPE scaling under the transformers v5 repair * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update RoPE drift guard for the recompute refactor and guard the v5 repair The drift guard's AST tripwire asserted the config-scaling call lived in the if config is not None branch of LlamaRotaryEmbedding.__init__. The fix moved that into _unsloth_recompute_inv_freq, so follow it there (with a fallback to the old inline branch) and add a guard that loader._fix_rope_inv_freq rebuilds inv_freq through the same helper. Also add a CPU functional check of the helper and drop the redundant standalone test. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/utils/test_rope_scaling_drift.py | 128 +++++++++++++++++++------ unsloth/models/llama.py | 41 ++++---- unsloth/models/loader.py | 19 ++-- 3 files changed, 132 insertions(+), 56 deletions(-) diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index b976654f87..98f7e2db62 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -31,6 +31,7 @@ requires_cuda = pytest.mark.skipif( REPO_ROOT = Path(__file__).resolve().parents[2] LLAMA_PY = REPO_ROOT / "unsloth" / "models" / "llama.py" +LOADER_PY = REPO_ROOT / "unsloth" / "models" / "loader.py" CLASS_NAME = "LlamaRotaryEmbedding" @@ -78,42 +79,88 @@ def _config_branch(init_fn): return None +def _iter_names_and_calls(node): + """(attribute/string names, bare-name calls, method-call attrs) under node.""" + names, calls, call_attrs = set(), set(), set() + for sub in ast.walk(node): + if isinstance(sub, ast.Attribute): + names.add(sub.attr) + elif isinstance(sub, ast.Constant) and isinstance(sub.value, str): + names.add(sub.value) + elif isinstance(sub, ast.Call): + if isinstance(sub.func, ast.Name): + calls.add(sub.func.id) + elif isinstance(sub.func, ast.Attribute): + call_attrs.add(sub.func.attr) + return names, calls, call_attrs + + +def _find_method(source_path, class_name, method_name): + for node in ast.walk(ast.parse(source_path.read_text())): + if isinstance(node, ast.ClassDef) and node.name == class_name: + for sub in node.body: + if isinstance(sub, ast.FunctionDef) and sub.name == method_name: + return sub + return None + + +def _find_function(source_path, function_name): + for node in ast.walk(ast.parse(source_path.read_text())): + if isinstance(node, ast.FunctionDef) and node.name == function_name: + return node + return None + + def test_config_path_inspects_rope_scaling(): init_fn = _load_class_init() - branch = _config_branch(init_fn) - assert branch is not None, ( - f"{CLASS_NAME}.__init__ no longer has an `if config is not None:` " - "branch; the config constructor path must read config.rope_scaling so " - "scaled models (llama3/linear/longrope) are not silently unscaled " - "(issue #2405)" - ) + # inv_freq is derived through the shared _unsloth_recompute_inv_freq helper + # (or still inlined in the config branch on older layouts); whichever scope + # holds the scaling must read config.rope_scaling and call + # _compute_config_rope_inv_freq, else scaled models run unscaled (#2405). + _, _, init_call_attrs = _iter_names_and_calls(init_fn) + scope = _find_method(LLAMA_PY, CLASS_NAME, "_unsloth_recompute_inv_freq") + if scope is not None: + assert "_unsloth_recompute_inv_freq" in init_call_attrs, ( + f"{CLASS_NAME}.__init__ no longer derives inv_freq via " + "_unsloth_recompute_inv_freq; keep the constructor wired to the " + "shared scaling helper or scaled configs silently lose RoPE scaling " + "(issue #2405)." + ) + else: + scope = _config_branch(init_fn) + assert scope is not None, ( + f"{CLASS_NAME}.__init__ has neither a _unsloth_recompute_inv_freq " + "helper nor an `if config is not None:` branch; the config path must " + "apply llama3/linear/longrope scaling (issue #2405)." + ) - names = set() - for stmt in branch.body: - for sub in ast.walk(stmt): - if isinstance(sub, ast.Attribute): - names.add(sub.attr) - elif isinstance(sub, ast.Constant) and isinstance(sub.value, str): - names.add(sub.value) + names, called, _ = _iter_names_and_calls(scope) assert "rope_scaling" in names, ( - f"{CLASS_NAME}.__init__ config path does not reference `rope_scaling`. " - "When a rotary class is built straight from a config (the path modern " - "transformers takes, since rotary moved to LlamaModel), the llama3 / " - "linear / longrope scaling must still be applied; otherwise long inputs " - "produce repeated-pattern gibberish (issue #2405)." + f"{CLASS_NAME} inv_freq computation does not reference `rope_scaling`; " + "scaled models (llama3/linear/longrope) would run unscaled and produce " + "repeated-pattern gibberish past the original context (issue #2405)." + ) + assert "_compute_config_rope_inv_freq" in called, ( + f"{CLASS_NAME} inv_freq computation no longer calls " + "_compute_config_rope_inv_freq; keep it wired or scaled configs silently " + "lose RoPE scaling again (issue #2405)." ) - called = { - sub.func.id - for stmt in branch.body - for sub in ast.walk(stmt) - if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name) - } - assert "_compute_config_rope_inv_freq" in called, ( - f"{CLASS_NAME}.__init__ config path no longer calls " - "_compute_config_rope_inv_freq; the CPU behavioral tests below cover " - "that helper directly, so the constructor must stay wired to it or " - "scaled configs silently lose RoPE scaling again (issue #2405)." + +def test_v5_repair_reuses_recompute(): + # transformers v5 blanks non-persistent buffers on load, so + # loader._fix_rope_inv_freq rebuilds inv_freq; it must reuse the scaled + # recompute, since an unscaled rebuild re-drops llama3 scaling (#2405). + fix_fn = _find_function(LOADER_PY, "_fix_rope_inv_freq") + assert fix_fn is not None, ( + "loader._fix_rope_inv_freq not found; if it was renamed, update this " + "guard so the v5 rope repair keeps applying config scaling (issue #2405)." + ) + _, _, call_attrs = _iter_names_and_calls(fix_fn) + assert "_unsloth_recompute_inv_freq" in call_attrs, ( + "loader._fix_rope_inv_freq no longer rebuilds inv_freq via " + "_unsloth_recompute_inv_freq; transformers v5 blanks the buffer on load " + "and an unscaled rebuild re-drops llama3 scaling (issue #2405)." ) @@ -189,6 +236,27 @@ def test_default_rope_type_matches_vanilla_inv_freq(): ) +def test_recompute_helper_scales_on_cpu(): + # Exercise the exact method loader._fix_rope_inv_freq calls, without CUDA. + from unsloth.models.llama import LlamaRotaryEmbedding, _get_rope_theta + + def recompute(config): + rot = object.__new__(LlamaRotaryEmbedding) + rot.attention_scaling = 1.0 + rot.base = _get_rope_theta(config, 10000.0) + rot.dim = config.head_dim + rot._unsloth_rope_config = config + return rot._unsloth_recompute_inv_freq().float().cpu() + + config = _make_config(LLAMA3_ROPE_SCALING) + assert torch.allclose( + recompute(config), _reference_inv_freq(config, "llama3"), rtol = 1e-4, atol = 1e-6 + ), "_unsloth_recompute_inv_freq dropped llama3 scaling (issue #2405)." + assert torch.allclose( + recompute(_make_config(None)), _vanilla_inv_freq(), rtol = 1e-4, atol = 1e-6 + ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled." + + def _cos_at_position(rot, position): """cos row at one position, built like _set_cos_sin_cache but CPU-only.""" inv_freq = rot.inv_freq.float().cpu() diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 564be09578..c25a031b82 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1756,7 +1756,6 @@ class LlamaRotaryEmbedding(torch.nn.Module): # Base-class-from-config path (modern transformers): derive inv_freq like # transformers so config.rope_scaling is not dropped (#2405). Scaled # subclasses are excluded to avoid double-scaling. - config_inv_freq = None if config is not None: # [TODO] Hack to pass in config - need to remove later base = _get_rope_theta(config, default = base) @@ -1769,32 +1768,17 @@ class LlamaRotaryEmbedding(torch.nn.Module): device = DEVICE_TYPE_TORCH max_position_embeddings = config.max_position_embeddings - rope_scaling = getattr(config, "rope_scaling", None) - if rope_scaling is not None and type(self) is LlamaRotaryEmbedding: - config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq( - config, - rope_scaling, - ) - self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base + # Kept so the v5 rope repair can rebuild the scaled inv_freq (#2405). + self._unsloth_rope_config = config # Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this self.current_rope_size = min(4 * 8192, self.max_position_embeddings) self.multi_gpu_cos_cached = [None] * DEVICE_COUNT self.multi_gpu_sin_cached = [None] * DEVICE_COUNT - if config_inv_freq is not None: - inv_freq = config_inv_freq # already scaled; skip subclass scaling - else: - # Normal Llama-3 RoPE - inv_freq = 1.0 / ( - self.base - ** ( - torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim - ) - ) - inv_freq = self._apply_inv_freq_scaling(inv_freq) + inv_freq = self._unsloth_recompute_inv_freq() self.register_buffer("inv_freq", inv_freq, persistent = False) # Build here to make `torch.jit.trace` work. @@ -1817,6 +1801,25 @@ class LlamaRotaryEmbedding(torch.nn.Module): """Override to apply custom inv_freq scaling (e.g., extended RoPE).""" return inv_freq + def _unsloth_recompute_inv_freq(self): + # Config scaling (llama3/yarn) first, else vanilla + subclass scaling. + # Shared by __init__ and the v5 rope repair so they cannot diverge. + config = getattr(self, "_unsloth_rope_config", None) + config_inv_freq = None + rope_scaling = getattr(config, "rope_scaling", None) if config is not None else None + if rope_scaling is not None and type(self) is LlamaRotaryEmbedding: + config_inv_freq, self.attention_scaling = _compute_config_rope_inv_freq( + config, + rope_scaling, + ) + if config_inv_freq is not None: + return config_inv_freq + inv_freq = 1.0 / ( + self.base + ** (torch.arange(0, self.dim, 2, dtype = torch.int64, device = "cpu").float() / self.dim) + ) + return self._apply_inv_freq_scaling(inv_freq) + def _apply_time_scaling(self, t): """Override to apply custom time scaling (e.g., linear scaling).""" return t diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index ba23197861..13342157b0 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -245,6 +245,7 @@ def _maybe_advise_fla_install(model_types): "transformers will use a slower pure PyTorch path." ) + def _fix_rope_inv_freq(model): """Fix inv_freq corruption caused by transformers v5 meta-device loading. @@ -268,14 +269,18 @@ def _fix_rope_inv_freq(model): and hasattr(module, "_apply_inv_freq_scaling") and hasattr(module, "multi_gpu_cos_cached") ): - inv_freq = 1.0 / ( - module.base - ** ( - torch.arange(0, module.dim, 2, dtype = torch.int64, device = "cpu").float() - / module.dim + if hasattr(module, "_unsloth_recompute_inv_freq"): + # Restore config scaling (llama3/yarn); unscaled here broke v5. + inv_freq = module._unsloth_recompute_inv_freq() + else: + inv_freq = 1.0 / ( + module.base + ** ( + torch.arange(0, module.dim, 2, dtype = torch.int64, device = "cpu").float() + / module.dim + ) ) - ) - inv_freq = module._apply_inv_freq_scaling(inv_freq) + inv_freq = module._apply_inv_freq_scaling(inv_freq) module.inv_freq = inv_freq for device_idx in range(len(module.multi_gpu_cos_cached)): if module.multi_gpu_cos_cached[device_idx] is not None: From cb9d90283000bab36f96a0926220e32c25874e21 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 09:13:29 -0700 Subject: [PATCH 018/113] Add the second blank line before _fix_rope_inv_freq (#6910) ruff-format requires two blank lines before a top-level function. loader.py carried only one, so the ruff-format-with-kwargs pre-commit hook reformats it and the run fails. This restores the expected spacing. From f0a5c52821e2f55b01c2a64d21e8fa9328d35250 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:06:06 -0700 Subject: [PATCH 019/113] studio: tool calling + healing parity for Llama-3, Mistral, Gemma 4 on safetensors + MLX (#5620) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * studio: tool calling for Llama-3, Mistral, Gemma 4 on safetensors + MLX (#5615) Adds tool calling for Llama-3, Mistral (pre-v11 + v11+ + [ARGS]), and Gemma 4 to the safetensors / transformers and MLX backends. Parser patched against llama.cpp / vLLM / SGLang per-family parsers and normalises to OpenAI shape. 96 targeted unit tests + cross-OS staging CI (ubuntu / macos-14 / windows) green on the multi-format probe. * studio: tool-call healing parity between safetensors / MLX and GGUF After the multi-format parser landed in #5615, the safetensors / MLX agentic loop and the GGUF loop still differed on healing behaviour. This commit closes the gaps in both directions so the two backends react the same way to identical model output. Changes: 1. core/inference/llama_cpp.py -- the GGUF BUFFERING state machine now wakes on every emission marker the shared parser knows. Was ("", " / Mistral [TOOL_CALLS] / Gemma 4 <|tool_call>). Stream cleanup is delegated to the same shared strip_tool_markup so leaked markup from any family is removed from assistant content. 2. core/inference/llama_cpp.py -- per-tool canonical heal key. When a tool arguments field is a bare string and JSON parsing fails, the GGUF path now heals to {"code": raw_args} for python, {"command": raw_args} for terminal, and {"query": raw_args} for everything else. Was hard-coded to {"query": raw_args}, which silently routed every python / terminal emission through web_search. Mirrors safetensors_agentic._CANONICAL_HEAL_ARG. 3. core/inference/safetensors_agentic.py -- re-prompt on plan- without-action. When the model emits a short forward-looking intent ("I'll search for that", "Let me check", "First, I will...") and no tool call, the loop nudges the model to act instead of silently returning a plan-only answer. Up to _MAX_REPROMPTS=3 (matches GGUF). The intent regex, character cap, and instruction text are byte-identical to the GGUF path. The buffer-end fall-through is unified so a buffered intent emission that never exits the BUFFERING state still triggers the re-prompt. 4. core/inference/safetensors_agentic.py -- extra iteration slots for re-prompts. The loop now budgets max_tool_iterations + _MAX_REPROMPTS + 1 total iterations and tracks the tool-call count separately, so a stalling model can be nudged 3x without eating the caller's tool-call budget. Mirrors the _extra slot reservation in the GGUF path. Tests (14 new safetensors-side units; 5 GGUF parity pins): TestLoopRePrompt -- intent-trigger, plain-answer, no-tools, cap-at-three, budget preserved, buffer-end intent. TestLoopCanonicalHealKey -- python / terminal / unknown. TestGGUFSafetensorsHealingParity -- shared markers used, shared strip used, canonical heal keys identical, intent regex matches same phrases, _MAX_REPROMPTS equal on both backends. All 110 targeted tests pass locally; the broader tool / inference / model-config / sandbox / anthropic / mlx suites stay green. Why this matters Without this parity, Llama-3.2 / Mistral / Gemma 4 emissions on Mac (MLX) and Linux-safetensors stop the agentic loop as soon as the model says "Let me...", because the GGUF re-prompt logic never existed on these backends. The two-marker GGUF BUFFERING tuple also let non-Qwen tool emissions stream out as plain prose when llama-server's structured channel did not pick them up. Both paths now drain the same way, heal the same way, and re-prompt the same way -- so a tool call that works on GGUF works identically on safetensors / MLX. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: fix tool-call parser bugs from gemini review on #5620 Three high-priority gemini findings on the tool-call parsing additions: 1. unicode_escape on UTF-8 bytes corrupts non-ASCII literals (e.g. ✨ becomes â\x9c¨). Replace with json.loads on a quoted string -- preserves emoji / CJK / RTL while still handling \n \t \uXXXX escapes. 2. Llama-3 sentinel stripping is order-dependent. A leading `<|eot_id|><|begin_of_text|>` left `<|begin_of_text|>` behind because the loop had already passed that sentinel. Loop until no sentinel matches at the start. 3. Mistral v11+ `[TOOL_CALLS] name { json }` regex uses non-greedy `\{.*?\}` which truncates at the first `}` of a nested JSON argument, leaking the tail (e.g. `}}`) into user-visible streamed text. Same problem for the v0.3 array pattern with nested brackets. Strip those with balanced brace/bracket scanning via a new `_strip_mistral_closed_calls` helper called from `strip_tool_markup`. Also fix the inference routes' parallel `_TOOL_XML_RE`: - Same nested-JSON truncation in the Mistral patterns; route the strip through the parser's balanced-scan helper via a thin `_strip_tool_xml` wrapper that all existing callers now use. - Llama-3 `<|python_tag|>[^\n<]*` stopped at any `<`, leaking the tail of any tool call whose argument contained a literal `<` (queries, code snippets). Relax to `[^\n]*` which keeps the strip confined to the actual end-of-line. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio/routes: make python_tag strip multi-line aware Earlier revisions of _TOOL_XML_RE in studio.backend.routes.inference oscillated between two bug shapes: 5615 r"<\|python_tag\|>[^\n<]*" -- stopped at any literal "<" so code='if x < 10: pass' leaked '< 10: pass)' to the user. 5620.1 r"<\|python_tag\|>[^\n]*" -- single-line only; the second line of python.call(code="a\nb") leaked. The full parser (_parse_llama3_python_tag) already handles both via balanced-brace scanning, so the parsing path was fine; the LEAK was in the streaming strip path that runs on every cumulative emission while content is still arriving. Switch to r"<\|python_tag\|>(?:[^<]|<(?!\|))*" so the strip consumes: * any character that is not a "<" (newlines, JSON, code, ...), * a "<" only when it is NOT followed by "|" (i.e. NOT a Llama-3 sentinel start like <|eot_id|>, <|eom_id|>, <|begin_of_text|>). This means: * code='if x < 10' stays inside the strip (5615 fix preserved), * multi-line code stays inside the strip (5620 round 2), * the strip terminates at the next Llama-3 sentinel so trailing assistant content survives. Tests: TestRoutesPythonTagStrip (8 cases) pytest test_safetensors_tool_loop.py test_safetensors_capability_advertise.py -> 118 passed in 1.81s (was 110). * [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: tighten verbose comments in tool-call parser sections Comments were narrating what the code already says. Cut historical "earlier revisions used X, then Y" narratives down to one-line WHY notes where the footgun still matters (canonical heal-key parity, balanced-brace vs non-greedy regex, ``(?:[^<]|<(?!\|))*`` over ``[^\n<]*``/``[^\n]*``). Drop section-header banners. No behaviour change. Re-ran: pytest studio/backend/tests/test_safetensors_tool_loop.py \ studio/backend/tests/test_safetensors_capability_advertise.py -q -> 118 passed. Regression replay (parser + _coerce_arguments on the 5 #5615 inputs) -> 21/21. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: parser robustness fixes for PR #5620 Three surgical extensions to the multi-format tool-call parser, each covering a real fine-tune / template emission shape that the current parser silently drops. No path narrows; all changes widen what is accepted. 1. `_parse_tool_call_json` now accepts both `arguments` and `parameters` keys. A Hermes / Qwen `{json}` wrapper around a Llama-3.2 fine-tune that emits the `parameters` key was extracting the tool name and silently discarding the args, producing a working-shaped call with an empty payload. The bare-JSON and python_tag paths already accepted both keys; this path now matches them. 2. `_TC_FUNC_START_RE`, `_TC_PARAM_START_RE`, and `_TC_PARAM_CLOSE_RE` now also match the attribute form `v` used by MiniCPM-5 and MiniMax-M2. Names land in either capture group, and `` is accepted as a short close. 3. `_parse_llama3_bare_json` sentinel-strip now consumes the role label inserted between `<|start_header_id|>` and `<|end_header_id|>` by Meta's official Llama-3.x chat template. Without this, every assistant turn re-fed through the template prefix `<|start_header_id|>assistant<|end_header_id|>\n\n{json}` parsed to zero calls, so any history-with-tool-call round-trip in production silently dropped. Tests in `studio/backend/tests/test_safetensors_tool_loop.py`: * `TestParserRobustness::test_tool_call_json_accepts_parameters_key` * `TestParserRobustness::test_function_xml_attribute_form` * `TestParserRobustness::test_function_xml_attribute_form_multi_param` * `TestParserRobustness::test_function_xml_legacy_equals_form_still_works` (regression guard for the existing `` syntax) * `TestParserRobustness::test_llama3_chat_template_round_trip` * `TestParserRobustness::test_llama3_round_trip_all_roles` * `TestParserRobustness::test_llama3_round_trip_with_eot_prefix` `pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py -q` goes from 118 to 125 passed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: terminate function-XML body at , not just `_parse_function_xml` was looking for `` (the Hermes wrapper) as the body terminator. When a model emits a standalone `v` followed by explanatory prose (which models routinely do), no `` is present, so the body extended to end-of-string and the trailing prose leaked into the LAST parameter value. Pre-existing on main (the legacy `` form had this bug too). Same affects PR #5620's new attribute-form `v` emission used by MiniCPM-5 / MiniMax-M2. Fix: `_TC_END_TAG_RE` now matches either `` OR ``. The existing `_TC_FUNC_CLOSE_RE` / `_TC_PARAM_CLOSE_RE` strips are unchanged. Multi-call inputs still bound each function at the next `` is preserved because the embedded close tag is ``, not ``). `pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py -q` goes from 125 to 127 passed. * Studio: tighten Llama-3.2 bare-JSON guard A fuzz pass on PR #5811 turned up that ``_parse_llama3_bare_json`` accepted ``parameters`` as a string, contradicting the docstring's "parameters or arguments is a dict" guard. Prose JSON like ``{"name":"foo","parameters":"a sentence"}`` would wrongly fire the parser, which the agentic loop would then heal into a real ``foo(query="a sentence")`` call. Same code lives on this branch, so the same fix applies here. Tightened guard: - ``parameters`` must be a dict (Llama-3 spec). - ``arguments`` may be a dict, or a JSON-encoded string that decodes to a dict (OpenAI shape, e.g. ``"arguments":"{\"q\":\"x\"}"``). Plain non-JSON strings or JSON-strings of lists / scalars / null no longer pass. Mirrors the fix landed in PR #5811 commit 615b8608. Adds the same 4 regression tests under TestParserMultiFormat. Existing test suite stays green: 127 -> 131 passing. * studio: fix safetensors tool-call parser gaps vs llama.cpp (Mistral CALL_ID / THINK, attribute-form signal) Three GGUF-parity fixes to the safetensors tool-call parser, each matching llama.cpp's reference behaviour: - Mistral Small 3.2 emits [TOOL_CALLS]name[CALL_ID][ARGS]{json}. The parser stopped after the name on seeing [CALL_ID] (neither [ARGS] nor {), dropping the call. Skip an optional [CALL_ID] segment in both the parse and strip paths. llama.cpp parses this (test-chat.cpp:4785). - Magistral wraps reasoning in [THINK]...[/THINK]. A [TOOL_CALLS] inside the reasoning was parsed as a real call, producing a phantom call. Strip a leading [THINK] block before scanning so only the post-reasoning call counts (test-chat.cpp:2285); a literal [THINK] inside a later argument is left intact. - The standalone MiniCPM-5 / MiniMax-M2 attribute form parsed correctly but was absent from TOOL_XML_SIGNALS and the markup strip patterns, so the streaming safety-net parse was gated off (dropping the call) and markup leaked into displayed text. Add the signal and broaden the strip regexes. Adds regression tests for all three. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: fire safetensors tool calls for the bare-JSON (Llama-3.2) form The agentic loop's streaming safety-net parse was gated on has_tool_signal(), which is False for the Llama-3.1 / 3.2 bare-JSON tool form {"name":..,"parameters":..} (no XML marker). Real tool calls were therefore dropped: the loop logged "model planned without calling tools", re-prompted three times, then gave up with zero tool calls, while GGUF's llama-server parses the same emission natively. Run parse_tool_calls_from_text() unconditionally in the safety net. The parser is strict (only fires on a valid tool-call shape) so plain answers are unaffected. Reproduced on a real unsloth/Llama-3.1-8B-Instruct run: the model emits {"name":"web_search","parameters":{...}} which now executes the tool instead of being re-prompted into a no-op. Adds a loop regression test for the bare-JSON form. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: complete strict-mode contract and fix parser import paths Address review findings on the multi-format tool-call parser: - Honor allow_incomplete=False in the remaining sub-parsers. The Llama-3 <|python_tag|>NAME.call(...) parser, the pre-v11 Mistral [TOOL_CALLS] array parser, and the Gemma 4 <|tool_call> parser ignored strict mode, so a truncated call (missing closing paren, ], or ) was still healed and executed with Auto-Heal disabled. Thread strictness through and reject the unclosed forms, matching the JSON and function-XML paths. - Drop the duplicate tool_call_parser import block in llama_cpp.py and the redundant un-aliased TOOL_XML_SIGNALS; only the _SHARED_TOOL_XML_SIGNALS alias is used as a value. - Import _strip_mistral_closed_calls from core.inference.tool_call_parser in routes/inference.py instead of studio.backend.core... The self-contained run.py launch mode only puts studio/backend on sys.path, so the absolute package path raised ModuleNotFoundError on the server-tool strip path. Add strict-mode regression tests for the truncated Llama-3 dot-call and the unclosed Mistral array. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: preserve XML param indentation and alias Mistral array parameters Two parser-correctness fixes found by auditing against the model chat templates and the SGLang / vLLM reference parsers: - Qwen3.5 XML parameter values lost their leading indentation. The chat template emits \nVALUE\n, but the parameter-start regex ate the wrapping newline AND the value's first-line indentation with a trailing \s*, then str.strip() removed the rest. Narrow the trailing class to horizontal whitespace only and trim exactly one wrapping newline (via _trim_param_value), preserving indentation in code/diff arguments. Matches SGLang's qwen3_coder detector. Applies to both _parse_function_xml (tool_call_parser.py) and the XML path in tool_healing.py. - Mistral pre-v11 array objects keyed on parameters dropped their payload. _consume_mistral_call read only the arguments key; alias parameters the same way the JSON/XML paths and SGLang's base detector do. Add regression tests for preserved multi-line indentation and the array parameters alias. * Studio: tighten tool-call parser comments Make the comments in the multi-format tool-call parser and its callers succinct: compress verbose docstrings/blocks to one or two lines, drop ones that restate the code, and trim the tiny balanced-scanner helpers. Correctness rationale and upstream provenance (SGLang/llama.cpp parity, the strict-mode / Auto-Heal contract, whitespace-preservation, and the Unicode / full-width-pipe notes) are kept in compact form. Comment-only: no code or behavior change (verified with comment_tools.py check --strip-docstrings; parser suite green). * Studio: make Llama-3 .call and Mistral-array healing parsing linear Two more O(n^2) ReDoS paths in the multi-format parser, both reachable from the agentic loop on a long truncated body with no length cap: - _LLAMA3_KV_RE.finditer over a .call(...) body retried at every offset of a long word run / unterminated quote (40K -> 14s). Replace with a hand-scan that reuses the same key/number/literal sub-regexes via anchored match and walks the string body by hand, so an unterminated quote is O(n). Verified byte-identical to the old regex over 200K fuzzed inputs. - _parse_mistral_array healing ran _balanced_brace_end from every { in the body (20K -> 17s). Walk top-level objects, advancing past each balanced {...}; this also drops the phantom call the old scan emitted from a nested argument object. Add adversarial-length linearity regressions plus positive .call kwargs and unclosed-array recovery coverage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: honor strict mode in safety-net, keep empty Gemma args, strip attribute-form function XML - safetensors safety-net parser now forwards allow_incomplete=auto_heal_tool_calls, matching the draining path, so a late incomplete tool call is not healed and executed when Auto-Heal is off. - Gemma empty bare value ({k:}) now serialises as "" instead of invalid {"k":}, which previously dropped the whole call. - Route _TOOL_XML_RE also strips the attribute form (MiniCPM-5 / MiniMax-M2) so it no longer leaks to the UI. * Studio: fix attribute-form function-XML literal close tag and zero-arg strict call Addresses Codex review of the attribute form in _parse_function_xml (MiniCPM-5 / MiniMax-M2): - End the call body at the LAST / within the call's window, so a literal close tag inside a code/search argument (e.g. print("")) is preserved instead of truncating the call. - Accept a closed call with no parameters as a valid zero-argument call in strict mode (the function close is already required), instead of rejecting it as a truncated call. - Tests for both, mirroring the legacy coverage. * Studio: fix tool-call parser/loop review findings on the multi-format path Address the live code-review findings on the safetensors/MLX + GGUF tool path: - routes: include the attribute form in the safetensors capability whitelist so MiniCPM-5 / MiniMax-M2 templates keep the tool pill (parser already handles the form; the post-filter wrongly suppressed it). - safetensors loop: build the plan-without-action re-prompt from the active tools instead of a hardcoded web_search/python string, and gate it on auto_heal_tool_calls, matching the GGUF loop. - safetensors loop: hold a leading bare-JSON object ({"name":..,"parameters":..}) during BUFFERING until it closes, then drain it as a tool call instead of streaming the raw JSON to clients. The DRAINING/STREAMING resolvers still recover a plain JSON answer, so this can never drop content. - parser: anchor the Llama-3 <|python_tag|>NAME.call(...) scan to the tag and chain ; -separated calls, so all semicolon-separated built-ins parse and a literal <|python_tag|>x.call(...) inside a JSON string argument no longer fires the wrong tool. - parser: consume the optional trailing after a named Mistral [TOOL_CALLS]name{json} call, mirroring the array shape. - GGUF streaming strip: use the shared parser patterns (which know [TOOL_CALLS] and <|python_tag|>) so a textual tool call entering DRAINING is stripped instead of leaking the marker to streaming clients. - routes: hoist the _strip_mistral_closed_calls import to module level. Adds regression tests covering each fix; existing parser suite stays green. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden multi-format tool-call detection from review findings Apply five targeted fixes from the review pass over the multi-format tool path: - routes: route display strip delegates to _strip_tool_xml so Mistral [TOOL_CALLS] blocks with nested JSON are removed from streamed display text, not just the XML forms. - tool_call_parser: skip function/parameter starts that fall inside an already-open parameter block (_inside_open_parameter) so nested example payloads are not mis-parsed as new calls; extract strip_llama3_leading_sentinels so the bare-JSON guard is shared. - safetensors_agentic: probe bare JSON through strip_llama3_leading_sentinels before the balanced-brace check so a leaked header sentinel does not defeat the guard. - tool_healing: allow dotted tool names in the Gemma wrapped start pattern. - llama_cpp (GGUF): buffer wrapper-less Llama-3.2 {"name":..} calls that carry no XML signal, drain a complete object silently and hold an incomplete one, and run the end-of-stream safety net unconditionally so markerless calls are detected and never leak the raw JSON (including truncated fragments). Adds regression tests for the GGUF bare-JSON streaming path and the Mistral display strip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: stop bare-JSON tool calls leaking at EOF, oversized, and into history The second review pass flagged that the Llama-3.2 bare-JSON tool-call handling still leaked raw JSON in several spots; ``strip_tool_markup`` only knows XML/bracket markup, so the bare-JSON form survived it. Fix them symmetrically across the safetensors and GGUF loops: - Safetensors stream-end resolver now routes a held bare-JSON fragment to DRAINING (mirroring GGUF) so a truncated ``{"name":..`` cut off by the end of the stream is dropped instead of flushed as assistant content. The 7/10 reviewer finding. - Both loops now drain (suppress) an oversized still-open bare-JSON call once it passes ``_MAX_BARE_JSON_BUFFER`` instead of streaming the raw prefix, gated on a ``"name"`` key so a giant plain JSON answer still streams; a complete oversized call still executes via the safety net. - Add a shared ``strip_leading_bare_json_call`` helper and apply it to the content kept for the assistant turn in both loops, so an executed bare-JSON call is not replayed as visible text or fed back as next-turn history. Plain JSON answers without a ``"name"`` key are untouched throughout. Adds regression tests for the EOF, oversized, and next-turn cases on both backends plus unit tests for the helper. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: bound the Llama-3 python_tag strip on real control sentinels The route display strip's <|python_tag|> arm ran to the next <| of any kind. A tool-call argument carrying a literal <|...|> token (for example <|cite|> inside a string value) truncated the strip early and leaked the call tail into the visible response. Narrow the stop condition to the genuine Llama control sentinels (eot_id, eom_id, python_tag, start/end_header_id, begin_of_text, finetune_right_pad_id) so embedded markup and JSON are consumed while real header/turn boundaries still bound the strip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: gate markerless bare JSON on enabled tools and close parser/strip asymmetries The Llama-3.2 custom_tools bare-JSON form has no marker, so any JSON object with a name key was read as a tool call. An ordinary JSON answer like {"name":"Alice","parameters":{"age":30}} was misclassified as a call to a disabled tool and dropped from the visible response. Gate the markerless form on the enabled tool names (threaded through parse_tool_calls_from_text and strip_leading_bare_json_call, supplied by both streaming loops): an object whose name is not an enabled tool is ordinary content. The marker-based forms keep their name-agnostic behaviour (an explicit signal is a real call attempt), and unrestricted mode stays ungated. Also fix two parser/strip asymmetries the parser already tolerated: - A literal inside a parameter value (print("")) truncated both the core and route strips at the first close, leaking the tail. Extend the strip to the call's real close (last before the next opener), mirroring the parser, without merging separate calls. - The single-object Mistral [TOOL_CALLS]{...} shape parsed but _strip_mistral_closed_calls left it, leaking the raw object into display. Strip the balanced object while keeping trailing prose, matching the array and name shapes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio tools: gate GGUF bare-JSON suppression on enabled tools and fix python-tag exponent parsing Pass-4 review follow-ups on the GGUF tool loop and Llama-3 parser: - The GGUF bare-JSON suppression sites still keyed off a raw "name" substring, so an ordinary JSON answer whose name is not an enabled tool was dropped when it was truncated, oversized, or reached the no-tool DRAINING fallback (the parser, helper, and safetensors paths were already gated). All three sites now use the shared enabled-name gate, and a held bare-JSON buffer that turns out not to be an enabled call is shown as the answer instead of dropped at stream end. - The Llama-3 python-tag numeric kwarg regex matched only the mantissa, so scientific notation was truncated to its leading digits (1e-3 parsed as 1) and a tool executed with the wrong value. The regex now accepts exponent and decimal forms, and the int/float classification keys off the exponent too. Adds regression tests for the truncated / oversized disabled-name JSON cases (and a counterpart that a truncated enabled call still does not leak) plus the scientific-notation kwargs. * Studio tools: gate safetensors bare-JSON drain, fix nested-name gate and function-XML strip Pass-4 review follow-ups on the shared parser / safetensors loop: - The safetensors oversized and end-of-stream bare-JSON drain branches keyed off a raw "name" substring, so a large or truncated ordinary JSON answer whose name is not an enabled tool was drained instead of streamed. Both now use the shared enabled-tool-name gate, matching the GGUF path. - strip_leading_bare_json_call matched the first "name" anywhere, so a plain JSON answer with a nested name equal to an enabled tool ({"result":{"name":"web_search"}}) was wrongly suppressed. It now extracts the TOP-LEVEL name only, walking past nested objects/arrays and keeping the text when a top-level value is truncated. - The function-XML display strip used a regex negative-lookahead that stopped at a literal opener inside a parameter value and then dropped the rest of the answer to EOF. A scan-based strip mirrors the parser (ignores openers inside an open via _inside_open_parameter) and closes each call at its real , so trailing assistant text after such a call survives. Adds regression tests for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tool parsing: 3.9 import safety, disabled-Auto-Heal contract, capability gate Round-2 review follow-ups on the multi-format tool-call parser: - tool_call_parser: add `from __future__ import annotations`. The module is dependency-light by design (external llama-server wrappers import it standalone) and the package targets python >=3.9, where its PEP 604 `int | None` return annotations would raise TypeError on import. - safetensors + GGUF drain fallback: gate the leading bare-JSON strip on auto_heal_tool_calls. With Auto-Heal off, a truncated enabled-name fragment that did not parse now stays visible, matching the XML strip in the same branch and the disabled-Auto-Heal contract. With Auto-Heal on it is still suppressed. - safetensors capability gate: match the bare-JSON `{"name":` template marker with a whitespace/escape-tolerant regex so a pretty-printed `{ "name" :` or JSON-escaped `{\"name\":` template is not mis-classified as tool-less. The parser already accepts that whitespace via raw_decode, so the gate must too. Regression tests added for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tool parsing: symmetric "function" bare-JSON alias and route strip parity Round-3 review follow-ups, all parser/strip symmetry fixes. - Bare-JSON "function" alias: the markerless parser accepts a call name via obj.get("name") or obj.get("function"), but the strip/gates only knew "name", so a {"function":} call executed while its raw JSON leaked. Teach _top_level_bare_json_name the alias (with "name" precedence and the same nested and truncated-name guards), and widen the guards in strip_leading_bare_json_call, the safetensors and GGUF _looks_like_enabled_bare_json gates, and the route capability marker regex. - Route display/history cleanup: strip a tail-only alias close (the parser accepts ...), and run the parser's guarded function-XML scan (_inside_open_parameter) before _TOOL_XML_RE so a literal nested inside an argument value does not truncate the strip and leak the tail. Regression tests added for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio tools: honor tool budget in GGUF loop and guard function-XML streaming strip Round 4 review fixes. Both are asymmetric-fix bugs where the final/steady path got a guard the analogous streaming/loop path did not. - GGUF tool-call budget: the safetensors loop counts real tool-call turns against max_tool_iterations (re-prompt stalls excepted), but the GGUF loop only bounded the turn count by the enlarged range (max_tool_iterations + _MAX_REPROMPTS). Since this PR raised _MAX_REPROMPTS from 1 to 3, a model that keeps making valid tool calls could run up to three extra tool rounds (with max_tool_iterations=1, four rounds instead of one). Add a _tool_iters_done counter that increments only when a tool actually executed in the turn, and stop once the caller's budget is spent so the post-loop final-answer nudge fires. A duplicate/disabled no-op turn is a correction turn (like a plan-without-action re-prompt) and does not consume budget, preserving the existing "already completed" re-prompt behavior. - Streaming display strip: the final strip runs the guarded _strip_function_xml_calls scanner (a literal inside a parameter value is data, not a nested call), but the GGUF and safetensors streaming strips still used only the open-ended regex arms. When a tool-call argument contained literal function markup, the regex tail ate everything to end-of-text and dropped the real trailing prose after the call's true . Run the guarded scanner (and the balanced Mistral strip) before the regex arms in both streaming paths so streaming and final display agree. Adds regression tests: GGUF valid tool calls respect max_tool_iterations, and the streaming strip keeps trailing prose after a function-XML call with a literal marker. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio tools: safetensors tool budget counts only executed turns (GGUF parity) Follow-up to the GGUF budget fix. The safetensors loop charged max_tool_iterations per non-re-prompt iteration (iteration + 1 - reprompt_count), so a duplicate/disabled no-op turn spent a budget slot even though no tool ran. With a small cap this dropped real work: for max_tool_iterations=2, a model that made a valid call, repeated it (an internal no-op correction turn), then made a distinct valid call executed only the first -- the third turn was sent with no tools and the distinct call was ignored. Track whether a turn actually executed a tool (set on record_result) and count only those turns against the cap, matching the GGUF loop. A duplicate/disabled no-op is a correction turn -- like a plan-without-action re-prompt -- and no longer consumes budget, so the model still gets its "already completed" nudge and another tool-enabled turn. Adds a regression test for the small-cap duplicate-then-distinct-call flow. * Studio: render the reasoning block for safetensors and MLX like GGUF enable_thinking chat templates (Qwen3/Qwen3.5/GLM) prefill an unclosed into the generation prompt, so the model emits only the closing then the answer. The safetensors/MLX chat stream emitted that as plain content, so the reasoning showed inline with no collapsible thinking block, while GGUF (which surfaces reasoning via reasoning_content) rendered one. This brings safetensors and MLX to parity. - _ResponsesReasoningExtractor gains a reasoning_prefilled mode that starts inside the reasoning block and splits on the first ; default False keeps GGUF and every existing caller byte-identical. It suppresses a stray re-emitted and holds partial markers back across chunk boundaries. - _sf_reasoning_prefill_mode gates the mode on reasoning being enabled for the request, an enable_thinking or enable_thinking_effort style, and the template actually using the standard / markers. Models with a bespoke reasoning channel (e.g. gemma's <|think|>/<|channel>) are excluded so their answer is never swallowed; gpt-oss (Harmony) and thinking-off requests are excluded too. - sf_tool_stream and stream_chunks (the latter also serves MLX) feed text through the extractor, emitting reasoning_content then content deltas, with a per-turn reset in the tool loop and a flush before each tool_start; only the visible delta reaches the monitor reply. The two non-streaming drains split reasoning_content the same way. - Tests: extractor prefilled mode (streaming and edge cases), the gate matrix including the gemma-style exclusion, and a route-replay of the tool-loop reasoning stream. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: don't force a tool re-prompt on a negated intent (safetensors parity) The safetensors _INTENT_SIGNAL claimed to mirror GGUF but was missing the negative lookahead, so a refusal like "I will not search the web for that" matched the "i will" intent and triggered the plan-without-action re-prompt (STOP... you MUST call a tool), overriding a valid no-tool answer. GGUF already excludes not/never. Add the same (?!\s+(?:not|never)\b) lookahead so both backends agree. Extends the intent parity test with negated refusals. * Studio: trim redundant comments (comment-only, AST-verified) * Studio: prevent Gemma tool-parser DoS on stray delimiters _gemma_parse_value returned the input index unchanged when text[i] was a stray delimiter (,}]), so the list and mapping caller loops that advance on the returned index spun forever at 100% CPU on malformed input such as [},]. Advance past the delimiter so parsing always terminates. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: strip Magistral [THINK] reasoning from final display/history strip_tool_markup removed [TOOL_CALLS] and markup but left a leading Magistral [THINK]...[/THINK] block intact, so its bracket-form reasoning (not the the reasoning channel renders) leaked into the safetensors display and conversation history while GGUF/llama.cpp routes it natively. Drop the leading reasoning block at end-of-turn (final=True) via the existing _strip_mistral_reasoning helper; streaming is untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honor reasoning_effort none in safetensors prefill; strip Magistral reasoning while streaming Two safetensors/MLX reasoning fixes surfaced in review: _sf_reasoning_prefill_mode only checked enable_thinking, so an enable_thinking_effort (GLM-5.2) request that disables thinking via reasoning_effort=none (without enable_thinking=False) still began in prefilled- mode. A plain answer with no was then swallowed whole into reasoning_content and the visible response came back empty. Thread reasoning_effort into the predicate and treat none as disabled, mirroring _request_reasoning_kwargs. strip_tool_markup_streaming stripped tool markup but not the leading Magistral [THINK]...[/THINK] bracket block, so the raw chain-of-thought leaked into the streamed safetensors content instead of the reasoning drawer (GGUF routes it natively). Apply _strip_mistral_reasoning first, matching the final strip; an unclosed [THINK] is held from the marker on so nothing flickers. * Mistral outer call wins over XML literals; align healer signals with its parser Two follow-ups on the shared-parser ordering after the healing-passthrough merge: - A well-formed [TOOL_CALLS] call whose JSON arguments quote tool XML parsed the literal instead of the outer call (executing the wrong tool). When the first XML signal sits inside a leading balanced Mistral body it is argument data, so the Mistral parser now runs first; an XML signal before the trigger keeps the normal order, so a [TOOL_CALLS] literal inside an XML call's arguments still stays data. - passthrough_healing buffered streams on the parser module's broadened signal list (now including <|python_tag|> and [TOOL_CALLS]) but promotes with core.tool_healing, which does not parse those forms: a streamed Mistral or Llama text call was held until finalization and flushed as prose. The healer keeps its own signal list limited to the formats it can promote, restoring immediate streaming for the rest. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: leading envelopes win over rehearsed literals - New _first_foreign_tool_signal shared by the leading-envelope guards adds <|python_tag|> to the protected signal set: the spelled-out literal inside a Mistral call's arguments (a query about Llama built-in tool syntax) executed the inner literal instead of the outer call. - New _xml_signal_inside_leading_bare_json guard, sibling of the Mistral one: a leading bare-JSON call whose string argument quotes tool XML (a code value citing ) had the literal promoted by the shared XML pass before the bare-JSON parser ran. - Magistral [THINK]...[/THINK] is dropped once at parse entry instead of only inside the Mistral parser, so a call rehearsed in the think block in a foreign format can no longer be promoted while the real call after the block is lost. Parse now agrees with the display strip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: a disabled leading bare-JSON object keeps its literals as data When the leading bare-JSON object is ordinary content (name not an enabled tool), the guard proved the first tool signal sits inside it, so falling through to the XML/python_tag passes promoted quoted string data as a real call. Drop the object and parse only the tail: a real call after the object still parses, nothing inside it can be promoted. * Address review: Mistral literals inside leading JSON, whitespace-tolerant wrapped Gemma opener - The leading bare-JSON guard now treats the [TOOL_CALLS] trigger as a foreign signal: the Mistral parser runs before the bare-JSON one, so a literal quoted inside the leading object's strings was promoted over the outer call (or over ordinary JSON content). - tool_healing's wrapped Gemma opener tolerates whitespace around call and the colon: sampling drift emits call: name{ and call : name{, and rejecting those lost the call entirely because no fallback re-parses the wrapped form. Strict mode still requires the closing tag. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: accept dotted Gemma argument keys in the key-quoting scanner The scanner quoted keys of [alnum_-] only, so a dotted key (user.name:...) was left unquoted, json.loads failed, and the whole wrapped call was lost (parse empty, strip wipes the markup). Dots now match the parser's own key/name charset. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: leading Mistral call owns the turn, dotted keys after bare values - A LEADING parseable [TOOL_CALLS] call now runs the Mistral parser first unconditionally: literal XML in trailing prose after the call was promoted by the earlier shared XML pass, executing the quoted example instead of the real leading call. XML leading keeps the normal order. - _GEMMA_NEXT_KEY_RE accepts dots so a dotted key after a bare value (query:foo,user.name:bob) ends the value at the comma instead of being swallowed into it, matching the round-earlier key-quoting charset. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: markup quoted inside a nameless leading JSON answer stays data The leading bare-JSON guard required a top-level name, so a structured JSON answer quoting tool markup in its strings (a response_format turn documenting a tool's syntax) had the literal promoted by the later passes. A nameless leading object that parses as real JSON now routes through the same decline-then-parse-the-tail path; non-JSON braced prose keeps the old behaviour, and a real call after the answer still parses. * Compress docstrings in the multi-format tool parser to their contract essence * verify_import_hoist: exempt __future__ imports and same-diff relocations Two false positives fired on this PR's refactor. A from __future__ import is a compiler directive whose name never appears as a runtime load, so HOISTED-IMPORT-UNUSED can never see it used, yet the file requires it for PEP 604 annotations on Python 3.9. TARGET-CHANGED flagged the deliberate move of the strip-pattern constants into core.inference.tool_call_parser as a silent re-point even though the old module-level target was removed and the new one added in the same diff. Both get narrow exemptions; a re-point to a pre-existing target is still caught, and the self-test negative controls all pass unchanged. * Leading bare-JSON calls own the turn; function calls end at the first balanced close The XML-signal guard for a leading bare-JSON call required the signal strictly inside the object, so a trailing XML example stole the turn from the leading call; it now applies the same inside-or-after rule as the Mistral guard. Function-XML calls also ended at the LAST close tag, which let prose after a closed call that mentions a literal close tag get swallowed into the final parameter value; calls now end at the first close tag that is not inside an open parameter, and the strip mirrors the same rule so parse and strip agree. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Attribute-form calls end at the first balanced close; bare-JSON strip requires the call shape The attribute form parser still kept the last close tag in the call window, folding prose after a closed call into the final parameter value. It now takes the first close not inside an open parameter, the same rule the equals form and the strip already use. The leading bare-JSON strip deleted any closed object whose top-level name matched an enabled tool, including plain JSON answers the parser correctly rejects as non-calls. The strip (and the drain gate that delegates to it) now requires the parser's exact call shape, so answers like {"name":"web_search","result":...} stream and display intact. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * False-alarm markers keep the answer; the bare-JSON strip consumes the whole chain The trailing strip arms dropped everything from a bare marker to EOF, so a normal answer that mentions [TOOL_CALLS] or another marker literally was truncated (or fully swallowed when it started with the literal) after the no-call drain fallback. Those arms now require a call-shaped lookahead or marker-at-EOF before dropping; truncated real calls still strip. Chained bare-JSON turns executed both calls but stripped only the first object, so the second call's raw JSON replayed into the next assistant history message alongside the structured tool_calls. The strip now consumes the entire chained run of call-shaped enabled objects while non-call answers, disabled names, and trailing prose stay intact. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Attribute-form containment, parameter-close-decides rule, preamble-tolerant Mistral guard, strict strip shape Four document-order and containment fixes. A leading attribute-form call now parses before the shared XML pass, so markup quoted in its parameter stays data. The open-parameter scan lets the parameter's own close tag decide, so any number of literal function closes inside one value stay data, restoring the pre-close-scan behavior for multi-close arguments. The leading-Mistral guard tolerates a visible preamble, with the leading-bare-JSON guard running first so a trigger quoted inside a leading JSON object stays data. The bare-JSON strip requires the parser's top-level name in every mode, so nested-name JSON answers survive name-agnostic stripping. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Let a leading <|python_tag|> call own the turn over quoted XML literals The leading-call ownership contract (a leading executable call owns the turn; foreign markup quoted in its string arguments or trailing prose stays data) was enforced for the bare-JSON, Mistral and attribute-form leading calls but not for the Llama-3 <|python_tag|> form. The shared tool_healing XML pass runs before _parse_llama3_python_tag and does not recognise <|python_tag|>, so a / / [TOOL_CALLS] literal quoted inside a <|python_tag|> .call(...) string argument (or its JSON parameters) was promoted and the wrong tool executed. Well-formed single-format examples: <|python_tag|>web_search.call(query="... ...") -> foo <|python_tag|>python.call(code="..") -> render_html both returned the phantom inner tool instead of the real leading call. Add a leading-<|python_tag|> guard mirroring the other leading-call guards: when the tag is the first tool signal, parse it before tool_healing so quoted foreign markup stays data. A foreign signal before the tag keeps normal document order. Added TestPythonTagOuterOverXmlLiteral (7 cases). * studio: tighten tool-calling comments to be shorter and clearer * studio: shorten tool-format comments in changed files --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: Daniel Han Co-authored-by: danielhanchen --- scripts/verify_import_hoist.py | 13 + studio/backend/core/inference/llama_cpp.py | 158 +- .../core/inference/passthrough_healing.py | 25 +- .../core/inference/safetensors_agentic.py | 198 ++- .../core/inference/tool_call_parser.py | 1434 ++++++++++++++++- studio/backend/core/tool_healing.py | 140 +- studio/backend/routes/inference.py | 237 ++- .../tests/test_gemma_tool_parse_edge_cases.py | 56 +- .../backend/tests/test_llama_cpp_tool_loop.py | 387 ++++- .../tests/test_responses_tool_passthrough.py | 118 ++ .../test_safetensors_capability_advertise.py | 223 ++- .../test_safetensors_reasoning_stream.py | 182 +++ .../tests/test_safetensors_tool_loop.py | 1224 ++++++++++++++ .../tests/test_tool_call_parser_strict.py | 866 ++++++++++ studio/backend/tests/test_tool_xml_strip.py | 155 +- 15 files changed, 5254 insertions(+), 162 deletions(-) create mode 100644 studio/backend/tests/test_safetensors_reasoning_stream.py diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b4c908b0cb..2d30265abe 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -564,6 +564,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] for n, tids in b["module_import_targets"].items(): if tids & after_used: continue # resolved -> fine + # `from __future__ import ...` is a compiler directive whose name is never loaded; skip it. + if all(t.startswith("from:__future__:") for t in tids): + continue newly_added = bool(tids - before_module_targets) was_used_before = bool(tids & before_used) if newly_added or was_used_before: @@ -588,9 +591,19 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] # package object and only *add* submodule attributes (e.g. adding # `import urllib.error` next to `import urllib.request`). Nothing the name # resolved to before is lost, so no reference is re-pointed -- skip it. + # + # A deliberate *relocation* is also benign: a name's import source moves A -> B in + # THIS diff (old `from A import x` removed, new `from B import x` added). Mirrors the + # TARGET-MISSING tolerance. Re-pointing to a pre-existing target (clash) is NOT exempted. + removed_module_targets = before_module_targets - after_module_targets for key, tafter in b["target_by_use"].items(): tbefore = a["target_by_use"].get(key) if tbefore and tbefore != tafter and (tbefore - tafter): + lost = tbefore - tafter + gained = tafter - tbefore + relocated = lost <= removed_module_targets and gained <= added_module_targets + if relocated: + continue findings.append( ( "BLOCKER", diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3ccfc5cdfe..5e67f6b484 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -38,9 +38,21 @@ from core.inference.llama_server_args import ( strip_shadowing_flags, strip_split_mode_only, ) -from core.tool_healing import ( + +# Share strip / signal constants with the multi-format parser so BUFFERING also +# catches Llama-3 / Mistral / Gemma 4. +from core.inference.tool_call_parser import ( _TOOL_ALL_PATS, - strip_tool_call_markup, + _balanced_brace_end, + _strip_function_xml_calls, + _strip_mistral_closed_calls, + TOOL_XML_SIGNALS as _SHARED_TOOL_XML_SIGNALS, + RAG_MAX_SEARCHES_PER_TURN, + RAG_SEARCH_CAP_NUDGE, + parse_tool_calls_from_text as _shared_parse_tool_calls_from_text, + strip_leading_bare_json_call, + strip_llama3_leading_sentinels, + strip_tool_markup as _shared_strip_tool_markup, ) from utils.native_path_leases import child_env_without_native_path_secret from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback @@ -48,12 +60,6 @@ from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs -from core.inference.tool_call_parser import ( - RAG_MAX_SEARCHES_PER_TURN, - RAG_SEARCH_CAP_NUDGE, - TOOL_XML_SIGNALS, - parse_tool_calls_from_text as _shared_parse_tool_calls_from_text, -) from core.inference.tool_loop_controller import ( ToolLoopController, tool_event_provenance, @@ -220,7 +226,7 @@ _INTENT_SIGNAL = re.compile( r"\b(?:now i|next i)\b" r")" ) -_MAX_REPROMPTS = 1 +_MAX_REPROMPTS = 3 # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. @@ -7881,12 +7887,17 @@ class LlamaCppBackend: # ── Message building (OpenAI format) ────────────────────────── @staticmethod - def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]: - """Thin wrapper around the shared parser in tool_call_parser - so safetensors and llama_cpp pick up the same fixes.""" + def _parse_tool_calls_from_text( + content: str, + *, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, + ) -> list[dict]: + """Wrapper around the shared parser; ``enabled_tool_names`` gates the markerless bare-JSON form.""" return _shared_parse_tool_calls_from_text( content, allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, ) @staticmethod @@ -8406,11 +8417,17 @@ class LlamaCppBackend: ) -> str: if not (auto_heal_tool_calls or force): return text - return strip_tool_call_markup(text, final = final) + return _shared_strip_tool_markup(text, final = final) def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text + # Shared patterns so a textual Mistral/Llama call entering DRAINING is stripped, not + # leaked. Mistral first; no final trim so incremental length comparisons hold. + text = _strip_mistral_closed_calls(text) + # Parser-accurate function-XML scan before the regex arms so a literal ```` + # in a value doesn't make the tail eat trailing prose after the real ````. + text = _strip_function_xml_calls(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text @@ -8456,6 +8473,13 @@ class LlamaCppBackend: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: + """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" + probe = strip_llama3_leading_sentinels(text.lstrip()) + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return False + return strip_leading_bare_json_call(probe, enabled_tool_names) != probe + tool_controller = ToolLoopController( tools = tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -8469,6 +8493,8 @@ class LlamaCppBackend: ) _MAX_BUFFER_CHARS = 32 + # Hold a leading ``{`` well past the 32-char XML cap until it balances (mirrors safetensors). + _MAX_BARE_JSON_BUFFER = 16384 _append_budget_exhausted_nudge = True # RAG: cap knowledge-base searches per assistant turn. The controller is # tool-agnostic, so this gate stays in the loop. @@ -8481,6 +8507,9 @@ class LlamaCppBackend: # "Hello!" won't match. Pattern compiled at module level # (_INTENT_SIGNAL). _reprompt_count = 0 + # Gates ``max_tool_iterations`` on real tool turns so reserved re-prompt slots don't + # extend the budget. Mirrors the safetensors guard. + _tool_iters_done = 0 _forced_tool_call_pending = False # Reserve extra iterations for re-prompts so they don't consume the @@ -8489,12 +8518,21 @@ class LlamaCppBackend: for iteration in range(max_tool_iterations + _extra): if cancel_event is not None and cancel_event.is_set(): return + # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget. + _turn_executed_real_tool = False active_tools = tool_controller.active_tools() if not active_tools: _append_budget_exhausted_nudge = False break - _tool_xml_signals = TOOL_XML_SIGNALS + # Gate the markerless bare-JSON form on enabled names so a JSON answer isn't misread as a call. + _enabled_tool_names = { + (tool.get("function") or {}).get("name") + for tool in active_tools + if (tool.get("function") or {}).get("name") + } + # Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows. + _tool_xml_signals = _SHARED_TOOL_XML_SIGNALS # Build payload -- stream: True so we detect tool signals # in the first 1-2 chunks without a non-streaming penalty. @@ -8777,7 +8815,36 @@ class LlamaCppBackend: is_prefix = True break - if is_match: + # Bare Llama-3.2 {"name":..} has no XML signal: hold an + # incomplete object, drain a complete one (mirrors safetensors). + _hold_buffer = False + # Whole buffer is the call (no visible prefix) -- drain silently. + _drain_silently = False + if not is_match and not is_prefix: + _bare = strip_llama3_leading_sentinels(stripped_buf) + if _bare.startswith("{"): + if _balanced_brace_end(_bare, 0) is None: + if len(stripped_buf) < _MAX_BARE_JSON_BUFFER: + _hold_buffer = True + elif _looks_like_enabled_bare_json( + _bare, _enabled_tool_names + ): + # Oversized still-open ENABLED-tool call: stop + # holding (memory bound) but DRAIN, not leak; + # a giant ordinary JSON answer still streams. + _drain_silently = True + elif self._parse_tool_calls_from_text( + content_buffer, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ): + _drain_silently = True + + if _drain_silently: + # No visible prefix -- the buffered text IS + # the call; drain without yielding it. + detect_state = _S_DRAINING + elif is_match: # Tool signal -- flush any visible # prefix before DRAINING so the # route sends it before tool_start. @@ -8794,7 +8861,9 @@ class LlamaCppBackend: "text": cleaned, } detect_state = _S_DRAINING - elif is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS: + elif _hold_buffer or ( + is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS + ): pass # keep buffering else: # Not a tool -- flush buffer @@ -8821,8 +8890,16 @@ class LlamaCppBackend: # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: stripped_buf = content_buffer.lstrip() + # A held bare-JSON fragment has no XML signal; route it to DRAINING. + _bare_eos = strip_llama3_leading_sentinels(stripped_buf) + # Gate on enabled names so a JSON answer isn't routed to DRAINING and dropped. + _is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json( + _bare_eos, _enabled_tool_names + ) if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals): detect_state = _S_DRAINING + elif _is_bare_tc: + detect_state = _S_DRAINING elif content_accum or reasoning_accum: detect_state = _S_STREAMING if content_buffer: @@ -8848,20 +8925,24 @@ class LlamaCppBackend: "text": cumulative_display, } else: + # No tool signal and no enabled bare-JSON call: a leading ``{`` is an ordinary + # JSON answer and must be shown; any other partial-markup prefix is dropped. + _held = strip_llama3_leading_sentinels(content_buffer.lstrip()) + if _held.startswith("{") and not _suppress_visible_output: + yield {"type": "content", "text": _held} return # ── STREAMING path: no tool call ── if detect_state == _S_STREAMING: - # Safety net: check for XML tool signals in content. The - # route layer resets prev_text on tool_start, so post-tool - # synthesis streams correctly even if content was emitted - # before the tool XML. - _safety_tc = None - if any(s in content_accum for s in _tool_xml_signals): - _safety_tc = self._parse_tool_calls_from_text( - content_accum, - allow_incomplete = auto_heal_tool_calls, - ) + # Safety net: re-parse the full content for tool calls. The route layer resets + # prev_text on tool_start, so post-tool synthesis streams correctly even if + # content was emitted before the tool XML. Unconditional (not gated on + # _tool_xml_signals): bare-JSON and Gemma wrapper-less calls carry no signal. + _safety_tc = self._parse_tool_calls_from_text( + content_accum, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ) if not _safety_tc: # ── Re-prompt on plan-without-action ── # If the model described its intent (forward-looking @@ -8978,10 +9059,13 @@ class LlamaCppBackend: for i in sorted(tool_calls_acc) if (tool_calls_acc[i].get("function", {}).get("name", "").strip()) ] or None - if not tool_calls and any(s in content_accum for s in _tool_xml_signals): + if not tool_calls: + # Unconditional re-parse: DRAINING means the buffer looked like a call, and + # bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on. tool_calls = self._parse_tool_calls_from_text( content_accum, allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) if tool_calls and not has_structured_tc: content_text = _strip_tool_markup( @@ -8989,6 +9073,11 @@ class LlamaCppBackend: final = True, force = True, ) + # ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call + # so the executed call isn't replayed as text or next-turn history. + content_text = strip_leading_bare_json_call( + content_text, _enabled_tool_names + ) if tool_calls: logger.info( f"Parsed {len(tool_calls)} tool call(s) from " @@ -9002,6 +9091,13 @@ class LlamaCppBackend: if content_accum: # Strip leaked tool-call XML before yielding. content_accum = _strip_tool_markup(content_accum, final = True) + # A truncated bare-JSON call has no XML to strip and didn't parse. With + # Auto-Heal on drop a leading ENABLED-tool fragment (plain JSON untouched); + # off keeps it visible per the strict contract. + if content_accum and active_tools and auto_heal_tool_calls: + content_accum = strip_leading_bare_json_call( + content_accum, _enabled_tool_names + ) if content_accum: yield {"type": "content", "text": content_accum} _meta = _build_metadata_event( @@ -9144,6 +9240,8 @@ class LlamaCppBackend: _kb_search_count += 1 completion = tool_controller.record_result(decision, result) resolved_provisional_tool_call_ids.add(decision.tool_call_id) + # A tool ran this turn, so it counts against the caller's budget. + _turn_executed_real_tool = True yield completion.tool_end_event() conversation.append(completion.tool_message()) @@ -9167,6 +9265,12 @@ class LlamaCppBackend: if tool_controller.force_final_answer or not tool_controller.active_tools(): _append_budget_exhausted_nudge = False break + # Count only real tool turns against the cap so reserved re-prompt slots can't + # become extra tool rounds; a no-op turn doesn't consume budget (GGUF parity). + if _turn_executed_real_tool: + _tool_iters_done += 1 + if _tool_iters_done >= max_tool_iterations: + break continue except httpx.ConnectError: diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index c73134b4a2..35855cc34d 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -29,10 +29,23 @@ import os from collections.abc import Mapping from typing import Any, Optional -from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal from core.inference.tool_loop_controller import coerce_tool_arguments from core.tool_healing import parse_tool_calls_from_text +# Only the formats this healer can promote. The parser's broader list adds Llama +# <|python_tag|> / Mistral [TOOL_CALLS], but buffering those here would flush a +# streamed call as prose, so keep a healer-aligned list. +_HEAL_SIGNALS = ( + "", + "<|tool_call>", + " bool: + return any(s in text for s in _HEAL_SIGNALS) + + # Read once at import (same convention as the other UNSLOTH_* switches). _HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1" # Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process @@ -44,7 +57,7 @@ def nudge_enabled(request_flag: Optional[bool]) -> bool: return _NUDGE_DEFAULT if request_flag is None else bool(request_flag) -_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS) +_MAX_SIGNAL_LEN = max(len(s) for s in _HEAL_SIGNALS) # A suspected-but-unclosed tool block larger than this is declared a false # alarm and flushed, bounding memory on a model rambling XML-lookalike text. _MAX_HOLD_CHARS = 64 * 1024 @@ -198,7 +211,7 @@ def heal_openai_message_events( if not isinstance(msg, dict) or msg.get("tool_calls"): return None content = msg.get("content") - if not isinstance(content, str) or not has_tool_signal(content): + if not isinstance(content, str) or not _has_heal_signal(content): return None parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True) tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None @@ -248,7 +261,7 @@ def heal_openai_message( def _earliest_signal(buffer: str) -> int: best = -1 - for signal in TOOL_XML_SIGNALS: + for signal in _HEAL_SIGNALS: index = buffer.find(signal) if index >= 0 and (best < 0 or index < best): best = index @@ -275,7 +288,7 @@ def _partial_signal_suffix(buffer: str) -> int: """Length of the longest buffer suffix that is a proper prefix of a signal.""" for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1): tail = buffer[-length:] - if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS): + if any(signal.startswith(tail) for signal in _HEAL_SIGNALS): return length return 0 @@ -508,7 +521,7 @@ def nudge_should_retry( if not message or message.get("tool_calls"): return False text = message.get("content") - if not isinstance(text, str) or not has_tool_signal(text): + if not isinstance(text, str) or not _has_heal_signal(text): return False return not _heal_would_promote(text, allowed_tools, tools) diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 0c96378d6c..b67c6cf7e7 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -22,11 +22,17 @@ from loggers import get_logger from core.inference.tool_call_parser import ( _TOOL_ALL_PATS, + _balanced_brace_end, + _strip_function_xml_calls, + _strip_mistral_closed_calls, + _strip_mistral_reasoning, BUDGET_EXHAUSTED_NUDGE, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, parse_tool_calls_from_text, + strip_leading_bare_json_call, + strip_llama3_leading_sentinels, strip_tool_markup, ) from core.inference.tool_loop_controller import ( @@ -50,6 +56,34 @@ logger = get_logger(__name__) # Buffer cap while disambiguating a possible tool-call prefix. _MAX_BUFFER_CHARS = 32 +# Memory bound for holding a leading bare-JSON object whose top-level "{" never balances. +_MAX_BARE_JSON_BUFFER = 16384 + +# Forward-looking intent ("I'll", "First,", "Step 1:") = planning; nudge a call. Negative +# lookahead drops negated forms ("I will not"). Mirrors GGUF. +_INTENT_SIGNAL = re.compile( + r"(?i)(" + r"\b(i['’](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" + r"|\b(?:first\b|step \d+:?|here['’]?s (?:my |the |a )?(?:plan|approach))" + r"|\b(?:now i|next i)\b" + r")" +) +_MAX_REPROMPTS = 3 +_REPROMPT_MAX_CHARS = 2000 +# Templated so the nudge names the caller's enabled tools. Mirrors GGUF tool_hint. +_REPROMPT_INSTRUCTION_TEMPLATE = ( + "STOP. Do NOT write code or explain. You MUST call a tool NOW. Call {tool_hint} immediately." +) + + +def _active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in active_tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + def strip_tool_markup_streaming( text: str, @@ -60,6 +94,12 @@ def strip_tool_markup_streaming( """Strip open-ended tool XML from display text without trimming whitespace.""" if not (auto_heal_tool_calls or tool_protocol_active): return text + # Mirror the final strip (no final trim): drop a leading Magistral ``[THINK]...[/THINK]`` + # block, then Mistral calls, then a parser-accurate function-XML scan before the regex + # arms. An unclosed ``[THINK]`` holds until ``[/THINK]`` so text stays monotonic. + text = _strip_mistral_reasoning(text) + text = _strip_mistral_closed_calls(text) + text = _strip_function_xml_calls(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text @@ -81,6 +121,14 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str: return status_for_tool(tool_name, arguments) +def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool: + """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" + probe = strip_llama3_leading_sentinels(text.lstrip()) + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return False + return strip_leading_bare_json_call(probe, enabled_tool_names) != probe + + _FUNCTION_SIGNAL_RE = re.compile(r"") _TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"') @@ -198,6 +246,10 @@ def run_safetensors_tool_loop( kb_search_count = 0 final_attempt_done = False next_call_id = 0 + reprompt_count = 0 + # Only turns that executed a tool count against ``max_tool_iterations``; a no-op or + # re-prompt turn must not consume budget (GGUF parity). + _executed_tool_iters = 0 def _tool_succeeded(tool_name: str) -> bool: key_prefix = f"{tool_name}:" @@ -215,9 +267,13 @@ def run_safetensors_tool_loop( _state_streaming = 1 _state_draining = 2 - for iteration in range(max_tool_iterations + 1): + # Reserve re-prompt slots so they don't eat the caller's tool budget. + _extra_iters = _MAX_REPROMPTS if max_tool_iterations > 0 else 0 + for iteration in range(max_tool_iterations + _extra_iters + 1): if cancel_event is not None and cancel_event.is_set(): return + # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget. + _turn_executed_real_tool = False if final_attempt_done: active_tools: list[dict] = [] @@ -229,6 +285,8 @@ def run_safetensors_tool_loop( tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools)) tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else () + # Gate the markerless bare-JSON form on enabled names so a JSON answer isn't misread as a call. + _enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools)) detect_state = _state_buffering content_buffer = "" @@ -367,6 +425,34 @@ def run_safetensors_tool_loop( is_prefix = True break + # Bare Llama-3.2 ``{"name":..,"parameters":..}`` carries no XML signal. Hold a leading + # ``{`` (after any sentinel) until it closes: drain if it parses as a call, else stream. + bare_probe = strip_llama3_leading_sentinels(stripped) + if ( + not is_match + and not is_prefix + and tool_protocol_active + and bare_probe.startswith("{") + ): + if _balanced_brace_end(bare_probe, 0) is None: + if len(stripped) < _MAX_BARE_JSON_BUFFER: + continue # object still open -- keep buffering + elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names): + # Oversized still-open ENABLED-tool call: stop holding (memory bound) but + # DRAIN, not leak; a giant ordinary JSON answer still streams. + detect_state = _state_draining + continue + elif parse_tool_calls_from_text( + content_buffer, + id_offset = next_call_id, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ): + # Closed object that parses as a bare-JSON call -- drain silently. + detect_state = _state_draining + continue + # Closed non-call object (or oversized non-call) -- stream as text. + if is_match: # Tool signal -- flush any visible prefix before DRAINING # so the route sends it before tool_start. @@ -419,44 +505,74 @@ def run_safetensors_tool_loop( if detect_state == _state_buffering: # Buffer never resolved -- tool XML or plain content? stripped = content_buffer.lstrip() + _bare_eos = strip_llama3_leading_sentinels(stripped) if ( stripped and tool_protocol_active and any(sig in stripped for sig in tool_xml_signals) ): detect_state = _state_draining + elif tool_protocol_active and _looks_like_enabled_bare_json( + _bare_eos, _enabled_tool_names + ): + # Held ENABLED-tool bare-JSON fragment has no XML signal; DRAIN it (a JSON answer + # falls through to the else and streams, GGUF parity). + detect_state = _state_draining else: + # Drain and fall through to STREAMING so the intent re-prompt + safety-net parser + # still fire on short emissions like "Let me search." that never exit BUFFERING. if content_buffer: cumulative_display += content_buffer - yield { - "type": "content", - "text": _strip_tool_markup_final( - cumulative_display, - auto_heal_tool_calls = auto_heal_tool_calls, - tool_protocol_active = False, - ), - } - yield {"type": "status", "text": ""} - return + cleaned = strip_tool_markup(cumulative_display, final = True) + if len(cleaned) > len(last_emitted): + last_emitted = cleaned + yield {"type": "content", "text": cleaned} + detect_state = _state_streaming if detect_state == _state_streaming: - # No tool detected mid-stream -- check for late tool XML. - safety_tc = None - saw_tool_signal = tool_protocol_active and any( - sig in content_accum for sig in tool_xml_signals + # Run the parser even with no XML signal (bare-JSON carries none); it's strict so + # plain answers stay untouched. Mirrors GGUF. + safety_tc = parse_tool_calls_from_text( + content_accum, + id_offset = next_call_id, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) - if saw_tool_signal: - safety_tc = parse_tool_calls_from_text( - content_accum, - id_offset = next_call_id, - allow_incomplete = auto_heal_tool_calls, - ) if not safety_tc: - # Final answer: if a literal tool marker in prose was stripped - # during streaming but did not parse as a real call, restore the - # raw cumulative text for core callers. Route-level cleanup can - # still apply the Auto-Heal display policy. - if saw_tool_signal and content_accum: + # Re-prompt only when the model planned without acting (intent signal); + # "4" / "Hello!" never trigger. Mirrors GGUF. + _stripped = content_accum.strip() + if ( + tools + and auto_heal_tool_calls + and reprompt_count < _MAX_REPROMPTS + and 0 < len(_stripped) < _REPROMPT_MAX_CHARS + and _INTENT_SIGNAL.search(_stripped) + and not final_attempt_done + ): + reprompt_count += 1 + logger.info( + "Safetensors re-prompt %d/%d: model planned without " + "calling tools (%d chars)", + reprompt_count, + _MAX_REPROMPTS, + len(_stripped), + ) + tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" + conversation.append({"role": "assistant", "content": _stripped}) + conversation.append( + { + "role": "user", + "content": _REPROMPT_INSTRUCTION_TEMPLATE.format(tool_hint = tool_hint), + } + ) + yield {"type": "status", "text": ""} + continue + + # Final answer. If a literal tool marker in prose was buffered but never + # parsed as a call, restore the raw text so the prose surfaces; route + # cleanup still applies the Auto-Heal policy. + if content_accum and any(sig in content_accum for sig in tool_xml_signals): yield {"type": "content", "text": content_accum} yield {"type": "status", "text": ""} return @@ -476,20 +592,24 @@ def run_safetensors_tool_loop( content_accum, id_offset = next_call_id, allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) if not tool_calls: # Parser found nothing. Auto-Heal-enabled display cleanup # strips unparseable tool XML; disabled Auto-Heal preserves # the raw text so literal/malformed markup stays visible. if content_accum: - yield { - "type": "content", - "text": _strip_tool_markup_final( - content_accum, - auto_heal_tool_calls = auto_heal_tool_calls, - tool_protocol_active = False, - ), - } + _drain_text = _strip_tool_markup_final( + content_accum, + auto_heal_tool_calls = auto_heal_tool_calls, + tool_protocol_active = False, + ) + # Drained bare-JSON call that didn't parse: with Auto-Heal on drop the fragment + # (plain JSON untouched); off keeps it visible per the strict contract. + if tool_protocol_active and auto_heal_tool_calls: + _drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names) + if _drain_text: + yield {"type": "content", "text": _drain_text} if provisional_render_html_started and not provisional_resolved: provisional_resolved = True yield { @@ -509,6 +629,9 @@ def run_safetensors_tool_loop( if tool_calls: next_call_id += len(tool_calls) + # Strip a leading bare-JSON call so it isn't replayed as text or next-turn history + # (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers. + content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names) if final_attempt_done: # Final-answer turn re-called a tool -- stop the loop. @@ -634,6 +757,8 @@ def run_safetensors_tool_loop( completion = tool_controller.record_result(decision, result) if provisional_match: provisional_resolved = True + # A tool ran this turn, so it counts against the caller's budget. + _turn_executed_real_tool = True yield completion.tool_end_event() conversation.append(completion.tool_message()) @@ -646,7 +771,10 @@ def run_safetensors_tool_loop( if not unrestricted_tools and not tool_controller.active_tools(): final_attempt_done = True continue - if iteration + 1 >= max_tool_iterations and not final_attempt_done: + # Count only real tool turns against the cap so a no-op turn doesn't consume budget (GGUF parity). + if _turn_executed_real_tool: + _executed_tool_iters += 1 + if _executed_tool_iters >= max_tool_iterations and not final_attempt_done: # Budget exhausted; nudge a final plain answer. final_attempt_done = True conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE}) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index ca3d1e4cbc..c31f4b272e 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -2,39 +2,74 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Backend-neutral tool-call XML parser shared by GGUF and safetensors. -Tolerates missing closing tags in either ``{json}`` -or ``v...`` shape. +Backend-neutral tool-call parser shared by GGUF, safetensors, and MLX, so the +safetensors + MLX agentic loop sees the same call shape llama-server gives GGUF: + + - ``{json}`` (Qwen / Hermes) + - ``v`` (Qwen3.5 xml) + - ``<|python_tag|>NAME.call(k="v", ...)`` (Llama-3 built-in tools) + - ``<|python_tag|>{"name":..., "parameters":...}`` (Llama-3 custom) + - ``{"name":..., "parameters":...}`` (Llama-3.2 bare JSON) + - ``[TOOL_CALLS] [{...}, ...]`` (Mistral v0.3 / Nemo / Small) + - ``[TOOL_CALLS]name{json}`` (Mistral v11+ / Magistral) + - ``[TOOL_CALLS]name[ARGS]{json}`` (Ministral / Mistral Large 3) + - ``<|tool_call>call:NAME{k:<|"|>v<|"|>}`` (Gemma 4) + +Missing closing tags / brackets are tolerated: models often truncate mid-stream. """ +# Keeps PEP 604 `X | None` lazy for python 3.9 (imported standalone by external servers). +from __future__ import annotations + +import json +import re +from typing import Any, Optional + +# Shared parser handles Qwen/Hermes, Qwen3.5 XML, Gemma 4; this module adds Llama-3, Mistral, bare JSON. from core import tool_healing as _tool_healing -_TOOL_ALL_PATS = _tool_healing._TOOL_ALL_PATS +# Flip the streaming buffer STREAMING->DRAINING so partial markup never leaks. +TOOL_XML_SIGNALS = ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", +) -def parse_tool_calls_from_text( - content: str, - *, - id_offset: int = 0, - allow_incomplete: bool = True, -) -> list[dict]: - return _tool_healing.parse_tool_calls_from_text( - content, - id_offset = id_offset, - allow_incomplete = allow_incomplete, - ) +# Closed pairs only (mid-stream); _TOOL_ALL_PATS eats unclosed tails at end-of-turn. +_TOOL_CLOSED_PATS = [ + re.compile(r".*?", re.DOTALL), + # Match to the real ```` (lookahead, not greedy ``.*``) so a literal + # ```` in a value doesn't truncate and each call stays separate. + re.compile( + r'' + r'(?:(?!).)*' + r"", + re.DOTALL, + ), + re.compile(r"<\|tool_call>.*?", re.DOTALL), +] +_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ + re.compile(r".*$", re.DOTALL), + re.compile(r'.*$', re.DOTALL), + # Bare-word markers drop a trailing truncated call only when the next chars look like + # a call start, so prose mentioning the marker is kept; a marker at end-of-text drops. + re.compile(r"<\|tool_call>(?=\s*call\s*:|\s*$).*$", re.DOTALL), + re.compile( + r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*[\[{])|\s*$).*$", + re.DOTALL, + ), + re.compile( + r"<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\()|\s*$).*$", + re.DOTALL, + ), +] -def strip_tool_markup(text: str, *, final: bool = False) -> str: - return _tool_healing.strip_tool_call_markup(text, final = final) - - -# Prefixes the streaming buffer watches for to gate in-progress text. -TOOL_XML_SIGNALS = ("", "<|tool_call>", "{json}``. +_TC_JSON_START_RE = re.compile(r"\s*\{") +# Qwen3.5 ```` plus attribute form ```` (MiniCPM-5, +# MiniMax-M2); name in group(1) or group(2). +_TC_FUNC_START_RE = re.compile(r'\s*') +# Body ends at ```` or ```` so trailing prose stays out of args. +_TC_END_TAG_RE = re.compile(r"") +_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") +# Horizontal whitespace only so the wrapping newline + indent survive (``_trim_param_value`` +# trims one newline), preserving code indent. +_TC_PARAM_START_RE = re.compile( + r'<(?:parameter|param)(?:=([\w\.\-]+)|\s+name="([\w\.\-]+)")>[^\S\n]*' +) +_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") + +# Llama-3 ``<|python_tag|>NAME.call(...)``. +_LLAMA3_PYTHON_TAG = "<|python_tag|>" +_LLAMA3_PY_CALL_RE = re.compile( + r"<\|python_tag\|>\s*([\w\.\-]+)\s*\.\s*call\s*\(", +) +# Anchored at the char after ``<|python_tag|>`` plus the ``; NAME.call(`` chain sep, so +# a ``.call(`` inside JSON args is ignored. +_LLAMA3_PY_CALL_HEAD_RE = re.compile(r"\s*([\w\.\-]+)\s*\.\s*call\s*\(") +_LLAMA3_CALL_CHAIN_RE = re.compile(r"\s*;\s*([\w\.\-]+)\s*\.\s*call\s*\(") +# ``.call(k=v)`` kwarg tokens, hand-scanned below (not finditer) to stay linear on a +# truncated body (ReDoS). +_LLAMA3_KEY_RE = re.compile(r"\w+") +_LLAMA3_WS_RE = re.compile(r"\s*") +# ints, decimals, sci notation; trailing ``(?![\w.])`` stops ``1.2.3`` truncating to ``1.2``. +_LLAMA3_NUM_RE = re.compile(r"-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?(?![\w.])") +_LLAMA3_LIT_RE = re.compile(r"true|false|null") + +# Mistral ``[TOOL_CALLS]`` trigger. v11+ chains ``name{json}`` (Magistral) or +# ``name[ARGS]{json}`` (Ministral / Large 3). +_MISTRAL_TRIGGER = "[TOOL_CALLS]" +_MISTRAL_ARGS_MARKER = "[ARGS]" +# Mistral Small 3.2 emits ``name[CALL_ID][ARGS]{json}`` (absent on Ministral / Magistral). +_MISTRAL_CALL_ID_MARKER = "[CALL_ID]" +# Magistral wraps reasoning in ``[THINK]...[/THINK]``; a ``[TOOL_CALLS]`` inside is not a real call. +_MISTRAL_THINK_OPEN = "[THINK]" +_MISTRAL_THINK_CLOSE = "[/THINK]" +_MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") + +# Gemma 4: ``<|tool_call>call:NAME{...}``, ``<|"|>`` wraps strings. +_GEMMA_TC_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w\.\-]+)\s*\{") +_GEMMA_STR_BEGIN = '<|"|>' +_GEMMA_STR_END = '<|"|>' +_GEMMA_TC_END = "" + + +def _balanced_bracket_end(text: str, start: int) -> int | None: + """Index of the ``]`` matching ``[`` at ``text[start]`` (ignores brackets in JSON strings).""" + if start >= len(text) or text[start] != "[": + return None + depth = 0 + in_string = False + esc = False + i = start + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _skip_mistral_call_id(text: str, pos: int) -> int: + """Skip an optional ``[CALL_ID]`` (Mistral Small 3.2); return the next token pos.""" + n = len(text) + i = pos + while i < n and text[i] in " \t\n\r": + i += 1 + if not text.startswith(_MISTRAL_CALL_ID_MARKER, i): + return pos + i += len(_MISTRAL_CALL_ID_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + # The id is a short opaque token; stop at whitespace or the next marker. + while i < n and text[i] not in " \t\n\r[{": + i += 1 + while i < n and text[i] in " \t\n\r": + i += 1 + return i + + +def _strip_mistral_reasoning(content: str) -> str: + """Drop a leading Magistral ``[THINK]`` block so rehearsed calls inside reasoning are not promoted; unclosed drops to EOF.""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if not content.startswith(_MISTRAL_THINK_OPEN, i): + return content + close = content.find(_MISTRAL_THINK_CLOSE, i + len(_MISTRAL_THINK_OPEN)) + if close == -1: + return content[:i] + return content[:i] + content[close + len(_MISTRAL_THINK_CLOSE) :] + + +def _strip_mistral_closed_calls(text: str) -> str: + """Strip cleanly-closed ``[TOOL_CALLS]`` blocks via balanced scanning (a non-greedy regex would truncate nested JSON); unclosed runs wait for ``final=True``.""" + n = len(text) + out = [] + cursor = 0 + while cursor < n: + idx = text.find(_MISTRAL_TRIGGER, cursor) + if idx == -1: + out.append(text[cursor:]) + break + out.append(text[cursor:idx]) + body_start = idx + len(_MISTRAL_TRIGGER) + i = body_start + while i < n and text[i] in " \t\n\r": + i += 1 + # Array shape: ``[TOOL_CALLS] [...]``. + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + if end is None: + # Truncated; let caller buffer / final-strip. + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Single-object shape ``[TOOL_CALLS] { json }``: the parser accepts it, so strip it too. + if i < n and text[i] == "{": + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Named shape: ``[TOOL_CALLS] name [ARGS]? { json }``. + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + out.append(text[idx:body_start]) + cursor = body_start + continue + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + i = _skip_mistral_call_id(text, i) + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + out.append(text[idx:i]) + cursor = i + continue + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + # Consume the optional EOS marker so ``...{json}`` doesn't leave ```` as content. + if text.startswith("", cursor): + cursor += len("") + return "".join(out) + + +_FUNC_CLOSE_TAG_RE = re.compile(r"") + + +def _strip_function_xml_calls(text: str, *, final: bool) -> str: + """Strip ```` calls by mirroring the parser: an opener inside an open ```` is data and each call closes at its first ```` that is not parameter data; ``final`` drops a trailing unclosed call.""" + starts = [ + m for m in _TC_FUNC_START_RE.finditer(text) if not _inside_open_parameter(text, m.start()) + ] + if not starts: + return text + out: list[str] = [] + pos = 0 + for idx, m in enumerate(starts): + if m.start() < pos: + continue # opener already inside a previously consumed call span + out.append(text[pos : m.start()]) + next_start = starts[idx + 1].start() if idx + 1 < len(starts) else len(text) + close = None + for cm in _FUNC_CLOSE_TAG_RE.finditer(text, m.end(), next_start): + if not _inside_open_parameter(text, cm.start()): + close = cm # first close that is not parameter data = the real close + break + if close is not None: + pos = close.end() + elif final: + pos = len(text) # trailing unclosed call -- drop to EOF + else: + out.append(text[m.start() :]) # keep the unclosed call buffered mid-stream + pos = len(text) + break + out.append(text[pos:]) + return "".join(out) + + +def strip_tool_markup(text: str, *, final: bool = False) -> str: + """Strip tool-call markup; ``final=True`` also drops trailing unclosed runs and trims.""" + if final: + # End-of-turn only: drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket form, + # not the ```` reasoning channel) so raw reasoning doesn't leak into display/history. + text = _strip_mistral_reasoning(text) + text = _strip_mistral_closed_calls(text) + # Scan-strip the function-XML form first (parser-accurate: a literal ```` in + # a value is data, not a call); the regex arms below cover the other formats. + text = _strip_function_xml_calls(text, final = final) + pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS + for pat in pats: + text = pat.sub("", text) + return text.strip() if final else text + + def has_tool_signal(text: str) -> bool: - """Return True if ``text`` contains any tool-call XML signal.""" return any(s in text for s in TOOL_XML_SIGNALS) + + +def _mistral_region_end(text: str, idx: int) -> int | None: + """Exclusive end of the balanced ``[TOOL_CALLS]`` call at ``idx``, or ``None`` when truncated (array, object, and named forms).""" + n = len(text) + i = idx + len(_MISTRAL_TRIGGER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + return None if end is None else end + 1 + if i < n and text[i] == "{": + end = _balanced_brace_end(text, i) + return None if end is None else end + 1 + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + return None + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + i = _skip_mistral_call_id(text, i) + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + return None + end = _balanced_brace_end(text, i) + return None if end is None else end + 1 + + +def _xml_signal_inside_leading_mistral(content: str) -> bool: + """True when a parseable Mistral call is the first tool emission in document order: it owns the turn, so later XML (quoted in its arguments or in trailing prose) is not promoted over it. A signal BEFORE the trigger keeps normal order.""" + trig = content.find(_MISTRAL_TRIGGER) + if trig < 0: + return False + first_xml = _first_foreign_tool_signal(content) + if first_xml is not None and first_xml < trig: + return False + # Only plain prose precedes the trigger (preamble-tolerant); prose merely mentioning + # the marker has no parseable region and keeps the normal order. + return _mistral_region_end(content, trig) is not None + + +_ATTR_FUNC_OPEN_RE = re.compile(r' int | None: + """Offset of the first signal a non-envelope parser would fire on (XML forms plus the Llama-3 ``<|python_tag|>`` marker).""" + first = None + for sig in ("", "<|tool_call>", ""): + p = content.find(sig) + if p >= 0 and (first is None or p < first): + first = p + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None and (first is None or attr.start() < first): + first = attr.start() + return first + + +def _xml_signal_inside_leading_bare_json(content: str) -> bool: + """True when the first foreign signal sits inside a LEADING bare-JSON call's balanced body: quoted argument data, so the bare-JSON parser takes the outer call first.""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if i >= n or content[i] != "{": + return False + end = _balanced_brace_end(content, i) + if end is None: + return False + if _top_level_bare_json_name(content[i : end + 1]) is None: + # Not a call object, but a nameless object that parses as real JSON is an envelope + # too (markup in its strings is data); non-JSON braced prose keeps the old behaviour. + try: + json.loads(content[i : end + 1]) + except ValueError: + return False + first_xml = _first_foreign_tool_signal(content) + # The Mistral trigger is foreign to a JSON envelope too, so fold it into first_xml. + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first_xml is None or trig < first_xml): + first_xml = trig + # Inside the balanced body the signal is quoted argument data, so the leading call owns + # the turn; a non-call object takes the decline path (dropped, only the tail parsed). + return first_xml is not None and i < first_xml + + +def parse_tool_calls_from_text( + content: str, + *, + id_offset: int = 0, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Return OpenAI-format tool calls, first-match wins. ``allow_incomplete`` heals truncated calls (``False`` = strict closed-only); ``enabled_tool_names`` gates the markerless bare-JSON form.""" + # Drop Magistral reasoning before any dispatch so a rehearsed call inside + # [THINK]...[/THINK] is not promoted; keeps the parse path aligned with the display strip. + content = _strip_mistral_reasoning(content) + + # A leading bare-JSON value is decided FIRST so markup quoted in its arguments stays + # data. Must precede the Mistral guard, whose preamble tolerance would else claim a + # trigger quoted inside the leading object. + if _xml_signal_inside_leading_bare_json(content): + calls = _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + if calls: + return calls + # Disabled/example name: the leading object is ordinary content. Drop it and parse + # only the tail -- a real call after it still parses, nothing inside it is promoted. + i = 0 + while i < len(content) and content[i] in " \t\n\r": + i += 1 + end = _balanced_brace_end(content, i) # guard guarantees a balanced object + return parse_tool_calls_from_text( + content[end + 1 :], + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + + # A [TOOL_CALLS] call that is the first tool emission owns the turn: XML quoted in its + # arguments or in trailing prose is not promoted, and a plain-prose preface keeps it. + if _xml_signal_inside_leading_mistral(content): + calls = _parse_mistral_tool_calls( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # A leading MiniCPM/MiniMax ```` call owns the turn: tool_healing + # does not know the wrapper, so gate it here. A signal before the opener keeps normal order. + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None: + first_other = None + for sig in ( + "", + "<|tool_call>", + "", + _MISTRAL_TRIGGER, + ): + p = content.find(sig) + if p >= 0 and (first_other is None or p < first_other): + first_other = p + if first_other is None or attr.start() < first_other: + calls = _parse_function_xml( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # A leading Llama-3 ``<|python_tag|>`` call owns the turn like the others: markup quoted + # in a ``.call(...)`` argument is not promoted. tool_healing does not know the tag, so + # gate it here. A foreign signal before the tag keeps normal order. + py_tag = content.find(_LLAMA3_PYTHON_TAG) + if py_tag >= 0: + first_other = None + for sig in ("", "<|tool_call>", "= 0 and (first_other is None or p < first_other): + first_other = p + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None and (first_other is None or attr.start() < first_other): + first_other = attr.start() + if first_other is None or py_tag < first_other: + calls = _parse_llama3_python_tag( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # Qwen/Hermes, Qwen3.5 XML, and Gemma 4 use the shared tool_healing parser (the + # strict/Auto-Heal + nested-marker + ``<|"|>`` handling GGUF relies on). + calls = _tool_healing.parse_tool_calls_from_text( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + ) + if calls: + return calls + + # Formats tool_healing does not cover: ```` (MiniCPM-5 / MiniMax-M2), + # Llama-3 and Mistral. Run only after tool_healing found nothing, so a strict-rejected + # call is never re-healed here. + for parser in ( + _parse_function_xml, # attribute form + _parse_llama3_python_tag, # Llama-3 <|python_tag|> + _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] + ): + calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # Llama-3.2 bare ``{"name":..., "parameters":...}``. Strict (starts with ``{`` + # and parses to the right shape) so plain prose stays untouched. + return _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + + +def _parse_tool_call_json( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + out: list[dict] = [] + for m in _TC_JSON_START_RE.finditer(content): + brace_start = m.end() - 1 + end = _balanced_brace_end(content, brace_start) + if end is None: + continue + # Strict mode: a balanced body that never closed its ```` is truncated + # (trailing prose after the close is still tolerated). + if not allow_incomplete and not content[end + 1 :].lstrip().startswith(""): + continue + try: + obj = json.loads(content[brace_start : end + 1]) + except (json.JSONDecodeError, ValueError): + continue + name = obj.get("name", "") + # Accept both ``arguments`` (Hermes/Qwen) and ``parameters`` (Llama-3 drift). + args = obj.get("arguments") + if args is None: + args = obj.get("parameters", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if not name: + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + return out + + +def _trim_param_value(val: str) -> str: + """Trim only the template's wrapping newline around an XML parameter value; ``str.strip()`` destroyed code/diff indentation.""" + if val.startswith("\n"): + val = val[1:] + if val.endswith("\n"): + val = val[:-1] + return val + + +def _inside_open_parameter(text: str, pos: int) -> bool: + """True if ``pos`` is inside an unclosed ```` block, i.e. the opener at ``pos`` is literal argument data, not a nested call.""" + last_param_open = -1 + for m in _TC_PARAM_START_RE.finditer(text, 0, pos): + last_param_open = m.start() + if last_param_open < 0: + return False + # The parameter's OWN close tag decides: if it closes after ``pos`` the position is + # argument data (even across literal ````); an unclosed one falls back to func close. + own_closes = [ + c + for c in ( + text.find("", last_param_open), + text.find("", last_param_open), + ) + if c >= 0 + ] + if own_closes: + return min(own_closes) > pos + func_closes = [ + c + for c in ( + text.find("", last_param_open), + text.find("", last_param_open), + ) + if c >= 0 + ] + return not func_closes or pos < min(func_closes) + + +def _parse_function_xml( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + out: list[dict] = [] + # Skip ```` openers that are literals inside an open parameter value, + # else the nested marker becomes a second call and truncates the real argument. + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + ] + for idx, fm in enumerate(func_starts): + # group(1) is ````, group(2) is ````. + func_name = fm.group(1) or fm.group(2) + body_start = fm.end() + next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + # The call ends at the FIRST / not inside an open parameter: + # a literal close in an argument is skipped as data, prose after the real close is not + # folded in (mirrors _strip_function_xml_calls). + close_match = None + for cm in _TC_END_TAG_RE.finditer(content, body_start, next_func): + if not _inside_open_parameter(content, cm.start()): + close_match = cm + break + has_close = close_match is not None + if has_close: + body_end = close_match.start() + else: + body_end = min(len(content), next_func) + # Strict mode: a call that never reached its close is truncated; do not heal it. + if not allow_incomplete and not has_close: + continue + body = _TC_FUNC_CLOSE_RE.sub("", content[body_start:body_end]) + + args: dict = {} + param_unclosed = False + # Same nested-literal guard: a ```` opener inside an open value is literal text. + param_starts = [ + pm + for pm in _TC_PARAM_START_RE.finditer(body) + if not _inside_open_parameter(body, pm.start()) + ] + if len(param_starts) == 1: + pm = param_starts[0] + raw_val = body[pm.end() :] + if not _TC_PARAM_CLOSE_RE.search(raw_val): + param_unclosed = True + val = _TC_PARAM_CLOSE_RE.sub("", raw_val) + args[pm.group(1) or pm.group(2)] = _trim_param_value(val) + else: + for pidx, pm in enumerate(param_starts): + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + ) + raw_val = body[val_start:next_param] + if not _TC_PARAM_CLOSE_RE.search(raw_val): + param_unclosed = True + val = _TC_PARAM_CLOSE_RE.sub("", raw_val) + args[pm.group(1) or pm.group(2)] = _trim_param_value(val) + + # Strict mode: every parameter must close; a dangling one means the call was cut off. + # A closed call with no parameters is a valid zero-argument call, so keep it. + if not allow_incomplete and param_unclosed: + continue + + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": func_name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _llama3_kv_value(body: str, p: int, n: int) -> tuple[Any, int | None]: + """One ``.call`` value at ``body[p:]``; returns ``(value, len)`` or ``(None, None)``.""" + if p >= n: + return None, None + if body[p] == '"': + # ``"((?:\\.|[^"\\])*)"`` by hand so an unterminated quote is O(n), not O(n^2). + j = p + 1 + while j < n: + c = body[j] + if c == "\\": + # ``\\.`` needs a following non-newline char; else the body can't match. + if j + 1 >= n or body[j + 1] == "\n": + return None, None + j += 2 + continue + if c == '"': + raw = body[p + 1 : j] + # json.loads keeps \n/\uXXXX escapes and literal UTF-8 (emoji/CJK) intact. + try: + return json.loads('"' + raw + '"'), j + 1 - p + except (json.JSONDecodeError, ValueError): + return raw, j + 1 - p + j += 1 + return None, None # unterminated + nm = _LLAMA3_NUM_RE.match(body, p) + if nm: + v = nm.group(0) + # Sci notation and decimals decode as float; a bare integer stays int. + return (float(v) if any(c in v for c in ".eE") else int(v)), nm.end() - p + lm = _LLAMA3_LIT_RE.match(body, p) + if lm: + return {"true": True, "false": False, "null": None}[lm.group(0)], lm.end() - p + return None, None + + +def _parse_llama3_kv_args(body: str) -> dict[str, Any]: + """Left-to-right ``k=v`` kwargs from a ``.call(...)`` body (linear scan; later keys win).""" + args: dict[str, Any] = {} + n = len(body) + i = 0 + while i < n: + km = _LLAMA3_KEY_RE.match(body, i) + if km is None: + i += 1 + continue + p = _LLAMA3_WS_RE.match(body, km.end()).end() + if p >= n or body[p] != "=": + i = km.end() + continue + p = _LLAMA3_WS_RE.match(body, p + 1).end() + val, length = _llama3_kv_value(body, p, n) + if length is None: + i = km.end() + continue + args[km.group(0)] = val + i = p + length + return args + + +def _parse_llama3_python_tag( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse Llama-3 ``<|python_tag|>`` emissions: ``NAME.call(...)``, bare JSON, ``; `` multi-call, ``parameters``/``arguments`` keys.""" + out: list[dict] = [] + if _LLAMA3_PYTHON_TAG not in content: + return out + + # 1. ``NAME.call(...)`` built-in form, anchored to ``<|python_tag|>`` (optionally + # ``; ``-chained) so a ``.call(...)`` inside a JSON string argument isn't mistaken for one. + pos = content.find(_LLAMA3_PYTHON_TAG) + truncated = False + while pos >= 0 and not truncated: + head = _LLAMA3_PY_CALL_HEAD_RE.match(content, pos + len(_LLAMA3_PYTHON_TAG)) + if head is None: + # Tag is the custom JSON form (``{...}``) or noise -- leave it to step 2. + break + name = head.group(1) + open_idx = head.end() + i = open_idx + while True: + i = open_idx + depth = 1 + in_string = False + esc = False + while i < len(content) and depth > 0: + ch = content[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + break + i += 1 + # Truncated ``.call(...)`` (no closing paren): reject in strict mode. + if not allow_incomplete and depth > 0: + truncated = True + break + body = content[open_idx:i] + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(_parse_llama3_kv_args(body)), + }, + } + ) + # ``)`` then optional ``; NAME.call(`` chains the next built-in call. + chain = _LLAMA3_CALL_CHAIN_RE.match(content, i + 1) + if chain is None: + break + name = chain.group(1) + open_idx = chain.end() + # Past the consumed region: a second ``<|python_tag|>`` may carry more calls. + pos = content.find(_LLAMA3_PYTHON_TAG, i + 1) + + # 2. ``<|python_tag|>{"name":.., "parameters":..}``; raw_decode peels ``; ``-separated objects. + if not out: + decoder = json.JSONDecoder() + idx = content.find(_LLAMA3_PYTHON_TAG) + while idx >= 0: + search_from = idx + len(_LLAMA3_PYTHON_TAG) + cursor = search_from + while cursor < len(content): + brace = content.find("{", cursor) + if brace < 0: + break + # Stop at the next ``<|python_tag|>``. + next_tag = content.find(_LLAMA3_PYTHON_TAG, search_from, brace) + if next_tag >= 0: + break + try: + obj, end_offset = decoder.raw_decode(content[brace:]) + except (json.JSONDecodeError, ValueError): + cursor = brace + 1 + continue + if not isinstance(obj, dict): + cursor = brace + end_offset + continue + name = obj.get("name") or obj.get("function") or "" + args = obj.get("parameters") if "parameters" in obj else obj.get("arguments", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor = brace + end_offset + idx = content.find(_LLAMA3_PYTHON_TAG, cursor) + return out + + +# Llama-3 special-token sentinels (chainable, any order) plus the header role label. +_LLAMA3_BARE_JSON_SENTINELS = ( + "<|begin_of_text|>", + "<|eot_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", +) +_LLAMA3_HEADER_ROLES = ("assistant", "user", "system", "tool", "ipython") + + +def strip_llama3_leading_sentinels(content: str) -> str: + """Strip leading Llama-3 sentinels leaked from a prior turn; shared by the parser and the streaming guards.""" + stripped = content.lstrip() + while True: + stripped = stripped.lstrip() + matched = False + for sentinel in _LLAMA3_BARE_JSON_SENTINELS: + if stripped.startswith(sentinel): + stripped = stripped[len(sentinel) :] + if sentinel == "<|start_header_id|>": + for role in _LLAMA3_HEADER_ROLES: + if stripped.startswith(role): + stripped = stripped[len(role) :] + break + matched = True + break + if not matched: + return stripped + + +def _parse_llama3_bare_json( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Llama-3.2 bare ``{"name":.., "parameters":..}`` (strict). ``enabled_tool_names`` keeps ordinary JSON answers from being misread; ``None`` is name-agnostic.""" + out: list[dict] = [] + stripped = strip_llama3_leading_sentinels(content) + if not stripped.startswith("{"): + return out + + decoder = json.JSONDecoder() + cursor = 0 + n = len(stripped) + while cursor < n: + # Skip whitespace and the Llama-3 ``;`` inter-call separator. + while cursor < n and stripped[cursor] in " \t\n\r;": + cursor += 1 + if cursor >= n or stripped[cursor] != "{": + break + try: + obj, end_offset = decoder.raw_decode(stripped[cursor:]) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(obj, dict): + break + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + break + # Markerless JSON is ambiguous: only a call when the name is an enabled tool. + if enabled_tool_names is not None and name not in enabled_tool_names: + break + # ``parameters`` must be a dict (Llama-3 spec); ``arguments`` may be a dict or a + # JSON-string of one (OpenAI). + if "parameters" in obj: + args = obj.get("parameters") + if not isinstance(args, dict): + break + args_str = json.dumps(args) + elif "arguments" in obj: + args = obj.get("arguments") + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + try: + parsed = json.loads(args) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(parsed, dict): + break + args_str = args + else: + break + else: + break + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor += end_offset + return out + + +def _parse_mistral_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse Mistral ``[TOOL_CALLS]`` emissions: pre-v11 array/object and v11+ named forms.""" + out: list[dict] = [] + content = _strip_mistral_reasoning(content) + idx = content.find(_MISTRAL_TRIGGER) + if idx < 0: + return out + + # Disambiguate the first occurrence: array / single object (pre-v11) or bare-name (v11+). + j = idx + len(_MISTRAL_TRIGGER) + k = j + while k < len(content) and content[k] in " \t\n\r": + k += 1 + if k >= len(content): + return out + + if content[k] == "[": + return _parse_mistral_array(content, k, id_offset, allow_incomplete = allow_incomplete) + + if content[k] == "{": + # Pre-v11 single ``{"name":...}``; fall through to v11+ if it carries no ``name``. + end = _balanced_brace_end(content, k) + if end is not None: + try: + obj = json.loads(content[k : end + 1]) + if isinstance(obj, dict) and obj.get("name"): + _consume_mistral_call(content[k : end + 1], out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + pass + + # v11+: walk every ``[TOOL_CALLS]``, parsing ``name{json}`` or ``name[ARGS]{json}``. + pos = idx + while pos >= 0: + cur = pos + len(_MISTRAL_TRIGGER) + nm = _MISTRAL_V11_NAME_RE.match(content, cur) + if not nm: + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + name = nm.group(1) + after_name = nm.end() + after_name = _skip_mistral_call_id(content, after_name) + if content.startswith(_MISTRAL_ARGS_MARKER, after_name): + after_name += len(_MISTRAL_ARGS_MARKER) + while after_name < len(content) and content[after_name] in " \t\n\r": + after_name += 1 + if after_name >= len(content) or content[after_name] != "{": + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + end = _balanced_brace_end(content, after_name) + if end is None: + break + try: + args = json.loads(content[after_name : end + 1]) + except (json.JSONDecodeError, ValueError): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + if not isinstance(args, dict): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = content.find(_MISTRAL_TRIGGER, end + 1) + return out + + +def _parse_mistral_array( + content: str, + start: int, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Pre-v11 ``[TOOL_CALLS] [{...}, ...]`` array form.""" + out: list[dict] = [] + j = start + depth = 0 + in_string = False + esc = False + while j < len(content): + ch = content[j] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + break + j += 1 + # An unclosed array (no matching ]) is truncated; reject in strict mode. + if not allow_incomplete and depth != 0: + return out + body = content[start : j + 1] if depth == 0 else content[start:] + + try: + arr = json.loads(body) + if isinstance(arr, list): + for obj in arr: + if isinstance(obj, dict): + _consume_mistral_call(json.dumps(obj), out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + if not allow_incomplete: + return out + + # Healing path for unclosed arrays: walk top-level objects, advancing past each + # balanced ``{...}`` (re-scanning from every ``{`` would be quadratic ReDoS). + pos = 0 + blen = len(body) + while pos < blen: + brace = body.find("{", pos) + if brace < 0: + break + end = _balanced_brace_end(body, brace) + if end is None: + break # truncated mid-object: nothing after it can balance + _consume_mistral_call(body[brace : end + 1], out, id_offset) + pos = end + 1 + return out + + +def _consume_mistral_call(obj_text: str, out: list[dict], id_offset: int) -> None: + try: + obj = json.loads(obj_text) + except (json.JSONDecodeError, ValueError): + return + if not isinstance(obj, dict): + return + name = obj.get("name") or "" + # Mistral uses ``arguments``; accept the ``parameters`` alias too. + args = obj.get("arguments") + if args is None: + args = obj.get("parameters", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if name: + out.append( + { + "id": obj.get("id") or f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + + +def _parse_gemma_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``.""" + out: list[dict] = [] + for m in _GEMMA_TC_RE.finditer(content): + name = m.group(1) + body_start = m.end() - 1 + end_marker = content.find(_GEMMA_TC_END, body_start) + # No closing tag: truncated call, reject in strict mode. + if not allow_incomplete and end_marker < 0: + continue + scan_end = end_marker if end_marker >= 0 else len(content) + end = _gemma_balanced_brace_end(content, body_start, scan_end) + if end is None: + continue + body = content[body_start + 1 : end] + try: + args = _gemma_parse_mapping_body(body) + except Exception: + args = {} + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _balanced_brace_end(text: str, brace_pos: int) -> int | None: + """Index of the ``}`` matching ``{`` at ``brace_pos`` (ignores braces in JSON strings).""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + in_string = False + esc = False + i = brace_pos + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +_BARE_JSON_NAME_RE = re.compile(r'"name"\s*:\s*"([^"]+)"') + + +def _top_level_bare_json_name(probe: str) -> Optional[str]: + """Top-level ``"name"`` (or ``"function"`` alias) of a bare-JSON object, else None; nested objects are skipped and truncated tails return None.""" + if not probe.startswith("{"): + return None + decoder = json.JSONDecoder() + function_value = None # the ``"function"`` alias, used only if no ``"name"`` key + i = 1 + n = len(probe) + while i < n: + while i < n and probe[i] in " \t\r\n,": + i += 1 + if i >= n or probe[i] == "}": + # End of object, no top-level ``"name"``: fall back to the ``"function"`` alias. + return function_value + if probe[i] != '"': + return None + try: + key, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(key, str): + return None + i += consumed + while i < n and probe[i] in " \t\r\n": + i += 1 + if i >= n or probe[i] != ":": + return None + i += 1 + while i < n and probe[i] in " \t\r\n": + i += 1 + if key == "name": + if i < n and probe[i] == '"': + try: + value, _consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + return value if isinstance(value, str) else None + return None + if key == "function" and function_value is None and i < n and probe[i] == '"': + # ``"function"`` is an alias; record it but keep scanning (``"name"`` wins). + try: + value, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + if isinstance(value, str): + function_value = value + i += consumed + continue + # Skip a non-name top-level value; a truncated one returns None (keep the text). + if i < n and probe[i] == "{": + end = _balanced_brace_end(probe, i) + if end is None: + return None + i = end + 1 + elif i < n and probe[i] == "[": + end = _balanced_bracket_end(probe, i) + if end is None: + return None + i = end + 1 + else: + try: + _value, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + i += consumed + # No top-level ``"name"`` key: fall back to the ``"function"`` alias if seen. + return function_value + + +def strip_leading_bare_json_call(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Remove leading Llama-3.2 bare-JSON calls (including a ``;``-chained run) + that ``strip_tool_markup`` misses; non-call text is unchanged and + ``enabled_tool_names`` gates like the parser. Consuming the whole chain + matters because the loops keep this text as next-turn assistant history: a + leftover executed call would be replayed alongside the structured + ``tool_calls``.""" + remainder = text + stripped_any = False + while True: + probe = strip_llama3_leading_sentinels(remainder.lstrip()) + # Skip the Llama-3 ``;`` inter-call separator between chained calls. + if stripped_any: + probe = probe.lstrip(" \t\n\r;") + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return probe.lstrip() if stripped_any else text + if enabled_tool_names is not None: + # Only suppress when the leading object's TOP-LEVEL name is an enabled tool + # (a nested ``"name"`` is data); an unknown name is kept. + name = _top_level_bare_json_name(probe) + if name not in enabled_tool_names: + return probe.lstrip() if stripped_any else text + end = _balanced_brace_end(probe, 0) + if end is None: + return "" # truncated bare-JSON call -- nothing recoverable + # A closed object must have the CALL SHAPE the parser accepts; an ordinary JSON + # answer it rejects is content, so keep it visible. + try: + obj = json.loads(probe[: end + 1]) + except (json.JSONDecodeError, ValueError): + return probe.lstrip() if stripped_any else text + if not _bare_json_call_shaped(obj): + return probe.lstrip() if stripped_any else text + remainder = probe[end + 1 :] + stripped_any = True + + +def _bare_json_call_shaped(obj) -> bool: + """The shape gate ``_parse_llama3_bare_json`` applies to a decoded object.""" + if not isinstance(obj, dict): + return False + # The parser requires a TOP-LEVEL name; a nested one is data, not the call name. + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + return False + if "parameters" in obj: + return isinstance(obj.get("parameters"), dict) + args = obj.get("arguments") + if isinstance(args, dict): + return True + if isinstance(args, str): + try: + return isinstance(json.loads(args), dict) + except (json.JSONDecodeError, ValueError): + return False + return False + + +def _gemma_balanced_brace_end(text: str, brace_pos: int, hard_stop: int) -> int | None: + """Like ``_balanced_brace_end`` but skips ``<|"|>`` strings and matches {}/[] symmetrically.""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + i = brace_pos + while i < hard_stop: + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return None + i = close + len(_GEMMA_STR_END) + continue + ch = text[i] + if ch == "{" or ch == "[": + depth += 1 + elif ch == "}" or ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _gemma_parse_value(text: str, i: int): + """Parse one Gemma arg value at ``i``; returns ``(value, next_index)``.""" + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return text[i + len(_GEMMA_STR_BEGIN) :], len(text) + return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END) + if text[i] == "{": + end = _gemma_balanced_brace_end(text, i, len(text)) + if end is None: + return {}, len(text) + return _gemma_parse_mapping_body(text[i + 1 : end]), end + 1 + if text[i] == "[": + j, depth = i, 0 + while j < len(text): + if text.startswith(_GEMMA_STR_BEGIN, j): + k = text.find(_GEMMA_STR_END, j + len(_GEMMA_STR_BEGIN)) + if k < 0: + j = len(text) + break + j = k + len(_GEMMA_STR_END) + continue + ch = text[j] + if ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + break + j += 1 + body = text[i + 1 : j] + items: list[Any] = [] + k = 0 + while k < len(body): + if body[k] in " \t\n\r,": + k += 1 + continue + v, k = _gemma_parse_value(body, k) + items.append(v) + return items, j + 1 + # Primitive: number / true/false/null / bare identifier. + end = i + while end < len(text) and text[end] not in ",}]" and not text.startswith(_GEMMA_STR_BEGIN, end): + end += 1 + if end == i: + # Stray delimiter, nothing consumed: advance past it so callers can't spin forever. + return "", i + 1 + raw = text[i:end].strip() + if raw == "true": + return True, end + if raw == "false": + return False, end + if raw == "null": + return None, end + try: + return int(raw), end + except ValueError: + pass + try: + return float(raw), end + except ValueError: + pass + return raw, end + + +def _gemma_parse_mapping_body(body: str) -> dict[str, Any]: + """Parse a Gemma argument mapping (content between `{` and `}`).""" + out: dict[str, Any] = {} + i = 0 + n = len(body) + while i < n: + while i < n and body[i] in " \t\n\r,": + i += 1 + if i >= n: + break + if body.startswith(_GEMMA_STR_BEGIN, i): + close = body.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + break + key = body[i + len(_GEMMA_STR_BEGIN) : close] + i = close + len(_GEMMA_STR_END) + else: + kstart = i + while i < n and body[i] != ":": + i += 1 + key = body[kstart:i].strip() + while i < n and body[i] in " \t\n\r": + i += 1 + if i < n and body[i] == ":": + i += 1 + while i < n and body[i] in " \t\n\r": + i += 1 + if i >= n: + out[key] = None + break + v, i = _gemma_parse_value(body, i) + out[key] = v + return out diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index e8367ad08c..ff8faf2308 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -27,12 +27,15 @@ _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") -_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{") +# Name class allows dots/hyphens for dotted Gemma names; whitespace-tolerant around +# ``call`` / ``:`` since drift emits ``call: name{`` and ``call : name{``. +_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") _TC_GEMMA_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -_TC_PARAM_START_RE = re.compile(r"\s*") +# Horizontal whitespace only so the newline + value indentation survive (_trim_param_value trims one newline). +_TC_PARAM_START_RE = re.compile(r"[^\S\n]*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _GEMMA_QUOTE = '<|"|>' _PARAM_CLOSE_TAG = "" @@ -43,7 +46,8 @@ _FUNC_CLOSE_TAG = "" # must be identifier-shaped (start with a letter or underscore); a comma # followed by digits-then-colon is value text such as a timestamp or ratio # (`meet at 10:00, 11:00 tomorrow`), not a new key. -_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:") +# Dots match the key-quoting scanner: a dotted key after a bare value must end the value at the comma. +_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") def _balanced_brace_end( @@ -223,7 +227,9 @@ def _quote_gemma_object_keys(src: str) -> str: while i < len(src) and src[i].isspace(): i += 1 key_name_start = i - while i < len(src) and (src[i].isalnum() or src[i] in "_-"): + # Dots match the parser's key/name charset: Gemma emits dotted argument keys + # (user.name:...) for namespaced schemas. + while i < len(src) and (src[i].isalnum() or src[i] in "_-."): i += 1 key_name = src[key_name_start:i] colon_pos = i @@ -267,7 +273,8 @@ def _quote_gemma_object_keys(src: str) -> str: json.loads(raw.strip()) parts.append(raw) except (json.JSONDecodeError, ValueError): - parts.append(json.dumps(raw.strip()) if raw.strip() else raw) + # Quote bare value; empty ({k:}) becomes "" so json.loads sees {"k":""} not invalid {"k":}. + parts.append(json.dumps(raw.strip())) else: parts.append(src[key_start:i]) return "".join(parts) @@ -291,9 +298,35 @@ def _inside_open_parameter(content: str, pos: int) -> bool: last_param_start = match.start() if last_param_start < 0: return False - last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos) - last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos) - return last_param_start > max(last_param_close, last_func_close) + # The parameter's OWN close tag decides: if it closes after ``pos`` the position is + # argument data (even across literal function closes); an unclosed one falls back to func close. + own_close = content.find(_PARAM_CLOSE_TAG, last_param_start) + if own_close >= 0: + return own_close > pos + func_close = content.find(_FUNC_CLOSE_TAG, last_param_start) + return func_close < 0 or pos < func_close + + +def _func_close_index(content: str, body_start: int, body: str) -> int: + """Index in ``body`` of the first ```` that is not argument + data (not inside an open parameter value); -1 when every close is data. + Taking the LAST close swallowed prose between the real close and a + literal ```` mentioned later in the answer.""" + idx = body.find(_FUNC_CLOSE_TAG) + while idx >= 0: + if not _inside_open_parameter(content, body_start + idx): + return idx + idx = body.find(_FUNC_CLOSE_TAG, idx + 1) + return -1 + + +def _trim_param_value(val: str) -> str: + """Trim only the wrapping newline (not str.strip) so code/diff argument indentation survives.""" + if val.startswith("\n"): + val = val[1:] + if val.endswith("\n"): + val = val[:-1] + return val def parse_tool_calls_from_text( @@ -349,7 +382,10 @@ def parse_tool_calls_from_text( if kind == "json": obj = json.loads(content[m.end() - 1 : end + 1]) name = obj.get("name", "") - arguments = obj.get("arguments", {}) + # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes ). + arguments = obj.get("arguments") + if arguments is None: + arguments = obj.get("parameters", {}) if isinstance(arguments, dict): arguments = json.dumps(arguments) else: @@ -382,7 +418,7 @@ def parse_tool_calls_from_text( body_end = len(content) body_end = min(body_end, next_func) body = content[body_start:body_end] - close_idx = body.rfind(_FUNC_CLOSE_TAG) + close_idx = _func_close_index(content, body_start, body) if close_idx >= 0: span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) body = body[:close_idx] @@ -404,7 +440,7 @@ def parse_tool_calls_from_text( val = stripped_val[: -len(_PARAM_CLOSE_TAG)] else: val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() + arguments[pm.group(1)] = _trim_param_value(val) else: valid_params = True for pidx, pm in enumerate(param_starts): @@ -422,7 +458,7 @@ def parse_tool_calls_from_text( val = stripped_val[: -len(_PARAM_CLOSE_TAG)] else: val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() + arguments[param_name] = _trim_param_value(val) if not valid_params: continue @@ -444,6 +480,86 @@ def parse_tool_calls_from_text( } ) call_spans.append((start, span_end)) + + if not tool_calls: + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + ] + for idx, fm in enumerate(func_starts): + func_name = fm.group(1) + body_start = fm.end() + next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + end_tag = _TC_END_TAG_RE.search(content[body_start:]) + if end_tag: + body_end = body_start + end_tag.start() + else: + body_end = len(content) + body_end = min(body_end, next_func) + body = content[body_start:body_end] + # Span for with_spans callers: through the close if present, else body end. + span_end = body_end + if not allow_incomplete: + close_idx = _func_close_index(content, body_start, body) + if close_idx < 0: + continue + body = body[:close_idx] + span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) + else: + # Terminate at the real close so trailing prose doesn't leak in; no close -> whole body. + close_idx = _func_close_index(content, body_start, body) + if close_idx >= 0: + body = body[:close_idx] + span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) + + arguments: dict = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + pm = param_starts[0] + val = body[pm.end() :] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + continue + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[pm.group(1)] = _trim_param_value(val) + else: + valid_params = True + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() + if pidx + 1 < len(param_starts) + else len(body) + ) + val = body[val_start:next_param] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + valid_params = False + break + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[param_name] = _trim_param_value(val) + if not valid_params: + continue + + tc = { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": func_name, + "arguments": json.dumps(arguments), + }, + } + tool_calls.append(tc) + call_spans.append((fm.start(), span_end)) + if with_spans: return tool_calls, call_spans return tool_calls diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 17be222d93..4393c1b304 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -603,6 +603,17 @@ def _chat_content_chunk(completion_id, created, model_name, text) -> str: ) +def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str: + """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block).""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(reasoning_content = text), + finish_reason = None, + ) + + def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: """Terminal stop chunk (empty delta) carrying the finish reason.""" return _chat_chunk_sse( @@ -1136,6 +1147,7 @@ from core.inference.key_exchange import decrypt_api_key from core.inference.model_ids import public_model_id from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client +from core.inference.tool_call_parser import _strip_function_xml_calls, _strip_mistral_closed_calls from core.inference.passthrough_healing import ( StreamToolCallHealer, heal_gate, @@ -1294,6 +1306,11 @@ async def artifact_preview_frame(allow_network: bool = False): ) +# Whitespace/escape-tolerant bare-JSON tool-template detector: matches pretty-printed and +# JSON-escaped ``{"name":`` plus the ``"function"`` alias. +_BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:') + + def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: """Classify reasoning/tool capabilities via the GGUF classifier so flags match across backends. gpt-oss is overridden: Harmony routes reasoning and @@ -1304,17 +1321,21 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: model_identifier = model_id, log_source = "safetensors", ) - # Our safetensors loop only parses {json}, - # ..., and Gemma native <|tool_call>.... - # Llama uses <|python_tag|>, Mistral uses [TOOL_CALLS]; advertising tools for - # those enables a pill the parser can't honour. GGUF is unaffected -- - # llama-server normalises every format into structured deltas. + # Markers the parser recognises; drop the pill if a template advertises tools but uses none. + # The bare-JSON ``{"name":`` form is matched whitespace-tolerantly below. + _PARSER_MARKERS = ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + ) if ( flags.get("supports_tools") and chat_template - and "" not in chat_template - and "" not in chat_template + and not any(m in chat_template for m in _PARSER_MARKERS) + and not _BARE_JSON_NAME_MARKER_RE.search(chat_template) ): logger.info( "safetensors: template advertises tools but uses an " @@ -1335,6 +1356,31 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: return flags +def _sf_reasoning_prefill_mode( + features: dict, + enable_thinking: Optional[bool], + template: Optional[str] = None, + reasoning_effort: Optional[str] = None, +) -> bool: + """Whether this request begins inside an unclosed ```` (Qwen3/GLM prefill it). Gated on the standard markers; bespoke channels, gpt-oss, and thinking-disabled requests are excluded. ``enable_thinking=None`` defaults ON.""" + if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"): + return False + tpl = template or "" + if "" not in tpl and "" not in tpl: + return False + if features.get("reasoning_always_on"): + return True + if not features.get("supports_reasoning"): + return False + if enable_thinking is False: + return False + # reasoning_effort="none" disables thinking on enable_thinking_effort (GLM-5.2) models like + # enable_thinking=False; without this the answer is swallowed into empty reasoning_content. + if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none": + return False + return True + + def _effective_enable_tools(payload) -> Optional[bool]: """Resolve `payload.enable_tools` against the process-level tool policy. @@ -1605,30 +1651,41 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str: return nudge + " " + _RAG_GROUNDING_NUDGE -# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py -# split across the visible/DRAIN boundary. Four leak shapes: -# 1. well-formed `...` / `...` -# 2. orphan opening to EOF (close was DRAINED) -# 3. bare orphan close (open was DRAINED) -# 4. tail-only `` (outer close truncated by EOS); anchored to -# `\Z` so mid-text `` in user code samples survives. +# Strip leaked tool-call markup: every shared-parser format plus the leak shapes +# ``llama_cpp.py``'s speculative buffer splits across the visible/DRAIN boundary. Mistral +# ``[TOOL_CALLS]`` uses the parser's balanced-brace helper (``\{.*?\}`` would truncate nested JSON). _TOOL_XML_RE = _re.compile( # Hyphen in the name char-class matches MCP tool names with dashes # (mcp__srv__list-issues) that would otherwise leak past this strip. - r"<(?:tool_call|function=[\w-]+)>.*?(?:|\Z)" + # The ``<|python_tag|>`` arm runs to the next REAL Llama sentinel or EOF, so a literal + # ``<|...|>`` token in an argument (e.g. ``<|cite|>``) doesn't truncate the strip. + # ```` plus the ```` attribute form; name class mirrors the parser. + # A CLOSED ``...`` extends to the last ```` before the next + # opener (so a literal ```` in a value can't truncate); this arm runs first. + r'(?:(?!).)*' + r'|<(?:tool_call|function(?:=[\w.\-]+|\s+name="[\w.\-]+"))>.*?(?:|\Z)' r"|<\|tool_call>.*?(?:|\Z)" r"|" r"|" - r"|\s*\Z", + r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*" + # ```` is the attribute-form alias of ````; strip a tail-only orphan. + r"|\s*\Z", _re.DOTALL, ) +def _strip_tool_xml(text: str) -> str: + """Mistral balanced-brace helper + guarded function-XML scan + ``_TOOL_XML_RE`` (skips openers inside an open ````).""" + return _TOOL_XML_RE.sub( + "", _strip_function_xml_calls(_strip_mistral_closed_calls(text), final = True) + ) + + def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str: - """Apply route-level XML leak cleanup only when Auto-Heal is enabled.""" + """Route-level tool-call leak cleanup (Auto-Heal only) via ``_strip_tool_xml``.""" if not auto_heal_tool_calls: return text - return _TOOL_XML_RE.sub("", text) + return _strip_tool_xml(text) logger = get_logger(__name__) @@ -6511,6 +6568,22 @@ async def openai_chat_completions( _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") _sf_features = _detect_safetensors_features(backend, _sf_tpl) + # Split prefilled-```` output into reasoning_content deltas (GGUF parity) so the UI + # renders the thinking block for safetensors and MLX. + _sf_parse_think = bool( + _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on") + ) + # Prefilled-open only for prefill styles with thinking on this request; gpt-oss excluded. + _sf_reasoning_prefilled = _sf_reasoning_prefill_mode( + _sf_features, payload.enable_thinking, _sf_tpl, payload.reasoning_effort + ) + + def _new_sf_reasoning_extractor(): + return _ResponsesReasoningExtractor( + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + cancel_event = threading.Event() completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) @@ -6652,6 +6725,19 @@ async def openai_chat_completions( gen = sf_generate_with_tools() prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() + + def _sf_flush_reasoning(): + # Drain the extractor at a turn boundary / stream end; only visible text reaches the monitor. + fr, fv = reasoning_extractor.finish() + out = [] + if fr: + out.append(_chat_reasoning_chunk(completion_id, created, model_name, fr)) + if fv: + api_monitor.append_reply(monitor_id, fv) + out.append(_chat_content_chunk(completion_id, created, model_name, fv)) + return out + while True: if cancel_event.is_set(): backend.reset_generation_state() @@ -6668,7 +6754,11 @@ async def openai_chat_completions( if event["type"] == "status": if not event["text"]: + # Turn boundary: flush reasoning, then start a fresh extractor. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() status_data = json.dumps( { "type": "tool_status", @@ -6680,7 +6770,11 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": + # Flush reasoning before tool_start so the thinking block closes ahead of the tool card. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() yield f"data: {json.dumps(event)}\n\n" continue @@ -6694,9 +6788,18 @@ async def openai_chat_completions( prev_text = clean_cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + # Split reasoning vs visible; only visible reaches the monitor. + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _chat_reasoning_chunk( + completion_id, created, model_name, reasoning_delta + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + for _c in _sf_flush_reasoning(): + yield _c yield _chat_final_chunk(completion_id, created, model_name, "stop") # Usage chunk from the last turn, same shape as the # GGUF tool loop's metadata. Request-scoped holder, so @@ -6774,18 +6877,27 @@ async def openai_chat_completions( return full_text content_text = await asyncio.to_thread(_drain_to_text) - api_monitor.set_reply(monitor_id, content_text) + # Split prefilled reasoning from the visible answer; monitor gets visible text only. + _reasoning_text, _visible_text = _extract_responses_reasoning( + content_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + api_monitor.set_reply(monitor_id, _visible_text) _stats = _sf_stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed") + _sf_msg_kwargs = {"content": _visible_text} + if _reasoning_text: + _sf_msg_kwargs["reasoning_content"] = _reasoning_text response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(content = content_text), + message = CompletionMessage(**_sf_msg_kwargs), finish_reason = "stop", ) ], @@ -6864,6 +6976,8 @@ async def openai_chat_completions( yield _chat_role_chunk(completion_id, created, model_name) prev_text = "" + # Split prefilled into reasoning_content deltas. Single turn (no per-turn reset); also MLX. + reasoning_extractor = _new_sf_reasoning_extractor() # Run the sync generator in a thread pool to avoid blocking the # event loop. Critical for compare mode: two SSE requests arrive # concurrently but the orchestrator serializes them via @@ -6892,9 +7006,21 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _chat_reasoning_chunk( + completion_id, created, model_name, reasoning_delta + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + final_reasoning, final_visible = reasoning_extractor.finish() + if final_reasoning: + yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning) + if final_visible: + api_monitor.append_reply(monitor_id, final_visible) + yield _chat_content_chunk(completion_id, created, model_name, final_visible) yield _chat_final_chunk(completion_id, created, model_name, "stop") # Usage chunk (choices=[], usage set), same shape as the # GGUF path so the speed popover works for MLX too. @@ -6956,18 +7082,27 @@ async def openai_chat_completions( for token in generate(): full_text = token + # Split prefilled reasoning from the visible answer; also covers MLX. + _reasoning_text, _visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + _plain_msg_kwargs = {"content": _visible_text} + if _reasoning_text: + _plain_msg_kwargs["reasoning_content"] = _reasoning_text response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(content = full_text), + message = CompletionMessage(**_plain_msg_kwargs), finish_reason = "stop", ) ], ) - api_monitor.set_reply(monitor_id, full_text) + api_monitor.set_reply(monitor_id, _visible_text) _stats = stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) @@ -7790,10 +7925,18 @@ def _responses_marker_holdback(text: str, markers: tuple[str, ...]) -> int: class _ResponsesReasoningExtractor: """Split local markup into Responses reasoning and visible text.""" - def __init__(self, *, parse_think_markers: bool = False) -> None: + def __init__( + self, + *, + parse_think_markers: bool = False, + reasoning_prefilled: bool = False, + ) -> None: self._buffer = "" - self._in_reasoning = False - self._parse_think_markers = parse_think_markers + # ``reasoning_prefilled``: output begins inside an unclosed ```` (Qwen3/GLM prefill), + # so start in reasoning to capture leading text until the first ````. + self._in_reasoning = reasoning_prefilled + # Splitting requires marker parsing; a prefilled open implies it. + self._parse_think_markers = parse_think_markers or reasoning_prefilled def feed( self, @@ -7816,14 +7959,21 @@ class _ResponsesReasoningExtractor: if self._in_reasoning: close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) if close_idx != -1: - reasoning_parts.append(self._buffer[:close_idx]) + reasoning_parts.append( + self._buffer[:close_idx].replace(_RESPONSES_THINK_OPEN, "") + ) self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] self._in_reasoning = False continue - keep = _responses_marker_holdback(self._buffer, (_RESPONSES_THINK_CLOSE,)) + # Hold back a trailing partial of either marker: the close (clean chunk-boundary split) + # and a stray open (so a re-emitted ```` isn't leaked into the reasoning drawer). + keep = _responses_marker_holdback( + self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN) + ) if keep == len(self._buffer): break - reasoning_parts.append(self._buffer[:-keep] if keep else self._buffer) + emit = self._buffer[:-keep] if keep else self._buffer + reasoning_parts.append(emit.replace(_RESPONSES_THINK_OPEN, "")) self._buffer = self._buffer[-keep:] if keep else "" break @@ -7860,7 +8010,7 @@ class _ResponsesReasoningExtractor: return "", remaining if self._in_reasoning: self._in_reasoning = False - return remaining, "" + return remaining.replace(_RESPONSES_THINK_OPEN, ""), "" return "", remaining.replace(_RESPONSES_THINK_CLOSE, "") @@ -7869,8 +8019,12 @@ def _extract_responses_reasoning( reasoning_content: Any = None, *, parse_think_markers: bool = False, + reasoning_prefilled: bool = False, ) -> tuple[str, str]: - extractor = _ResponsesReasoningExtractor(parse_think_markers = parse_think_markers) + extractor = _ResponsesReasoningExtractor( + parse_think_markers = parse_think_markers, + reasoning_prefilled = reasoning_prefilled, + ) reasoning, visible = extractor.feed(text, reasoning_content) final_reasoning, final_visible = extractor.finish() return reasoning + final_reasoning, visible + final_visible @@ -9700,7 +9854,7 @@ async def anthropic_messages( # Strip stale tool-call XML from conversation for _msg in openai_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): - _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + _msg["content"] = _strip_tool_xml(_msg["content"]).strip() def _run_tool_gen(): return llama_backend.generate_chat_completion_with_tools( @@ -9854,7 +10008,7 @@ async def _anthropic_tool_stream( # content event that was purely tool XML doesn't count as text. if etype == "content": event = dict(event) - event["text"] = _TOOL_XML_RE.sub("", event["text"]) + event["text"] = _strip_tool_xml(event["text"]) # disable_parallel_tool_use: keep only the first tool_use block, # dropping every later tool_start and its paired tool_end (robust # to empty tool-call ids — tracked by state, not id matching). @@ -10040,7 +10194,7 @@ async def _anthropic_tool_non_streaming( etype = event.get("type", "") if etype == "content": # Strip leaked tool-call XML - clean = _TOOL_XML_RE.sub("", event["text"]) + clean = _strip_tool_xml(event["text"]) new = clean[len(prev_text) :] prev_text = clean if new: @@ -10509,10 +10663,11 @@ async def _anthropic_passthrough_non_streaming( else: text = message.get("content") or "" if text: - # Keep unpromoted bytes when healing is active; legacy stripping is - # only for opted-out or no-client-tool requests. + # Keep unpromoted bytes when healing is active; legacy stripping is only for opted-out + # or no-client-tool requests. _strip_tool_xml also cleans Mistral [TOOL_CALLS] and + # guarded function-XML, not just _TOOL_XML_RE. if not healing_active: - text = _TOOL_XML_RE.sub("", text) + text = _strip_tool_xml(text) text = text.strip() if text: content_blocks.append(AnthropicResponseTextBlock(text = text)) diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index 8df8d37a52..63df86ec17 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -21,7 +21,10 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -from core.inference.tool_call_parser import parse_tool_calls_from_text +from core.inference.tool_call_parser import ( + _gemma_parse_value, + parse_tool_calls_from_text, +) def _args(call: dict) -> dict: @@ -45,6 +48,17 @@ def test_normal_multi_key_arguments_still_split(): assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} +def test_empty_bare_value_becomes_empty_string_not_dropped(): + # An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON). + calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}") + assert len(calls) == 1, calls + assert _args(calls[0]) == {"query": "", "unit": "celsius"} + + only = parse_tool_calls_from_text("<|tool_call>call:get{q:}") + assert len(only) == 1, only + assert _args(only[0]) == {"q": ""} + + def test_bare_value_with_timestamps_after_comma_is_kept(): # A comma followed by digits-then-colon (a timestamp/ratio) is value text, # not a new key, so the whole query must be preserved as one argument. @@ -159,3 +173,43 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): ) calls = parse_tool_calls_from_text(content) assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_gemma_parse_value_always_advances_on_stray_delimiter(): + # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the + # parser, or a looping caller spins forever (DoS). + for delim in (",", "}", "]"): + text = delim + "rest" + value, nxt = _gemma_parse_value(text, 0) + assert nxt > 0, (delim, value, nxt) + + +def test_malformed_gemma_array_does_not_hang(): + # ``[},]`` (stray ``}`` in a list body) hung the buggy parser; the timeout fails + # the regression loudly instead of blocking CI forever. + import threading + + result: dict = {} + + def _run(): + result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:[},]}") + + t = threading.Thread(target = _run, daemon = True) + t.start() + t.join(timeout = 10.0) + assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed array input" + + +def test_malformed_gemma_mapping_value_does_not_hang(): + # A stray ``}`` where a mapping value is expected must also terminate. + import threading + + result: dict = {} + + def _run(): + result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:}},b:1}") + + t = threading.Thread(target = _run, daemon = True) + t.start() + t.join(timeout = 10.0) + assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed mapping input" diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 05d2a0b80a..8977d6e92a 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -20,7 +20,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -from core.inference.llama_cpp import _PROVISIONAL_ARGS_MIN_CHARS, LlamaCppBackend +from core.inference.llama_cpp import ( + _MAX_REPROMPTS, + _PROVISIONAL_ARGS_MIN_CHARS, + LlamaCppBackend, +) from state import tool_approvals from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision @@ -1036,9 +1040,11 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch): def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): """No-tool re-prompt attempts should not concatenate into the UI.""" - streams = [ - [_sse({"content": "I will use render_html now."}), _done()], - [_sse({"content": "Understood. I will use render_html now."}), _done()], + # One initial response plus one stream per re-prompt (count from the shared cap). + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + streams += [ + [_sse({"content": "Understood. I will use render_html now."}), _done()] + for _ in range(_MAX_REPROMPTS) ] payloads: list[dict] = [] backend = _make_backend(monkeypatch, streams, payloads) @@ -1073,7 +1079,7 @@ def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] assert content_texts == ["I will use render_html now."] - assert len(payloads) == 2 + assert len(payloads) == _MAX_REPROMPTS + 1 def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch): @@ -1200,6 +1206,66 @@ def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatc ) +def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch): + # Inline Mistral ``[TOOL_CALLS]`` after a visible preface: the DRAINING flush must use the + # shared parser patterns (the legacy set leaked the marker to clients). + streams = [ + [_sse({"content": 'Let me search. [TOOL_CALLS]web_search{"query":"cats"}'}), _done()], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("[TOOL_CALLS]" not in t for t in content_texts), content_texts + assert any("Let me search." in t for t in content_texts) + + +def test_textual_llama_python_tag_marker_not_leaked(monkeypatch): + # Same leak class for the Llama-3 built-in ``<|python_tag|>NAME.call(...)`` form. + streams = [ + [_sse({"content": '<|python_tag|>web_search.call(query="cats")'}), _done()], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("<|python_tag|>" not in t for t in content_texts), content_texts + + def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): """Suppression ends once a forced re-prompt actually calls a tool.""" @@ -1738,6 +1804,189 @@ def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch): assert calls == [("python", {"code": big_code})] +def _streamed_content(text: str, frag: int = 4) -> list[str]: + """Stream content token-by-token like llama-server; ``frag`` sets the chunk size.""" + chunks = [_sse({"content": text[i : i + frag]}) for i in range(0, len(text), frag)] + chunks.append(_done()) + return chunks + + +def test_bare_json_tool_call_streamed_is_not_leaked_and_executes(monkeypatch): + """A wrapper-less bare-JSON call must be held while incomplete, drained silently, and executed with nothing leaking.""" + + bare_call = '{"name": "web_search", "parameters": {"query": "weather in Sydney"}}' + first_stream = _streamed_content(bare_call) + final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Weather: sunny, 22C." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather in Sydney?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "weather in Sydney"})] + assert any( + event.get("type") == "tool_end" and event.get("tool_name") == "web_search" + for event in events + ) + + # The bare JSON never leaked to the user-visible stream. + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all('"name"' not in t for t in content_texts), content_texts + assert all("web_search" not in t for t in content_texts), content_texts + # The post-tool synthesis is still streamed. + assert any("sunny in Sydney" in t for t in content_texts), content_texts + + +def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(monkeypatch): + """Markerless JSON with a non-enabled name is the answer, not a phantom call.""" + + answer = '{"name": "Alice", "parameters": {"age": 30}}' + first_stream = _streamed_content(answer) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "give me a person record"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_incomplete_bare_json_truncation_is_not_leaked(monkeypatch): + """If generation is cut off mid bare-JSON object (no closing brace), the held + fragment must be stripped at stream end rather than dumped to the user.""" + + truncated = '{"name": "web_search", "parameters": {"query": "weather in S' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all('{"name"' not in t for t in content_texts), content_texts + + +def test_gguf_truncated_disabled_name_json_is_preserved_when_tools_active(monkeypatch): + """A truncated JSON answer with a non-enabled name must still be shown (resolvers are gated on enabled names).""" + + truncated = '{"name": "Alice", "parameters": {"age": 30' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "give json"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_gguf_truncated_enabled_name_json_is_still_suppressed(monkeypatch): + """Counterpart guard: a truncated ENABLED-tool bare call (``web_search``) cut off + mid-JSON still must NOT leak -- the gate only spares disabled / non-tool names.""" + + truncated = '{"name": "web_search", "parameters": {"query": "weather in S' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all('{"name"' not in t for t in content_texts), content_texts + + +def test_gguf_oversized_disabled_name_json_is_preserved(monkeypatch): + """An oversized still-open JSON answer with a non-enabled name streams as content, not a phantom drain.""" + + cap = 16384 + big = "A" * (cap + 5000) + answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes + first_stream = [_sse({"content": answer[i : i + 2000]}) for i in range(0, len(answer), 2000)] + first_stream.append(_done()) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "long json"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts[:1] + + def _usage_done(usage: dict, finish_reason: str = "stop") -> str: """A terminal SSE chunk carrying llama-server's ``usage`` block, the way the real server reports it on the final chunk of a completion.""" @@ -1813,3 +2062,131 @@ def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch): metadata = [e for e in events if e.get("type") == "metadata"] assert metadata, "expected a metadata event" assert "prompt_tokens_details" not in metadata[-1]["usage"] + + +def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): + """An oversized bare-JSON call drains rather than streams, and still executes via the safety net.""" + + cap = 16384 + big = "A" * (cap + 5000) + full = '{"name":"python","parameters":{"code":"' + big + '"}}' + first_stream = [_sse({"content": full[i : i + 2000]}) for i in range(0, len(full), 2000)] + first_stream.append(_done()) + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert not any(t.lstrip().startswith('{"name') for t in content_texts), content_texts[:1] + assert calls and calls[0][0] == "python" + assert len(calls[0][1].get("code", "")) > cap + + +def test_gguf_bare_json_call_not_replayed_in_next_turn_content(monkeypatch): + """After a bare-JSON call executes, the kept assistant message must not carry the raw call as content.""" + + import copy + + first_stream = [ + _sse({"content": '{"name":"web_search","parameters":{"query":"cats"}}'}), + _done(), + ] + final_stream = [_sse({"content": "Found."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "RESULT") + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert len(payloads) >= 2 + asst = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"] + assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst + + +def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(monkeypatch): + """Auto-Heal OFF keeps a truncated enabled-name fragment visible; ON suppresses it (strip gated on auto_heal_tool_calls).""" + + trunc = '{"name":"web_search","parameters":{"query":"weather' + + def _run(auto_heal): + stream = [_sse({"content": trunc}), _done()] + backend = _make_backend(monkeypatch, [stream], []) + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + auto_heal_tool_calls = auto_heal, + ) + ) + contents = "".join(e.get("text", "") for e in events if e.get("type") == "content") + return calls, contents + + calls_off, contents_off = _run(False) + assert calls_off == [], calls_off + assert "web_search" in contents_off, contents_off + + calls_on, contents_on = _run(True) + assert calls_on == [], calls_on + assert "web_search" not in contents_on, contents_on + + +def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch): + """Re-prompt slots must not extend the tool budget: stop after ``max_tool_iterations`` executed rounds.""" + # More tool-call streams than the budget: leaked re-prompt slots would run 2+3=5 rounds; + # honouring the budget stops after 2, then a tool-less final-answer pass. + streams = [ + _structured_tool_call("web_search", {"query": f"q{i}"}, f"call_{i}") for i in range(6) + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search repeatedly"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + # Exactly two executed tool rounds, then one final-answer pass. + assert len(calls) == 2, calls + assert len(payloads) == 3, len(payloads) + # The final pass is the budget-exhausted nudge and carries no tools. + assert _tool_names(payloads[2]) == [], _tool_names(payloads[2]) + assert any( + m.get("role") == "user" and "used all available tool calls" in m.get("content", "") + for m in payloads[2]["messages"] + ), payloads[2]["messages"] diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index a7ceb49ed9..ce5688be3e 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -59,6 +59,7 @@ from models.inference import ( ResponsesUsage, ) from routes.inference import ( + _ResponsesReasoningExtractor, _SameTaskStreamingResponse, _build_chat_request, _chat_tool_calls_to_responses_output, @@ -795,6 +796,7 @@ class TestResponsesNonStreamingAdapter: def test_monitor_records_translated_visible_text(self, monkeypatch): import routes.inference as inf_mod + import routes.inference as inf_mod async def fake_chat_completions(chat_req, request): assert request.state.skip_api_monitor is True @@ -1988,6 +1990,122 @@ class TestTranslatedMessagesValidate: ChatMessage(**m.model_dump(exclude_none = True)) +# reasoning_prefilled: Qwen3/GLM enable_thinking templates prefill an unclosed , so generation +# begins inside the think block and emits only the closing ; extractor starts in reasoning. +class TestReasoningPrefilledExtractor: + def test_prefilled_single_feed_splits_lone_close(self): + # T1: reasoning...answer with a prefilled (unseen) open tag. + reasoning, visible = _extract_responses_reasoning( + "plananswer", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "plan" + assert visible == "answer" + + def test_prefilled_never_closed_is_all_reasoning(self): + # T2: truncated mid-thought (no ) -> all reasoning (GGUF parity). + reasoning, visible = _extract_responses_reasoning( + "still thinking with no close", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "still thinking with no close" + assert visible == "" + + def test_prefilled_close_split_across_feeds(self): + # T3: straddles two feed() calls; holdback resolves it. + ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + r1, v1 = ex.feed("planans") + fr, fv = ex.finish() + assert (r1 + r2 + fr) == "plan" + assert (v1 + v2 + fv) == "ans" + + def test_prefilled_close_split_one_char_per_feed(self): + # T4: every char in its own feed still splits correctly. + ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + reasoning, visible = "", "" + for ch in "planx": + r, v = ex.feed(ch) + reasoning += r + visible += v + fr, fv = ex.finish() + assert (reasoning + fr) == "plan" + assert (visible + fv) == "x" + + def test_prefilled_empty_generation(self): + # T5: nothing generated. + reasoning, visible = _extract_responses_reasoning( + "", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "" + assert visible == "" + + def test_prefilled_whitespace_after_close_is_visible(self): + # T6: Qwen commonly emits \n\n before the answer. + reasoning, visible = _extract_responses_reasoning( + "plan\n\nanswer", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "plan" + assert visible == "\n\nanswer" + + def test_prefilled_stray_open_tag_is_suppressed(self): + # T7: a re-emitted literal inside prefilled reasoning is dropped, not leaked. + reasoning, visible = _extract_responses_reasoning( + "abc", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "ab" + assert visible == "c" + assert "" not in reasoning + + def test_prefilled_close_at_start_empty_reasoning(self): + # T8: model closed immediately (empty reasoning) then answered. + reasoning, visible = _extract_responses_reasoning( + "hi", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "" + assert visible == "hi" + + def test_not_prefilled_lone_close_preserves_current_behavior(self): + # T9: without prefilled, a lone keeps pre-fix behavior (reasoning stays visible, tag dropped). + reasoning, visible = _extract_responses_reasoning( + "reasoningans", + parse_think_markers = True, + reasoning_prefilled = False, + ) + assert reasoning == "" + assert visible == "reasoningans" + + def test_not_prefilled_full_pair_still_splits(self): + # T10: normal explicit .. (GGUF / Harmony) unchanged. + reasoning, visible = _extract_responses_reasoning( + "rv", + parse_think_markers = True, + reasoning_prefilled = False, + ) + assert reasoning == "r" + assert visible == "v" + + def test_prefilled_ignored_when_markers_not_parsed(self): + # T11: a non-reasoning model (parse_think_markers False) passes text straight through. + reasoning, visible = _extract_responses_reasoning( + "just an answer", + parse_think_markers = False, + reasoning_prefilled = False, + ) + assert reasoning == "" + assert visible == "just an answer" + + # ===================================================================== # Streaming passthrough healing — text-form calls promoted in order # ===================================================================== diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 671af93708..643d64af7a 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -127,9 +127,8 @@ def test_detect_safetensors_features_gptoss_disables_tools(): assert flags["supports_tools"] is False -# Llama-3 / Mistral advertise tools but emit <|python_tag|> / [TOOL_CALLS], -# which our parser can't read. The route helper must not flip supports_tools=True -# for them, else the UI enables a pill the agentic loop can't honour. +# Llama-3 / Mistral / Gemma 4 tool-call formats are parser-supported, so supports_tools stays True; +# only templates matching none of the known markers are suppressed. LLAMA3_TEMPLATE = """ {%- if tools %} @@ -161,27 +160,106 @@ MISTRAL_TEMPLATE = """ {%- endfor %} """ +GEMMA4_TEMPLATE = """ +{%- if tools %} + {{- 'Tools available. Emit calls as ' }} + {{- '<|tool_call>call:NAME{key:<|"|>val<|"|>}' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" -def test_detect_safetensors_features_llama3_template_suppresses_tools(): - """Llama-3 emits <|python_tag|>; safetensors loop cannot parse it.""" + +def test_detect_safetensors_features_llama3_template_keeps_tools_on(): + """Llama-3 emits <|python_tag|>; parser now supports it.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE) - assert flags["supports_tools"] is False + assert flags["supports_tools"] is True -def test_detect_safetensors_features_mistral_template_suppresses_tools(): - """Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it.""" +def test_detect_safetensors_features_mistral_template_keeps_tools_on(): + """Mistral emits [TOOL_CALLS]; parser now supports it.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3") flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_gemma4_template_keeps_tools_on(): + """Gemma 4 emits <|tool_call>; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-E2B-it-UD-MLX-4bit") + flags = _detect_safetensors_features(backend, GEMMA4_TEMPLATE) + assert flags["supports_tools"] is True + + +LLAMA3_2_BARE_JSON_TEMPLATE = """ +{%- if tools %} + {{- 'Given the following functions, respond with JSON for a function call.' }} + {{- 'Respond in the format {"name": function name, "parameters": dictionary}.' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if 'tool_calls' in message %} + {{- '{"name": "' + message.tool_calls[0].function.name + '", '}} + {{- '"parameters": ' + (message.tool_calls[0].function.arguments | tojson) + '}' }} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_llama3_2_bare_json_keeps_tools_on(): + """Llama-3.2 bare JSON is supported, so the pill stays enabled.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, LLAMA3_2_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +MINICPM5_ATTRIBUTE_TEMPLATE = """ +{%- if tools %} + {{- 'Available tools. Emit calls as ' }} + {{- 'value' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" + + +def test_detect_safetensors_features_attribute_function_form_keeps_tools_on(): + """The attribute form ```` must be whitelisted or the pill is wrongly suppressed.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "openbmb/MiniCPM-5") + flags = _detect_safetensors_features(backend, MINICPM5_ATTRIBUTE_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_unknown_format_suppresses_tools(): + """Tools advertised with no known marker must be suppressed.""" + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}<|im_start|>system\n" + "Emit tool calls as JSON-RPC notifications inside the response." + "<|im_end|>{%- endif %}" + ) + backend = SimpleNamespace(active_model_name = "custom/unknown-tool-format") + flags = _detect_safetensors_features(backend, tpl) assert flags["supports_tools"] is False def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on(): - """Sanity check: gate only suppresses non-Qwen formats.""" + """Sanity check: Qwen marker still flips supports_tools.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B") @@ -454,3 +532,130 @@ def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors(): assert flags["supports_tools"] is True assert flags["supports_reasoning"] is True assert flags["supports_preserve_thinking"] is True + + +# Templates advertising tools whose ``{"name":`` example is pretty-printed or JSON-escaped. +_WHITESPACE_BARE_JSON_TEMPLATE = ( + "{%- if tools %}\n" + "To call a tool, output JSON of the form:\n" + '{ "name" : "function_name", "parameters": { } }\n' + "{%- endif %}\n" + "{{ messages }}" +) +_ESCAPED_BARE_JSON_TEMPLATE = ( + "{%- if tools %}\n" + 'Respond with {\\"name\\": \\"fn\\", \\"parameters\\": {}}\n' + "{%- endif %}\n" + "{{ messages }}" +) +_TOOLS_ADVERTISED_NO_PARSEABLE_FORM = ( + "{%- if tools %}\nYou may use the available tools.\n{%- endif %}\n{{ messages }}" +) + + +def test_detect_safetensors_features_keeps_tools_for_pretty_printed_bare_json(): + # Pretty-printed bare-JSON (``{ "name" :``) keeps supports_tools: parser accepts the whitespace. + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _WHITESPACE_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_keeps_tools_for_escaped_bare_json(): + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _ESCAPED_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_drops_tools_when_no_parseable_form(): + # Negative control: tools advertised but no parser-recognised emission form -> pill dropped. + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _TOOLS_ADVERTISED_NO_PARSEABLE_FORM) + assert flags["supports_tools"] is False + + +def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json(): + # The {"function":...} bare-JSON alias keeps supports_tools, mirroring {"name":...}. + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}\n" + 'Respond with {"function": "fn", "parameters": {}}\n' + "{%- endif %}\n" + "{{ messages }}" + ) + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, tpl) + assert flags["supports_tools"] is True + + +# _sf_reasoning_prefill_mode gates the prefilled- extractor for enable_thinking models. +class TestSafetensorsReasoningPrefillGate: + # Qwen3-style template with the standard / markers. + _QWEN_TPL = "{% if enable_thinking %}{% endif %}......" + # gemma-style bespoke reasoning channel -- no standard markers. + _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" + + def _features(self, **over): + base = { + "supports_reasoning": True, + "reasoning_always_on": False, + "reasoning_style": "enable_thinking", + } + base.update(over) + return base + + def test_g1_enable_thinking_true(self): + # G1: Qwen3.5 template + explicit enable_thinking=True -> prefilled. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True + + def test_g2_enable_thinking_none_defaults_on(self): + # G2: default request (None) -> prefilled (Qwen3/GLM templates default on). + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True + + def test_g3_enable_thinking_false(self): + # G3: thinking explicitly off -> not prefilled. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False + + def test_g4_gpt_oss_reasoning_effort_excluded(self): + # G4: gpt-oss uses explicit tags via HarmonyTextStreamer -> normal mode. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_style = "reasoning_effort") + assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False + + def test_g5_enable_thinking_effort_included(self): + # G5: GLM-style enable_thinking_effort also prefills. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_style = "enable_thinking_effort") + assert _sf_reasoning_prefill_mode(feats, None, self._QWEN_TPL) is True + + def test_g6_non_reasoning_model(self): + # G6: no reasoning capability -> never prefilled. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(supports_reasoning = False, reasoning_style = None) + assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False + + def test_g7_reasoning_always_on(self): + # G7: hardcoded- template -> prefilled regardless of the flag. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_always_on = True) + assert _sf_reasoning_prefill_mode(feats, False, self._QWEN_TPL) is True + + def test_g8_gemma_bespoke_channel_excluded(self): + # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled (else the + # whole answer is swallowed as reasoning). Regression guard. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False + + def test_g9_missing_template_not_prefilled(self): + # G9: no template available -> conservative (not prefilled). + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, None) is False diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py new file mode 100644 index 0000000000..9158d1ad5e --- /dev/null +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Safetensors/MLX reasoning-block parity with GGUF. + +enable_thinking templates prefill an unclosed ````, so the stream must split the leading +text into ``reasoning_content`` deltas (per turn, monitor gets visible text only). Replays a copy +of ``sf_tool_stream``'s reasoning loop from routes/inference.py against synthetic events. +""" + +from __future__ import annotations + +import sys +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) + +from routes.inference import ( + _ResponsesReasoningExtractor, + _sf_reasoning_prefill_mode, + _strip_tool_xml_for_display, +) + + +def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict: + """Mirror sf_tool_stream's reasoning loop: diff cumulative snapshots, reset (flushing) on turn end.""" + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + reasoning_deltas: list[str] = [] + visible_deltas: list[str] = [] + monitor: list[str] = [] + tool_starts: list[dict] = [] + order: list[str] = [] # "reasoning" | "visible" | "tool_start" sequence + + def _flush(): + fr, fv = extractor.finish() + if fr: + reasoning_deltas.append(fr) + order.append("reasoning") + if fv: + visible_deltas.append(fv) + monitor.append(fv) + order.append("visible") + + for event in events: + etype = event["type"] + if etype == "status": + if not event["text"]: + _flush() + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + continue + if etype in ("tool_start", "tool_end"): + if etype == "tool_start": + _flush() + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + tool_starts.append(event) + order.append("tool_start") + continue + clean = _strip_tool_xml_for_display(event.get("text", ""), auto_heal_tool_calls = True) + new_text = clean[len(prev_text) :] + prev_text = clean + if not new_text: + continue + r, v = extractor.feed(new_text) + if r: + reasoning_deltas.append(r) + order.append("reasoning") + if v: + visible_deltas.append(v) + monitor.append(v) + order.append("visible") + _flush() + return { + "reasoning": "".join(reasoning_deltas), + "visible": "".join(visible_deltas), + "monitor": "".join(monitor), + "tool_starts": tool_starts, + "order": order, + } + + +def test_s1_plain_stream_splits_prefilled_reasoning(): + # S1: plain/MLX single turn -> reasoning delta + visible delta; monitor visible-only. + events = [ + {"type": "content", "text": "Let me compute 17*23"}, + {"type": "content", "text": "Let me compute 17*23 = 391The answer is 391."}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + assert out["reasoning"] == "Let me compute 17*23 = 391" + assert out["visible"] == "The answer is 391." + assert out["monitor"] == "The answer is 391." + assert "" not in out["reasoning"] and "" not in out["visible"] + + +def test_s2_reasoning_flushed_before_tool_start(): + # S2: reasoning streamed as reasoning_content, then flushed BEFORE tool_start. + events = [ + {"type": "content", "text": "I should search"}, + {"type": "content", "text": "I should search Sydney weather"}, + {"type": "tool_start", "tool_name": "web_search", "tool_call_id": "c0"}, + {"type": "tool_end", "tool_name": "web_search", "tool_call_id": "c0"}, + {"type": "status", "text": ""}, + {"type": "content", "text": "Found itSydney is 21C today."}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + # Both turns' reasoning surfaced, answer only from turn 2. + assert "I should search Sydney weather" in out["reasoning"] + assert "Found it" in out["reasoning"] + assert out["visible"] == "Sydney is 21C today." + assert out["monitor"] == "Sydney is 21C today." + # Ordering: the pre-tool reasoning is emitted before the tool_start. + assert out["order"].index("reasoning") < out["order"].index("tool_start") + + +def test_s3_extractor_resets_each_turn(): + # S3: multi-turn -> the two turns' reasoning are distinct (fresh extractor each). + events = [ + {"type": "content", "text": "turn1 thoughtspartial"}, + {"type": "status", "text": ""}, + {"type": "content", "text": "turn2 thoughtsfinal answer"}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + assert out["reasoning"] == "turn1 thoughtsturn2 thoughts" + assert out["visible"] == "partialfinal answer" + + +def test_s4_harmony_full_tags_normal_mode(): + # S4: gpt-oss / explicit-tag models use normal mode (prefilled=False). + events = [{"type": "content", "text": "reasoning herevisible answer"}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["reasoning"] == "reasoning here" + assert out["visible"] == "visible answer" + + +def test_s5_thinking_off_no_reasoning_deltas(): + # S5: thinking disabled -> not prefilled, no , all content is visible. + events = [{"type": "content", "text": "Just the plain answer, no thinking."}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["reasoning"] == "" + assert out["visible"] == "Just the plain answer, no thinking." + assert out["monitor"] == "Just the plain answer, no thinking." + + +_THINK_TPL = "...{% if enable_thinking %}{% endif %}......" + + +def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): + # GLM-5.2 enable_thinking_effort + reasoning_effort="none" disables thinking like + # enable_thinking=False, so prefilled must be OFF (else the answer is swallowed into reasoning). + feats = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "none") is False + # Thinking on (effort level or default) still prefills. + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "high") is True + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, None) is True + # An explicit enable_thinking=False also disables (unchanged). + assert _sf_reasoning_prefill_mode(feats, False, _THINK_TPL, "high") is False + # reasoning_always_on wins regardless of reasoning_effort. + always = {**feats, "reasoning_always_on": True} + assert _sf_reasoning_prefill_mode(always, None, _THINK_TPL, "none") is True + # Plain enable_thinking models (Qwen) have no "none" sentinel; unaffected. + plain = {"reasoning_style": "enable_thinking", "supports_reasoning": True} + assert _sf_reasoning_prefill_mode(plain, None, _THINK_TPL, "none") is True + + # End-to-end: with prefilled=False, a plain no- answer stays visible. + events = [{"type": "content", "text": "The capital of France is Paris."}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["visible"] == "The capital of France is Paris." + assert out["reasoning"] == "" + # The buggy prefilled=True path is what swallowed the whole answer (guard the delta). + swallowed = _replay_sf_reasoning_stream(events, prefilled = True) + assert swallowed["visible"] == "" + assert swallowed["reasoning"] == "The capital of France is Paris." diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 3f2d49f0dd..984d5f8ae9 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -138,6 +138,20 @@ class TestParser: assert len(result) == 1 assert "print('hi')" in result[0]["function"]["arguments"] + def test_xml_param_preserves_leading_indentation(self): + # Only the wrapping newline is trimmed, so code indentation survives. + text = ( + "\n" + " indented = 1\n" + " more\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "code": " indented = 1\n more" + } + def test_function_signal_inside_parameter_is_literal(self): text = ( "" @@ -189,6 +203,13 @@ class TestParser: text = 'before <|tool_call>call:terminal{command:"ls"} after' assert strip_tool_markup(text) == "before after" + def test_strip_named_mistral_call_consumes_trailing_eos(self): + # The named [TOOL_CALLS]name{json} shape must eat the optional trailing . + text = '[TOOL_CALLS]web_search{"query":"cats"}' + assert strip_tool_markup(text) == "" + text = '[TOOL_CALLS]web_search{"query":"cats"} and then' + assert strip_tool_markup(text) == " and then" + def test_strip_markup_unclosed_final(self): text = "before {partial" # final=True drops the trailing run. @@ -214,6 +235,376 @@ class TestParser: == "before " ) + def test_streaming_strip_keeps_prose_after_function_xml_with_literal_marker(self): + # A literal in a value is data: the strip closes at the REAL , keeping prose. + raw = ( + "pref " + 'print("") tail' + ) + assert strip_tool_markup_streaming(raw) == "pref tail" + # Streaming and final strip agree on the visible text (final also trims). + assert strip_tool_markup_streaming(raw) == strip_tool_markup(raw, final = True) + + def test_streaming_strip_drops_leading_magistral_reasoning(self): + # Magistral reasoning is a leading [THINK]...[/THINK] block; the streaming strip must drop it. + closed = "[THINK]Let me think. 2+2 is 4.[/THINK]The answer is 4." + assert strip_tool_markup_streaming(closed) == "The answer is 4." + assert strip_tool_markup_streaming(closed) == strip_tool_markup(closed, final = True) + # Unclosed mid-stream reasoning is held; cleaned text grows only after [/THINK]. + assert strip_tool_markup_streaming("[THINK]still thinking") == "" + assert strip_tool_markup_streaming("[THINK]r[/THINK]The") == "The" + assert strip_tool_markup_streaming("[THINK]r[/THINK]The answer") == "The answer" + # A non-leading [THINK] is ordinary prose, left untouched. + assert strip_tool_markup_streaming("hi [THINK] later") == "hi [THINK] later" + + +class TestParserMultiFormat: + """Shared-parser coverage: every family's emission maps to the same OpenAI shape.""" + + # Llama-3 + + def test_llama3_python_tag_dot_call(self): + # Llama-3 built-in tools: <|python_tag|>NAME.call(k="v", ...). + import json + + text = '<|python_tag|>brave_search.call(query="weather in Tokyo")' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "brave_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "weather in Tokyo"} + + def test_llama3_python_tag_dot_call_multi_arg(self): + import json + + text = "<|python_tag|>get_weather.call(" 'location="Tokyo", units="celsius", days=5)' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius", "days": 5} + + def test_llama3_python_tag_json_form(self): + import json + + text = '<|python_tag|>{"name":"web_search","parameters":{"query":"hi","n":5}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "hi", "n": 5} + + def test_llama3_python_tag_json_form_with_eom(self): + # Llama-3 emits <|eom_id|> after the JSON; must not break parsing. + import json + + text = '<|python_tag|>{"name":"python","parameters":{"code":"print(2+2)"}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"code": "print(2+2)"} + + def test_llama3_strip_markup_final(self): + text = '<|python_tag|>brave_search.call(query="x")' + assert strip_tool_markup(text, final = True) == "" + + # Llama-3.2 bare JSON ``custom_tools`` + + def test_llama3_2_bare_json_parameters(self): + # Llama-3.2-Instruct emits bare JSON directly as content, no <|python_tag|> prefix. + import json + + text = '{"name":"web_search","parameters":{"query":"Tokyo weather"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "Tokyo weather"} + + def test_llama3_2_bare_json_arguments_key(self): + import json + + text = '{"name":"add","arguments":{"a":1,"b":2}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"a": 1, "b": 2} + + def test_llama3_2_bare_json_multi_call(self): + # Llama-3 may chain calls with "; " per training template. + text = '{"name":"a","parameters":{}}; {"name":"b","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_llama3_2_bare_json_with_eom_sentinel(self): + text = '{"name":"x","parameters":{"y":1}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_leading_sentinel_skipped(self): + # Sometimes prior <|eot_id|> leaks into the next turn. + text = '<|eot_id|>{"name":"x","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_plain_prose_does_not_fire(self): + # Defensive: must NOT fire on plain assistant prose. + text = "Hello world, how are you today?" + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_embedded_in_prose_does_not_fire(self): + # Defensive: JSON embedded in prose must NOT fire (content must START with `{`). + text = 'The tool result was: {"name":"foo"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_name_does_not_fire(self): + text = '{"result":"ok","data":[1,2,3]}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_args_does_not_fire(self): + text = '{"name":"x"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_args_not_dict_does_not_fire(self): + text = '{"name":"x","parameters":42}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_parameters_does_not_fire(self): + # Llama-3 spec: parameters must be a dict; a string value must NOT trigger. + text = '{"name":"foo","parameters":"this is a sentence"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_not_json_does_not_fire(self): + # OpenAI arguments may be a JSON-string of a dict, but a plain non-JSON string must not pass. + text = '{"name":"foo","arguments":"not json"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_json_dict_fires(self): + # OpenAI shape: arguments is a JSON-encoded string of a dict. + text = '{"name":"foo","arguments":"{\\"q\\":\\"x\\"}"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + # arguments stays as the original JSON-string. + assert result[0]["function"]["arguments"] == '{"q":"x"}' + + def test_llama3_2_bare_json_string_arguments_json_non_dict_does_not_fire(self): + # JSON-string that parses to a list / scalar / null must NOT fire. + for bad in ( + '{"name":"foo","arguments":"[1,2,3]"}', + '{"name":"foo","arguments":"\\"plain\\""}', + '{"name":"foo","arguments":"null"}', + '{"name":"foo","arguments":"42"}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + + # Mistral pre-v11 + + def test_mistral_pre_v11_array(self): + import json + + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"query":"hello"},"id":"abc"}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + # Mistral provides its own id; preserve it. + assert result[0]["id"] == "abc" + assert json.loads(result[0]["function"]["arguments"]) == {"query": "hello"} + + def test_mistral_array_parameters_key_alias(self): + import json + + # Array object keyed on parameters (not arguments) must keep its payload. + text = '[TOOL_CALLS] [{"name":"get_weather","parameters":{"city":"Paris"}}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Paris"} + + def test_mistral_pre_v11_array_multi(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"x":1},"id":"id1"},' + '{"name":"b","arguments":{"y":2},"id":"id2"}]' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_mistral_pre_v11_unclosed_array(self): + # Closing ] truncated: parser must heal off individual objects. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"},"id":"id"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # Mistral v11+ + + def test_mistral_v11_single(self): + # Magistral / Mistral Small 3.1: bare name{json} after trigger. + import json + + text = '[TOOL_CALLS]add{"a":3.5,"b":4}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 3.5, "b": 4} + + def test_mistral_v11_parallel(self): + # v11+ parallel: [TOOL_CALLS]a{...}[TOOL_CALLS]b{...}. + text = '[TOOL_CALLS]add{"a":1}[TOOL_CALLS]sub{"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "add" + assert result[1]["function"]["name"] == "sub" + + def test_mistral_v11_with_args_marker(self): + # Ministral / Mistral Large 3: [TOOL_CALLS]name[ARGS]{json}. + import json + + text = '[TOOL_CALLS]add[ARGS]{"a":1,"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 1, "b": 2} + + def test_mistral_strip_markup_v11(self): + text = '[TOOL_CALLS]add{"a":1}' + assert strip_tool_markup(text, final = True) == "" + + def test_mistral_call_id_form(self): + # Mistral Small 3.2: the [CALL_ID] segment must be skipped, not treated as a stop (llama.cpp test-chat.cpp:4785). + import json + + text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "special_function" + assert json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_mistral_call_id_form_parallel(self): + text = ( + '[TOOL_CALLS]special_function[CALL_ID]000000001[ARGS]{"arg1": 1}' + "[TOOL_CALLS]special_function_with_opt[CALL_ID]000000002" + '[ARGS]{"arg1": 1, "arg2": 2}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "special_function" + assert result[1]["function"]["name"] == "special_function_with_opt" + + def test_mistral_call_id_form_stripped(self): + text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' + assert strip_tool_markup(text, final = True) == "" + + def test_mistral_think_reasoning_ignored(self): + # A [TOOL_CALLS] inside [THINK]...[/THINK] is reasoning; only the call after [/THINK] counts (llama.cpp test-chat.cpp:2285). + import json + + text = ( + '[THINK]Let me think about [TOOL_CALLS]fake[ARGS]{"x":1} ' + 'and more[/THINK][TOOL_CALLS]real_fn[ARGS]{"y":2}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "real_fn" + assert json.loads(result[0]["function"]["arguments"]) == {"y": 2} + + def test_mistral_think_reasoning_no_real_call(self): + # Reasoning that mentions a call but emits none after [/THINK] yields no calls. + text = '[THINK]I might call [TOOL_CALLS]fake[ARGS]{"x":1}[/THINK]Done.' + assert parse_tool_calls_from_text(text) == [] + + def test_mistral_think_literal_in_argument_preserved(self): + # A literal [THINK] inside a real tool argument must not be stripped or corrupt the parse. + import json + + text = '[TOOL_CALLS]search[ARGS]{"q":"explain the [THINK] token"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"q": "explain the [THINK] token"} + + # Gemma 4 + + def test_gemma4_simple_call(self): + import json + + text = ( + "<|tool_call>call:get_weather{" + 'location:<|"|>Tokyo<|"|>,units:<|"|>celsius<|"|>}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius"} + + def test_gemma4_with_primitives(self): + import json + + text = ( + "<|tool_call>call:set_pref{" + "enabled:true,attempts:5,threshold:1.5,nickname:null}" + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"enabled": True, "attempts": 5, "threshold": 1.5, "nickname": None} + + def test_gemma4_nested_args(self): + # Gemma 4 nests dicts / lists with bare keys and <|"|> strings. + import json + + text = ( + "<|tool_call>call:search{" + 'query:<|"|>foo<|"|>,filters:{site:<|"|>example.com<|"|>,recent:true},' + 'tags:[<|"|>a<|"|>,<|"|>b<|"|>]}' + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "foo" + assert args["filters"] == {"site": "example.com", "recent": True} + assert args["tags"] == ["a", "b"] + + def test_gemma4_multi_call(self): + text = "<|tool_call>call:a{x:1}<|tool_call>call:b{y:2}" + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_gemma4_unclosed_does_not_raise(self): + # Truncated mid-stream; must not raise. + text = '<|tool_call>call:foo{x:<|"|>bar<|"|>' + result = parse_tool_calls_from_text(text) + assert isinstance(result, list) + + def test_gemma4_strip_markup_final(self): + text = "<|tool_call>call:foo{x:1}" + assert strip_tool_markup(text, final = True) == "" + + # Cross-format sentinels + + def test_all_markers_in_tool_xml_signals(self): + # Streaming buffer wakes up on every emission marker. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + for marker in ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + ): + assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + + def test_has_tool_signal_for_all_formats(self): + assert has_tool_signal('<|python_tag|>brave_search.call(q="x")') + assert has_tool_signal('[TOOL_CALLS] [{"name":"x"}]') + assert has_tool_signal('[TOOL_CALLS]add{"a":1}') + assert has_tool_signal("<|tool_call>call:foo{}") + # ──────────────────────────────────────────────────────────────────── # run_safetensors_tool_loop @@ -347,6 +738,130 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success(): assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) +def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): + # A late unclosed heals only with Auto-Heal on; off, it must not execute. + prose = "Sure, let me look that up for you right now. " + incomplete = '{"name":"web_search","arguments":{"query":"weather in Sydney"}}' + + loop_off, exec_off = _make_loop( + turns = [[prose, incomplete], ["Final answer."]], + exec_results = ["RESULT"], + auto_heal_tool_calls = False, + max_tool_iterations = 3, + ) + events_off = _collect_events(loop_off) + assert exec_off.calls == [], "disabled Auto-Heal must not execute a healed incomplete call" + assert not [e for e in events_off if e.get("type") == "tool_start"] + + loop_on, exec_on = _make_loop( + turns = [[prose, incomplete], ["Final answer."]], + exec_results = ["RESULT"], + auto_heal_tool_calls = True, + max_tool_iterations = 3, + ) + _collect_events(loop_on) + assert exec_on.calls == [("web_search", {"query": "weather in Sydney"})], exec_on.calls + + +def test_bare_json_tool_call_is_not_streamed_as_content(): + # Llama-3.2 bare form carries no XML signal: BUFFER until the object closes, never leak the JSON. + bare = '{"name":"web_search","parameters":{"query":"cats"}}' + loop, exec_fn = _make_loop( + turns = [[bare], ["Here are the results."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t or "web_search" in t for t in contents), contents + assert any("Here are the results." in t for t in contents) + + +def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(): + # Markerless JSON whose "name" is not an enabled tool must be shown, not dropped. + answer = '{"name":"Alice","parameters":{"age":30}}' + loop, exec_fn = _make_loop(turns = [[answer]], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents + + +def test_bare_json_tool_call_split_across_chunks_is_not_streamed(): + # Same as above but the bare object arrives split mid-key, held across chunks until it balances. + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_', 'search","parameters":{"query":"cats"}}'], + ["Done."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t or "web_search" in t for t in contents), contents + + +def test_leading_json_answer_is_not_dropped(): + # A leading {...} that is NOT a call must still surface; the hold only delays it. + obj = '{"answer": 42, "note": "done"}' + loop, exec_fn = _make_loop( + turns = [[obj]], + exec_results = [], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert any('"answer"' in t for t in contents), contents + + +def _reprompt_loop(*, auto_heal_tool_calls): + """Drive one restricted tool with an intent-only first turn to exercise the nudge; returns conversations and events.""" + captured: list[list] = [] + + def fake_single_turn(messages, active_tools = None): + captured.append(list(messages)) + if len(captured) == 1: + yield "I'll search for that now." # forward-looking intent, no call + else: + yield "Final answer." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "find X"}], + tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], + execute_tool = exec_fn, + auto_heal_tool_calls = auto_heal_tool_calls, + max_tool_iterations = 3, + ) + ) + return captured, events + + +def test_reprompt_names_only_active_tools_not_hardcoded(): + # The nudge must name the tools actually enabled, not hardcoded web_search/python. + captured, _events = _reprompt_loop(auto_heal_tool_calls = True) + assert len(captured) >= 2, "intent prose should have triggered a re-prompt turn" + reprompt = captured[1][-1] + assert reprompt["role"] == "user" + assert "search_knowledge_base" in reprompt["content"] + assert "web_search" not in reprompt["content"] + assert "python" not in reprompt["content"] + + +def test_reprompt_suppressed_when_auto_heal_disabled(): + # With Auto-Heal off the nudge stays silent for GGUF parity, so only the initial generation runs. + captured, events = _reprompt_loop(auto_heal_tool_calls = False) + assert len(captured) == 1, captured + contents = [e["text"] for e in events if e["type"] == "content"] + assert any("search for that" in t for t in contents) + + class TestLoopBasic: def test_plain_answer(self): # No tool XML; loop should yield content then status="". @@ -406,6 +921,85 @@ class TestLoopBasic: contents = [e for e in events if e["type"] == "content"] assert "Result: 1" in contents[-1]["text"] + def test_llama3_python_tag_form(self): + # The loop must recognise Llama-3's <|python_tag|> marker, drain the turn, and execute the call. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|python_tag|>web_search.call(", + 'query="weather in Tokyo"', + ")", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in Tokyo"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_llama3_bare_json_form_fires_tool(self): + # Llama-3.1/3.2 bare-JSON calls carry no XML signal; the safety-net parse must still fire + # the tool. Regression for the has_tool_signal gate that dropped these. + loop, exec_fn = _make_loop( + turns = [ + ['{"name": "web_search", "parameters": {"query": "weather in SF"}}'], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 18C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in SF"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_mistral_pre_v11_form(self): + # Pre-v11 Mistral emission: [TOOL_CALLS] [{...}]. + loop, exec_fn = _make_loop( + turns = [ + [ + '[TOOL_CALLS] [{"name":"web_search",', + '"arguments":{"query":"hi"},"id":"abc"}]', + ], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + # Mistral-provided ids must propagate to tool_start events. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "abc" + + def test_mistral_v11_form(self): + # v11+ Mistral emission: bare name{json} after the trigger. + loop, exec_fn = _make_loop( + turns = [ + ['[TOOL_CALLS]web_search{"query":"hi"}'], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + + def test_gemma4_form(self): + # Gemma 4 emission: <|tool_call>call:NAME{...}. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_call>call:web_search{", + 'query:<|"|>weather<|"|>', + "}", + ], + ["sunny"], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + def test_render_html_emits_provisional_tool_start(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -765,6 +1359,55 @@ class TestLoopBehaviour: assert len(duplicate_nudges) == 1 assert captured_tool_names[2] == ["web_search", "python"] + def test_duplicate_noop_does_not_consume_budget_at_small_cap(self): + # A duplicate no-op turn must NOT spend the tool budget: only turns that execute a tool + # count (GGUF parity), so a distinct call can still follow at max_tool_iterations=2. + captured_tool_names: list[list[str]] = [] + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"python","arguments":{"code":"print(1)"}}'], + ["final"], + ] + ) + + def fake_single_turn(messages, active_tools = None): + captured_tool_names.append( + [ + tool["function"]["name"] + for tool in (active_tools or []) + if tool.get("function", {}).get("name") + ] + ) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-result", "python-result"]) + _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + + # Both distinct tools execute; the repeated call in between did not cost a slot. + assert exec_fn.calls == [ + ("web_search", {"query": "x"}), + ("python", {"code": "print(1)"}), + ] + # The turn after the duplicate still offered tools (budget not yet spent). + assert captured_tool_names[2] == ["web_search", "python"] + def test_repeated_duplicate_noop_transitions_to_final_attempt(self): captured_tool_names: list[list[str]] = [] turns = iter( @@ -953,6 +1596,234 @@ class TestLoopBehaviour: assert "boom" in tool_end["result"] +class TestLoopRePrompt: + """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``_MAX_REPROMPTS`` extra slots.""" + + def test_intent_signal_triggers_reprompt(self): + # Turn 1: intent signal, no tool call. + # Turn 2 (re-prompt): proper tool call -> executes. + # Turn 3: final answer. + loop, exec_fn = _make_loop( + turns = [ + ["Let me search for that."], + [ + '{"name":"web_search","arguments":' + '{"query":"sky color"}}' + ], + ["The sky is blue."], + ], + exec_results = ["Blue (Rayleigh scattering)"], + ) + events = _collect_events(loop) + # web_search must have been called once (after the re-prompt). + assert exec_fn.calls == [("web_search", {"query": "sky color"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "blue" in contents[-1]["text"].lower() + + def test_intent_signal_without_tools_does_not_reprompt(self): + # Same intent signal but no tools enabled -- must NOT re-prompt. + loop, exec_fn = _make_loop( + turns = [["Let me think about that for a moment."]], + exec_results = [], + ) + # _make_loop hard-codes three tools; rebuild without tools. + from core.inference.safetensors_agentic import run_safetensors_tool_loop + + def _gen(_messages): + yield "Let me think about that for a moment." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = exec_fn, + ) + ) + assert exec_fn.calls == [] + contents = [e for e in events if e["type"] == "content"] + assert contents and "think" in contents[-1]["text"].lower() + + def test_direct_answer_does_not_trigger_reprompt(self): + # Plain answer with no intent words: do NOT re-prompt. + loop, exec_fn = _make_loop( + turns = [["4"]], + exec_results = [], + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = [e for e in events if e["type"] == "content"] + assert contents and contents[-1]["text"].strip() == "4" + + def test_max_reprompts_capped_at_three(self): + # Model keeps stalling with intent -- after 3 re-prompts the loop must give up. + turns = [["Let me search for that."]] * 6 # well over the cap + loop, exec_fn = _make_loop( + turns = turns, + exec_results = [], + ) + events = _collect_events(loop, max_events = 500) + # No tool ever ran, but the loop terminated cleanly. + assert exec_fn.calls == [] + statuses = [e for e in events if e["type"] == "status"] + assert statuses and statuses[-1]["text"] == "" + + def test_short_intent_below_buffer_threshold_triggers_reprompt(self): + # Short emission that never exits BUFFERING must still trigger the intent re-prompt. + loop, exec_fn = _make_loop( + turns = [ + ["Let me check."], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["found"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "x"})] + + def test_reprompt_does_not_consume_tool_budget(self): + # max_tool_iterations=1: the re-prompt must not eat the slot, so the real call still runs. + loop, exec_fn = _make_loop( + turns = [ + # 1. Intent stall (re-prompt 1/3). + ["Let me search for that."], + # 2. Real tool call (uses the budget slot). + ['{"name":"web_search","arguments":{"query":"weather"}}'], + # 3. Budget exhausted -> nudged final answer. + ["Final: it is sunny"], + ], + exec_results = ["sunny"], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + +class TestLoopCanonicalHealKey: + """Per-tool canonical heal key (``code``/``command``/``query``), mirroring GGUF.""" + + def test_python_bare_string_heals_to_code(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"python","arguments":"print(1)"}' ""], + ["done"], + ], + exec_results = ["1\n"], + ) + events = _collect_events(loop) + # The bare string must heal to {"code": ...}, not {"query": ...}, so the python sandbox runs it. + assert exec_fn.calls == [("python", {"code": "print(1)"})] + + def test_terminal_bare_string_heals_to_command(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"terminal","arguments":"ls -la"}' ""], + ["done"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("terminal", {"command": "ls -la"})] + + def test_unknown_tool_bare_string_heals_to_query(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":"hello"}' ""], + ["ok"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hello"})] + + +class TestGGUFSafetensorsHealingParity: + """Pin GGUF vs safetensors/MLX loop parity so a regression on either side breaks CI.""" + + def test_gguf_imports_shared_signal_markers(self): + # The GGUF BUFFERING machine must wake on every shared emission marker, else calls slip past as prose. + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert "_SHARED_TOOL_XML_SIGNALS" in src, ( + "GGUF agentic loop must reuse the shared TOOL_XML_SIGNALS " + "tuple so it wakes on all five emission formats" + ) + + def test_gguf_uses_shared_strip_helper(self): + # The GGUF stream-cleanup must delegate to the shared strip_tool_markup for every family. + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert ( + "_shared_strip_tool_markup" in src + ), "GGUF stream cleanup must delegate to the shared strip_tool_markup helper" + + def test_gguf_uses_canonical_heal_keys(self): + # GGUF and safetensors heal a bare-string argument to the same canonical key via the shared coerce_tool_arguments. + from core.inference.tool_loop_controller import ( + _CANONICAL_HEAL_ARG, + coerce_tool_arguments, + ) + + assert _CANONICAL_HEAL_ARG["python"] == "code" + assert _CANONICAL_HEAL_ARG["terminal"] == "command" + assert coerce_tool_arguments("print(1)", heal = True, tool_name = "python").arguments == { + "code": "print(1)" + } + assert coerce_tool_arguments("ls -la", heal = True, tool_name = "terminal").arguments == { + "command": "ls -la" + } + assert coerce_tool_arguments("weather", heal = True, tool_name = "web_search").arguments == { + "query": "weather" + } + + def test_intent_regex_matches_same_phrases_as_gguf(self): + # The intent re-prompt regex must match the SAME phrases on both backends. + from core.inference.llama_cpp import _INTENT_SIGNAL as gguf_re + from core.inference.safetensors_agentic import ( + _INTENT_SIGNAL as sf_re, + ) + + for phrase in ( + "I'll search for that", + "I will look it up", + "Let me check", + "I am going to call the tool", + "First, I will explore", + "Here's my plan", + "Now I need to call web_search", + ): + assert gguf_re.search(phrase), f"GGUF missed {phrase!r}" + assert sf_re.search(phrase), f"safetensors missed {phrase!r}" + + for plain in ( + "4", + "Hello!", + "The sky is blue.", + "I can help with that.", + "I should mention", + "Let's go.", + # Negated intent is a refusal, not a plan: neither backend may re-prompt on it. + "I will not search the web for that.", + "I'll never call that tool.", + ): + assert not gguf_re.search(plain), f"GGUF wrongly fired on {plain!r}" + assert not sf_re.search(plain), f"safetensors wrongly fired on {plain!r}" + + def test_max_reprompts_equal_on_both_backends(self): + from core.inference.llama_cpp import _MAX_REPROMPTS as gguf_cap + from core.inference.safetensors_agentic import _MAX_REPROMPTS as sf_cap + assert gguf_cap == sf_cap == 3 + + class TestLoopControl: def test_cancel_event_breaks_loop(self): cancel = threading.Event() @@ -1407,5 +2278,358 @@ class TestGptOssNameDetection: assert is_gpt_oss_model_name(cast(str, None)) is False +# Routes-level python_tag strip (multi-line; stop on next sentinel) +class TestRoutesPythonTagStrip: + """``_TOOL_XML_RE`` must consume multi-line code, embedded JSON, and bare ``<`` (earlier ``[^\n<]*`` / ``[^\n]*`` revisions leaked tails); the streaming route-level strip is the regression-prone path.""" + + def _strip(self, text: str) -> str: + # Import inside the test so a routes-module import error doesn't fail collection. + from routes.inference import _strip_tool_xml + return _strip_tool_xml(text) + + def test_single_line_python_tag_stripped(self): + # Floor: the original 5620 single-line behaviour still works. + text = '<|python_tag|>brave_search.call(query="weather")' + assert self._strip(text) == "" + + def test_python_tag_with_less_than_in_code(self): + # 5615 regression: a literal < inside code must NOT terminate the strip early. + text = '<|python_tag|>python.call(code="if x < 10: pass")' + assert self._strip(text) == "" + + def test_python_tag_multiline_code_stripped(self): + # 5620 round-1 regression: multi-line code's second line leaked. + text = '<|python_tag|>python.call(code="line1\nline2\nline3")' + assert self._strip(text) == "" + + def test_python_tag_multiline_with_less_than(self): + # Combined: multi-line code AND literal < in code. + text = ( + '<|python_tag|>python.call(code="for i in range(10):\n' + " if i < 5:\n" + ' print(i)")' + ) + assert self._strip(text) == "" + + def test_python_tag_stops_at_eom_sentinel(self): + # Strip stops at the next Llama-3 <| sentinel so trailing assistant content survives. + text = '<|python_tag|>python.call(code="multi\nline")' "<|eom_id|>final answer text" + assert self._strip(text) == "<|eom_id|>final answer text" + + def test_python_tag_stops_at_eot_sentinel(self): + text = '<|python_tag|>brave_search.call(query="x")' "<|eot_id|>after" + assert self._strip(text) == "<|eot_id|>after" + + def test_python_tag_json_form_multiline_stripped(self): + # The JSON form of python_tag with newlines inside string args. + text = '<|python_tag|>{"name":"python","parameters":{"code":"a = 1\nb = 2\nprint(a+b)"}}' + assert self._strip(text) == "" + + def test_python_tag_with_eom_then_trailing_python_tag(self): + # Two python_tag emissions back-to-back across a sentinel: both strip independently. + text = ( + '<|python_tag|>brave_search.call(query="a")' + "<|eom_id|>" + '<|python_tag|>python.call(code="x=1")' + ) + # <|eom_id|> between the two strips remains; both python_tag blocks are consumed. + assert self._strip(text) == "<|eom_id|>" + + +# Robustness fixes uncovered while validating against vLLM / sglang. +class TestParserRobustness: + def test_tool_call_json_accepts_parameters_key(self): + # Hermes wrapper using parameters instead of arguments; this path now accepts both keys. + import json + + text = "\n" '{"name": "search", "parameters": {"q": "ramen"}}\n' "" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "search" + assert json.loads(result[0]["function"]["arguments"]) == {"q": "ramen"} + + def test_function_xml_attribute_form(self): + # MiniCPM-5 / MiniMax-M2 attribute syntax: v. + import json + + text = '' 'Tokyo' "" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_xml_attribute_form_multi_param(self): + import json + + text = ( + '' + 'Tokyo' + 'celsius' + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"city": "Tokyo", "unit": "celsius"} + + def test_function_xml_legacy_equals_form_still_works(self): + # Regression guard: the old v syntax must keep parsing after the regex broadening. + import json + + text = "Tokyo" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_attribute_form_has_tool_signal(self): + # The standalone form must flip the streaming buffer, else the call is dropped. + assert has_tool_signal('') is True + + def test_function_attribute_form_strip_markup(self): + # The attribute form must also be stripped from displayed text, like . + text = 'result X' + assert strip_tool_markup(text, final = True) == "result" + + def test_llama3_chat_template_round_trip(self): + # Llama-3.x prefixes assistant turns with <|start_header_id|>...<|end_header_id|>; the + # sentinel-strip must reach past the role label to the JSON body, else history calls drop. + import json + + text = ( + "<|start_header_id|>assistant<|end_header_id|>\n\n" + '{"name": "get_weather", "parameters": {"city": "Tokyo"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_llama3_round_trip_all_roles(self): + # Same logic must work for every role the chat template inserts. + import json + for role in ("assistant", "user", "system", "tool", "ipython"): + text = ( + f"<|start_header_id|>{role}<|end_header_id|>\n\n" + '{"name": "f", "parameters": {"x": 1}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1, f"failed for role={role}" + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + + def test_llama3_round_trip_with_eot_prefix(self): + # Prior turn closes with <|eot_id|>, then the new header opens; both sentinels + role must be consumed. + import json + + text = ( + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + '{"name": "f", "parameters": {}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "f" + + def test_function_xml_followed_by_prose(self): + # Body must terminate at even without a wrapper, else prose leaks into the value. + import json + + text = ( + "" + "Tokyo" + "\n\nHere is what I found." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_attribute_xml_followed_by_prose(self): + # Same expectation for the MiniCPM-5 attribute form. + import json + + text = ( + '' + 'Tokyo' + "\n\nLet me know if you need anything else." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + +def test_truncated_bare_json_at_eof_is_not_leaked(): + # Stream ends mid bare-JSON: the held fragment must be dropped at EOF, not flushed as content. + loop, _exec = _make_loop( + turns = [['{"name":"web_search","parameters":{"query":"weather in S']], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t for t in contents), contents + + +def test_oversized_bare_json_call_is_not_leaked_and_executes(): + # A bare-JSON call exceeding _MAX_BARE_JSON_BUFFER must DRAIN, not stream the prefix, and still execute. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + full = '{"name":"python","parameters":{"code":"' + big + '"}}' + chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] + loop, exec_fn = _make_loop(turns = [chunks, ["done"]], exec_results = ["OK"], max_tool_iterations = 2) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any(t.lstrip().startswith('{"name') for t in contents), contents[:1] + assert exec_fn.calls and exec_fn.calls[0][0] == "python" + assert len(exec_fn.calls[0][1].get("code", "")) > _MAX_BARE_JSON_BUFFER + + +def test_oversized_plain_json_answer_still_streams(): + # A giant plain JSON answer (no "name" key) is NOT a call and must still stream. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + full = '{"result":"' + big + '"}' + chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] + loop, _exec = _make_loop(turns = [chunks], max_tool_iterations = 1) + events = _collect_events(loop) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert '"result"' in contents + + +def test_oversized_disabled_name_json_answer_still_streams(): + # A giant still-open JSON answer whose "name" is NOT an enabled tool must stream, not drain. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes + chunks = [answer[i : i + 2000] for i in range(0, len(answer), 2000)] + loop, exec_fn = _make_loop(turns = [chunks], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents[:80] + + +def test_truncated_disabled_name_json_is_shown_at_eof(): + # A truncated JSON answer whose name is not an enabled tool must be shown at EOF. + truncated = '{"name":"Alice","parameters":{"age":' + loop, exec_fn = _make_loop(turns = [[truncated]], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents + + +def test_truncated_plain_json_with_nested_enabled_name_is_visible(): + # A truncated answer with only a NESTED "name" must be shown: the gate uses the TOP-LEVEL name. + loop, exec_fn = _make_loop( + turns = [['{"result":{"name":"web_search","age":']], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert '"result"' in contents and "web_search" in contents, contents + + +def test_bare_json_call_not_replayed_in_next_turn_content(): + # After a bare-JSON call executes, the next-turn assistant content must not contain the raw call. + captured: list[list[dict]] = [] + exec_fn = FakeExecuteTool(["RESULT"]) + + def st(messages, active_tools = None): + captured.append([dict(m) for m in messages]) + if len(captured) == 1: + yield '{"name":"web_search","parameters":{"query":"cats"}}' + else: + yield "Found." + + _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + ) + assert len(captured) >= 2, captured + asst = [m for m in captured[1] if m.get("role") == "assistant"] + assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst + + if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): + # With Auto-Heal OFF a truncated enabled-name bare-JSON fragment stays visible; with it ON, suppressed. + trunc = '{"name":"web_search","parameters":{"query":"weather' + off, exec_off = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = False) + events_off = _collect_events(off) + assert exec_off.calls == [], exec_off.calls + contents_off = "".join(e["text"] for e in events_off if e["type"] == "content") + assert "web_search" in contents_off, contents_off + + on, exec_on = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = True) + events_on = _collect_events(on) + assert exec_on.calls == [], exec_on.calls + contents_on = "".join(e["text"] for e in events_on if e["type"] == "content") + assert "web_search" not in contents_on, contents_on + + +def test_looks_like_enabled_bare_json_accepts_function_alias(): + # The buffering gate must recognise the "function" bare-JSON alias, so it is buffered, not streamed. + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + + enabled = {"web_search"} + assert _looks_like_enabled_bare_json( + '{"function":"web_search","parameters":{"q":"x"}}', enabled + ) + # A non-tool "function" value is an ordinary JSON answer -> not gated. + assert not _looks_like_enabled_bare_json('{"function":"Alice","parameters":{}}', enabled) + + +class TestFalseAlarmMarkerProse: + def test_leading_marker_prose_streams_intact(self): + # An answer starting with a literal marker is a false alarm: the full prose must reach the client. + text = "[TOOL_CALLS] is the Mistral tool marker. More prose after." + loop, exec_fn = _make_loop(turns = [[text]]) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert texts and texts[-1] == text + + def test_chained_bare_json_calls_not_replayed_in_history(self): + # Both chained calls execute; the next-turn history must not contain the second call's raw JSON. + chained = ( + '{"name":"web_search","parameters":{"q":"first"}};' + '{"name":"python","parameters":{"code":"x"}}' + ) + convs = [] + turn_iter = iter([[chained], ["Final answer."]]) + + def gen(messages, active_tools = None): + convs.append([dict(m) for m in messages]) + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(["r1", "r2"]) + loop = run_safetensors_tool_loop( + single_turn = gen, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + ) + _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search", "python"] + assistant = next(m for m in convs[1] if m["role"] == "assistant") + assert '"python"' not in (assistant.get("content") or "") diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 39fdd151be..7664126d91 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -102,6 +102,22 @@ class TestFunctionStyleTrailingText: text = "weather london" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + def test_attribute_form_literal_close_tag_is_preserved(self): + # Attribute form ends at the LAST , so a literal close inside code survives. + text = ( + '' + 'print("")' + " all done" + ) + call = _only(text) + assert call == {"name": "python", "arguments": {"code": 'print("")'}} + + def test_closed_zero_param_attribute_call_is_accepted_in_strict_mode(self): + # A closed zero-param call is valid; strict mode must not treat it as truncated. + assert _only('') == {"name": "ping", "arguments": {}} + # A no-arg call that never closes is still rejected as truncated. + assert parse_tool_calls_from_text('', allow_incomplete = False) == [] + class TestParityWithJsonStyle: def test_json_tool_call_with_trailing_prose_is_accepted(self): @@ -176,6 +192,37 @@ class TestGemmaNativeStyle: } +class TestLlama3PythonTagStrict: + def test_closed_dot_call_is_accepted(self): + text = '<|python_tag|>get_weather.call(location="Tokyo")' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + assert json.loads(calls[0]["function"]["arguments"]) == {"location": "Tokyo"} + + def test_truncated_dot_call_is_rejected(self): + # No closing paren (depth > 0 at EOF): truncated, reject in strict mode. + text = '<|python_tag|>get_weather.call(location="Tokyo"' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # Auto-Heal still recovers it. + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestMistralArrayStrict: + def test_closed_array_is_accepted(self): + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"}}]' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + def test_unclosed_array_is_rejected(self): + # Missing the closing ]; strict mode must not heal it. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"}}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # Auto-Heal still recovers the object by hand. + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + class TestHealingPathUnaffected: def test_auto_heal_still_repairs_unclosed_function(self): text = "cats" @@ -197,3 +244,822 @@ class TestHealingPathUnaffected: assert text[span[0] : span[1]] == ( "cats" ) + + def test_wrapperless_fallback_calls_carry_spans(self): + # The wrapperless fallback must report spans so consumers strip exactly the markup. + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + closed = "before cats after" + calls, spans = parse_with_spans(closed, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "cats"} + (span,) = spans + assert closed[span[0] : span[1]] == ( + "cats" + ) + + healed = "x dogs" + calls, spans = parse_with_spans(healed, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "dogs"} + (span,) = spans + assert healed[span[0] : span[1]] == "dogs" + + +class TestParserLinearity: + """Llama-3 ``.call`` kwargs and Mistral-array healing must stay linear (a regex-per-offset blew up on long truncated bodies).""" + + def test_llama3_unterminated_call_arg_is_linear(self): + import time + + text = '<|python_tag|>upload.call(data="' + "A" * 200_000 # no closing quote/paren + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_llama3_huge_wordrun_call_arg_is_linear(self): + import time + + text = "<|python_tag|>upload.call(" + "a" * 200_000 # giant word run, no '=' + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_mistral_unclosed_array_open_braces_is_linear(self): + import time + + text = "[TOOL_CALLS] [" + "{" * 200_000 # unclosed array, all open braces + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_llama3_call_kwargs_still_parse(self): + text = '<|python_tag|>do.call(s="hi 😀", n=42, f=1.5, b=true, z=null)' + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == { + "s": "hi 😀", + "n": 42, + "f": 1.5, + "b": True, + "z": None, + } + + def test_llama3_call_scientific_notation_args_parse(self): + # Scientific notation must decode as float (the old regex truncated 1e-3 -> 1). + text = "<|python_tag|>calc.call(x=1e-3, y=-2E+4, z=0.5e2, n=42)" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"x": 1e-3, "y": -2e4, "z": 50.0, "n": 42} + assert isinstance(args["n"], int) and isinstance(args["x"], float) + + def test_mistral_unclosed_array_recovers_top_level_objects(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"k":1}},' + '{"name":"b","arguments":{"j":2}}' # missing closing ] + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["a", "b"] + + +class TestLlamaBuiltinChainAndNesting: + """Llama-3 ``.call`` built-ins: ``; `` chaining and nested-tag isolation.""" + + def test_semicolon_chained_builtin_calls_all_parse(self): + # Only the first call is anchored to <|python_tag|>; the rest chain via ';'. + text = "<|python_tag|>alpha.call(x=1); beta.call(y=2); gamma.call(z=3)" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["alpha", "beta", "gamma"] + assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} + + def test_nested_python_tag_in_json_string_arg_is_not_a_call(self): + # A <|python_tag|> literal inside a code arg is data: the outer "python" call wins. + text = ( + '<|python_tag|>{"name":"python","parameters":' + '{"code":"<|python_tag|>os.call(\'rm -rf /\')"}}' + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "python" + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "<|python_tag|>os.call('rm -rf /')" + + def test_single_builtin_call_unchanged(self): + text = '<|python_tag|>web_search.call(query="cats")' + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + +def test_strip_leading_bare_json_call_drops_complete_call(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A complete Llama-3.2 bare-JSON call is removed; trailing prose is kept. + assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"cats"}}') == "" + assert ( + strip_leading_bare_json_call('{"name":"python","parameters":{"code":"x"}} done') == "done" + ) + + +def test_strip_leading_bare_json_call_drops_truncated_call(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A truncated call (no closing brace) collapses to "" -- nothing recoverable. + assert ( + strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"weather in S') + == "" + ) + + +def test_strip_leading_bare_json_call_preserves_plain_json_and_prose(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # No "name" key -> plain JSON answer, left untouched. + assert ( + strip_leading_bare_json_call('{"result": 42, "ok": true}') == '{"result": 42, "ok": true}' + ) + # Prose before the brace -> not a leading bare call, untouched. + assert strip_leading_bare_json_call('here is {"name":"x"}') == 'here is {"name":"x"}' + # Ordinary text untouched. + assert strip_leading_bare_json_call("just a sentence.") == "just a sentence." + + +def test_bare_json_gated_on_enabled_tool_names(): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + alice = '{"name":"Alice","parameters":{"age":30}}' + real = '{"name":"web_search","parameters":{"query":"cats"}}' + # With an enabled set, markerless JSON whose name is not a tool is NOT a call. + assert parse_tool_calls_from_text(alice, enabled_tool_names = {"web_search"}) == [] + # A real call (enabled name) still parses. + got = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in got] == ["web_search"] + # No enabled set (None) keeps the name-agnostic behaviour for direct callers. + assert [c["function"]["name"] for c in parse_tool_calls_from_text(alice)] == ["Alice"] + # Marker-based forms are NOT gated (an explicit signal is a real call attempt). + xml = '{"name":"Alice","arguments":{}}' + assert parse_tool_calls_from_text(xml, enabled_tool_names = {"web_search"}) + + +def test_strip_leading_bare_json_call_gated_on_enabled_tool_names(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + alice = '{"name":"Alice","parameters":{"age":30}}' + # Not an enabled tool -> ordinary JSON answer, kept verbatim. + assert strip_leading_bare_json_call(alice, {"web_search"}) == alice + # Enabled tool -> a real call, stripped (trailing prose kept). + assert ( + strip_leading_bare_json_call( + '{"name":"web_search","parameters":{"q":1}} hi', {"web_search"} + ) + == "hi" + ) + + +def test_function_xml_strip_keeps_literal_close_tag_in_param_value(): + from core.inference.tool_call_parser import strip_tool_markup + + # Strip uses the LAST so a literal in a value survives; calls strip independently. + text = 'print("") done' + assert strip_tool_markup(text, final = True) == "done" + two = ( + "a 1 mid " + "2 end" + ) + assert strip_tool_markup(two, final = True) == "a mid end" + + +def test_function_xml_strip_keeps_trailing_text_after_literal_open_tag(): + from core.inference.tool_call_parser import parse_tool_calls_from_text, strip_tool_markup + + # A literal opener inside a value is data: the strip keeps " done". + text = 'print("") done' + assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python" + assert strip_tool_markup(text, final = True) == "done" + # Non-final (streaming) keeps an unclosed call buffered, does not eat prose early. + open_text = 'pre print("")' + assert strip_tool_markup(open_text, final = False) == open_text + + +def test_final_strip_removes_magistral_think_reasoning(): + from core.inference.tool_call_parser import strip_tool_markup + + # Magistral reasoning is [THINK]...[/THINK]; end-of-turn must drop it. + text = "[THINK]The user greeted me, I should say hi.[/THINK]Hello! How can I help?" + assert strip_tool_markup(text, final = True) == "Hello! How can I help?" + # A [TOOL_CALLS] living inside the reasoning goes with it. + with_call = '[THINK]Maybe I should search.[/THINK][TOOL_CALLS]search{"q":"x"}' + assert strip_tool_markup(with_call, final = True) == "" + + +def test_streaming_strip_keeps_magistral_think_buffered(): + from core.inference.tool_call_parser import strip_tool_markup + + # Mid-stream (final=False) leaves the reasoning block intact; only end-of-turn removes it. + text = "[THINK]still thinking" + assert strip_tool_markup(text, final = False) == text + + +def test_final_strip_leaves_non_magistral_bracket_text_untouched(): + from core.inference.tool_call_parser import strip_tool_markup + + # Only a LEADING [THINK] block is reasoning; unrelated bracketed prose stays. + text = "See [THINK about it] later" + assert strip_tool_markup(text, final = True) == "See [THINK about it] later" + + +def test_strip_leading_bare_json_call_ignores_nested_name(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A nested "name" must NOT gate the strip; the JSON answer is kept verbatim. + nested_trunc = '{"result":{"name":"web_search","age":' + nested_full = '{"result":{"name":"web_search","age":1}}' + assert strip_leading_bare_json_call(nested_trunc, {"web_search"}) == nested_trunc + assert strip_leading_bare_json_call(nested_full, {"web_search"}) == nested_full + # A real top-level call (even with a top-level array before the name) still strips. + assert ( + strip_leading_bare_json_call( + '{"data":[1,2],"name":"web_search","parameters":{}}', {"web_search"} + ) + == "" + ) + + +def test_mistral_single_object_call_is_stripped_for_display(): + from core.inference.tool_call_parser import ( + _strip_mistral_closed_calls, + parse_tool_calls_from_text, + ) + + # The parser accepts single-object [TOOL_CALLS]{...}, so the strip must remove it too. + text = '[TOOL_CALLS]{"name":"web_search","arguments":{"filters":{"date":"2024"}}} tail' + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == ["web_search"] + assert _strip_mistral_closed_calls(text) == " tail" + # A literal [TOOL_CALLS] in prose (no following object) is left untouched. + assert _strip_mistral_closed_calls("See the [TOOL_CALLS] docs") == "See the [TOOL_CALLS] docs" + + +def test_tool_call_parser_declares_future_annotations_for_py39_import(): + # PEP 604 X | None annotations need `from __future__ import annotations` on py3.9; guard it stays. + from pathlib import Path + src = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "tool_call_parser.py" + ).read_text() + assert "from __future__ import annotations" in src + + +def test_bare_json_function_alias_parses_and_strips_symmetrically(): + # The "function" alias for the call name must parse and strip symmetrically. + from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_leading_bare_json_call, + _top_level_bare_json_name, + ) + + enabled = {"web_search"} + text = '{"function":"web_search","parameters":{"query":"cats"}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = enabled) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert strip_leading_bare_json_call(text, enabled) == "" + + # "name" still takes precedence when both are present; nested aliases are data. + assert _top_level_bare_json_name('{"function":"foo","name":"web_search"}') == "web_search" + assert _top_level_bare_json_name('{"function":"web_search"}') == "web_search" + assert _top_level_bare_json_name('{"result":{"function":"web_search"}}') is None + # A non-enabled function-alias object is ordinary content and is preserved. + assert ( + strip_leading_bare_json_call('{"function":"not_a_tool","parameters":{}}', enabled) + == '{"function":"not_a_tool","parameters":{}}' + ) + + +class TestMistralOuterOverXmlLiteral: + """Quoted tool XML inside a [TOOL_CALLS] call's arguments is data; the outer call executes. Reverse order keeps the XML.""" + + def test_mistral_v11_arg_quoting_function_xml(self): + text = ( + '[TOOL_CALLS]web_search[ARGS]{"query":"literal ' + '1"}' + ) + for strict in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = not strict) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert "" in json.loads(calls[0]["function"]["arguments"])["query"] + + def test_mistral_array_arg_quoting_tool_call_json(self): + text = ( + '[TOOL_CALLS][{"name":"web_search","arguments":{"query":' + '"see {\\"name\\":\\"evil\\"}"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_outer_keeps_winning_over_mistral_literal(self): + text = ( + '{"name":"web_search","arguments":' + '{"query":"docs say [TOOL_CALLS]evil[ARGS]{}"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestHealerSignalAlignment: + """The healer buffers only promotable formats; Mistral/Llama text calls stream through.""" + + def test_heal_signals_subset_of_promotable_formats(self): + from core.inference.passthrough_healing import _HEAL_SIGNALS + assert set(_HEAL_SIGNALS) == {"", "<|tool_call>", "evil.call(x=1)"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what is <|python_tag|>evil.call(x=1)" + + +class TestPythonTagOuterOverXmlLiteral: + """A leading Llama-3 ``<|python_tag|>`` call owns the turn: tool XML/Mistral + markup quoted in a ``.call(...)`` string argument (or in trailing prose) is + data, so the outer call executes -- parity with the bare-JSON / Mistral / + attribute-form leading-ownership rules. XML before the tag keeps normal order.""" + + def test_call_arg_quoting_complete_function_xml(self): + # A closed in a .call() code arg must not beat the leading python_tag call. + text = ( + '<|python_tag|>python.call(code="' + '1")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "1" + + def test_call_arg_quoting_bare_function_tag_in_query(self): + # A query mentioning must search, not execute a phantom tool. + text = '<|python_tag|>web_search.call(query="how do I use in llama")' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "how do I use in llama" + + def test_call_arg_quoting_tool_call_json(self): + text = ( + "<|python_tag|>save_file.call(content=" + '"{\\"name\\": \\"delete\\", \\"arguments\\": {}}")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["save_file"] + + def test_json_form_code_arg_quoting_function_xml(self): + # JSON emission: a in the code arg is data; the outer "python" call runs. + text = ( + '<|python_tag|>{"name":"python","parameters":' + '{"code":"ls"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "ls" + + def test_call_arg_quoting_mistral_trigger(self): + text = '<|python_tag|>web_search.call(query="see [TOOL_CALLS]evil[ARGS]{}")' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_call_wins_over_trailing_xml(self): + # A leading python_tag call owns the turn even when a real XML literal follows. + text = ( + '<|python_tag|>web_search.call(query="cats") ' + "1" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_before_python_tag_keeps_xml_order(self): + # A foreign signal BEFORE the tag keeps normal document order (XML wins). + text = ( + "x " + '<|python_tag|>python.call(code="y")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestBareJsonOuterOverXmlLiteral: + """Quoted tool XML inside a leading bare-JSON call is data; XML before the JSON keeps normal order.""" + + def test_bare_json_code_arg_quoting_function_xml(self): + text = ( + '{"name": "python", "arguments": ' + '{"code": "run() # ls"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "run() # ls" + + def test_bare_json_outer_unrestricted_mode(self): + text = '{"name": "python", "parameters": {"code": "ls"}}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + + def test_xml_before_json_keeps_xml_order(self): + text = ( + "cats" + ' {"name": "python", "arguments": {"code": "x"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestMagistralThinkRehearsal: + """A call rehearsed inside [THINK]...[/THINK] is reasoning; the real call after wins, and parse agrees with strip.""" + + def test_function_xml_rehearsal_in_think_is_not_promoted(self): + text = ( + '[THINK]I could emit {"query":"x"}' + ' here[/THINK][TOOL_CALLS] [{"name":"terminal","arguments":{"cmd":"ls"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"] + + def test_hermes_rehearsal_in_think_is_not_promoted(self): + text = ( + '[THINK]maybe {"name":"web_search","arguments":' + '{"query":"x"}}[/THINK]' + '[TOOL_CALLS] [{"name":"terminal","arguments":{"cmd":"ls"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"] + + def test_unclosed_think_parses_nothing(self): + text = '[THINK]let me try {"query":"x"}' + assert parse_tool_calls_from_text(text) == [] + + +class TestDisabledBareJsonLiteralNotPromoted: + """A leading non-enabled-name object is content: nothing inside promotes, and a call after it still parses.""" + + def test_literal_inside_disabled_json_stays_data(self): + text = ( + '{"name": "Alice", "note": "try ' + 'x"}' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_python_tag_literal_inside_disabled_json_stays_data(self): + text = '{"name": "Alice", "note": "<|python_tag|>web_search.call(query=1)"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_real_call_after_disabled_json_still_parses(self): + text = ( + '{"name": "Alice", "note": "x"} ' + '{"name": "web_search", "arguments": {"query": "cats"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestMistralLiteralInsideLeadingJson: + """A [TOOL_CALLS] literal quoted inside a leading JSON object must not be promoted over it.""" + + def test_outer_json_call_wins_over_mistral_literal(self): + text = '{"name": "python", "arguments": {"code": "[TOOL_CALLS]web_search{}"}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "[TOOL_CALLS]web_search{}" + + def test_disabled_outer_json_keeps_mistral_literal_as_data(self): + text = '{"name": "Alice", "note": "[TOOL_CALLS]web_search{}"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaWrappedWhitespace: + """Whitespace drift around ``call``/``:`` in wrapped Gemma calls must still parse (no fallback exists).""" + + def test_space_after_call_colon_parses(self): + text = '<|tool_call>call: web_search{query:<|"|>cats<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_space_around_colon_parses(self): + text = '<|tool_call>call : web_search{query:<|"|>cats<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_strict_mode_still_requires_the_closing_tag(self): + text = '<|tool_call>call: web_search{query:<|"|>cats<|"|>}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + +class TestGemmaDottedArgumentKeys: + """Dotted Gemma keys (namespaced schemas) must survive key-quoting or the call is lost.""" + + def test_dotted_key_parses(self): + text = '<|tool_call>call:web_search{user.name:<|"|>bob<|"|>, query:<|"|>x<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"user.name": "bob", "query": "x"} + + +class TestLeadingMistralCallOwnsTheTurn: + """A leading Mistral call wins in document order over literal XML in trailing prose.""" + + def test_leading_mistral_wins_over_trailing_xml_literal(self): + text = ( + '[TOOL_CALLS]web_search[ARGS]{"query":"cats"} ' + "Note: 1" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_leading_keeps_normal_order(self): + text = ( + "x " + "[TOOL_CALLS]evil[ARGS]{}" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestGemmaDottedKeyAfterBareValue: + def test_dotted_key_after_bare_value_is_a_boundary(self): + text = "<|tool_call>call:web_search{query:foo,user.name:bob}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "foo", "user.name": "bob"} + + +class TestNamelessLeadingJsonAnswerIsData: + """A nameless leading JSON answer is an envelope: quoted markup stays data, and a call after it parses.""" + + def test_xml_literal_inside_json_answer_stays_data(self): + text = '{"answer": "use x"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_real_call_after_json_answer_still_parses(self): + text = ( + '{"answer": "docs"} {"name": "web_search", ' + '"arguments": {"query": "cats"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestLeadingBareJsonOwnsTurnOverTrailingXml: + """Document order: a leading closed bare-JSON call owns the turn even when + tool XML appears AFTER it (inside-or-after, mirroring the Mistral rule).""" + + def test_leading_call_wins_over_trailing_xml(self): + text = ( + '{"name":"lookup","parameters":{"q":"first"}} Example: ' + '{"name":"delete_all","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "first"} + + def test_chained_leading_calls_win_over_trailing_xml(self): + text = ( + '{"name":"lookup","parameters":{"q":"first"}};' + '{"name":"lookup","parameters":{"q":"second"}} ' + '{"name":"delete_all","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls + + def test_non_call_leading_object_defers_to_trailing_real_call(self): + # Nameless/disabled-name objects decline: dropped, and the real trailing call still parses. + for lead in ('{"answer": 42}', '{"name":"draft","parameters":{}}'): + text = lead + ' {"name":"delete_all","arguments":{}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"delete_all"}) + assert [c["function"]["name"] for c in calls] == ["delete_all"], (lead, calls) + + def test_leading_xml_call_still_wins_over_trailing_bare_json(self): + text = ( + '{"name":"delete_all","arguments":{}} ' + 'Example: {"name":"lookup","parameters":{"q":"x"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["delete_all"], calls + + +class TestProseCloseTagAfterClosedFunctionCall: + """A literal in prose after a closed call is data: the call + ends at its first close that is not parameter data, so arguments never + swallow the prose between the real close and the literal.""" + + def test_arguments_do_not_swallow_prose(self): + text = ( + "cats" + " Done. The tag closes a call." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_literal_close_inside_open_parameter_stays_data(self): + text = 'print("")' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + def test_attribute_form_arguments_do_not_swallow_prose(self): + # The attribute form shares the first-balanced-close rule: prose closes never fold in. + text = ( + 'cats' + " Done. The tag closes a call." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_attribute_form_literal_close_in_open_parameter_stays_data(self): + text = 'print("")' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + def test_attribute_form_two_calls_both_parse(self): + text = ( + 'cats' + 'x=1' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "python"}) + assert [c["function"]["name"] for c in calls] == ["web_search", "python"], calls + + +class TestEnabledNameJsonAnswerIsContent: + """A JSON answer whose top-level name matches an enabled tool but has no + call shape is content: the parser rejects it, so the strip and the drain + gate must keep it visible too.""" + + def test_answer_survives_strip(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + ans = '{"name":"web_search","result":"no call"}' + assert strip_leading_bare_json_call(ans, {"web_search"}) == ans + + def test_answer_does_not_route_to_draining(self): + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + assert not _looks_like_enabled_bare_json( + '{"name":"web_search","result":"no call"}', {"web_search"} + ) + + def test_real_call_still_strips_and_drains(self): + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + from core.inference.tool_call_parser import strip_leading_bare_json_call + + real = '{"name":"web_search","parameters":{"q":"x"}}' + assert strip_leading_bare_json_call(real, {"web_search"}) == "" + assert _looks_like_enabled_bare_json(real, {"web_search"}) + + def test_arguments_string_call_still_strips(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + call = '{"name":"web_search","arguments":"{\\"q\\":\\"x\\"}"} tail' + assert strip_leading_bare_json_call(call, {"web_search"}) == "tail" + + +class TestAttributeFormLeadingContainment: + """A leading attribute-form call owns the turn: markup quoted inside its + parameter is data, not a call for the shared XML parser to promote.""" + + def test_quoted_tool_call_inside_param_stays_data(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'find ' + '{"name":"delete","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert "delete" in json.loads(calls[0]["function"]["arguments"])["query"] + + def test_real_xml_call_before_attribute_form_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + '{"name":"delete","arguments":{}} Example: ' + 'x' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + assert calls[0]["function"]["name"] == "delete" + + +class TestParameterKeepsMultipleLiteralCloses: + """A parameter that provably closes with its own tag keeps every literal + function close inside it as data (regression: the first literal close was + treated as ending the parameter, truncating the value).""" + + def test_two_literal_closes_in_one_parameter(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + '' + "a b c " + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "a b c" + } + + def test_strip_removes_the_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + text = ( + '' + "a b c after" + ) + assert strip_tool_markup(text, final = True) == "after" + + def test_unclosed_parameter_still_heals_at_function_close(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + calls = parse_tool_calls_from_text( + "val", + enabled_tool_names = {"web_search"}, + ) + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "val"} + + +class TestMistralPreambleOwnership: + """A visible preface before the first Mistral call must not hand the turn + to a later XML literal: the Mistral call is first in document order.""" + + def test_v11_named_form_after_preface(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'pref [TOOL_CALLS]web_search[ARGS]{"query":"cats"} Note ' + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_array_form_after_preface(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'pref [TOOL_CALLS][{"name":"web_search","arguments":{"query":"cats"}}] Note ' + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_call_before_trigger_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + "1 then " + '[TOOL_CALLS][{"name":"web_search","arguments":{}}]' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert calls[0]["function"]["name"] == "evil" + + def test_prose_mention_without_call_shape_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + "See [TOOL_CALLS] docs for details. " + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"evil"}) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestBareJsonStripRequiresTopLevelName: + """The strip's shape gate requires the parser's TOP-LEVEL name in every + mode: a JSON answer with only a nested name is content, even name-agnostic.""" + + def test_nested_name_answer_survives_name_agnostic_strip(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + ans = '{"parameters":{},"result":{"name":"web_search"}}' + assert strip_leading_bare_json_call(ans) == ans + assert strip_leading_bare_json_call(ans, {"web_search"}) == ans + + def test_real_call_still_strips_name_agnostic(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') == "" diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index c2dc1fe8db..7fe52a664d 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -24,15 +24,34 @@ import re as _re _src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() _m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) assert _m, "could not extract _TOOL_XML_RE source" -_ns = {"_re": _re} +# Provide both helpers so the extracted _strip_tool_xml_for_display resolves. +from core.inference.tool_call_parser import _strip_function_xml_calls, _strip_mistral_closed_calls + +_ns = { + "_re": _re, + "_strip_mistral_closed_calls": _strip_mistral_closed_calls, + "_strip_function_xml_calls": _strip_function_xml_calls, +} exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) _TOOL_XML_RE = _ns["_TOOL_XML_RE"] + +_xml_helper = _re.search( + r"def _strip_tool_xml\(text: str\) -> str:\n(?: .+\n)+", + _src, +) +assert _xml_helper, "could not extract _strip_tool_xml source" +assert "_strip_mistral_closed_calls" in _xml_helper.group( + 0 +), "extracted _strip_tool_xml no longer runs the Mistral balanced strip" +exec(_xml_helper.group(0), _ns) + _helper = _re.search( r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n" r"(?: .+\n)+", _src, ) assert _helper, "could not extract _strip_tool_xml_for_display source" +assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates" exec(_helper.group(0), _ns) _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] @@ -46,6 +65,15 @@ def test_route_display_strip_respects_disabled_auto_heal_contract(): assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) +def test_route_display_strip_removes_mistral_tool_calls_with_nested_json(): + # [TOOL_CALLS] with nested JSON needs the Mistral balanced-brace strip, not the regex. + text = 'ok [TOOL_CALLS]web_search{"filters":{"date":"2024"},"query":"cats"} tail' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "web_search" not in out, out + assert out == "ok tail" + + def test_strips_well_formed_tool_call(): text = ( "Let me search.\n" @@ -73,6 +101,25 @@ def test_strips_function_only_well_formed(): assert "Done." in cleaned +def test_strips_function_attribute_form(): + # Attribute form must strip from the route too; dotted/hyphenated names included. + text = ( + 'Sure.\n\n' + "\nSydney\n\n\nDone." + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "" not in cleaned + assert "Sure." in cleaned and "Done." in cleaned + + dotted = 'A x B' + assert _TOOL_XML_RE.sub("", dotted) == "A B" + + # Auto-Heal-disabled display contract still preserves literal markup. + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + assert "" not in cleaned + + +# Llama-3 <|python_tag|> arm bounds on REAL sentinels only +def test_python_tag_strip_consumes_literal_sentinel_in_arg(): + # A literal <|...|> token inside the arg must not end the strip early. + text = '<|python_tag|>{"name": "send", "parameters": {"text": "use <|cite|> here"}}' + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "", f"python_tag call leaked at literal sentinel: {cleaned!r}" + + +@pytest.mark.parametrize( + "sentinel", + [ + "<|eot_id|>", + "<|eom_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + ], +) +def test_python_tag_strip_stops_at_real_sentinel(sentinel): + # A real control sentinel bounds the strip so following text survives. + text = f'<|python_tag|>{{"name": "x", "parameters": {{}}}}{sentinel}visible answer' + cleaned = _TOOL_XML_RE.sub("", text) + assert ( + cleaned == f"{sentinel}visible answer" + ), f"strip did not stop at real sentinel {sentinel!r}: {cleaned!r}" + + +def test_python_tag_strip_restarts_on_second_python_tag(): + # A second <|python_tag|> opens a new region; both are stripped. + text = '<|python_tag|>{"name": "a"}<|python_tag|>{"name": "b"}' + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "", f"second python_tag region leaked: {cleaned!r}" + + +def test_route_strip_removes_param_alias_close_tag(): + # Orphan (attribute-form alias of ) must strip too. + assert _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + assert ( + _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + ) + + +def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup(): + # A literal in a value must not truncate the strip. + text = " tail" + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail" + + +def test_strip_keeps_prose_after_closed_function_call_with_literal_close(): + # The call ends at its first non-data close; prose after (even a literal ) survives. + from core.inference.tool_call_parser import strip_tool_markup + text = ( + "cats" + " Done. The tag closes a call." + ) + assert strip_tool_markup(text, final = True) == "Done. The tag closes a call." + + +def test_final_strip_keeps_prose_mentioning_bare_markers(): + # A false-alarm marker in prose must not drop trailing text; only call-start-shaped text drops. + from core.inference.tool_call_parser import strip_tool_markup + for text in ( + "See [TOOL_CALLS] docs for details. More prose after.", + "<|python_tag|> is the Llama marker. Explanation continues.", + "The <|tool_call> opener wraps Gemma calls.", + ): + assert strip_tool_markup(text, final = True) == text + # A bare marker at end-of-text is a fragment and still drops. + assert strip_tool_markup("Answer text [TOOL_CALLS]", final = True) == "Answer text" + + +def test_final_strip_still_drops_truncated_marker_calls(): + from core.inference.tool_call_parser import strip_tool_markup + for text in ( + '[TOOL_CALLS][{"name":"web_search","argu', + '[TOOL_CALLS]web_search[ARGS]{"q":"x', + '<|python_tag|>{"name":"web_search","par', + '<|python_tag|>foo.call(items=["a', + "<|tool_call>call:web_search{query:tru", + ): + assert strip_tool_markup(text, final = True) == "" + + +def test_chained_bare_json_strip_consumes_all_calls(): + # Next-turn history must not keep an executed call, else it replays. + from core.inference.tool_call_parser import strip_leading_bare_json_call + + enabled = {"web_search", "python"} + chained = ( + '{"name":"web_search","parameters":{"q":"first"}};' + '{"name":"python","parameters":{"code":"x"}}' + ) + assert strip_leading_bare_json_call(chained, enabled_tool_names = enabled) == "" + assert ( + strip_leading_bare_json_call(chained + " trailing prose", enabled_tool_names = enabled) + == "trailing prose" + ) + # The chain stops at a non-call answer object, which stays visible. + call_then_answer = ( + '{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}' + ) + assert ( + strip_leading_bare_json_call(call_then_answer, enabled_tool_names = enabled) + == '{"name":"web_search","result":"data"}' + ) From f38672da65e420a11323f0e5aa4649449a20e66c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:07:56 -0700 Subject: [PATCH 020/113] Studio: stop chat generation on the assistant-turn-end token (fixes Qwen3.5 loop) (#6804) * Studio: stop chat generation on the assistant-turn-end token A small chat model (e.g. Qwen3.5-0.8B) looped on the safetensors path: it emitted a valid response or tool call, then ran past its turn and re-emitted the call, hallucinating <|im_start|>user turns. Root cause: the model's tokenizer.eos_token is synced to the config document terminator (<|endoftext|>, 248044) while chat turns actually end with <|im_end|> (248046), so generate_stream's single eos_token_id never stopped at the turn boundary. Stop on every assistant-turn-end marker the vocab defines (tokenizer.eos plus <|im_end|>, <|eot_id|>, , ...). Verified on the real weights: the single-eos control loops (400 tokens) while the fixed set yields a clean 38-token tool call and a clean answer from the tool result. No-op when eos is already the turn-ender (the id just dedups). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: repair chat generation_config.eos_token_id at load time Qwen3.5 / Qwen3.6 small chat checkpoints declare the chat turn-end as tokenizer.eos_token (<|im_end|>) but ship config.eos_token_id = <|endoftext|> and no generation_config.json (upstream shipped generation_config only on the large chat models). So every .generate() path that reads generation_config -- the vision path and tool loops, not just generate_stream -- never stops at the turn boundary and loops. At load time, when the tokenizer's own eos is a chat turn-end marker but generation_config.eos_token_id omits it, add it. This fixes the config once for all generation paths and complements the generate_stream turn-end stop. No-op for base models (eos is a plain document terminator) and already-correct configs. Verified on unsloth/Qwen3.5-0.8B: 248044 -> [248044, 248046]. * Studio: derive chat turn-end eos from the template, resolve once at load Address PR review of the turn-end stop handling: - Do not call tokenizer.get_vocab() per generation request (serializes the whole 100k+ vocab). Resolve the turn-end tokens once at load and cache them on model_info; generate_stream reads the cache. - Derive turn-end markers from the chat_template the model actually uses, not raw vocab membership, so a base/coder model that merely carries ChatML control tokens in a shared vocab is not stopped early, and a loader that synced tokenizer.eos to the document terminator is still covered. - Skip harmony/gpt-oss templates: <|end|> there is an intra-message channel delimiter, not the turn end (dropped <|return|> from the marker list too). - Move the logic to a dependency-light module (core.inference.chat_eos) so the unit test does not import the full unsloth/torch inference stack. Verified on unsloth/Qwen3.5-0.8B (gen_config 248044 -> [248044, 248046], clean 38-token tool call with generation_config-only stopping), Phi-3.5 (adds <|end|>), Llama-3 / Qwen3 (unchanged), and a harmony template (left untouched). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: refresh turn-end eos after the mapper installs its template For a MODEL_TO_TEMPLATE_MAPPER model whose own tokenizer ships no chat_template, the effective template is applied at generate time via get_chat_template, but the turn-end eos ids were resolved once at load when the template was still empty, so only the document eos was cached. Qwen2.5 / Yi base checkpoints (eos <|endoftext|>, ChatML turns end with <|im_end|>) then run past the assistant boundary in generate_stream and loop. Re-resolve the turn-end eos from the now-templated tokenizer and refresh the cached ids right after applying the mapper template, so generate_stream stops at the ChatML turn end. Add a regression test. * Studio: union turn-end eos refresh into load-time cache instead of overwriting get_chat_template can return a different tokenizer whose vocab was remapped (Gemma folds onto the eos id), while generate_stream re-reads the original model_info tokenizer. Overwriting the cache with the refreshed set dropped a valid load-time id (e.g. =107) and let generation run past the real turn marker. Union the refresh into the existing cache so it can only add ids, never drop a valid one. Add a regression test covering the destructive-swap case the prior test missed. * Studio: resolve refreshed turn-end ids on the generation tokenizer, add Gemma-4 marker Two residual gaps in the turn-end eos refresh: - For map_eos_token=True mapped templates (e.g. chatml on a Yi-6B base), get_chat_template returns a tokenizer whose vocab folds the turn-end token onto the document eos id, while generate_stream re-reads the original tokenizer. The refresh resolved ids on the returned tokenizer, so it stored the doc eos and missed the real turn-end id, and generation ran past the boundary. Read the turn-end marker strings from the mapped template but resolve their ids on the original generation tokenizer (new resolve_chat_turn_end_eos_ids_using). - Add Gemma-4's turn terminator to the marker allowlist; those templates keep a document eos so resolve otherwise missed the real turn marker. Add regression tests for both. * Fix turn-end detection for Starling, multi-variant and vision templates; keep tests collectable The turn-end marker set missed OpenChat/Starling's barred <|end_of_turn|> (distinct from Gemma's unbarred form), so Starling generations ran past the assistant boundary. A dict/list chat_template (Hermes-3 style default+tool_use variants) hit an early non-string return and skipped detection; flatten and scan every variant. Vision models carry the chat_template on the ProcessorMixin, not the unwrapped inner tokenizer, so read markers from the template-carrying container while resolving ids on the generation tokenizer. The refresh test constructs the real backend, so it is guarded with a module-level skip when unsloth/unsloth_zoo is absent (the lightweight pytest matrix), and core.inference package init is made lazy so the dependency-light chat_eos tests collect without the heavy stack. * Studio: tighten chat turn-end eos comments * Studio: condense chat turn-end eos comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/__init__.py | 43 +++- studio/backend/core/inference/chat_eos.py | 109 ++++++++++ studio/backend/core/inference/inference.py | 68 +++++- .../tests/test_chat_eos_template_refresh.py | 194 ++++++++++++++++++ .../backend/tests/test_chat_turn_end_eos.py | 150 ++++++++++++++ 5 files changed, 558 insertions(+), 6 deletions(-) create mode 100644 studio/backend/core/inference/chat_eos.py create mode 100644 studio/backend/tests/test_chat_eos_template_refresh.py create mode 100644 studio/backend/tests/test_chat_turn_end_eos.py diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 2faf70bb79..ad78157418 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation. The default get_inference_backend() returns an InferenceOrchestrator that delegates to a subprocess. The original InferenceBackend runs inside the subprocess and can be imported directly from .inference when needed. + +Public names are resolved lazily (PEP 562): importing this package -- or a +dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull +the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML +backend and its Studio dependencies). Those load only when a public name is +actually accessed, so standalone helpers stay unit-testable without the full +inference stack. """ -from .orchestrator import InferenceOrchestrator, get_inference_backend -from .llama_cpp import LlamaCppBackend - -# Expose InferenceOrchestrator as InferenceBackend for backward compat. -InferenceBackend = InferenceOrchestrator +from typing import TYPE_CHECKING __all__ = [ "InferenceBackend", @@ -21,3 +24,33 @@ __all__ = [ "get_inference_backend", "LlamaCppBackend", ] + +# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator. +_LAZY_ATTRS = { + "InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"), + "InferenceBackend": ("orchestrator", "InferenceOrchestrator"), + "get_inference_backend": ("orchestrator", "get_inference_backend"), + "LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"), +} + + +def __getattr__(name): + try: + submodule, attr = _LAZY_ATTRS[name] + except KeyError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + from importlib import import_module + + value = getattr(import_module(f"{__name__}.{submodule}"), attr) + globals()[name] = value # cache so later access skips __getattr__ + return value + + +def __dir__(): + return sorted(set(globals()) | set(__all__)) + + +if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names + from .llama_cpp import LlamaCppBackend + from .orchestrator import InferenceOrchestrator, get_inference_backend + InferenceBackend = InferenceOrchestrator diff --git a/studio/backend/core/inference/chat_eos.py b/studio/backend/core/inference/chat_eos.py new file mode 100644 index 0000000000..2a5d0db228 --- /dev/null +++ b/studio/backend/core/inference/chat_eos.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve a chat model's assistant-turn-end stop tokens. + +Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships +config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its +small chat variants ship no generation_config), so generation runs past the turn +and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns. + +Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it +actually uses to end a turn), not raw vocab membership: a base/coder model can +carry ChatML control tokens in a shared vocab without using them, and a loader +may have synced ``eos_token`` to the document terminator. Dependency-light (no +torch / unsloth) so it is unit-testable without the full inference stack. +""" + +from typing import Optional + +# Canonical assistant-turn-end markers per chat family. +_CHAT_TURN_END_TOKENS = ( + "<|im_end|>", # ChatML: Qwen, Yi + "<|eot_id|>", # Llama 3.x + "<|eom_id|>", # Llama 3.x tool turns + "", # Gemma + "", # Gemma-4 + "<|end|>", # Phi + "<|end_of_turn|>", # OpenChat / Starling (barred, distinct from Gemma's) +) +# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has +# its own streamer, so its eos is left untouched. +_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>") + + +def _eos_id_set(eos_token_id) -> set: + if isinstance(eos_token_id, (list, tuple)): + return {int(t) for t in eos_token_id if t is not None} + if eos_token_id is not None: + return {int(eos_token_id)} + return set() + + +def _collect_template_text(chat_template) -> str: + """Flatten a tokenizer ``chat_template`` into one scannable string. + + Usually the template is a single jinja string, but multi-variant models + (e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a + ``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list + of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case + would skip turn-end detection for those valid models, so gather every string + leaf (variant names are harmless: they never contain the markers). + """ + if isinstance(chat_template, str): + return chat_template + if isinstance(chat_template, dict): + values = chat_template.values() + elif isinstance(chat_template, (list, tuple)): + values = chat_template + else: + return "" + parts = [_collect_template_text(v) for v in values] + return "\n".join(p for p in parts if p) + + +def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list: + """eos of ``id_tokenizer`` plus any canonical turn-end marker the + ``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` -- + the tokenizer generation actually uses. + + Pass the same tokenizer for both at load time. After a mapped ``get_chat_template`` + pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective + template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped + template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab + folds the turn-end token onto the doc-eos id, and generate_stream re-reads the + original tokenizer, so resolving ids on the mapped tokenizer would store the wrong + (doc-eos) id and let generation run past the real turn marker.""" + ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None)) + template = _collect_template_text(getattr(template_tokenizer, "chat_template", None)) + if not template or any(h in template for h in _HARMONY_MARKERS): + return sorted(ids) + unk = getattr(id_tokenizer, "unk_token_id", None) + for marker in _CHAT_TURN_END_TOKENS: + if marker in template: + try: + tid = id_tokenizer.convert_tokens_to_ids(marker) + except Exception: + tid = None + if tid is not None and tid != unk and int(tid) >= 0: + ids.add(int(tid)) + return sorted(ids) + + +def resolve_chat_turn_end_eos_ids(tokenizer) -> list: + """tokenizer.eos plus any canonical turn-end marker the model's chat_template + actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended + to be resolved once at load. Returns eos unchanged for harmony templates.""" + return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer) + + +def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]: + """Merged eos_token_id list, or None if ``current_eos`` already covers every + resolved turn-end id. Used to repair a model's generation_config at load so + every ``.generate()`` path (vision, tool loops) stops at the turn boundary.""" + if not turn_end_ids: + return None + current_set = _eos_id_set(current_eos) + if set(turn_end_ids) <= current_set: + return None + return sorted(current_set | set(turn_end_ids)) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 4dca4db768..eaee5a213a 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -27,6 +27,10 @@ from utils.hardware import ( from core.inference.audio_codecs import AudioCodecManager from core.inference.runtime_context import runtime_context_length from core.inference.message_content import content_to_text +from core.inference.chat_eos import ( + chat_eos_repair, + resolve_chat_turn_end_eos_ids_using, +) from io import StringIO import structlog from loggers import get_logger @@ -210,6 +214,50 @@ class InferenceBackend: # API uses -1 to disable top-k; transformers uses 0. return 0 if top_k < 0 else top_k + def _resolve_chat_eos(self, model_name: str) -> None: + """Resolve this chat model's assistant-turn-end stop tokens once at load, + cache them in model_info, and repair generation_config so every + ``.generate()`` path stops at the turn boundary. + + Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with + ``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no + ``generation_config.json``, so paths that read ``generation_config`` (the + vision path, tool loops) run past the turn and loop. Turn-end markers are + derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids), + so base/coder models and harmony templates are left untouched. + """ + info = self.models.get(model_name) or {} + model = info.get("model") + container = info.get("tokenizer") + tokenizer = getattr(container, "tokenizer", container) # unwrap processors + if model is None or tokenizer is None: + return + # Vision models carry the chat_template on the processor, not the inner + # tokenizer. Read markers from whichever has one, but resolve ids on the + # generation tokenizer, else the vision path misses the turn-end token. + template_source = container if getattr(container, "chat_template", None) else tokenizer + try: + turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer) + except Exception as e: # never block a load on eos resolution + logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e) + return + info["chat_turn_end_eos_ids"] = turn_end_ids + + gen = getattr(model, "generation_config", None) + if gen is None: + return + repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids) + if repaired is None: + return + previous = gen.eos_token_id + gen.eos_token_id = repaired + logger.info( + "Repaired generation_config.eos_token_id for %s: %s -> %s", + model_name, + previous, + repaired, + ) + def load_model( self, config: ModelConfig, @@ -496,6 +544,7 @@ class InferenceBackend: max_seq_length, ) + self._resolve_chat_eos(model_name) self._load_chat_template_info(model_name) self.active_model_name = model_name @@ -946,6 +995,22 @@ class InferenceBackend: tokenizer, chat_template = template_name, ) + # The mapper installs the effective template only now, at generate + # time, so re-resolve and UNION into the load-time cache (never + # overwrite). get_chat_template can return a remapped tokenizer + # (turn-end folded onto doc-eos) while generate_stream reads the + # original, so take marker strings from the mapped template but + # resolve their ids on the original. + try: + _gen_tok = model_info.get("tokenizer") or tokenizer + refreshed = resolve_chat_turn_end_eos_ids_using( + getattr(tokenizer, "tokenizer", tokenizer), + getattr(_gen_tok, "tokenizer", _gen_tok), + ) + existing = model_info.get("chat_turn_end_eos_ids") or [] + model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed)) + except Exception as e: + logger.warning(f"Could not refresh chat turn-end eos after template: {e}") else: logger.info( f"No registered Unsloth template for {self.active_model_name}, using tokenizer default" @@ -1382,7 +1447,8 @@ class InferenceBackend: min_p = min_p, repetition_penalty = repetition_penalty, do_sample = temperature > 0, - eos_token_id = tokenizer.eos_token_id, + # Resolved once at load (chat_template-derived turn-end tokens). + eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id, pad_token_id = tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id, diff --git a/studio/backend/tests/test_chat_eos_template_refresh.py b/studio/backend/tests/test_chat_eos_template_refresh.py new file mode 100644 index 0000000000..75d0117015 --- /dev/null +++ b/studio/backend/tests/test_chat_eos_template_refresh.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Mapper models whose own tokenizer ships no chat_template have their turn-end +eos resolved at LOAD from an empty template (document eos only). The effective +template is installed later, at generate time, via get_chat_template, so the +turn-end-eos cache must be refreshed then; otherwise generate_stream runs past +the ChatML <|im_end|> boundary and loops (the exact bug this PR fixes). +""" + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +# These tests construct InferenceBackend, pulling the full stack. CI may lack +# unsloth/unsloth_zoo (ImportError) or have a broken CUDA/bitsandbytes setup +# (RuntimeError); skip at module level so collection is not aborted (exit 2). +try: + from core.inference import inference as inf_mod # noqa: E402 + from core.inference.inference import InferenceBackend # noqa: E402 +except (ImportError, RuntimeError) as exc: # pragma: no cover - env-dependent + pytest.skip( + f"full inference backend unavailable ({type(exc).__name__}: {exc})", + allow_module_level = True, + ) + +_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" +_GEMMA = "{% for m in messages %}{{m.role}}\n{{m.content}}{% endfor %}" + + +class _FakeTokenizer: + def __init__( + self, + eos_id, + chat_template = "", + token_ids = None, + ): + self.eos_token_id = eos_id + self.chat_template = chat_template + self.pad_token_id = eos_id + self.unk_token_id = None + self._ids = dict(token_ids or {}) + + def convert_tokens_to_ids(self, tok): + return self._ids.get(tok) + + +def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch): + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/qwen2.5-0.5b" + + # No chat_template at load, so the cache stored only the document eos, though + # <|im_end|> is atomic in the vocab (unused until the mapper installs a template). + bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645}) + model_info = { + "tokenizer": bare_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [151643], + } + backend.models = {backend.active_model_name: model_info} + + # The mapper installs a ChatML template (turns end with <|im_end|>) at generate time. + templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False + ) + + # Stub the tail so the generator runs through the refresh without a real model. + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # After the template is applied the cache must include the ChatML turn-end id. + assert model_info["chat_turn_end_eos_ids"] == [151643, 151645] + + +def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkeypatch): + # Regression: get_chat_template can return a remapped tokenizer (Gemma: + # folded onto the eos id) while generate_stream re-reads the original. Resolving on + # the swap yields a narrower set, so the refresh must UNION, never overwrite. + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/gemma-2b-it" + + # Original tokenizer (used by generate_stream): =107 distinct from + # eos=1, so the load-time cache resolved to [1, 107]. + orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 107}) + model_info = { + "tokenizer": orig_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [1, 107], + } + backend.models = {backend.active_model_name: model_info} + + # Destructively-swapped tokenizer: now maps onto eos id 1, so + # resolving on it yields only [1] (drops 107). + swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 1}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False + ) + + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # The load-time =107 must survive: overwriting with the swapped + # [1] would regress and loop past the turn. + assert model_info["chat_turn_end_eos_ids"] == [1, 107] + + +def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkeypatch): + # Yi-style map_eos_token=True: the original carries <|im_end|> at its own id, but + # get_chat_template folds it onto the doc-eos id. generate_stream uses the original, + # so read marker strings from the mapped template but ids from the original. + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "01-ai/yi-6b" + + # Original: no template of its own, doc eos = 2, <|im_end|> atomic = 7. + orig_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7}) + model_info = { + "tokenizer": orig_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [2], + } + backend.models = {backend.active_model_name: model_info} + + # Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2. + remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False + ) + + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2. + assert model_info["chat_turn_end_eos_ids"] == [2, 7] + + +class _FakeProcessor: + """A ProcessorMixin-like container: carries the chat_template itself and + wraps the real text tokenizer as ``.tokenizer`` (the vision layout).""" + + def __init__(self, chat_template, tokenizer): + self.chat_template = chat_template + self.tokenizer = tokenizer + + +def test_resolve_chat_eos_reads_vision_processor_template(): + # Vision model: the chat_template lives on the processor while the inner tokenizer + # ships none. _resolve_chat_eos must read the marker from the processor but resolve + # its id on the inner tokenizer, and repair generation_config. + from types import SimpleNamespace + + inner_tok = _FakeTokenizer(1, chat_template = "", token_ids = {"": 107}) + processor = _FakeProcessor(_GEMMA, inner_tok) + model = SimpleNamespace(generation_config = SimpleNamespace(eos_token_id = 1)) + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/gemma-3-4b-it" + model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True} + backend.models = {backend.active_model_name: model_info} + + backend._resolve_chat_eos(backend.active_model_name) + + assert model_info["chat_turn_end_eos_ids"] == [1, 107] + # generation_config repaired so the vision .generate() path stops at the turn. + assert model.generation_config.eos_token_id == [1, 107] diff --git a/studio/backend/tests/test_chat_turn_end_eos.py b/studio/backend/tests/test_chat_turn_end_eos.py new file mode 100644 index 0000000000..c49e39f8fe --- /dev/null +++ b/studio/backend/tests/test_chat_turn_end_eos.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""chat_eos: resolve assistant-turn-end stop tokens from the chat_template and +repair generation_config so a chat model whose eos is a bare document terminator +(Qwen3.5: config eos <|endoftext|>, turns end with <|im_end|>) stops at the turn +boundary instead of running past it and looping. Dependency-light: imported here +without the full inference stack. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.chat_eos import ( # noqa: E402 + chat_eos_repair, + resolve_chat_turn_end_eos_ids, + resolve_chat_turn_end_eos_ids_using, +) + + +class _FakeTokenizer: + def __init__( + self, + eos_id, + chat_template = "", + token_ids = None, + unk_token_id = None, + ): + self.eos_token_id = eos_id + self.chat_template = chat_template + self.unk_token_id = unk_token_id + self._ids = dict(token_ids or {}) + + def convert_tokens_to_ids(self, tok): + return self._ids.get(tok, self.unk_token_id) + + +# ---- resolve_chat_turn_end_eos_ids --------------------------------------- + +_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" + + +def test_qwen35_adds_im_end_from_template(): + # eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046). + tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046}) + assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046] + + +def test_marker_in_vocab_but_not_in_template_is_ignored(): + # Base/coder model: <|im_end|> is in the vocab but the template does not use + # it, so it must not become a stop token. + tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046}) + assert resolve_chat_turn_end_eos_ids(tok) == [248044] + + +def test_harmony_template_is_left_untouched(): + # gpt-oss/harmony: <|end|> is a channel delimiter, not the turn end. + harmony = "<|start|>assistant<|channel|>analysis<|message|>...<|end|>" + tok = _FakeTokenizer(200002, chat_template = harmony, token_ids = {"<|end|>": 200007}) + assert resolve_chat_turn_end_eos_ids(tok) == [200002] + + +def test_llama3_eot_id_from_template(): + tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009}) + assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009] + + +def test_gemma4_turn_marker_from_template(): + # Gemma-4 ends turns with while keeping a document eos, so must + # be added as a stop token. + tok = _FakeTokenizer( + 1, chat_template = ".........", token_ids = {"": 106} + ) + assert resolve_chat_turn_end_eos_ids(tok) == [1, 106] + + +def test_resolve_using_reads_markers_from_template_but_ids_from_generation_tokenizer(): + # map_eos_token=True: the mapped template remaps <|im_end|> onto the doc-eos id, + # but the original keeps it atomic. Reading marker STRINGS from the template but + # IDS on the original recovers the real turn-end id (7), not the doc-eos id (2). + template_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + id_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7}) + assert resolve_chat_turn_end_eos_ids_using(template_tok, id_tok) == [2, 7] + # Same tokenizer for both reproduces the plain resolve (load-time behaviour). + assert resolve_chat_turn_end_eos_ids_using(template_tok, template_tok) == [2] + + +def test_list_eos_preserved(): + tok = _FakeTokenizer([1, 2], chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + assert resolve_chat_turn_end_eos_ids(tok) == [1, 2] + + +def test_missing_marker_maps_to_unk_and_is_skipped(): + tok = _FakeTokenizer(7, chat_template = _CHATML, token_ids = {}, unk_token_id = 0) + assert resolve_chat_turn_end_eos_ids(tok) == [7] + + +def test_starling_barred_end_of_turn_from_template(): + # OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from + # Gemma's ). eos synced to =2, turn marker at 32000. + starling = "GPT4 Correct Assistant: hi<|end_of_turn|>" + tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000] + + +def test_dict_chat_template_scans_all_variants(): + # Hermes-3 style: chat_template is a {name: template} dict. Detection must scan + # every variant, not bail because the container is not a plain str. + tmpl = {"default": "{{ messages }}", "tool_use": _CHATML} + tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 5] + + +def test_list_of_dicts_chat_template_scans_all_variants(): + # tokenizer_config.json stores multi-templates as a list of {name, template}. + tmpl = [{"name": "default", "template": _CHATML}] + tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 5] + + +def test_dict_harmony_template_left_untouched(): + # A multi-variant container whose variant is harmony must still be left alone. + tmpl = {"default": "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"} + tok = _FakeTokenizer(200002, chat_template = tmpl, token_ids = {"<|end|>": 200007}) + assert resolve_chat_turn_end_eos_ids(tok) == [200002] + + +# ---- chat_eos_repair ------------------------------------------------------ + + +def test_repair_adds_missing_turn_end(): + assert chat_eos_repair(248044, [248044, 248046]) == [248044, 248046] + + +def test_repair_from_missing_generation_config_eos(): + assert chat_eos_repair(None, [248046]) == [248046] + + +def test_repair_noop_when_already_covered(): + assert chat_eos_repair([248046, 248044], [248046]) is None + + +def test_repair_noop_when_no_turn_end_ids(): + assert chat_eos_repair(248044, []) is None From e9f49c62dd078f59421be788ad56ff10ae4b8a01 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:08:39 -0700 Subject: [PATCH 021/113] studio: deterministic backend tool-calling wiring test (#6836) * studio: deterministic backend tool-calling wiring test Add a deterministic, download-free test that exercises the shared tool-calling seam both inference backends use. InferenceBackend (transformers) and MLXInferenceBackend both render the prompt through apply_chat_template_for_generation(..., tools=...) and stream cumulative text into run_safetensors_tool_loop. The existing test_safetensors_tool_loop.py covers the parser and the loop state machine with fake generators but does not cover the backend's own tool-injection seam, so a regression that drops the tool schema before the tokenizer, or fails to feed a tool result back into generation, would slip through. The test drives that seam with fakes: a tokenizer that records the tools it is handed, a canned tool-call generation, and a stub executor. It asserts the full chain: tools reach the chat template, the loop parses the call, the tool is dispatched once with the parsed arguments, the result is fed back, generation re-enters, and the final answer streams after the tool result. It also guards that the raw tool-call markup never leaks to the client as content. The test imports no torch, unsloth, or mlx, so it runs in the portable Backend CI alongside the tool-call parser tests and stays sub-second. Follow-up to the parser test PRs #5620 and #5704. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: assert the tool result is fed back before the final turn Strengthen the wiring test so single_turn records each turn's conversation and the test asserts the tool result message is present in the conversation handed to the final generation turn. Event ordering alone did not catch a loop that stops appending the tool output before re-entering generation, because the fake generation ignores the conversation; this closes that gap. * studio: tighten comments in tool-calling wiring test * studio: shorten comments in tool-calling wiring test --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../tests/test_safetensors_toolcall_wiring.py | 179 ++++++++++++++++++ 1 file changed, 179 insertions(+) create mode 100644 studio/backend/tests/test_safetensors_toolcall_wiring.py diff --git a/studio/backend/tests/test_safetensors_toolcall_wiring.py b/studio/backend/tests/test_safetensors_toolcall_wiring.py new file mode 100644 index 0000000000..5c298a7966 --- /dev/null +++ b/studio/backend/tests/test_safetensors_toolcall_wiring.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Deterministic backend-wiring test for the safetensors / MLX tool-calling path. + +The parser and the cumulative-text state machine are already covered exhaustively by +``test_safetensors_tool_loop.py`` with fake generators. What that suite does not touch is the +*backend's own tool-injection seam*: both ``InferenceBackend`` (transformers) and +``MLXInferenceBackend`` render the prompt through the shared +``apply_chat_template_for_generation(..., tools=...)`` helper and stream cumulative text into the +shared ``run_safetensors_tool_loop`` (see ``core/inference/inference.py`` and +``core/inference/mlx_inference.py`` -- both call the same helper and the same loop, so a single CPU +test of that seam covers the macOS MLX path too). + +This test drives that exact seam with deterministic fakes -- a fake tokenizer that records the +``tools`` it is handed, a canned tool-call generation, and a stub executor -- and asserts the full +agentic chain end to end: + + tools injected into the template -> loop parses the call -> tool dispatched once -> + tool result fed back -> generation re-entered -> final answer streamed. + +It is the deterministic, download-free stand-in for the real-model MLX / GGUF browser tool-calling +end-to-end: it imports no torch / unsloth / mlx, so it runs in the portable Backend CI alongside the +tool-call parser tests. Follow-up to the parser test PRs (#5620 / #5704). +""" + +from core.inference.chat_template_helpers import apply_chat_template_for_generation +from core.inference.safetensors_agentic import run_safetensors_tool_loop + +TOOL_NAME = "get_weather" +TOOL_ARGS = {"city": "Paris"} +FAKE_TOOL = { + "type": "function", + "function": { + "name": TOOL_NAME, + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} +# Full parser matrix lives in test_safetensors_tool_loop.py. +TOOL_CALL_TEXT = '{"name": "get_weather", "arguments": {"city": "Paris"}}' +FINAL_ANSWER = "The weather in Paris is sunny and 22C." +TOOL_RESULT = "Paris: sunny, 22C" + + +class RecordingTokenizer: + """Fake tokenizer that records the ``tools`` handed to ``apply_chat_template``. + + Modelled on ``TestChatTemplateHelper._Tok`` in ``test_safetensors_tool_loop.py``: it accepts the + real helper's kwargs and returns a canned prompt, so the test can assert the backend seam actually + forwarded the tool schema -- a silent drop on a chat-template fallback would leave ``tools_seen`` + holding ``None``. + """ + + def __init__(self): + self.tools_seen: list = [] + self.call_count = 0 + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ): + self.call_count += 1 + self.tools_seen.append(kwargs.get("tools")) + return "PROMPT" + + +class StubExecutor: + """Stand-in for ``core.inference.tools.execute_tool``: records calls, returns a fixed result. + + A fake tool name plus this stub means no real python / terminal / web / RAG side effect can run. + """ + + def __init__(self, result: str): + self.result = result + self.calls: list[tuple[str, dict]] = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + rag_scope = None, + disable_sandbox = False, + ): + self.calls.append((name, arguments)) + return self.result + + +def _collect(generator, max_events = 200): + events = [] + for ev in generator: + events.append(ev) + if len(events) >= max_events: + break + return events + + +def _tool_names(tools): + return [(t.get("function") or {}).get("name") for t in (tools or [])] + + +def test_backend_seam_injects_tools_and_drives_full_tool_loop(): + """The shared backend seam forwards tools into the chat template, and the loop parses the call, + dispatches it once, feeds the result back, and re-enters generation for the final answer.""" + tok = RecordingTokenizer() + executor = StubExecutor(TOOL_RESULT) + turns = iter([TOOL_CALL_TEXT, FINAL_ANSWER]) + active_tools_seen: list = [] + conversations_seen: list = [] + + def single_turn(conversation, *, active_tools = None): + # Mirror the real _single_turn: render via the shared helper, then yield cumulative snapshots. + active_tools_seen.append(active_tools) + conversations_seen.append([dict(m) for m in conversation]) + apply_chat_template_for_generation(tok, conversation, tools = active_tools) + text = next(turns) + mid = len(text) // 2 + acc = "" + for chunk in (text[:mid], text[mid:]): + acc += chunk + yield acc + + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "What is the weather in Paris?"}], + tools = [FAKE_TOOL], + execute_tool = executor, + max_tool_iterations = 3, + ) + ) + + # 1. Helper forwarded the tool schema to the tokenizer (seam does not drop tools). + assert tok.tools_seen, "tokenizer.apply_chat_template was never called" + assert tok.tools_seen[0], "tool schema was dropped before reaching the tokenizer" + assert TOOL_NAME in _tool_names(tok.tools_seen[0]) + + # 2. Loop offered the tool to the first generation turn. + assert active_tools_seen and active_tools_seen[0] is not None + assert TOOL_NAME in _tool_names(active_tools_seen[0]) + + # 3 / 4 / 5. Exactly one tool_start, one dispatch with parsed args, one tool_end with the result. + tool_starts = [e for e in events if e["type"] == "tool_start"] + tool_ends = [e for e in events if e["type"] == "tool_end"] + assert len(tool_starts) == 1 and tool_starts[0]["tool_name"] == TOOL_NAME + assert executor.calls == [(TOOL_NAME, TOOL_ARGS)], executor.calls + assert len(tool_ends) == 1 and tool_ends[0]["result"] == TOOL_RESULT + + # 6. Final answer streams after the tool result: loop appended it and re-entered generation. + contents = [e for e in events if e["type"] == "content"] + assert contents and FINAL_ANSWER in contents[-1]["text"] + last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end") + last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content") + assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result" + + # 6b. Tool result fed back into the conversation before the final turn (6 alone misses this: + # the fake generation ignores the conversation). + assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call" + final_turn_convo = conversations_seen[1] + assert any( + TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo + ), "tool result was not fed back into the conversation before the final generation turn" + + # 7. Guard: raw tool-call markup never leaked to the client as content. + for e in contents: + assert "" not in e["text"] + assert TOOL_NAME not in e["text"] From e9ea45b6a51776cff1185c79c6a9dc0991712080 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:12:22 -0700 Subject: [PATCH 022/113] Studio: coerce tool_call arguments to dict before chat templating (fixes MLX tool follow-up error) (#6807) * Studio: coerce tool_call arguments to dict before chat templating Strict tool chat templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and raise "TypeError: Can only get item pairs from a mapping" when a prior assistant tool call is re-rendered on the next turn. The agentic loop stores arguments in the OpenAI JSON-string form (as_assistant_tool_call), which is correct on the wire and for llama-server, but the transformers / MLX paths apply_chat_template directly and hit the strict Jinja templates. Normalize each assistant tool_call's function.arguments from a JSON string to a dict inside apply_chat_template_for_generation (shared by both the MLX and safetensors paths). A dict renders on strict and lenient templates alike; non-JSON / non-dict values are left untouched, and the OpenAI-format as_assistant_tool_call (used by the GGUF path + API responses) is unchanged. Verified against the real mlx-community/Qwen3.5-2B-8bit template: string args raised the tester's error, the fix renders cleanly, and the lenient unsloth/Qwen3.5-0.8B template still works. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make tool-arg coercion a string-first fallback (non-regressive) Render the original OpenAI string-arg form first and only coerce arguments to a dict when the template raises the mapping TypeError, instead of always coercing. Any template that already renders is now byte-identical (a template that emits arguments verbatim keeps the JSON string, not a Python dict repr). Verified across Llama-3, Qwen2.5, Qwen3, Qwen3.5, Phi-3.5 (byte-identical) and mlx-community/Qwen3.5-2B-8bit (strict -> fixed). Gemma-3 / Mistral tool-template errors are unrelated (role alternation / tool-id length) and identical with or without the change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make core.inference package init lazy so dependency-light helpers import standalone Importing any core.inference submodule ran the package __init__, which eagerly imported orchestrator and llama_cpp; both pull loggers -> structlog (and httpx), so a dependency-light helper like chat_template_helpers dragged in the full heavy stack and its unit test failed to collect in a backend env without structlog. Defer those imports to attribute access via PEP 562 __getattr__, mirroring the lazy pattern already in core/__init__.py. The re-exports resolve unchanged on first access. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Retry dict-coercion for strict templates that raise non-TypeError apply_chat_template_for_generation only retried the OpenAI JSON-string arguments coercion when the first render raised TypeError (the arguments.items() form). The bundled gemma-4.jinja instead rejects string arguments with raise_exception, which surfaces as a Jinja error, so a second tool turn with string function.arguments propagated and failed rather than retrying with the parsed dict. Broaden the outer catch to Exception, still gated on there being a string arg to normalize (normalized is messages -> re-raise), so unrelated template errors and templates that already render are unaffected. * Tighten comments in tool-call argument coercion helper and tests * Tighten tool-call argument coercion comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/chat_template_helpers.py | 90 ++++++++-- .../test_chat_template_tool_arguments.py | 157 ++++++++++++++++++ 2 files changed, 229 insertions(+), 18 deletions(-) create mode 100644 studio/backend/tests/test_chat_template_tool_arguments.py diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index b85e9c348a..f58c93b7fe 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -6,9 +6,51 @@ Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg fallback for templates that reject reasoning/tools args. """ +import json from typing import Optional +def _normalize_tool_call_arguments(messages: list) -> list: + """Coerce each assistant ``tool_calls[].function.arguments`` from a JSON + string to a dict. + + The OpenAI wire format carries ``arguments`` as a JSON string, but some chat + templates (e.g. the stricter Qwen tool templates shipped with mlx-community + checkpoints) iterate ``arguments.items()`` and raise + ``TypeError: Can only get item pairs from a mapping.`` on the string form + when a prior tool call is re-rendered on the next turn. A dict works on both + strict and lenient templates, so parse the string; leave non-JSON or non-dict + values untouched. Returns the original list unchanged when nothing needed + coercing (no copy).""" + mutated = False + out: list = [] + for msg in messages: + tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if not tool_calls: + out.append(msg) + continue + new_calls = [] + msg_changed = False + for call in tool_calls: + fn = call.get("function") if isinstance(call, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + try: + parsed = json.loads(args) + except (ValueError, TypeError): + parsed = None + if isinstance(parsed, dict): + call = {**call, "function": {**fn, "arguments": parsed}} + msg_changed = True + new_calls.append(call) + if msg_changed: + out.append({**msg, "tool_calls": new_calls}) + mutated = True + else: + out.append(msg) + return out if mutated else messages + + def apply_chat_template_for_generation( tokenizer, messages: list, @@ -38,21 +80,33 @@ def apply_chat_template_for_generation( attempts.append(dict(reasoning_kwargs)) attempts.append({}) - last_exc: Optional[Exception] = None - for kwargs in attempts: - try: - return tokenizer.apply_chat_template( - messages, - tokenize = False, - add_generation_prompt = True, - **kwargs, - ) - except TypeError as e: - last_exc = e - continue - except Exception as e: - last_exc = e - break - if last_exc is not None: - raise last_exc - raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + def _render(msgs: list) -> str: + last_exc: Optional[Exception] = None + for kwargs in attempts: + try: + return tokenizer.apply_chat_template( + msgs, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ) + except TypeError as e: + last_exc = e + continue + except Exception as e: + last_exc = e + break + if last_exc is not None: + raise last_exc + raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + + try: + return _render(messages) + except Exception: + # Strict tool templates reject the JSON-string ``arguments`` form via + # TypeError or a broad Jinja raise_exception, so retry with dicts coerced. + # Original messages render first, so working templates stay byte-identical. + normalized = _normalize_tool_call_arguments(messages) + if normalized is messages: + raise + return _render(normalized) diff --git a/studio/backend/tests/test_chat_template_tool_arguments.py b/studio/backend/tests/test_chat_template_tool_arguments.py new file mode 100644 index 0000000000..13d1ecabaa --- /dev/null +++ b/studio/backend/tests/test_chat_template_tool_arguments.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""apply_chat_template_for_generation must coerce assistant tool_call arguments +from the OpenAI JSON-string form to a dict before rendering. Strict tool +templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and +raise "Can only get item pairs from a mapping." on the string form when a prior +tool call is re-rendered on the next turn (MLX + transformers paths). +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.chat_template_helpers import ( # noqa: E402 + _normalize_tool_call_arguments, + apply_chat_template_for_generation, +) + + +def _conv(arguments): + return [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "id": "c1", + "function": {"name": "web_search", "arguments": arguments}, + } + ], + }, + {"role": "tool", "name": "web_search", "content": "21C sunny"}, + ] + + +class _StrictTemplateTokenizer: + """Mimics a strict Qwen tool template: rejects string tool_call arguments.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + args = call.get("function", {}).get("arguments") + if isinstance(args, str): + raise TypeError("Can only get item pairs from a mapping.") + return "RENDERED" + + +def test_string_arguments_are_parsed_to_dict(): + out = _normalize_tool_call_arguments(_conv('{"query": "sweden"}')) + args = out[1]["tool_calls"][0]["function"]["arguments"] + assert args == {"query": "sweden"} + + +def test_dict_arguments_untouched_and_no_copy(): + conv = _conv({"query": "sweden"}) + assert _normalize_tool_call_arguments(conv) is conv + + +def test_non_json_string_left_as_is(): + out = _normalize_tool_call_arguments(_conv("not json")) + assert out[1]["tool_calls"][0]["function"]["arguments"] == "not json" + + +def test_render_succeeds_on_strict_template_with_string_arguments(): + # Regression: strict template + string args used to raise. + result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}')) + assert result == "RENDERED" + + +class _RecordingTokenizer: + """Lenient template: renders whatever arguments it is given (string or dict).""" + + def __init__(self): + self.seen_arguments = None + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + self.seen_arguments = call.get("function", {}).get("arguments") + return "RENDERED" + + +def test_lenient_template_receives_original_string_untouched(): + # Lenient template must see the exact original string, not a coerced dict. + tok = _RecordingTokenizer() + apply_chat_template_for_generation(tok, _conv('{"query": "x"}')) + assert tok.seen_arguments == '{"query": "x"}' + + +def test_messages_without_tool_calls_pass_through_unchanged(): + conv = [{"role": "user", "content": "hi"}] + assert _normalize_tool_call_arguments(conv) is conv + + +class _RaiseExceptionTemplateTokenizer: + """Mimics the bundled gemma-4.jinja: rejects string tool_call arguments via + ``raise_exception(...)``, which surfaces as a Jinja error, NOT a TypeError.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + args = call.get("function", {}).get("arguments") + if isinstance(args, str): + raise ValueError( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string." + ) + return "RENDERED" + + +def test_render_succeeds_on_raise_exception_template_with_string_arguments(): + # Regression: gemma-4.jinja rejects string args via a non-TypeError; retry must still coerce. + result = apply_chat_template_for_generation( + _RaiseExceptionTemplateTokenizer(), _conv('{"query": "x"}') + ) + assert result == "RENDERED" + + +def test_unrelated_template_error_still_propagates_with_dict_args(): + # Failure unrelated to string args (dict args, nothing to coerce) must propagate. + class _AlwaysRaises: + def apply_chat_template(self, messages, **kw): + raise ValueError("template is broken") + + with pytest.raises(ValueError, match = "broken"): + apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"})) From eb1ef44255e4a409c70343611e97c15c2ba197d3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 10:39:37 -0700 Subject: [PATCH 023/113] Studio: Gemma tool-call streaming follow-ups + nested-XML escape fix (#6476) (#6611) * Quote-aware Gemma strip, symmetric unstarted cleanup, ReDoS anchor Address review findings on the tool-strip and streaming paths: - strip_tool_call_markup stripped Gemma-native spans with a plain regex that stops at the first , so a literal close marker inside a <|"|>-quoted argument truncated the span and leaked its suffix into visible text. A brace/quote-aware _strip_gemma_native_spans now removes complete spans (keeping an incomplete one unless final), matching the parser's own balance logic. - The Gemma close pattern this PR added (<\|tool_call>.*?) had no \Z fallback, so a run of unclosed markers backtracked from every open position (quadratic, and the streaming stripper re-scans per token). It is now anchored to (?:|\Z) like routes/inference.py's _TOOL_XML_RE, linear with identical output on well-formed input. - _SameTaskStreamingResponse added unstarted_cleanup for the OpenAI passthrough, but the local GGUF/safetensors streams that enter _TrackedCancel before returning only unregister in the generator finally, which never runs if the client disconnects before the body iterator starts, leaking cancel-registry entries. Each such stream now passes unstarted_cleanup to exit its tracker. - __call__ reads _unstarted_cleanup via getattr so a response built through __new__ (the cancel-timing test) without __init__ does not raise AttributeError; the test also sets the attribute explicitly. - Document that the verbatim /v1/chat/completions passthrough delegates /<|tool_call> splitting to llama-server (--jinja, --reasoning-format auto) and is intentionally not re-parsed locally, noting the llama.cpp dependency. Adds a regression test for the close-marker-inside-quoted-argument strip. * Tighten comments on the tool-strip and streaming paths Compress the verbose comment blocks added with the Gemma tool-call / streaming work to crisp one or two liners, drop restatements of obvious code, and shorten docstrings, keeping the load-bearing rationale (ReDoS anchor, quote-aware strip, unstarted-cleanup, llama.cpp passthrough dependency). Code is unchanged (verified comment-only via AST/ast signature, docstrings stripped). * Harden Gemma parse/strip: span-aware XML fallback and quote-aware streaming - Security: the XML fallback in parse_tool_calls_from_text scanned the whole content for markers and only skipped those inside an open XML parameter, not those inside a collected JSON/Gemma candidate span. A balanced but unparsable Gemma call whose argument data contained XML tool markup (<|tool_call>call:outer{code:...}) therefore fell through to the fallback and returned an executable terminal call. The fallback now also excludes markers inside any candidate span, including ones that failed to parse. - strip_tool_call_markup no longer skips the generic Gemma regex after running the quote-aware _strip_gemma_native_spans, so a closed Gemma span the helper cannot match (malformed, e.g. <|tool_call>{"name":"x"}) is still stripped instead of leaking its opener and payload into visible text. - _strip_gemma_native_spans stops at the first unbalanced start instead of re-scanning every later start to EOF, keeping it linear on a run of unclosed markers rather than quadratic. - The GGUF and safetensors streaming strippers run _strip_gemma_native_spans before the regex patterns, so a well-formed streamed call whose quoted argument contains a literal close marker no longer leaks its suffix into incremental display. Adds regression tests for the nested-XML escape and the malformed-span strip. * Avoid remainder copy in _strip_gemma_native_spans Match the Gemma close marker with re pos directly on the buffer instead of slicing tail = text[brace_end + 1:] on every span. The streaming strippers re-scan a growing cumulative buffer per token, so the per-span remainder copy was quadratic. Behavior is unchanged. * Exclude unclosed Gemma/JSON starts from the XML tool-call fallback The nested-XML guard only skipped markers inside recorded candidate spans, but a span is recorded only when the braces balance. An unbalanced call such as <|tool_call>call:outer{code:... recorded no span, so the fallback still promoted the inner to an executable terminal call. Treat unclosed JSON/Gemma starts as exclusion spans through EOF before scanning. Standalone calls with no preceding unclosed start still parse. Regression tests added. * Skip doomed tool-strip passes to avoid quadratic rescans The lazy closed-pair strip patterns (.*?, .*?) rescan to EOF from every opener when their close token is absent, which is O(n^2) and re-runs per streamed token. Add strip_tool_patterns, which skips a pass whose close token is not present in the text; output is identical to the per-pattern loop (verified by fuzz), and a degenerate run drops from ~minutes to milliseconds. Used by strip_tool_call_markup and the GGUF/safetensors streaming strippers. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use full tool-call envelopes to close nested-XML escape variants Key the parser and stripper off the full <|tool_call>... / ... envelope (start to close marker, searched after the braces; EOF if unclosed) instead of just the braces: - XML between the closing brace and the close marker (call:outer{broken:{x}}...) is now inside the envelope, so the fallback no longer promotes it to a tool call. - A balanced inner call inside an unclosed outer (call:outer{code:<|tool_call>call:terminal{...}) is skipped via the envelope nested check, not just the XML fallback. - strip_tool_call_markup searches for the close marker after the braces, so junk before is stripped through the close and text after it is preserved instead of truncated to EOF; a no-close run stops early (linear). Regression tests added; standalone XML and well-formed calls unaffected. * Fix non-final Gemma strip and missing-close recovery for PR #6611 Split the nested-skip from the XML fallback exclusion: nesting is decided by each marker's brace region, so a balanced call after one with a missing close marker is recovered instead of being swallowed to EOF. Only the XML fallback keeps the search-to-close envelope, so trailing nested markup still cannot escape as an executable call. Use a closed-only Gemma pattern in the non-final strip list so an incomplete block is preserved (matching the JSON and function paths); the final list keeps the close-or-EOF Gemma pattern in its original position, so streaming display output is byte-for-byte unchanged. Add regression tests for both cases. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Block gap-nested tool markers and fix XML strip order for PR #6611 Decide candidate nesting by a per-marker coverage region paired with a per-format stack (a close after the braces pops the nearest still-open marker of that format). A closed outer call now covers up to its own close marker, so a JSON or Gemma tool marker smuggled between the outer braces and that close is treated as data instead of being executed. An outer that balances but has no close of its own covers only its brace region, so a later sibling after an omitted close marker is still recovered (adjacent calls use an exclusive end bound so the next call is not misread as nested). Strip every closed pair (JSON, Gemma, function) before any to-EOF sweep, so a closed function call whose parameter text contains a bare Gemma opener is removed as a unit and the to-EOF sweep can no longer drop the visible text after the close. Add regression tests for both. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Strip closed tool blocks before the Gemma final sweep for PR #6611 The final display strip ran the quote-aware Gemma helper before the closed JSON/function patterns. A closed ... or ... block whose argument data held a call-form Gemma opener (e.g. a "<|tool_call>call:t{" string) was read as an incomplete Gemma span and truncated to EOF, dropping the block's close and any visible text after it. Strip closed JSON/function blocks first, so such a block is removed as a unit before the helper runs. Centralize the final strip order in a shared strip_tool_markup_final so strip_tool_call_markup and both streaming display wrappers (safetensors, llama_cpp) stay in sync, and apply the same closed-block pre-pass to the non-final path. Add regression tests for the JSON and function variants. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Recover XML/JSON siblings after a close-less tool marker for PR #6611 Two fixes so the XML fallback and marker coverage recover a later valid call after an earlier marker omits its close, matching the candidate loop: Reuse the candidate marker-coverage in the XML fallback instead of a separate search-to-close-or-EOF envelope. A balanced but close-less marker now covers only its brace region there too, so a following sibling is recovered rather than filtered as nested data; an unbalanced marker still covers to EOF and a closed one still covers through its close, so nested XML stays blocked. Ignore a close token that falls inside another call's balanced braces when pairing closes in _marker_coverage. Such a token is that call's quoted argument data, so it no longer pops an earlier close-less marker and extends its coverage over a later valid sibling. Add regression tests for both. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make the closed-block strip pre-pass Gemma-span-aware The final display strip ran the closed JSON/function regex pre-pass before removing Gemma-native spans, so a literal quoted inside a Gemma argument plus any later (a real call's close or even prose) was deleted across the Gemma boundary. That mangled the Gemma close marker, the quote-aware helper then saw an unclosed opener, and the whole visible tail after the call was truncated. The pre-pass now skips matches that start inside a complete Gemma span (that text is the span's argument data) and resumes scanning at the end of the covering span, so a real function-XML call after the Gemma call is still stripped. The original ordering rationale is preserved: a Gemma opener inside a JSON or function argument still cannot truncate that block, covered by regression tests for both directions. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim comments in the Gemma streaming and strip pipeline to essentials * Tighten comments in the Gemma strip and streaming disconnect paths * Fold marker-collection comment to two lines --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../core/inference/tool_call_parser.py | 14 +- studio/backend/core/tool_healing.py | 366 +++++++++++------- studio/backend/routes/inference.py | 87 +++-- .../tests/test_gemma_tool_parse_edge_cases.py | 195 ++++++++-- .../tests/test_tool_call_parser_strict.py | 41 ++ studio/backend/tests/test_tool_strip_guard.py | 76 ++++ 6 files changed, 562 insertions(+), 217 deletions(-) create mode 100644 studio/backend/tests/test_tool_strip_guard.py diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index c31f4b272e..9e82e40de2 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -530,13 +530,23 @@ def parse_tool_calls_from_text( # Formats tool_healing does not cover: ```` (MiniCPM-5 / MiniMax-M2), # Llama-3 and Mistral. Run only after tool_healing found nothing, so a strict-rejected - # call is never re-healed here. + # call is never re-healed here. Blank any JSON/Gemma marker coverage first: markup inside + # a marker's span (even one that failed to parse) is that call's data, not a sibling, so + # a nested ```` / ``<|python_tag|>`` / ``[TOOL_CALLS]`` must not be promoted. + fallback_content = content + coverage = _tool_healing.marker_coverage(content) + if coverage: + chars = list(content) + for cov_start, cov_end in coverage: + for i in range(cov_start, min(cov_end, len(chars))): + chars[i] = " " + fallback_content = "".join(chars) for parser in ( _parse_function_xml, # attribute form _parse_llama3_python_tag, # Llama-3 <|python_tag|> _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] ): - calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + calls = parser(fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete) if calls: return calls diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index ff8faf2308..b91403ed57 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -10,20 +10,45 @@ orchestrator, structlog, httpx, or the rest of the studio backend. import json import re -# Pre-compiled patterns for tool XML stripping. The hyphen in the name -# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues, -# issue-number) parse alongside the built-ins. +# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs +# strip first so a closed call goes as a unit before any to-EOF sweep reaches +# nested markup; only the final list adds the .*$ EOF sweeps. +_TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) +_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) +_TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) +_TC_GEMMA_END_PAT = re.compile(r"") _TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), - re.compile(r"<\|tool_call>.*?", re.DOTALL), - re.compile(r""), - re.compile(r".*?", re.DOTALL), + _TC_JSON_CLOSED_PAT, + _TC_GEMMA_CLOSED_PAT, + _TC_FUNC_CLOSED_PAT, + _TC_GEMMA_END_PAT, ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ - re.compile(r".*$", re.DOTALL), re.compile(r"<\|tool_call>.*$", re.DOTALL), + re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] +# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in +# their argument data cannot make the helper truncate the block and its tail. +_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT] +# A lazy closed-pair pattern whose close token is absent rescans to EOF from +# every opener (quadratic, re-run per streamed token); skip that doomed pass. +_PAT_REQUIRED_TOKEN = { + _TC_JSON_CLOSED_PAT: "", + _TC_GEMMA_CLOSED_PAT: "", + _TC_FUNC_CLOSED_PAT: "", +} + + +def strip_tool_patterns(text: str, patterns) -> str: + """Apply ``patterns`` in order, skipping closed-pair passes with no close token.""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + text = pat.sub("", text) + return text + # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") @@ -40,13 +65,9 @@ _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _GEMMA_QUOTE = '<|"|>' _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" -# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next -# `key:` pair. A comma NOT followed by a key token is part of the value (e.g. -# `location:New York, NY`), so it must not terminate the value. The key token -# must be identifier-shaped (start with a letter or underscore); a comma -# followed by digits-then-colon is value text such as a timestamp or ratio -# (`meet at 10:00, 11:00 tomorrow`), not a new key. -# Dots match the key-quoting scanner: a dotted key after a bare value must end the value at the comma. +# A bare (unquoted) Gemma value ends at `}` or at a comma beginning the next +# identifier-shaped `key:` pair; a comma before a non-key (`New York, NY`, +# `10:00, 11:00`) stays in the value. Dots let a dotted key end the value. _GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") @@ -143,14 +164,8 @@ def _split_top_level_commas(src: str) -> list: def _quote_gemma_array_elements(body: str) -> str: - """Normalise the elements of a Gemma array value so json.loads succeeds. - - Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of - objects (``items:[{path:a}]``) whose keys/values also lack quotes; left - as-is json.loads fails and the whole call is dropped. Bare string elements - are quoted, object and nested-array elements are normalised recursively, and - quoted strings (already normalised from ``<|"|>``), numbers, and JSON - literals are preserved.""" + """Normalise a Gemma array value (``labels:[bug,ui]``) so json.loads succeeds: + quote bare strings, recurse into objects/arrays, keep quoted/JSON literals.""" out: list[str] = [] for element in _split_top_level_commas(body): stripped = element.strip() @@ -158,11 +173,9 @@ def _quote_gemma_array_elements(body: str) -> str: out.append(element) continue if stripped[0] == "{": - # Object element: quote its keys/bare values like a top-level object. out.append(_quote_gemma_object_keys(stripped)) continue if stripped[0] == "[": - # Nested array: normalise its elements too. inner_end = _balanced_bracket_end(stripped, 0) if inner_end == len(stripped) - 1: out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]") @@ -241,15 +254,12 @@ def _quote_gemma_object_keys(src: str) -> str: parts.append(src[i:colon_pos]) parts.append(":") i = colon_pos + 1 - # Gemma may emit bare string values ({unit:celsius}); quote them so - # json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is. + # Quote bare string values ({unit:celsius}); JSON stays as-is. ws = i while i < len(src) and src[i].isspace(): i += 1 parts.append(src[ws:i]) if i < len(src) and src[i] == "[": - # Array value: quote bare string elements (e.g. labels:[bug,ui]) - # so json.loads succeeds instead of dropping the call. arr_end = _balanced_bracket_end(src, i) if arr_end < 0: parts.append(src[i:]) @@ -259,9 +269,7 @@ def _quote_gemma_object_keys(src: str) -> str: i = arr_end + 1 elif i < len(src) and src[i] not in '"{': v_start = i - # Consume the bare value up to `}` or a comma that starts the - # next key:value pair; a comma inside the value (e.g. - # `New York, NY`) does not terminate it. + # Bare value: up to `}` or a comma that starts the next key:pair. while i < len(src): if src[i] == "}": break @@ -329,6 +337,68 @@ def _trim_param_value(val: str) -> str: return val +def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: + """Coverage ``[start, end]`` per marker, used to skip markers that are another + call's data. Closes pair to markers via a per-format stack so an inner close + is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with + a paired close cover through it (markers before the close are data); balanced + without one cover only the braces, so a later sibling is still recovered.""" + n = len(content) + brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0] + events = [] # (position, order) with order 0 = braces-done, 1 = close marker + for idx, (_start, brace_end, _kind, _m) in enumerate(markers): + if brace_end >= 0: + events.append((brace_end, 0, _kind, idx)) + for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)): + for cm in close_re.finditer(content): + # A close inside another call's balanced braces is quoted data; it + # must not pop an earlier close-less marker and swallow a sibling. + if any(s < cm.start() < be for s, be in brace_regions): + continue + events.append((cm.start(), 1, kind, cm.end())) + events.sort(key = lambda e: (e[0], e[1])) + waiting = {"json": [], "gemma": []} + close_end_for: dict[int, int] = {} + for _pos, order, kind, payload in events: + if order == 0: + waiting[kind].append(payload) # marker index, now awaiting its close + elif waiting[kind]: + close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here + coverage = [] + for idx, (start, brace_end, _kind, _m) in enumerate(markers): + if brace_end < 0: + coverage.append((start, n)) + elif idx in close_end_for: + coverage.append((start, close_end_for[idx])) + else: + coverage.append((start, brace_end)) + return coverage + + +def _build_markers(content: str): + """JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document + order; ``brace_end < 0`` marks an unbalanced (to-EOF) open.""" + markers = [] + for start_re, gemma, kind in ( + (_TC_JSON_START_RE, False, "json"), + (_TC_GEMMA_START_RE, True, "gemma"), + ): + for m in start_re.finditer(content): + if _inside_open_parameter(content, m.start()): + continue + brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma) + markers.append((m.start(), brace_end, kind, m)) + markers.sort(key = lambda c: c[0]) + return markers + + +def marker_coverage(content: str) -> list[tuple[int, int]]: + """Coverage spans of JSON/Gemma tool markers so other parsers can treat markup + inside a marker's coverage (even a marker that failed to parse) as that call's + data rather than a sibling call.""" + return _marker_coverage(content, _build_markers(content)) + + def parse_tool_calls_from_text( content: str, *, @@ -350,37 +420,26 @@ def parse_tool_calls_from_text( """ tool_calls: list[dict] = [] call_spans: list[tuple] = [] - # Collect every supported call format with spans, then emit in document - # order. A marker inside another call's argument string is data, not a - # separate executable call. - parsed_items = [] # (start, span_end, name, arguments) - candidates = [] # (start, brace_end, kind, match) - for m in _TC_JSON_START_RE.finditer(content): - if _inside_open_parameter(content, m.start()): - continue - end = _balanced_brace_end(content, m.end() - 1) - if end >= 0: - candidates.append((m.start(), end, "json", m)) - for m in _TC_GEMMA_START_RE.finditer(content): - if _inside_open_parameter(content, m.start()): - continue - end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True) - if end >= 0: - candidates.append((m.start(), end, "gemma", m)) - candidates.sort(key = lambda c: c[0]) - - candidate_spans = [(s, e) for s, e, _kind, _m in candidates] - for idx, (start, end, kind, m) in enumerate(candidates): - if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx): + # Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside + # another call's coverage, or an open value, is data not executed. + markers = _build_markers(content) + coverage = _marker_coverage(content, markers) + parsed_items = [] # (start, span_end, name, arguments) in document order + for idx, (start, brace_end, kind, m) in enumerate(markers): + # A marker starting inside another's coverage is that call's data. The + # end is exclusive so a marker at a close's end is an adjacent sibling. + if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx): continue + if brace_end < 0: + continue # unclosed: not parseable; the fallback still excludes its XML if not allow_incomplete: - tail = content[end + 1 :].lstrip() + tail = content[brace_end + 1 :].lstrip() close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE if close_re.match(tail) is None: continue try: if kind == "json": - obj = json.loads(content[m.end() - 1 : end + 1]) + obj = json.loads(content[m.end() - 1 : brace_end + 1]) name = obj.get("name", "") # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes ). arguments = obj.get("arguments") @@ -390,10 +449,11 @@ def parse_tool_calls_from_text( arguments = json.dumps(arguments) else: name = m.group(1) - arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end])) + arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) except (json.JSONDecodeError, ValueError): continue - span_end = end + 1 + # Span reaches through the close tag when present, else just the braces. + span_end = brace_end + 1 close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE ws = len(content[span_end:]) - len(content[span_end:].lstrip()) close_m = close_re.match(content, span_end + ws) @@ -401,11 +461,15 @@ def parse_tool_calls_from_text( span_end = close_m.end() parsed_items.append((start, span_end, name, arguments)) + # Function-XML calls promote in document order alongside marker calls (the + # #6801 contract). A inside any marker's coverage is excluded -- + # even if that marker failed to parse -- so nested XML cannot escape; one + # after a balanced close-less marker is a sibling, not swallowed to EOF. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) - and not any(s <= fm.start() <= e for s, e in candidate_spans) + and not any(s <= fm.start() < e for s, e in coverage) ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) @@ -481,90 +545,106 @@ def parse_tool_calls_from_text( ) call_spans.append((start, span_end)) - if not tool_calls: - func_starts = [ - fm - for fm in _TC_FUNC_START_RE.finditer(content) - if not _inside_open_parameter(content, fm.start()) - ] - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - # Span for with_spans callers: through the close if present, else body end. - span_end = body_end - if not allow_incomplete: - close_idx = _func_close_index(content, body_start, body) - if close_idx < 0: - continue - body = body[:close_idx] - span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) - else: - # Terminate at the real close so trailing prose doesn't leak in; no close -> whole body. - close_idx = _func_close_index(content, body_start, body) - if close_idx >= 0: - body = body[:close_idx] - span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - pm = param_starts[0] - val = body[pm.end() :] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - continue - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = _trim_param_value(val) - else: - valid_params = True - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - valid_params = False - break - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = _trim_param_value(val) - if not valid_params: - continue - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) - call_spans.append((fm.start(), span_end)) - if with_spans: return tool_calls, call_spans return tool_calls +def _strip_gemma_native_spans(text: str, *, final: bool) -> str: + """Remove complete Gemma-native spans, brace/quote-balanced so a literal + ```` in a quoted argument cannot truncate the span. An incomplete + span is dropped to EOF when ``final``, else kept (still streaming).""" + out: list[str] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + # Unbalanced: nothing completes from here on. Drop the rest if final, + # else keep it; stop either way (rescanning would be quadratic). + if final: + out.append(text[cursor:start]) + cursor = len(text) + break + # Junk between } and is malformed-call markup: strip through + # the close, keep text after it. No close anywhere means stop (linear). + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + if final: + out.append(text[cursor:start]) + cursor = len(text) + break + out.append(text[cursor:start]) + cursor = close.end() + out.append(text[cursor:]) + return "".join(out) + + +def _gemma_span_ranges(text: str) -> list: + """``(start, end)`` of each complete Gemma-native span; same walk as + ``_strip_gemma_native_spans`` without stripping.""" + ranges: list[tuple] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + break + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + break + ranges.append((start, close.end())) + cursor = close.end() + return ranges + + +def _strip_closed_blocks_outside_gemma(text: str) -> str: + """Closed JSON/function pre-pass that skips matches starting inside a complete + Gemma span: deleting across the span boundary would mangle the Gemma close and + truncate the tail. A skipped match resumes at the covering span's end, so a + real function-XML call after the span is still stripped.""" + ranges = _gemma_span_ranges(text) + if not ranges: + return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + for pat in _TOOL_CLOSED_BLOCK_PATS: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + out: list[str] = [] + pos = 0 + while True: + m = pat.search(text, pos) + if m is None: + out.append(text[pos:]) + break + covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None) + if covering is not None: + out.append(text[pos : covering[1]]) + pos = covering[1] + continue + out.append(text[pos : m.start()]) + pos = m.end() + new_text = "".join(out) + if new_text != text: + text = new_text + ranges = _gemma_span_ranges(text) + return text + + +def strip_tool_markup_final(text: str) -> str: + """Final display strip, shared with the streaming wrappers so all paths order + the passes identically: Gemma-aware closed JSON/function blocks first, then + well-formed Gemma spans (quote-aware), then the regex sweeps mop up malformed + spans and drop any unclosed remainder to EOF. Whitespace is kept.""" + text = _strip_closed_blocks_outside_gemma(text) + text = _strip_gemma_native_spans(text, final = True) + return strip_tool_patterns(text, _TOOL_ALL_PATS) + + def strip_tool_call_markup(text: str, *, final: bool = False) -> str: """Strip tool-call XML markup from text. @@ -572,7 +652,9 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str: When ``final`` is True, trailing incomplete tool-call blocks are removed too, and the result is stripped of surrounding whitespace. """ - patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in patterns: - text = pat.sub("", text) - return text.strip() if final else text + if final: + return strip_tool_markup_final(text).strip() + # Non-final: same ordering as the final path, but incomplete blocks are kept. + text = _strip_closed_blocks_outside_gemma(text) + text = _strip_gemma_native_spans(text, final = False) + return strip_tool_patterns(text, _TOOL_CLOSED_PATS) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4393c1b304..1a1a934009 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -852,17 +852,14 @@ class _SameTaskStreamingResponse(StreamingResponse): **kwargs, ) -> None: super().__init__(*args, **kwargs) - # Async callable invoked when the client disconnects before the body - # iterator is ever advanced. A generator that never started cannot run - # its own try/finally, so a stream that acquires resources before its - # first yield (the passthrough opens an upstream httpx stream eagerly) - # passes this to release them. + # Released when the client disconnects before the body iterator starts: + # its try/finally never runs, so a stream that opens resources before the + # first yield (the passthrough's upstream httpx stream) passes this. self._unstarted_cleanup = unstarted_cleanup async def __call__(self, scope, receive, send) -> None: - # Track whether the body iterator was ever advanced: send() only emits a - # body message after the generator yields its first chunk, so a failure - # before then means it never entered its try/finally. + # send() emits a body message only after the first chunk, so no body + # message means the generator never entered its try/finally. body_started = False async def _tracking_send(message) -> None: @@ -873,15 +870,11 @@ class _SameTaskStreamingResponse(StreamingResponse): try: await self.stream_response(_tracking_send) - except OSError: - # Client disconnected mid-send. + except OSError: # client disconnected mid-send if body_started: - # The generator produced at least one chunk and is suspended in - # its try/finally. Throw CancelledError into it (not aclose's - # GeneratorExit) so its `except asyncio.CancelledError` handler - # runs and finishes any api_monitor entry; GeneratorExit would - # skip it and only run `finally`. Fall back to aclose() without - # athrow. + # Generator is suspended in its try/finally: throw CancelledError + # (not aclose's GeneratorExit) so its handler finishes the + # api_monitor entry. Fall back to aclose() without athrow. athrow = getattr(self.body_iterator, "athrow", None) if athrow is not None: try: @@ -893,16 +886,16 @@ class _SameTaskStreamingResponse(StreamingResponse): if aclose is not None: await aclose() else: - # http.response.start failed before the body iterator advanced, - # so its try/finally never armed and aclose()/athrow() are no-ops - # on an unstarted generator. Release any resources acquired - # before the first yield via the explicit cleanup hook. + # Generator never started; aclose()/athrow() are no-ops on it, so + # release eager resources via the hook. getattr guards a response + # built through __new__ without __init__ (tests, pickling). aclose = getattr(self.body_iterator, "aclose", None) if aclose is not None: await aclose() - if self._unstarted_cleanup is not None: + cleanup = getattr(self, "_unstarted_cleanup", None) + if cleanup is not None: try: - await self._unstarted_cleanup() + await cleanup() except Exception: pass raise ClientDisconnect() @@ -910,6 +903,16 @@ class _SameTaskStreamingResponse(StreamingResponse): await self.background() +def _tracked_cancel_unstarted_cleanup(tracker): + """unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when + the generator's finally (which normally exits it) never runs.""" + + async def _cleanup() -> None: + tracker.__exit__(None, None, None) + + return _cleanup + + async def _aclose_stream_resources( *, watchers = (), @@ -4069,12 +4072,9 @@ async def generate_stream( _DONE = object() while True: if cancel_event.is_set(): - # The disconnect watcher set cancel_event between chunks. - # Reset the backend here: closing the Python generator does - # not signal a subprocess backend, so without this it keeps - # decoding after the client is gone. The finally's reset is - # guarded on cancel_event being unset, so it will not run - # again for this path. + # Watcher set cancel_event between chunks. Reset here: closing + # the generator does not signal a subprocess backend, so it would + # keep decoding. The finally's reset is guarded, so no double-run. backend.reset_generation_state() break chunk = await asyncio.to_thread(next, gen, _DONE) @@ -5684,6 +5684,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( audio_input_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6163,6 +6164,7 @@ async def openai_chat_completions( if payload.stream: return _SameTaskStreamingResponse( gguf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6419,6 +6421,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( gguf_stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6852,6 +6855,7 @@ async def openai_chat_completions( if payload.stream: return _SameTaskStreamingResponse( sf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -7067,6 +7071,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -9977,11 +9982,8 @@ async def _anthropic_tool_stream( drop_until_tool_end = False gen = run_gen() - # Concurrent disconnect watcher: the loop only polls is_disconnected() - # between events, so a client disconnect during a long prefill or - # generation step would otherwise hold the decode slot until the next - # event or a failed send. The watcher sets cancel_event so the backend - # stops promptly. + # Watcher to cancel on disconnect: the in-loop poll fires only between + # events, so a mid-prefill disconnect would otherwise hold the decode slot. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -10073,11 +10075,8 @@ async def _anthropic_plain_stream( captured_finish_reason = None gen = run_gen() - # Concurrent disconnect watcher: the loop only polls is_disconnected() - # between chunks, so a client disconnect during a long prefill or - # generation step would otherwise hold the decode slot until the next - # chunk or a failed send. The watcher sets cancel_event so the backend - # stops promptly. + # Watcher to cancel on disconnect: the in-loop poll fires only between + # chunks, so a mid-prefill disconnect would otherwise hold the decode slot. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -11030,6 +11029,10 @@ async def _openai_passthrough_stream( response ``id``, ``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so the client sees a standard OpenAI response. + + Reasoning/tool-call splitting is delegated to llama-server (``--jinja + --reasoning-format auto``), so ``delta.content`` carries no raw markup and is + deliberately not re-parsed locally, unlike the ``/completion`` paths. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" body = _build_openai_passthrough_body( @@ -11446,11 +11449,9 @@ async def _openai_passthrough_stream( delta = choice.get("delta") if isinstance(delta, dict) and delta.get("tool_calls"): saw_tool_call_delta = True - # Detect an upstream error chunk independently of API - # monitoring: when monitor_id is None (skip_api_monitor), - # _monitor_openai_sse_line returns before inspecting the - # error, so without this the synthetic-finish guard would - # emit a successful finish_reason after a failed stream. + # Detect an error chunk independently of API monitoring + # (skip_api_monitor returns early), else the synthetic + # finish would fire after a failed stream. if _monitor_openai_error_message(chunk_data): saw_stream_error = True # With healing active, a content-bearing line may be replaced by diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index 63df86ec17..fff6b240c5 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -1,15 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Edge cases in Gemma-native tool-call parsing. - -Covers two failure modes: - 1. A bare (unquoted) string argument that contains a comma, e.g. - ``location:New York, NY`` -- the comma must not be treated as the next - key boundary, or the whole call is dropped. - 2. A tool-call marker that appears INSIDE another call's argument string is - data, not a real call, so it must not be promoted to a second tool call. -""" +"""Gemma-native tool-call parsing edge cases: commas inside bare string values, +and markers inside another call's argument data staying data.""" from __future__ import annotations @@ -25,6 +18,7 @@ from core.inference.tool_call_parser import ( _gemma_parse_value, parse_tool_calls_from_text, ) +from core.tool_healing import strip_tool_call_markup def _args(call: dict) -> dict: @@ -43,8 +37,6 @@ def test_bare_string_argument_with_comma_is_kept(): def test_normal_multi_key_arguments_still_split(): calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}') assert len(calls) == 1, calls - # Numbers stay numeric, bare strings get quoted, an explicit quoted comma - # stays inside its value. assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} @@ -60,8 +52,7 @@ def test_empty_bare_value_becomes_empty_string_not_dropped(): def test_bare_value_with_timestamps_after_comma_is_kept(): - # A comma followed by digits-then-colon (a timestamp/ratio) is value text, - # not a new key, so the whole query must be preserved as one argument. + # A comma before digits-then-colon (timestamp/ratio) is value text, not a key. calls = parse_tool_calls_from_text( "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" ) @@ -70,8 +61,6 @@ def test_bare_value_with_timestamps_after_comma_is_kept(): def test_marker_inside_json_argument_is_not_a_second_call(): - # A python call whose `code` argument contains a Gemma marker string. The - # marker is data and must not execute as a second `terminal` call. content = ( '{"name":"python","arguments":{"code":' '"x = 1 # <|tool_call>call:terminal{command:ls}"}}' @@ -89,8 +78,6 @@ def test_two_separate_gemma_calls_both_parse(): def test_mixed_format_calls_preserve_document_order(): - # A Gemma-native call precedes a JSON-format call in the text; tools execute - # in returned order, so `create` must come before `read`. content = ( "<|tool_call>call:create{path:a} then " '{"name":"read","arguments":{"path":"a"}}' @@ -100,8 +87,6 @@ def test_mixed_format_calls_preserve_document_order(): def test_json_marker_inside_gemma_argument_is_not_a_second_call(): - # The reverse of the JSON-outer case: a JSON-style marker inside a Gemma - # call's quoted argument is code text, not a second `terminal` call. content = ( '<|tool_call>call:python{code:<|"|>' 'print({"name":"terminal","arguments":{"command":"ls"}})' @@ -112,18 +97,14 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call(): def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call(): - # An UNQUOTED Gemma value containing a literal marker: the outer object fails - # to normalize (the inner braces/marker break the JSON), but the inner marker - # is nested in the outer candidate span, so it must not be promoted to a - # standalone `terminal` call. The safe outcome is no executed tool call. + # The outer object fails to normalize, but the nested marker is covered by + # its span; safe outcome is no executed call at all. content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}}" calls = parse_tool_calls_from_text(content) assert "terminal" not in [c["function"]["name"] for c in calls], calls def test_bare_string_array_argument_is_quoted(): - # Gemma may emit an array of bare strings without per-element quotes; they - # must be quoted so the call is not dropped. calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}") assert len(calls) == 1, calls assert _args(calls[0]) == {"labels": ["bug", "ui"]} @@ -137,8 +118,6 @@ def test_array_keeps_numbers_and_quoted_elements(): def test_array_of_objects_is_normalised(): - # Arrays of objects are a common tool-schema shape; their (unquoted) keys and - # bare values must be normalised too, not left verbatim, or the call drops. calls = parse_tool_calls_from_text( "<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}" ) @@ -152,9 +131,6 @@ def test_nested_array_elements_are_normalised(): def test_gemma_marker_inside_xml_parameter_is_not_a_second_call(): - # An XML-style call whose value contains a - # Gemma marker: the marker is the parameter's data, not a separate terminal - # call, so only the python call must be returned. content = ( "" "x = 1 # <|tool_call>call:terminal{command:ls}" @@ -175,6 +151,165 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): assert [c["function"]["name"] for c in calls] == ["python"], calls +def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): + # Parse keeps the quoted close marker as data; strip removes the whole span. + text = '<|tool_call>call:python{code:<|"|>print("")<|"|>}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1, calls + assert _args(calls[0]) == {"code": 'print("")'} + assert strip_tool_call_markup("before " + text + " after") == "before after" + assert strip_tool_call_markup("before " + text + " after", final = True) == "before after" + + +def test_nested_xml_in_malformed_gemma_call_does_not_execute(): + # The failed Gemma candidate's span still covers its nested . + text = ( + "<|tool_call>call:outer{code:id" + ", broken:{x}}" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_unbalanced_gemma_call_with_xml_does_not_execute(): + # Unclosed braces cover to EOF, so the trailing is excluded. + text = ( + "<|tool_call>call:outer{code:" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_standalone_function_xml_still_parses(): + text = "id" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_xml_between_braces_and_close_marker_does_not_execute(): + # Coverage runs to the close marker, so in the gap is data. + text = ( + "<|tool_call>call:outer{broken:{x}}" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_balanced_inner_call_inside_unclosed_outer_does_not_execute(): + text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_preserves_text_after_malformed_gemma_close(): + # Junk before the close is a malformed span: strip through it, keep the tail. + text = "pre <|tool_call>call:t{a:1} note post" + assert strip_tool_call_markup(text) == "pre post" + assert strip_tool_call_markup(text, final = True) == "pre post" + + +def test_malformed_closed_gemma_span_is_stripped(): + assert ( + strip_tool_call_markup('before <|tool_call>{"name":"x"} after') + == "before after" + ) + + +def test_valid_call_after_missing_close_is_recovered(): + # A close-less call covers only its braces, so the later call is recovered. + text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}" + names_inc = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True) + ] + assert "b" in names_inc, names_inc + names_strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] + assert names_strict == ["b"], names_strict + + +def test_strip_non_final_keeps_incomplete_gemma_block(): + text = "before <|tool_call>call:t{" + assert strip_tool_call_markup(text) == text + assert strip_tool_call_markup(text, final = True) == "before" + + +def test_json_call_between_gemma_braces_and_close_does_not_execute(): + # A JSON call between the outer's braces and its close is covered data. + text = ( + "<|tool_call>call:outer{broken:{x}}" + '{"name":"terminal","arguments":{"command":"id"}}' + "" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_gemma_call_between_gemma_braces_and_close_does_not_execute(): + # Same escape with a Gemma-native inner marker. + text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): + # The to-EOF Gemma sweep must not eat visible text after . + text = ( + 'before print("<|tool_call>") after' + ) + assert strip_tool_call_markup(text, final = True) == "before after" + assert strip_tool_call_markup(text) == "before after" + + +def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener(): + # A call-form Gemma opener quoted in a closed block must not truncate it. + xml = "<|tool_call>call:t{" + json_block = ( + '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' + ) + for block in (xml, json_block): + text = "before " + block + " after" + assert strip_tool_call_markup(text, final = True) == "before after", block + assert strip_tool_call_markup(text) == "before after", block + + +def test_function_sibling_after_close_less_gemma_marker_is_recovered(): + # The close-less marker covers only its braces; the XML sibling is recovered. + text = ( + "<|tool_call>call:bad{broken:{x}} " + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered(): + # A close token quoted in the later call must not extend the earlier + # close-less marker's coverage over that call. + gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|><|"|>}' + names = [ + c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False) + ] + assert names == ["b"], names + json_text = ( + '{"name":"a","arguments":{}} ' + '{"name":"b","arguments":{"x":""}}' + ) + names_j = [ + c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) + ] + assert "b" in names_j, names_j + + def test_gemma_parse_value_always_advances_on_stray_delimiter(): # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the # parser, or a looping caller spins forever (DoS). diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 7664126d91..fded2a8443 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -1063,3 +1063,44 @@ class TestBareJsonStripRequiresTopLevelName: def test_real_call_still_strips_name_agnostic(self): from core.inference.tool_call_parser import strip_leading_bare_json_call assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') == "" + + +class TestGemmaAwareClosedBlockPrePass: + """The closed JSON/function strip pre-pass must not delete across a complete + Gemma span (a quoted plus a later real ).""" + + def test_literal_function_in_gemma_arg_with_later_real_call(self): + from core.tool_healing import strip_tool_call_markup + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " ls" + " after" + ) + assert strip_tool_call_markup(text, final = True) == "before after" + + def test_literal_function_in_gemma_arg_with_prose_closer(self): + from core.tool_healing import strip_tool_call_markup + + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " then use to close. after" + ) + out = strip_tool_call_markup(text, final = True) + assert out.startswith("before") + assert out.endswith("after") + assert "call:python" not in out + + def test_gemma_opener_inside_json_arg_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + text = ( + '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' + ) + assert strip_tool_call_markup(text, final = True) == "after" + + def test_gemma_opener_inside_function_param_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + text = ( + 'x = "<|tool_call>call:t{"' + " after" + ) + assert strip_tool_call_markup(text, final = True) == "after" diff --git a/studio/backend/tests/test_tool_strip_guard.py b/studio/backend/tests/test_tool_strip_guard.py new file mode 100644 index 0000000000..dfa3101882 --- /dev/null +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""strip_tool_patterns must match the plain per-pattern loop while skipping the +quadratic no-match rescan of a closed-pair sweep whose close token is absent.""" + +import random +import sys +import time +from pathlib import Path + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from core.tool_healing import ( + _TOOL_ALL_PATS, + _TOOL_CLOSED_PATS, + strip_tool_call_markup, + strip_tool_patterns, +) + + +def _naive(text, patterns): + for pat in patterns: + text = pat.sub("", text) + return text + + +_TOKENS = [ + "", + "", + "<|tool_call>", + "", + "", + "", + "", + "", + "", + "call:fn{", + "}", + "{", + '<|"|>', + "A", + " ", + "\n", + "id", + "x:1", + "", +] + + +def test_guard_matches_plain_loop_on_fuzz(): + rng = random.Random(1234) + for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS): + for _ in range(20000): + s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10))) + assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns) + + +def test_strip_markup_representative_cases_unchanged(): + assert strip_tool_call_markup("a {} b") == "a b" + assert strip_tool_call_markup("a 1 b") == "a b" + # Non-final keeps an unclosed block; final strips it to EOF. + assert strip_tool_call_markup("a {partial") == "a {partial" + assert strip_tool_call_markup("a {partial", final = True) == "a" + + +def test_no_quadratic_blowup_on_unclosed_markers(): + # Unguarded, this took minutes. + big = "" * 20000 + "" * 20000 + t0 = time.perf_counter() + out = strip_tool_call_markup(big, final = True) + assert time.perf_counter() - t0 < 2.0 + assert out == "" From c00c1e70c8a9f5a4cdbac61fc73b55e89c52be08 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 15:40:46 -0700 Subject: [PATCH 024/113] studio: tool calling for DeepSeek (R1/V3/V3.1), GLM 4.x, Kimi K2 on safetensors + MLX (#5624) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * studio: tool calling for Llama-3, Mistral, Gemma 4 on safetensors + MLX (#5615) Adds tool calling for Llama-3, Mistral (pre-v11 + v11+ + [ARGS]), and Gemma 4 to the safetensors / transformers and MLX backends. Parser patched against llama.cpp / vLLM / SGLang per-family parsers and normalises to OpenAI shape. 96 targeted unit tests + cross-OS staging CI (ubuntu / macos-14 / windows) green on the multi-format probe. * studio: tool-call healing parity between safetensors / MLX and GGUF After the multi-format parser landed in #5615, the safetensors / MLX agentic loop and the GGUF loop still differed on healing behaviour. This commit closes the gaps in both directions so the two backends react the same way to identical model output. Changes: 1. core/inference/llama_cpp.py -- the GGUF BUFFERING state machine now wakes on every emission marker the shared parser knows. Was ("", " / Mistral [TOOL_CALLS] / Gemma 4 <|tool_call>). Stream cleanup is delegated to the same shared strip_tool_markup so leaked markup from any family is removed from assistant content. 2. core/inference/llama_cpp.py -- per-tool canonical heal key. When a tool arguments field is a bare string and JSON parsing fails, the GGUF path now heals to {"code": raw_args} for python, {"command": raw_args} for terminal, and {"query": raw_args} for everything else. Was hard-coded to {"query": raw_args}, which silently routed every python / terminal emission through web_search. Mirrors safetensors_agentic._CANONICAL_HEAL_ARG. 3. core/inference/safetensors_agentic.py -- re-prompt on plan- without-action. When the model emits a short forward-looking intent ("I'll search for that", "Let me check", "First, I will...") and no tool call, the loop nudges the model to act instead of silently returning a plan-only answer. Up to _MAX_REPROMPTS=3 (matches GGUF). The intent regex, character cap, and instruction text are byte-identical to the GGUF path. The buffer-end fall-through is unified so a buffered intent emission that never exits the BUFFERING state still triggers the re-prompt. 4. core/inference/safetensors_agentic.py -- extra iteration slots for re-prompts. The loop now budgets max_tool_iterations + _MAX_REPROMPTS + 1 total iterations and tracks the tool-call count separately, so a stalling model can be nudged 3x without eating the caller's tool-call budget. Mirrors the _extra slot reservation in the GGUF path. Tests (14 new safetensors-side units; 5 GGUF parity pins): TestLoopRePrompt -- intent-trigger, plain-answer, no-tools, cap-at-three, budget preserved, buffer-end intent. TestLoopCanonicalHealKey -- python / terminal / unknown. TestGGUFSafetensorsHealingParity -- shared markers used, shared strip used, canonical heal keys identical, intent regex matches same phrases, _MAX_REPROMPTS equal on both backends. All 110 targeted tests pass locally; the broader tool / inference / model-config / sandbox / anthropic / mlx suites stay green. Why this matters Without this parity, Llama-3.2 / Mistral / Gemma 4 emissions on Mac (MLX) and Linux-safetensors stop the agentic loop as soon as the model says "Let me...", because the GGUF re-prompt logic never existed on these backends. The two-marker GGUF BUFFERING tuple also let non-Qwen tool emissions stream out as plain prose when llama-server's structured channel did not pick them up. Both paths now drain the same way, heal the same way, and re-prompt the same way -- so a tool call that works on GGUF works identically on safetensors / MLX. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: fix tool-call parser bugs from gemini review on #5620 Three high-priority gemini findings on the tool-call parsing additions: 1. unicode_escape on UTF-8 bytes corrupts non-ASCII literals (e.g. ✨ becomes â\x9c¨). Replace with json.loads on a quoted string -- preserves emoji / CJK / RTL while still handling \n \t \uXXXX escapes. 2. Llama-3 sentinel stripping is order-dependent. A leading `<|eot_id|><|begin_of_text|>` left `<|begin_of_text|>` behind because the loop had already passed that sentinel. Loop until no sentinel matches at the start. 3. Mistral v11+ `[TOOL_CALLS] name { json }` regex uses non-greedy `\{.*?\}` which truncates at the first `}` of a nested JSON argument, leaking the tail (e.g. `}}`) into user-visible streamed text. Same problem for the v0.3 array pattern with nested brackets. Strip those with balanced brace/bracket scanning via a new `_strip_mistral_closed_calls` helper called from `strip_tool_markup`. Also fix the inference routes' parallel `_TOOL_XML_RE`: - Same nested-JSON truncation in the Mistral patterns; route the strip through the parser's balanced-scan helper via a thin `_strip_tool_xml` wrapper that all existing callers now use. - Llama-3 `<|python_tag|>[^\n<]*` stopped at any `<`, leaking the tail of any tool call whose argument contained a literal `<` (queries, code snippets). Relax to `[^\n]*` which keeps the strip confined to the actual end-of-line. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: tool calling for DeepSeek (R1/V3/V3.1), GLM 4.x, Kimi K2 Adds three more emission-family parsers to tool_call_parser.py so the shared safetensors / MLX / GGUF agentic loop covers the major open- weight reasoning families. Patterns ported from llama.cpp (common/chat-parser.cpp legacy pre-PEG branch), vLLM (tool_parsers/deepseekv3*, glm4_moe, kimi_k2), and SGLang (function_call/deepseekv31_detector, glm4_moe_detector, kimik2_detector). All three references are MIT (llama.cpp) or Apache-2.0 (vLLM, SGLang). Formats covered: DeepSeek R1 <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function <｜tool▁sep｜>NAME\n```json\n{...}\n```<｜tool▁call▁end｜> <｜tool▁calls▁end｜> -- args wrapped in a Markdown json fence, ``function`` literal prefix per llama.cpp common_chat_parse_ deepseek_r1 (chat-parser.cpp:801-820) DeepSeek V3/V3.1 <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>NAME <｜tool▁sep｜>{json}<｜tool▁call▁end｜><｜tool▁calls▁end｜> -- bare JSON, no code fence, no ``function`` prefix per llama.cpp common_chat_parse_deepseek_v3_1 (chat-parser.cpp:822-879) GLM 4.5/4.6/4.7 NAME\nk1 \nv1... -- strings raw, non-strings JSON-encoded per chat_template.jinja; multi-call is back-to-back blocks. Per llama.cpp common_chat_parse_glm_4_5 (chat-parser.cpp:1040-1052) Kimi K2 <|tool_calls_section_begin|><|tool_call_begin|> functions.NAME:IDX<|tool_call_argument_begin|>{json} <|tool_call_end|><|tool_calls_section_end|> -- bare name recovered by stripping ``functions.`` prefix and ``:IDX`` suffix; full id preserved as tool_calls[i].id so the roundtrip replays verbatim. Per llama.cpp common_chat_parse_kimi_k2 (chat-parser.cpp:896-913) Marker collisions GLM uses the same ```` opener as Qwen but with a bare function name + ```` body (Qwen has ``\s*{`` after the tag). The dispatch keeps Qwen first; Qwen's _TC_JSON_START_RE returns no matches on a GLM emission, so the fall-through to _parse_glm_tool_ calls handles it correctly. Existing Qwen tests confirm zero regression. Streaming buffer TOOL_XML_SIGNALS extended from 5 markers to 12 so the BUFFERING state machine wakes on every new family's section opener. Added the DeepSeek alternative markers (ASCII underscores, short ``<｜tool▁calls｜>`` form) because real checkpoints emit those variants. Strip patterns _TOOL_CLOSED_PATS adds DeepSeek envelope (``<｜tool▁calls▁begin｜>... <｜tool▁calls▁end｜>``) and Kimi section (``<|tool_calls_section_begin|> ...<|tool_calls_section_end|>``). _TOOL_ALL_PATS adds the same plus the unclosed-tail variants so a truncated stream does not leak markup. Route gate _detect_safetensors_features._PARSER_MARKERS grows to include DeepSeek and Kimi markers plus ```` (the unique GLM signal). _TOOL_XML_RE (the route-layer markup-strip regex) gets DeepSeek and Kimi closed-pair patterns. _TOOL_TEMPLATE_MARKERS in llama_cpp.py adds ``message['role'] == 'tool'``, ``message['tool_calls']``, and ``tool_calls is defined`` so the classifier recognises DeepSeek's subscripted-access template style (it has no top-level ``{% if tools %}`` block). Tests (39 new): TestParserDeepSeek (7) -- R1 fence, short-form opener, V3.1 bare, multi-call, with-reasoning, strip, signal-wakes-streaming TestParserGLM (6) -- single, mixed types, multi-call, unclosed-heal, no-Qwen-regression, strip TestParserKimi (6) -- single, multi-call, dotted-name, unclosed, strip, signal-wakes-streaming TestParserCrossFormatRouting (2) -- dispatch routing, signal coverage TestLoopBasic loop integration (3) -- DeepSeek / GLM / Kimi end-to-end Capability advertise (3) -- DeepSeek / GLM / Kimi templates flip supports_tools=True All 398 targeted tests pass locally (115 safetensors + 27 capability + rest of tool / inference / sandbox / model-config suites). Builds on PR #5620 (parser + healing parity for Llama-3 / Mistral / Gemma 4); will rebase cleanly onto main once #5620 lands. PR opened as draft - do not merge until validated against real models for each family. Sources - llama.cpp common/chat-parser.cpp lines 801-913, 1040-1052 (MIT) - vLLM vllm/tool_parsers/deepseekv31_tool_parser.py (Apache-2.0) - vLLM vllm/tool_parsers/glm4_moe_tool_parser.py (Apache-2.0) - vLLM vllm/tool_parsers/kimi_k2_tool_parser.py (Apache-2.0) - SGLang python/sglang/srt/function_call/{deepseekv31,glm4_moe,kimik2}_ detector.py (Apache-2.0) - Live chat templates: deepseek-ai/DeepSeek-V3.1, zai-org/GLM-4.6, moonshotai/Kimi-K2-Instruct, unsloth/DeepSeek-V3-0324, unsloth/GLM-4.5-Air, unsloth/Kimi-K2-Instruct * studio/routes: make python_tag strip multi-line aware Earlier revisions of _TOOL_XML_RE in studio.backend.routes.inference oscillated between two bug shapes: 5615 r"<\|python_tag\|>[^\n<]*" -- stopped at any literal "<" so code='if x < 10: pass' leaked '< 10: pass)' to the user. 5620.1 r"<\|python_tag\|>[^\n]*" -- single-line only; the second line of python.call(code="a\nb") leaked. The full parser (_parse_llama3_python_tag) already handles both via balanced-brace scanning, so the parsing path was fine; the LEAK was in the streaming strip path that runs on every cumulative emission while content is still arriving. Switch to r"<\|python_tag\|>(?:[^<]|<(?!\|))*" so the strip consumes: * any character that is not a "<" (newlines, JSON, code, ...), * a "<" only when it is NOT followed by "|" (i.e. NOT a Llama-3 sentinel start like <|eot_id|>, <|eom_id|>, <|begin_of_text|>). This means: * code='if x < 10' stays inside the strip (5615 fix preserved), * multi-line code stays inside the strip (5620 round 2), * the strip terminates at the next Llama-3 sentinel so trailing assistant content survives. Tests: TestRoutesPythonTagStrip (8 cases) pytest test_safetensors_tool_loop.py test_safetensors_capability_advertise.py -> 118 passed in 1.81s (was 110). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: review follow-ups for DeepSeek / GLM / Kimi tool calling Four fixes addressing review of the parent commit: 1. GLM coercion: tighten the json.loads -> ast.literal_eval -> raw cascade to only deserialize when the body unambiguously looks like a JSON literal (object, array, JSON-encoded string, true/false/null, or numeric). Strings like ``True`` / ``None`` (Python literals, not JSON) and arbitrary prose now stay raw. The bare-numeric / bare-boolean ambiguity with string args remains an inherent limitation of the template without schema access -- documented in the new comment. Drops the ast import entirely (closes Gemini's :1036 suggestion). 2. Kimi K2 bare-counter ids (e.g. ``<|tool_call_begin|>3``) are now dropped rather than surfaced as a tool literally named "3". Matches vLLM behaviour; SGLang's schema-infer fallback is out of scope at the parse site. Real Kimi K2 emissions use ``functions.NAME:IDX`` so this is the exception path. 3. Restore the elaborate ``<|python_tag|>(?:[^<]|<(?!\|))*`` clause in routes.inference._TOOL_XML_RE -- the simpler ``[^\n<]*`` form regressed PR #5620's multi-line / literal-``<`` python_tag fix. Restore ``TestRoutesPythonTagStrip`` (8 tests) adapted to call ``_TOOL_XML_RE.sub`` directly since the ``_strip_tool_xml`` helper was inlined this PR. 4. Add the spaced and backslash-escaped DeepSeek opener variants (``<｜tool calls begin｜>``, ``<｜tool\_calls\_begin｜>``) to ``TOOL_XML_SIGNALS`` for streaming-gate parity with ``_DEEPSEEK_BEGIN_RE``. Also updates the llama.cpp / vLLM citations in the parser docstrings: ``common/chat-parser.cpp`` was split into ``common/chat.cpp`` + ``common/chat-peg-parser.cpp`` by llama.cpp PR #18675, and vLLM moved the tool parsers from ``vllm/entrypoints/openai/tool_parsers/`` to ``vllm/tool_parsers/``. Pin to pre-refactor commit ``51fa458a92d6`` where the cited line numbers still resolve. New regression tests in ``test_pr5624_regressions.py`` cover the GLM coercion heuristic shapes, GLM literal-``<`` in arg_value, Kimi K2 dotted name, Kimi K2 bare-counter drop, DeepSeek V3.1 truncated mid-stream, and routes-layer strip across all three new families. Tests: pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py studio/backend/tests/test_pr5624_regressions.py -q -> 170 passed in 1.91s * [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: tighten verbose comments in tool-call parser sections Comments were narrating what the code already says. Cut historical "earlier revisions used X, then Y" narratives down to one-line WHY notes where the footgun still matters (canonical heal-key parity, balanced-brace vs non-greedy regex, ``(?:[^<]|<(?!\|))*`` over ``[^\n<]*``/``[^\n]*``). Drop section-header banners. No behaviour change. Re-ran: pytest studio/backend/tests/test_safetensors_tool_loop.py \ studio/backend/tests/test_safetensors_capability_advertise.py -q -> 118 passed. Regression replay (parser + _coerce_arguments on the 5 #5615 inputs) -> 21/21. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: GLM 4.7 no-newline emission + Kimi multi-section parity Two fixes surfaced by triple-confirm verification against the live HF chat templates and upstream llama.cpp / vLLM / SGLang parsers. 1. GLM 4.7 silent drop ``zai-org/GLM-4.7/chat_template.jinja`` line 65 uses ``{{- '' + tc.name -}}`` which Jinja strips trailing whitespace from, so the first ```` follows the function name with NO ``\n`` between them. Real emissions look like ``get_weathercityLondon ``. The previous ``_GLM_TC_OPEN_RE`` ended the name with ``\n`` so GLM-4.7 calls were silently dropped (parser returned ``[]``). Fix: relax the name terminator to a lookahead that accepts EITHER ``\n`` OR the next ````: _GLM_TC_OPEN_RE = re.compile( r"\s*([^\n<{][^\n<]*?)\s*(?=\n|)" ) The first-char restriction ``[^\n<{]`` still excludes Qwen's ``{json}`` form so the Qwen-vs-GLM dispatch remains mutually exclusive. 2. Kimi multi-section parity with vLLM / SGLang ``vllm/tool_parsers/kimi_k2_tool_parser.py`` and SGLang's ``kimik2_detector.py`` both use ``re.findall`` and so collect every ``<|tool_calls_section_begin|>...<|tool_calls_section_end|>`` block in a single stream. The previous implementation stopped at the first ``<|tool_calls_section_end|>``. Kimi K2 doesn't emit multi-section in practice, but parity is cheap. Fix: wrap the existing per-call body parser in an outer loop that advances past each ``<|tool_calls_section_end|>`` and continues to the next ``<|tool_calls_section_begin|>``. Body parsing extracted to ``_parse_kimi_section_body`` for clarity. Truncated final section is still surfaced via the existing in-body balanced-brace walk. Verified independently against the live HF templates: * GLM-4.7 emission constructed from the live template parses to the expected ``{name, arguments}`` shape. * GLM-4.5 / 4.6 newline shape continues to parse (the lookahead also matches ``\n``). * Qwen ``{json}`` still dispatches to the Qwen path -- the first-char restriction stops the GLM regex from biting JSON bodies. * Kimi two-section stream surfaces both calls in order with full ids preserved. * Bare-counter Kimi ids still drop. Tests added in ``test_pr5624_regressions.py``: * ``test_glm_4_7_no_newlines_between_name_and_arg_key`` * ``test_glm_4_7_no_newlines_multi_call`` * ``test_glm_4_7_does_not_break_qwen_path`` * ``test_kimi_two_sections_in_one_stream_both_parse`` pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py studio/backend/tests/test_pr5624_regressions.py -q -> 174 passed in 1.93s pytest studio/backend/tests/ -q -k 'not gpu and not llama_cpp_integration' -> 2038 passed, 15 failed (pre-existing CI gaps). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: parser robustness fixes for PR #5620 Three surgical extensions to the multi-format tool-call parser, each covering a real fine-tune / template emission shape that the current parser silently drops. No path narrows; all changes widen what is accepted. 1. `_parse_tool_call_json` now accepts both `arguments` and `parameters` keys. A Hermes / Qwen `{json}` wrapper around a Llama-3.2 fine-tune that emits the `parameters` key was extracting the tool name and silently discarding the args, producing a working-shaped call with an empty payload. The bare-JSON and python_tag paths already accepted both keys; this path now matches them. 2. `_TC_FUNC_START_RE`, `_TC_PARAM_START_RE`, and `_TC_PARAM_CLOSE_RE` now also match the attribute form `v` used by MiniCPM-5 and MiniMax-M2. Names land in either capture group, and `` is accepted as a short close. 3. `_parse_llama3_bare_json` sentinel-strip now consumes the role label inserted between `<|start_header_id|>` and `<|end_header_id|>` by Meta's official Llama-3.x chat template. Without this, every assistant turn re-fed through the template prefix `<|start_header_id|>assistant<|end_header_id|>\n\n{json}` parsed to zero calls, so any history-with-tool-call round-trip in production silently dropped. Tests in `studio/backend/tests/test_safetensors_tool_loop.py`: * `TestParserRobustness::test_tool_call_json_accepts_parameters_key` * `TestParserRobustness::test_function_xml_attribute_form` * `TestParserRobustness::test_function_xml_attribute_form_multi_param` * `TestParserRobustness::test_function_xml_legacy_equals_form_still_works` (regression guard for the existing `` syntax) * `TestParserRobustness::test_llama3_chat_template_round_trip` * `TestParserRobustness::test_llama3_round_trip_all_roles` * `TestParserRobustness::test_llama3_round_trip_with_eot_prefix` `pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py -q` goes from 118 to 125 passed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim verbose comments in tool-call parser sections for PR #5624 Pure comment / docstring tightening on top of the GLM 4.7 + Kimi multi-section fixes. No behavioural change. * Drop multi-paragraph prelude and post-refactor citation chatter in the DeepSeek, GLM and Kimi parser docstrings; keep the shape and upstream-commit pin. * Collapse ``parse_tool_calls_from_text``'s 9 per-family blocks into a single ordered loop with one combined comment. * Tighten the GLM coercion, Kimi bare-counter and ``_TOOL_XML_RE`` comments to one or two lines each. * Same trim pass on ``_PARSER_MARKERS`` and the regression-test docstrings. Tests: pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py studio/backend/tests/test_pr5624_regressions.py -q -> 174 passed in 2.00s * Fix O(N^2) DeepSeek V3.1 backtracking for PR #5624 Adversarial input ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>fn<｜tool▁sep｜>`` followed by a long body that does NOT contain a closing brace caused the V3 path's ``([^\n<]+?)<｜tool▁sep｜>`` regex to backtrack quadratically: at each position the lazy quantifier extends one char at a time looking for a sep that isn't there, taking ~19s on 50k chars. Replace the regex search with ``str.find`` on the sep marker plus a left-walk to recover the name. ``str.find`` is O(N); the walk stops on ``\n`` (turn boundary), ``<`` (start of a tag), or ``>`` (end of an optional ``<｜tool▁call▁begin｜>`` prefix). Same observable behaviour as the regex on every canonical input. Tests: test_deepseek_v3_1_huge_truncated_body_is_linear (new) -- 50k chars must parse in < 1s. pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py studio/backend/tests/test_pr5624_regressions.py -q -> 175 passed in 1.97s pytest studio/backend/tests/ -q -k 'not gpu and not llama_cpp_integration' -> 2038 passed, 15 pre-existing failures unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: terminate function-XML body at , not just `_parse_function_xml` was looking for `` (the Hermes wrapper) as the body terminator. When a model emits a standalone `v` followed by explanatory prose (which models routinely do), no `` is present, so the body extended to end-of-string and the trailing prose leaked into the LAST parameter value. Pre-existing on main (the legacy `` form had this bug too). Same affects PR #5620's new attribute-form `v` emission used by MiniCPM-5 / MiniMax-M2. Fix: `_TC_END_TAG_RE` now matches either `` OR ``. The existing `_TC_FUNC_CLOSE_RE` / `_TC_PARAM_CLOSE_RE` strips are unchanged. Multi-call inputs still bound each function at the next `` is preserved because the embedded close tag is ``, not ``). `pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py -q` goes from 125 to 127 passed. * Studio: tighten Llama-3.2 bare-JSON guard A fuzz pass on PR #5811 turned up that ``_parse_llama3_bare_json`` accepted ``parameters`` as a string, contradicting the docstring's "parameters or arguments is a dict" guard. Prose JSON like ``{"name":"foo","parameters":"a sentence"}`` would wrongly fire the parser, which the agentic loop would then heal into a real ``foo(query="a sentence")`` call. Same code lives on this branch, so the same fix applies here. Tightened guard: - ``parameters`` must be a dict (Llama-3 spec). - ``arguments`` may be a dict, or a JSON-encoded string that decodes to a dict (OpenAI shape, e.g. ``"arguments":"{\"q\":\"x\"}"``). Plain non-JSON strings or JSON-strings of lists / scalars / null no longer pass. Mirrors the fix landed in PR #5811 commit 615b8608. Adds the same 4 regression tests under TestParserMultiFormat. Existing test suite stays green: 127 -> 131 passing. * Studio: skip non-scalar args in python_tag JSON form The JSON sub-path of ``_parse_llama3_python_tag`` was fabricating ``{"value": args}`` when the model emitted a non-dict / non-string ``arguments`` value (e.g. ``42``, ``[1,2,3]``, ``null``, ``true``). This silently turned a malformed emission into a real tool call, which the agentic loop would then execute with arguments the model never intended. Tightened: skip the call instead of fabricating. The same behaviour now matches the bare-JSON guard tightened earlier (strict-guard merge from PR #5620, inherited via merge here). Added a regression test covering the four non-scalar shapes. Pass count on this branch: 158 -> 159. Sites in ``_parse_tool_call_json`` and ``_consume_mistral_call`` keep the existing looser behaviour for now; both are reached only after explicit ```` / ``[TOOL_CALLS]`` markers so the false-positive surface there is much narrower. * studio: fix safetensors tool-call parser gaps vs llama.cpp (Mistral CALL_ID / THINK, attribute-form signal) Three GGUF-parity fixes to the safetensors tool-call parser, each matching llama.cpp's reference behaviour: - Mistral Small 3.2 emits [TOOL_CALLS]name[CALL_ID][ARGS]{json}. The parser stopped after the name on seeing [CALL_ID] (neither [ARGS] nor {), dropping the call. Skip an optional [CALL_ID] segment in both the parse and strip paths. llama.cpp parses this (test-chat.cpp:4785). - Magistral wraps reasoning in [THINK]...[/THINK]. A [TOOL_CALLS] inside the reasoning was parsed as a real call, producing a phantom call. Strip a leading [THINK] block before scanning so only the post-reasoning call counts (test-chat.cpp:2285); a literal [THINK] inside a later argument is left intact. - The standalone MiniCPM-5 / MiniMax-M2 attribute form parsed correctly but was absent from TOOL_XML_SIGNALS and the markup strip patterns, so the streaming safety-net parse was gated off (dropping the call) and markup leaked into displayed text. Add the signal and broaden the strip regexes. Adds regression tests for all three. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: fix GLM and Kimi K2 safetensors tool-call parser gaps vs llama.cpp Four GGUF-parity fixes for the GLM and Kimi K2 families: - GLM 4.7 zero-argument inline call name was dropped: the open-tag lookahead only allowed \n or after the name. Allow too so a no-arg call parses to empty args (vLLM / SGLang / llama.cpp all parse it). - GLM string argument values were stripped, losing significant leading / trailing whitespace in code / diff arguments. Keep the raw value for the string fallback and only strip the copy used to probe for a JSON literal, matching vLLM glm4_moe which never strips string args. - Kimi K2 calls emitted without the <|tool_calls_section_begin|> wrapper were dropped. llama.cpp makes the section optional (Kimi can call a tool straight after reasoning without opening a section); parse a bare <|tool_call_begin|> when no section is present. - Kimi K2 malformed / truncated JSON in one call dropped every later call in the section. Skip the bad call and keep parsing so valid subsequent calls are recovered (vLLM parity). Adds regression tests for all four. * studio: fire safetensors tool calls for the bare-JSON (Llama-3.2) form The agentic loop's streaming safety-net parse was gated on has_tool_signal(), which is False for the Llama-3.1 / 3.2 bare-JSON tool form {"name":..,"parameters":..} (no XML marker). Real tool calls were therefore dropped: the loop logged "model planned without calling tools", re-prompted three times, then gave up with zero tool calls, while GGUF's llama-server parses the same emission natively. Run parse_tool_calls_from_text() unconditionally in the safety net. The parser is strict (only fires on a valid tool-call shape) so plain answers are unaffected. Reproduced on a real unsloth/Llama-3.1-8B-Instruct run: the model emits {"name":"web_search","parameters":{...}} which now executes the tool instead of being re-prompted into a no-op. Adds a loop regression test for the bare-JSON form. * studio: fire safetensors tool calls for Gemma 4 (native template + stripped parser) Gemma-4 safetensors fired no tools while its GGUF fired reliably. Three gaps: - The Studio swaps in the Unsloth "gemma-4" chat template, which does not render the tools schema (the model's native template does), so the model never saw the tools. Fall back to the model's native template when the override template renders identically with and without tools. Same fix helps any family whose override template drops tools. - skip_special_tokens strips the <|tool_call> wrapper and <|"|> string markers, so a streamed Gemma-4 call arrives as a bare call:NAME{k:v, ...} with unquoted values. Parse that form, keeping commas/braces inside a code or command value, normalising surrounding quotes, and stripping the leaked markup from the final answer. - Without a grammar a small model can loop, repeating one call for the whole tool budget. Collapse exact-duplicate calls within a turn and force a final answer after a turn that made no new tool progress (llama-server's lazy grammar prevents this loop on the GGUF side). Adds parser tests for the bare/stripped Gemma-4 form. * [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: complete strict-mode contract and fix parser import paths Address review findings on the multi-format tool-call parser: - Honor allow_incomplete=False in the remaining sub-parsers. The Llama-3 <|python_tag|>NAME.call(...) parser, the pre-v11 Mistral [TOOL_CALLS] array parser, and the Gemma 4 <|tool_call> parser ignored strict mode, so a truncated call (missing closing paren, ], or ) was still healed and executed with Auto-Heal disabled. Thread strictness through and reject the unclosed forms, matching the JSON and function-XML paths. - Drop the duplicate tool_call_parser import block in llama_cpp.py and the redundant un-aliased TOOL_XML_SIGNALS; only the _SHARED_TOOL_XML_SIGNALS alias is used as a value. - Import _strip_mistral_closed_calls from core.inference.tool_call_parser in routes/inference.py instead of studio.backend.core... The self-contained run.py launch mode only puts studio/backend on sys.path, so the absolute package path raised ModuleNotFoundError on the server-tool strip path. Add strict-mode regression tests for the truncated Llama-3 dot-call and the unclosed Mistral array. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden DeepSeek/Kimi tool-call parsing and strip Address review findings on the DeepSeek and Kimi parsers: - Honor allow_incomplete=False for DeepSeek. An envelope with no closing <｜tool▁calls▁end｜> is truncated mid-stream; reject it in strict mode instead of healing the body out to EOF, matching the strict XML and Mistral paths. - Do not skip a following tool call when the current call's end marker is missing. The DeepSeek V3 and Kimi loops advanced by searching forward for the next <｜tool▁call▁end｜> / <|tool_call_end|>, which could land on a later call's end marker and drop the call in between. Advance by the JSON end; the loop re-locates the next call marker from there. - Strip truncated DeepSeek and Kimi section blocks in the route-level display regex. The patterns required the closing marker; add the end-of-text alternative so a block truncated by EOS does not leak raw markup to the UI. Add regression tests for the truncated DeepSeek envelope, and for DeepSeek and Kimi multi-call recovery when the first call's end marker is missing. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: preserve XML param indentation and alias Mistral array parameters Two parser-correctness fixes found by auditing against the model chat templates and the SGLang / vLLM reference parsers: - Qwen3.5 XML parameter values lost their leading indentation. The chat template emits \nVALUE\n, but the parameter-start regex ate the wrapping newline AND the value's first-line indentation with a trailing \s*, then str.strip() removed the rest. Narrow the trailing class to horizontal whitespace only and trim exactly one wrapping newline (via _trim_param_value), preserving indentation in code/diff arguments. Matches SGLang's qwen3_coder detector. Applies to both _parse_function_xml (tool_call_parser.py) and the XML path in tool_healing.py. - Mistral pre-v11 array objects keyed on parameters dropped their payload. _consume_mistral_call read only the arguments key; alias parameters the same way the JSON/XML paths and SGLang's base detector do. Add regression tests for preserved multi-line indentation and the array parameters alias. * Studio: DeepSeek strip sync, Gemma nested args, GLM/Kimi strict mode Parser-correctness fixes found by auditing DeepSeek/GLM/Kimi against vLLM, SGLang, and the model chat templates: - DeepSeek: the short <｜tool▁calls｜> opener (and the space / escaped-underscore spellings) was parsed but never stripped, so a short-opener envelope leaked raw markup to the UI. Share one opener alternation between _DEEPSEEK_BEGIN_RE and the strip patterns (and the route-level display regex) so a signal we parse can never be left un-stripped. - Gemma wrapper-less stream: a nested object/array argument (loc:{city:NYC}, labels:[bug,ui]) was kept as a literal string. Parse it recursively when the bare value is a balanced {} / [], falling back to the raw string for a truncated value. - GLM and Kimi ignored allow_incomplete. With Auto-Heal off, a GLM block with no , a Kimi section with no <|tool_calls_section_end|>, or a Kimi call with no <|tool_call_end|> are truncated and must be rejected, matching the strict behavior of the JSON/XML/Mistral/DeepSeek paths and vLLM/SGLang. Add regression tests for the short-opener strip, the Gemma nested args, and GLM / Kimi strict-mode rejection. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten tool-call parser comments Make the comments in the multi-format tool-call parser and its callers succinct: compress verbose docstrings/blocks to one or two lines, drop ones that restate the code, and trim the tiny balanced-scanner helpers. Correctness rationale and upstream provenance (SGLang/llama.cpp parity, the strict-mode / Auto-Heal contract, whitespace-preservation, and the Unicode / full-width-pipe notes) are kept in compact form. Comment-only: no code or behavior change (verified with comment_tools.py check --strip-docstrings; parser suite green). * Studio: tighten DeepSeek/GLM/Kimi parser comments Compress the comments added for the DeepSeek/GLM/Kimi parsers and the Gemma wrapper-less helpers to one or two lines, keeping the upstream provenance (llama.cpp 51fa458a92d6), the O(N^2) / strict-mode rationale, and the vLLM parity notes intact. Comment-only: no code or behavior change (verified with comment_tools.py check --strip-docstrings; parser suite green). * Studio: make DeepSeek R1 / GLM parsing linear and close routes strip gaps Review follow-up for the DeepSeek/GLM/Kimi parser: - DeepSeek R1 detection used a greedy ``([^\n]+)\n```json`` regex that backtracks O(N^2) on a fence-less truncated body; scan with str.find instead (mirrors the V3 path). - GLM arg pairs used a lazy-group finditer that rescanned to EOF from each bare in an unclosed body (O(N^2)); walk pairs with str.find. - The route display strip (_TOOL_XML_RE) accepted fewer DeepSeek openers than the parser (missed the space / escaped-underscore spellings) and missed bare section-less Kimi calls, so a call we parse could leak raw markup to the UI. Reuse the parser's shared _DEEPSEEK_OPEN_RE_SRC and add a bare-Kimi arm. Add ReDoS-linearity regressions for the R1 and GLM paths, a positive R1 fenced-json parse test, and routes-strip tests for the space/escaped DeepSeek openers and the bare Kimi call. * Studio: fix test_mcp_servers _TOOL_XML_RE reconstruction after _DS_OPEN_SRC reuse The routes strip fix made _TOOL_XML_RE reference the module-level _DS_OPEN_SRC variable. test_mcp_servers reconstructs the regex by exec-ing the extracted compile() source in a namespace that only defined _re, so it raised NameError. Inject _DS_OPEN_SRC into that namespace, matching the same fix already applied in test_tool_xml_strip. * Studio: make Llama-3 .call and Mistral-array healing parsing linear Two more O(n^2) ReDoS paths in the multi-format parser, both reachable from the agentic loop on a long truncated body with no length cap: - _LLAMA3_KV_RE.finditer over a .call(...) body retried at every offset of a long word run / unterminated quote (40K -> 14s). Replace with a hand-scan that reuses the same key/number/literal sub-regexes via anchored match and walks the string body by hand, so an unterminated quote is O(n). Verified byte-identical to the old regex over 200K fuzzed inputs. - _parse_mistral_array healing ran _balanced_brace_end from every { in the body (20K -> 17s). Walk top-level objects, advancing past each balanced {...}; this also drops the phantom call the old scan emitted from a nested argument object. Add adversarial-length linearity regressions plus positive .call kwargs and unclosed-array recovery coverage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: strengthen #5624 regression assertions and strip-test harness guards - test_strip_tool_markup_handles_deepseek_envelope used `A or B` where B was the preservation property the next line already asserts, masking the real check. Replace with an explicit assertion that the call name and args are stripped. - The test_tool_xml_strip source-extraction harness reconstructs _TOOL_XML_RE and _strip_tool_xml_for_display from routes/inference.py via lazy regexes that could silently grab a shorter slice. Assert the extracted regex carries the DeepSeek / bare-Kimi arms and the helper body reached the _TOOL_XML_RE.sub call. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: honor strict mode in safety-net, keep empty Gemma args, strip attribute-form function XML - safetensors safety-net parser now forwards allow_incomplete=auto_heal_tool_calls, matching the draining path, so a late incomplete tool call is not healed and executed when Auto-Heal is off. - Gemma empty bare value ({k:}) now serialises as "" instead of invalid {"k":}, which previously dropped the whole call. - Route _TOOL_XML_RE also strips the attribute form (MiniCPM-5 / MiniMax-M2) so it no longer leaks to the UI. * Studio: linearize wrapper-less Gemma nested-arg parsing and correct parser provenance - _gemma_parse_value/_gemma_parse_mapping/_gemma_parse_array now parse nested {}/[] in a single forward pass instead of pre-scanning each subtree with a balanced-brace walk and re-parsing it. Deeply nested wrapper-less Gemma args were O(n^2); they are now ~linear (and ~40x faster at depth 400). - Correct the DeepSeek/GLM/Kimi provenance comments: the cited commit 51fa458a92d6 is unrelated, and GLM/Kimi were never standalone common_chat_parse_* functions (llama.cpp uses common_chat_params_init_glm_4_5 plus a generalized XML parser, PRs #15904 / #16932). - Add tests: Gemma deep-nesting linearity, nested object/array preservation, same-turn distinct-call cap, and the native-template tool-render fallback. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: guard Gemma value parser against non-advancement and missing tokenizer Addresses Gemini review: - _gemma_parse_value now consumes one character when a stray }/]/, sits where a value is expected, so _gemma_parse_array can never stall at the same index on malformed input (a latent infinite loop). - _render_with_native_template returns None when neither a tokenizer nor a processor is present instead of raising AttributeError. - Tests for both. * Studio: fix attribute-form function-XML literal close tag and zero-arg strict call Addresses Codex review of the attribute form in _parse_function_xml (MiniCPM-5 / MiniMax-M2): - End the call body at the LAST / within the call's window, so a literal close tag inside a code/search argument (e.g. print("")) is preserved instead of truncating the call. - Accept a closed call with no parameters as a valid zero-argument call in strict mode (the function close is already required), instead of rejecting it as a truncated call. - Tests for both, mirroring the legacy coverage. * Studio: drop scratch review/planning artifacts from the branch * Studio: fix tool-call parser/loop review findings on the multi-format path Address the live code-review findings on the safetensors/MLX + GGUF tool path: - routes: include the attribute form in the safetensors capability whitelist so MiniCPM-5 / MiniMax-M2 templates keep the tool pill (parser already handles the form; the post-filter wrongly suppressed it). - safetensors loop: build the plan-without-action re-prompt from the active tools instead of a hardcoded web_search/python string, and gate it on auto_heal_tool_calls, matching the GGUF loop. - safetensors loop: hold a leading bare-JSON object ({"name":..,"parameters":..}) during BUFFERING until it closes, then drain it as a tool call instead of streaming the raw JSON to clients. The DRAINING/STREAMING resolvers still recover a plain JSON answer, so this can never drop content. - parser: anchor the Llama-3 <|python_tag|>NAME.call(...) scan to the tag and chain ; -separated calls, so all semicolon-separated built-ins parse and a literal <|python_tag|>x.call(...) inside a JSON string argument no longer fires the wrong tool. - parser: consume the optional trailing after a named Mistral [TOOL_CALLS]name{json} call, mirroring the array shape. - GGUF streaming strip: use the shared parser patterns (which know [TOOL_CALLS] and <|python_tag|>) so a textual tool call entering DRAINING is stripped instead of leaking the marker to streaming clients. - routes: hoist the _strip_mistral_closed_calls import to module level. Adds regression tests covering each fix; existing parser suite stays green. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix DeepSeek/GLM/Gemma tool-call review findings Address the live code-review findings specific to the DeepSeek / GLM / Kimi and native-template additions: - parser: in strict mode (Auto-Heal off) require the per-call <｜tool▁call｜end｜> terminator for DeepSeek V3 calls instead of executing on a bare balanced object closed only by the envelope end. - parser: keep GLM string arguments that begin with a quote verbatim (drop the leading-quote case from the JSON-decode probe) so a quoted search query is not decoded down to its inner text. - parser: reject a GLM call with an unclosed in strict mode, and under Auto-Heal keep the partial value rather than dropping it to a no-arg call. - parser: add a balanced wrapper-less Gemma strip (call:NAME{...}) so a nested object/array argument is removed whole instead of leaving a trailing brace; run the balanced Mistral and Gemma strips on the streaming display paths too. - safetensors loop: buffer a leading wrapper-less Gemma call:NAME{...} so it drains and executes instead of streaming the raw call text. - inference: render the native-template fallback on a shallow tokenizer copy instead of mutating the shared tokenizer outside the generation lock, and load the native template from base_model for LoRA adapters. Adds regression tests for each; existing parser suite stays green. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden multi-format tool-call detection from review findings Apply five targeted fixes from the review pass over the multi-format tool path: - routes: route display strip delegates to _strip_tool_xml so Mistral [TOOL_CALLS] blocks with nested JSON are removed from streamed display text, not just the XML forms. - tool_call_parser: skip function/parameter starts that fall inside an already-open parameter block (_inside_open_parameter) so nested example payloads are not mis-parsed as new calls; extract strip_llama3_leading_sentinels so the bare-JSON guard is shared. - safetensors_agentic: probe bare JSON through strip_llama3_leading_sentinels before the balanced-brace check so a leaked header sentinel does not defeat the guard. - tool_healing: allow dotted tool names in the Gemma wrapped start pattern. - llama_cpp (GGUF): buffer wrapper-less Llama-3.2 {"name":..} calls that carry no XML signal, drain a complete object silently and hold an incomplete one, and run the end-of-stream safety net unconditionally so markerless calls are detected and never leak the raw JSON (including truncated fragments). Adds regression tests for the GGUF bare-JSON streaming path and the Mistral display strip. * [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: stop bare-JSON tool calls leaking at EOF, oversized, and into history The second review pass flagged that the Llama-3.2 bare-JSON tool-call handling still leaked raw JSON in several spots; ``strip_tool_markup`` only knows XML/bracket markup, so the bare-JSON form survived it. Fix them symmetrically across the safetensors and GGUF loops: - Safetensors stream-end resolver now routes a held bare-JSON fragment to DRAINING (mirroring GGUF) so a truncated ``{"name":..`` cut off by the end of the stream is dropped instead of flushed as assistant content. The 7/10 reviewer finding. - Both loops now drain (suppress) an oversized still-open bare-JSON call once it passes ``_MAX_BARE_JSON_BUFFER`` instead of streaming the raw prefix, gated on a ``"name"`` key so a giant plain JSON answer still streams; a complete oversized call still executes via the safety net. - Add a shared ``strip_leading_bare_json_call`` helper and apply it to the content kept for the assistant turn in both loops, so an executed bare-JSON call is not replayed as visible text or fed back as next-turn history. Plain JSON answers without a ``"name"`` key are untouched throughout. Adds regression tests for the EOF, oversized, and next-turn cases on both backends plus unit tests for the helper. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: bound the Llama-3 python_tag strip on real control sentinels The route display strip's <|python_tag|> arm ran to the next <| of any kind. A tool-call argument carrying a literal <|...|> token (for example <|cite|> inside a string value) truncated the strip early and leaked the call tail into the visible response. Narrow the stop condition to the genuine Llama control sentinels (eot_id, eom_id, python_tag, start/end_header_id, begin_of_text, finetune_right_pad_id) so embedded markup and JSON are consumed while real header/turn boundaries still bound the strip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden GLM/Gemma parsing, cap GGUF textual calls, share native-template fallback GLM 4.x parser walked a body pre-bounded by the first , so a string argument containing a literal (e.g. code that prints it) was truncated. Walk arg_key/arg_value pairs against the full content instead, since each is delimited by its own and the call's real close is the that precedes the next . Add a truncated wrapper-less Gemma pattern (call:NAME{... with no closing brace) to the markup strip so a call cut off mid-arguments does not leak raw into the visible stream. It runs after the closed form, so a complete call keeps trailing prose. Cap and dedup tool calls parsed from the GGUF TEXTUAL fallback at _MAX_TOOL_CALLS_PER_TURN, mirroring the safetensors loop. Structured delta.tool_calls are grammar-bounded by llama-server, but text parsed straight from content is not, so one runaway turn could fan out into dozens of executions. Extract the native-chat-template fallback into chat_template_helpers (render_native_template / render_with_native_template_fallback) so the transformers and MLX text backends share one implementation. The MLX text path now applies it too, so an Unsloth override template that drops the tools schema no longer silently stops MLX from advertising tools. The MLX VLM path renders via the processor for image tokens and is intentionally left on its own render. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: gate markerless bare JSON on enabled tools and close parser/strip asymmetries The Llama-3.2 custom_tools bare-JSON form has no marker, so any JSON object with a name key was read as a tool call. An ordinary JSON answer like {"name":"Alice","parameters":{"age":30}} was misclassified as a call to a disabled tool and dropped from the visible response. Gate the markerless form on the enabled tool names (threaded through parse_tool_calls_from_text and strip_leading_bare_json_call, supplied by both streaming loops): an object whose name is not an enabled tool is ordinary content. The marker-based forms keep their name-agnostic behaviour (an explicit signal is a real call attempt), and unrestricted mode stays ungated. Also fix two parser/strip asymmetries the parser already tolerated: - A literal inside a parameter value (print("")) truncated both the core and route strips at the first close, leaking the tail. Extend the strip to the call's real close (last before the next opener), mirroring the parser, without merging separate calls. - The single-object Mistral [TOOL_CALLS]{...} shape parsed but _strip_mistral_closed_calls left it, leaking the raw object into display. Strip the balanced object while keeping trailing prose, matching the array and name shapes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio tools: fix strip/parse symmetry and native-template token for DeepSeek/GLM/Kimi Pass-3 review follow-ups on the multi-format tool parser: - Bare Kimi call (<|tool_call_begin|>...<|tool_call_end|> with no section wrapper) is accepted by the parser, so add it to the closed strip patterns so the streaming (non-final) display strip removes it instead of leaking the markup mid-generation. - Route display strip now also runs the wrapper-less Gemma cleanup, so a Gemma 4 call:NAME{..} no longer leaks into the visible answer. - MLX model record carries base_model for a LoRA adapter so the native-template fallback loads the base repo template rather than the adapter's (often template-less) tokenizer. - Native-template reload forwards the load-time HF token so a gated/private model's repo template can still be fetched (transformers and MLX text paths). - GGUF end-of-stream bare-call heuristic is gated on the enabled tool names so a truncated ordinary JSON object ({"name":"Alice","age":) streams as the answer instead of being dropped as a tool call. Adds regression tests for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio tools: gate GGUF bare-JSON suppression on enabled tools and fix python-tag exponent parsing Pass-4 review follow-ups on the GGUF tool loop and Llama-3 parser: - The GGUF bare-JSON suppression sites still keyed off a raw "name" substring, so an ordinary JSON answer whose name is not an enabled tool was dropped when it was truncated, oversized, or reached the no-tool DRAINING fallback (the parser, helper, and safetensors paths were already gated). All three sites now use the shared enabled-name gate, and a held bare-JSON buffer that turns out not to be an enabled call is shown as the answer instead of dropped at stream end. - The Llama-3 python-tag numeric kwarg regex matched only the mantissa, so scientific notation was truncated to its leading digits (1e-3 parsed as 1) and a tool executed with the wrong value. The regex now accepts exponent and decimal forms, and the int/float classification keys off the exponent too. Adds regression tests for the truncated / oversized disabled-name JSON cases (and a counterpart that a truncated enabled call still does not leak) plus the scientific-notation kwargs. * Studio: drop accidentally committed async worker transcripts Eight generated reviewer / async-worker transcripts were committed under studio/backend/async_task_outputs/. They are not imported or referenced by any code and carry only internal task state, so they should never ship in the repo. Remove them and gitignore the directory so they cannot be re-added. * Studio tools: gate safetensors bare-JSON drain, fix nested-name gate and function-XML strip Pass-4 review follow-ups on the shared parser / safetensors loop: - The safetensors oversized and end-of-stream bare-JSON drain branches keyed off a raw "name" substring, so a large or truncated ordinary JSON answer whose name is not an enabled tool was drained instead of streamed. Both now use the shared enabled-tool-name gate, matching the GGUF path. - strip_leading_bare_json_call matched the first "name" anywhere, so a plain JSON answer with a nested name equal to an enabled tool ({"result":{"name":"web_search"}}) was wrongly suppressed. It now extracts the TOP-LEVEL name only, walking past nested objects/arrays and keeping the text when a top-level value is truncated. - The function-XML display strip used a regex negative-lookahead that stopped at a literal opener inside a parameter value and then dropped the rest of the answer to EOF. A scan-based strip mirrors the parser (ignores openers inside an open via _inside_open_parameter) and closes each call at its real , so trailing assistant text after such a call survives. Adds regression tests for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep tools prompt when native-template probe raises; make helper tests hermetic Pass-4 review follow-ups on the native-template fallback: - render_with_native_template_fallback re-renders the live template with tools=None to detect whether it dropped the schema. A template that requires tools can raise on that probe; that must not discard the already-valid tools prompt. The probe is now wrapped so any error returns the original formatted_prompt (transformers would otherwise fall back to manual formatting and lose the schema; MLX would let the exception escape). - The native-template helper tests imported InferenceBackend just to reach the thin wrapper, which pulls in unsloth and its optional vllm package metadata. They now call the dependency-light render_native_template helper directly so they pass in a backend/test environment without vllm. Adds a probe-raises regression test. * Tool parsing: 3.9 import safety, disabled-Auto-Heal contract, capability gate Round-2 review follow-ups on the multi-format tool-call parser: - tool_call_parser: add `from __future__ import annotations`. The module is dependency-light by design (external llama-server wrappers import it standalone) and the package targets python >=3.9, where its PEP 604 `int | None` return annotations would raise TypeError on import. - safetensors + GGUF drain fallback: gate the leading bare-JSON strip on auto_heal_tool_calls. With Auto-Heal off, a truncated enabled-name fragment that did not parse now stays visible, matching the XML strip in the same branch and the disabled-Auto-Heal contract. With Auto-Heal on it is still suppressed. - safetensors capability gate: match the bare-JSON `{"name":` template marker with a whitespace/escape-tolerant regex so a pretty-printed `{ "name" :` or JSON-escaped `{\"name\":` template is not mis-classified as tool-less. The parser already accepts that whitespace via raw_decode, so the gate must too. Regression tests added for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GLM tool-call display strip: treat literal close tag in arg value as data Round-2 review follow-up on the GLM 4.x tool-call format. The GLM call shape is NAMEkv .... The parser was hardened to walk arg_key / arg_value pairs so a literal inside an argument value (e.g. print("")) is treated as data and the call's real close is the that precedes the next . The display strips still used a non-greedy .*? regex, which stopped at the literal and leaked the call's tail into visible content and stale history. Add _strip_glm_calls, a scan that mirrors the parser's close detection, and run it before the regex arms in every strip pipeline: the core strip_tool_markup, the route _strip_tool_xml display/history cleanup, and the safetensors + GGUF streaming strips. Qwen / Hermes {json} has no NAME token after the opener, so it is left to the regex arms unchanged. Regression tests cover the literal-close-tag leak (core + route), normal GLM calls, back-to-back GLM calls, zero-arg GLM, truncated GLM, and untouched Qwen. * Tool parsing: symmetric "function" bare-JSON alias and route strip parity Round-3 review follow-ups, all parser/strip symmetry fixes. - Bare-JSON "function" alias: the markerless parser accepts a call name via obj.get("name") or obj.get("function"), but the strip/gates only knew "name", so a {"function":} call executed while its raw JSON leaked. Teach _top_level_bare_json_name the alias (with "name" precedence and the same nested and truncated-name guards), and widen the guards in strip_leading_bare_json_call, the safetensors and GGUF _looks_like_enabled_bare_json gates, and the route capability marker regex. - Route display/history cleanup: strip a tail-only alias close (the parser accepts ...), and run the parser's guarded function-XML scan (_inside_open_parameter) before _TOOL_XML_RE so a literal nested inside an argument value does not truncate the strip and leak the tail. Regression tests added for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio tools: fix DeepSeek strict recovery, Kimi dotted names, Gemma spaced streaming Round 3 review fixes for the DeepSeek / GLM / Kimi tool-call parsing path. - DeepSeek R1 and V3/V3.1 strict parsing (Auto-Heal off): when a call is truncated (missing closing fence or terminator), skip it and keep scanning for later well-formed calls instead of breaking out and dropping the rest of the envelope. This matches the Kimi strict parser's recovery behaviour. - Kimi dotted tool names: keep the full name after stripping only the functions. prefix and :idx suffix, e.g. functions.mcp.server-list:0 stays mcp.server-list. The previous split on "." truncated dotted MCP names to their last segment. This matches current vLLM (tool_id.split(":")[0].removeprefix("functions.")) and SGLang (^(?:functions\.)?(?P[\w.\-]+):(?P\d+)$). - Gemma wrapper-less call streaming: hold the whitespace-tolerant prefix (call : NAME) in the streaming suppression buffer, matching the parser's _GEMMA_BARE_TC_RE, so the spaced spelling split across chunks is buffered instead of leaking as visible text. Applied to both the safetensors and llama.cpp streaming paths. - Remove dead _render_with_native_template method and the now-unused copy import from inference.py; the live path uses render_with_native_template_fallback. Adds regression tests for DeepSeek R1/V3 strict recovery, Kimi full dotted name preservation, and the Gemma spaced-call streaming suppression. * Studio tools: honor tool budget in GGUF loop and guard function-XML streaming strip Round 4 review fixes. Both are asymmetric-fix bugs where the final/steady path got a guard the analogous streaming/loop path did not. - GGUF tool-call budget: the safetensors loop counts real tool-call turns against max_tool_iterations (re-prompt stalls excepted), but the GGUF loop only bounded the turn count by the enlarged range (max_tool_iterations + _MAX_REPROMPTS). Since this PR raised _MAX_REPROMPTS from 1 to 3, a model that keeps making valid tool calls could run up to three extra tool rounds (with max_tool_iterations=1, four rounds instead of one). Add a _tool_iters_done counter that increments only when a tool actually executed in the turn, and stop once the caller's budget is spent so the post-loop final-answer nudge fires. A duplicate/disabled no-op turn is a correction turn (like a plan-without-action re-prompt) and does not consume budget, preserving the existing "already completed" re-prompt behavior. - Streaming display strip: the final strip runs the guarded _strip_function_xml_calls scanner (a literal inside a parameter value is data, not a nested call), but the GGUF and safetensors streaming strips still used only the open-ended regex arms. When a tool-call argument contained literal function markup, the regex tail ate everything to end-of-text and dropped the real trailing prose after the call's true . Run the guarded scanner (and the balanced Mistral strip) before the regex arms in both streaming paths so streaming and final display agree. Adds regression tests: GGUF valid tool calls respect max_tool_iterations, and the streaming strip keeps trailing prose after a function-XML call with a literal marker. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio tools: safetensors tool budget counts only executed turns (GGUF parity) Follow-up to the GGUF budget fix. The safetensors loop charged max_tool_iterations per non-re-prompt iteration (iteration + 1 - reprompt_count), so a duplicate/disabled no-op turn spent a budget slot even though no tool ran. With a small cap this dropped real work: for max_tool_iterations=2, a model that made a valid call, repeated it (an internal no-op correction turn), then made a distinct valid call executed only the first -- the third turn was sent with no tools and the distinct call was ignored. Track whether a turn actually executed a tool (set on record_result) and count only those turns against the cap, matching the GGUF loop. A duplicate/disabled no-op is a correction turn -- like a plan-without-action re-prompt -- and no longer consumes budget, so the model still gets its "already completed" nudge and another tool-enabled turn. Adds a regression test for the small-cap duplicate-then-distinct-call flow. * Studio tools: fix stale Kimi dotted-name regression test test_pr5624_regressions.py still expected functions.my.tool:0 to resolve to the last segment (tool). The parser now preserves the full dotted name (my.tool) after removing only the functions. prefix and :idx suffix, matching current vLLM/SGLang so dotted MCP names like mcp.server-list survive. Update the assertion, name, and module docstring to the corrected contract (the raw id is still preserved on the call). * Studio: render the reasoning block for safetensors and MLX like GGUF enable_thinking chat templates (Qwen3/Qwen3.5/GLM) prefill an unclosed into the generation prompt, so the model emits only the closing then the answer. The safetensors/MLX chat stream emitted that as plain content, so the reasoning showed inline with no collapsible thinking block, while GGUF (which surfaces reasoning via reasoning_content) rendered one. This brings safetensors and MLX to parity. - _ResponsesReasoningExtractor gains a reasoning_prefilled mode that starts inside the reasoning block and splits on the first ; default False keeps GGUF and every existing caller byte-identical. It suppresses a stray re-emitted and holds partial markers back across chunk boundaries. - _sf_reasoning_prefill_mode gates the mode on reasoning being enabled for the request, an enable_thinking or enable_thinking_effort style, and the template actually using the standard / markers. Models with a bespoke reasoning channel (e.g. gemma's <|think|>/<|channel>) are excluded so their answer is never swallowed; gpt-oss (Harmony) and thinking-off requests are excluded too. - sf_tool_stream and stream_chunks (the latter also serves MLX) feed text through the extractor, emitting reasoning_content then content deltas, with a per-turn reset in the tool loop and a flush before each tool_start; only the visible delta reaches the monitor reply. The two non-streaming drains split reasoning_content the same way. - Tests: extractor prefilled mode (streaming and edge cases), the gate matrix including the gemma-style exclusion, and a route-replay of the tool-loop reasoning stream. * Studio: render the reasoning block for safetensors and MLX like GGUF enable_thinking chat templates (Qwen3/Qwen3.5/GLM) prefill an unclosed into the generation prompt, so the model emits only the closing then the answer. The safetensors/MLX chat stream emitted that as plain content, so the reasoning showed inline with no collapsible thinking block, while GGUF (which surfaces reasoning via reasoning_content) rendered one. This brings safetensors and MLX to parity. - _ResponsesReasoningExtractor gains a reasoning_prefilled mode that starts inside the reasoning block and splits on the first ; default False keeps GGUF and every existing caller byte-identical. It suppresses a stray re-emitted and holds partial markers back across chunk boundaries. - _sf_reasoning_prefill_mode gates the mode on reasoning being enabled for the request, an enable_thinking or enable_thinking_effort style, and the template actually using the standard / markers. Models with a bespoke reasoning channel (e.g. gemma's <|think|>/<|channel>) are excluded so their answer is never swallowed; gpt-oss (Harmony) and thinking-off requests are excluded too. - sf_tool_stream and stream_chunks (the latter also serves MLX) feed text through the extractor, emitting reasoning_content then content deltas, with a per-turn reset in the tool loop and a flush before each tool_start; only the visible delta reaches the monitor reply. The two non-streaming drains split reasoning_content the same way. - Tests: extractor prefilled mode (streaming and edge cases), the gate matrix including the gemma-style exclusion, and a route-replay of the tool-loop reasoning stream. * [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: don't force a tool re-prompt on a negated intent (safetensors parity) The safetensors _INTENT_SIGNAL claimed to mirror GGUF but was missing the negative lookahead, so a refusal like "I will not search the web for that" matched the "i will" intent and triggered the plan-without-action re-prompt (STOP... you MUST call a tool), overriding a valid no-tool answer. GGUF already excludes not/never. Add the same (?!\s+(?:not|never)\b) lookahead so both backends agree. Extends the intent parity test with negated refusals. * studio: parse the outer envelope before DeepSeek/Kimi markers embedded in its args parse_tool_calls_from_text ran the DeepSeek/Kimi marker pre-pass before the shared / parser. When a Qwen/Hermes call's argument contained literal Kimi/DeepSeek markup (for example a user asking the model to explain that syntax), the pre-pass matched the embedded marker and returned it, executing the wrong tool and dropping the real call. Skip the pre-pass when a or envelope opens before the first DeepSeek/Kimi marker, so the shared parser takes the outer call; a genuine marker-led call (no leading envelope) still goes through the pre-pass. Tests for the embedded-marker case and the control. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: trim redundant comments (comment-only, AST-verified) * Studio: trim redundant comments (comment-only, AST-verified) * Studio: prevent Gemma tool-parser DoS on stray delimiters _gemma_parse_value returned the input index unchanged when text[i] was a stray delimiter (,}]), so the list and mapping caller loops that advance on the returned index spun forever at 100% CPU on malformed input such as [},]. Advance past the delimiter so parsing always terminates. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: strip Magistral [THINK] reasoning from final display/history strip_tool_markup removed [TOOL_CALLS] and markup but left a leading Magistral [THINK]...[/THINK] block intact, so its bracket-form reasoning (not the the reasoning channel renders) leaked into the safetensors display and conversation history while GGUF/llama.cpp routes it natively. Drop the leading reasoning block at end-of-turn (final=True) via the existing _strip_mistral_reasoning helper; streaming is untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep times in wrapper-less Gemma tool arguments The wrapper-less Gemma value scanner used _GEMMA_KEY_RE = [\w.\-]+ for keys, which also matches a digit-leading token, so a comma followed by a time or ratio inside a value (call:web_search{query:meet at 10:00, 11:00 tomorrow}) was misread as a new 11: key, truncating the query and injecting a bogus argument. Require keys to start with a letter or underscore, matching the identifier-start rule the wrapped path already uses (_GEMMA_NEXT_KEY_RE). Add a regression test. * Studio: treat markers/close-tags inside tool-call arguments as data Four parser correctness fixes where a valid argument string was mistaken for structure: - DeepSeek: find the envelope-end token outside JSON strings, so a query/code argument containing the literal token no longer truncates the body and drops the whole call. - GLM: locate the real as the one whose next token is / / end, so a value containing a literal (or ) is kept instead of executing the tool with corrupted arguments. - Attribute-form envelopes now count in the embedded-marker guard, so a DeepSeek/Kimi marker inside a parameter value does not hijack the outer call and run the wrong tool. - Wrapper-less Gemma call:NAME{...} is gated on the enabled tool names (parse and display strip), mirroring the Llama bare-JSON gate, so a disabled/example name in prose is not stolen as a call and the real answer is preserved. Add regression tests for each. * Gate route Gemma wrapperless strip by enabled tools; make Kimi section-end search string-aware Route-level display stripping now threads the enabled tool-name set into the Gemma wrapperless-call strip, so prose that mentions a disabled tool (call:foo{...}) is preserved while active tool calls are still stripped. This mirrors the parser-level gate already used in tool_call_parser. The Kimi section-end lookup now searches outside JSON string literals, so a section-end marker appearing inside an argument string no longer triggers a false truncation that drops a valid tool call. * Run DeepSeek/Kimi pre-pass when a closed tool-call example precedes a real block The marker pre-pass was skipped whenever any / opener appeared before the first DeepSeek/Kimi marker, even when that opener was a CLOSED syntax example in prose that ends before the real block. In that case parse_tool_calls_from_text skipped the DeepSeek/Kimi parsers and the genuine tool call was dropped while a phantom tool named in the example ran instead. Only treat a marker as embedded in a leading envelope when removing the closed outer / envelopes also removes every marker (the marker actually sat inside one). A marker left standing is a real call, so the pre-pass runs. The legitimate case of a marker inside a closed outer envelope's arguments is preserved. * Honor reasoning_effort none in safetensors prefill; strip Magistral reasoning while streaming Two safetensors/MLX reasoning fixes surfaced in review: _sf_reasoning_prefill_mode only checked enable_thinking, so an enable_thinking_effort (GLM-5.2) request that disables thinking via reasoning_effort=none (without enable_thinking=False) still began in prefilled- mode. A plain answer with no was then swallowed whole into reasoning_content and the visible response came back empty. Thread reasoning_effort into the predicate and treat none as disabled, mirroring _request_reasoning_kwargs. strip_tool_markup_streaming stripped tool markup but not the leading Magistral [THINK]...[/THINK] bracket block, so the raw chain-of-thought leaked into the streamed safetensors content instead of the reasoning drawer (GGUF routes it natively). Apply _strip_mistral_reasoning first, matching the final strip; an unclosed [THINK] is held from the marker on so nothing flickers. * Heal truncated outer tool envelopes and keep quoted Gemma args intact Two follow-ups from review of the marker pre-pass and Gemma parsing: The leading-envelope guard only removed CLOSED outer / envelopes before deciding whether a DeepSeek/Kimi marker was embedded, so a truncated outer call missing its close tag (whose argument embeds a marker) was treated as a standalone marker and the embedded sample ran instead of the intended outer call being Auto-Healed. Decide on the last outer opener before the marker and whether it closed before the marker instead, so a closed syntax example still runs the pre-pass while a real closed-or-truncated outer call keeps it. The wrapper-less Gemma argument scan tracked bracket depth but not quotes, so a quoted value containing a comma followed by a key-like token (a search query such as "weather, location: Boston") was split mid-string, truncating the value and fabricating an extra argument. Track quote state (with escapes) so the top-level comma boundary is only taken outside quoted spans. * Span outer envelopes to their real close when locating embedded markers Locating the DeepSeek/Kimi marker relative to a leading outer envelope used the FIRST close tag after the opener, so a literal or inside an argument value (for example python code that contains the text) was mistaken for the envelope boundary. The marker after it was then treated as a standalone call and the embedded sample ran instead of the intended outer call. Match the closed outer envelopes with the shared patterns that already extend to the real final close (a literal close inside a value is data), and treat a marker that survives their removal as embedded only when a still-open (truncated) outer opener precedes it, so Auto-Heal still repairs a truncated outer call. A closed syntax example before a genuine block still runs the pre-pass. * Span the tool_call outer envelope to its real close in the marker guard The leading-envelope check reused the lazy .*? strip pattern, so a Qwen/Hermes JSON argument containing a literal ended the span early. A DeepSeek/Kimi sample later in that same string then survived the closed-envelope removal, and the pre-pass executed the embedded call instead of the outer . The arm already spanned to its real close; give the same real-close pattern (with the negative lookahead that keeps back-to-back calls separate) so a literal close inside a value is data. * Preserve no-tool Gemma prose and keep later R1 calls when healing a close Two review follow-ups: _gemma_strip_gate returned None when no tools were enabled, and None means strip every markerless call:NAME{...} block, so a no-tool answer that documents the syntax (or the Anthropic display path, which passes an empty tool list as None) had that prose deleted. It is a display/history gate, so return the enabled-name set instead -- an empty set when no tool is enabled, which strips nothing because every call:NAME{...} is then prose. The DeepSeek R1 heal path located the close fence with an unbounded forward search, so when a first call had balanced JSON but omitted its fence the search landed on a LATER call's terminator and pos advanced past that valid call, dropping it. Match the close immediately after the JSON (whitespace-skipped) like the strict path, and advance by just the JSON when it is absent, so a multi-call turn keeps its later well-formed calls (heal is now a superset of strict). * Resume wrapper-less Gemma scan past a consumed call's balanced body The markerless call:NAME{...} scan used finditer, which resumes right after the opening call: token, so a nested call:OTHER{...} mentioned inside the first call's own quoted string argument (for example a web_search query that quotes the Gemma tool syntax) was re-matched and returned as a spurious second tool call, executing an unintended tool. Walk with a manual cursor that resumes after the outer call's balanced body (brace matching already skips quoted braces), so a call's arguments are never rescanned. Genuinely separate back-to-back calls and disabled/example prose are unaffected. * Mistral outer call wins over XML literals; align healer signals with its parser Two follow-ups on the shared-parser ordering after the healing-passthrough merge: - A well-formed [TOOL_CALLS] call whose JSON arguments quote tool XML parsed the literal instead of the outer call (executing the wrong tool). When the first XML signal sits inside a leading balanced Mistral body it is argument data, so the Mistral parser now runs first; an XML signal before the trigger keeps the normal order, so a [TOOL_CALLS] literal inside an XML call's arguments still stays data. - passthrough_healing buffered streams on the parser module's broadened signal list (now including <|python_tag|> and [TOOL_CALLS]) but promotes with core.tool_healing, which does not parse those forms: a streamed Mistral or Llama text call was held until finalization and flushed as prose. The healer keeps its own signal list limited to the formats it can promote, restoring immediate streaming for the rest. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: Gemma wrapper-less marker literals and quotes, GLM embedded close pair - The Gemma fallback deferral now keys on an actual wrapped opener (_GEMMA_TC_RE), not the wrapper literal anywhere in content: a wrapper-less call whose argument merely mentions <|tool_call> has nothing tool_healing can parse, and deferring it lost the call entirely (not executed and stripped from display). - New _gemma_body_brace_end boundary scanner honors single- and double-quoted strings like _gemma_parse_stripped_body, shared by parse and strip, so a quoted brace in a code argument (code:print('}')) no longer truncates the executed arguments or the strip span. - _glm_value_close now requires a structural to sit at balanced quote state: the full pair embedded inside a string literal is data, not an early close. When no candidate balances, the first token-valid close wins as before. * Address review: leading envelopes win over rehearsed literals - New _first_foreign_tool_signal shared by the leading-envelope guards adds <|python_tag|> to the protected signal set: the spelled-out literal inside a Mistral call's arguments (a query about Llama built-in tool syntax) executed the inner literal instead of the outer call. - New _xml_signal_inside_leading_bare_json guard, sibling of the Mistral one: a leading bare-JSON call whose string argument quotes tool XML (a code value citing ) had the literal promoted by the shared XML pass before the bare-JSON parser ran. - Magistral [THINK]...[/THINK] is dropped once at parse entry instead of only inside the Mistral parser, so a call rehearsed in the think block in a foreign format can no longer be promoted while the real call after the block is lost. Parse now agrees with the display strip. * [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 * Address review: a disabled leading bare-JSON object keeps its literals as data When the leading bare-JSON object is ordinary content (name not an enabled tool), the guard proved the first tool signal sits inside it, so falling through to the XML/python_tag passes promoted quoted string data as a real call. Drop the object and parse only the tail: a real call after the object still parses, nothing inside it can be promoted. * Address review: apostrophes in raw Gemma values, GLM strict key contract, per-model template token - Quote openers in the wrapper-less Gemma boundary and body scanners now require value-start context (after : { [ ( , =): an apostrophe inside an unquoted value (query:what's the weather) opened quote mode, swallowed the real closing brace, and lost the whole call on common contraction queries. Quoted values keep hiding delimiters as before. - A GLM with no tag now rejects the call in strict mode, matching the unclosed-value contract, instead of executing the tool with the argument silently dropped; Auto-Heal keeps the lenient skip. - The native-template fallback reads the hf_token stored on the model record instead of the instance-wide last-load token, so a later token-less load cannot break template fetches for a previously loaded gated model (both the transformers and MLX backends). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: Mistral literals inside leading JSON, whitespace-tolerant wrapped Gemma opener - The leading bare-JSON guard now treats the [TOOL_CALLS] trigger as a foreign signal: the Mistral parser runs before the bare-JSON one, so a literal quoted inside the leading object's strings was promoted over the outer call (or over ordinary JSON content). - tool_healing's wrapped Gemma opener tolerates whitespace around call and the colon: sampling drift emits call: name{ and call : name{, and rejecting those lost the call entirely because no fallback re-parses the wrapped form. Strict mode still requires the closing tag. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: DeepSeek/Kimi markers inside leading JSON and Mistral envelopes stay data The DeepSeek/Kimi pre-pass runs before the outer-call parsers, and _marker_inside_leading_envelope only protected XML envelopes: a marker quoted inside a leading bare-JSON or Mistral call's argument strings was promoted as a separate no-arg call and the real outer call dropped. The guard now recognizes those two leading envelopes as well; standalone DeepSeek/Kimi calls keep parsing. * Address review: accept dotted Gemma argument keys in the key-quoting scanner The scanner quoted keys of [alnum_-] only, so a dotted key (user.name:...) was left unquoted, json.loads failed, and the whole wrapped call was lost (parse empty, strip wipes the markup). Dots now match the parser's own key/name charset. * Address review: a real DeepSeek/Kimi call after a disabled leading JSON object still parses DeepSeek/Kimi markers are foreign signals for the leading bare-JSON guard too: a marker literal inside a disabled leading object made the envelope guard skip the pre-pass for the whole message, so a real DeepSeek/Kimi call after the object was dropped. Routing the case through the guard's drop-and-parse-the-tail recursion reaches the real call while the literal inside the object stays data. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: leading Mistral call owns the turn, dotted keys after bare values - A LEADING parseable [TOOL_CALLS] call now runs the Mistral parser first unconditionally: literal XML in trailing prose after the call was promoted by the earlier shared XML pass, executing the quoted example instead of the real leading call. XML leading keeps the normal order. - _GEMMA_NEXT_KEY_RE accepts dots so a dotted key after a bare value (query:foo,user.name:bob) ends the value at the comma instead of being swallowed into it, matching the round-earlier key-quoting charset. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: a leading wrapper-less Gemma call owns the turn A quoted foreign literal inside a leading wrapper-less Gemma call's argument (a query citing another tool syntax) was promoted by tool_healing before the Gemma fallback ran, executing the quoted example and dropping the outer call. New leading guard, sibling of the Mistral and bare-JSON ones, gated on an enabled name since the form is markerless. Foreign markup leading keeps the normal order. * Fix merge resolution: restore both leading-guard test classes intact * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: markup quoted inside a nameless leading JSON answer stays data The leading bare-JSON guard required a top-level name, so a structured JSON answer quoting tool markup in its strings (a response_format turn documenting a tool's syntax) had the literal promoted by the later passes. A nameless leading object that parses as real JSON now routes through the same decline-then-parse-the-tail path; non-JSON braced prose keeps the old behaviour, and a real call after the answer still parses. * Address review: JSON answers stay data, nested Gemma quotes, earliest envelope, no failure caching - A whole-content JSON value is a structured answer: the markerless Gemma scan and its strip no longer promote or strip a quoted example of an enabled tool's syntax inside it. - Nested stripped-stream Gemma values now unquote quoted string leaves recursively, so {loc:{city:"New York"}} hands the tool New York, matching the top-level coercion. - The DeepSeek/Kimi pre-pass dispatches by earliest envelope opener, so a leading real call wins over a trailing example of the sibling format in either direction. - A failed native-template fetch is no longer cached as no-template: the next call retries after the model record's token is fixed or a transient Hub error clears; only definitive loads are cached. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: closed calls precede the marker pre-pass, truncated Gemma scan stops, quoted nested delimiters - A closed non-DeepSeek/Kimi call preceding the first DS/Kimi marker owns the turn: a trailing syntax example, or one quoted inside a wrapped Gemma argument, was promoted by the pre-pass and dropped the real leading call. Wrapped Gemma joins the outer-envelope pattern sets. - An unbalanced wrapper-less Gemma call now stops the scan (mirroring the strip contract) instead of resuming inside its own argument text, where a quoted enabled call would be promoted. - Raw-quoted strings in nested stripped-stream Gemma values hide delimiters, so {city:"New, York"} is one value instead of a split pair, returned unquoted like the top-level coercion. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: string-marker literals in wrapper-less args, mid-value quoted phrases - The wrapper-less deferral guard no longer keys on the <|"|> literal: a real call whose argument merely mentions the string marker was deferred to tool_healing, which has no wrapped opener to parse, losing the call. The wrapped-opener check alone owns the deferral. - Double quotes now also open at the start of a word, so a quoted phrase mid-value (query:find "weather, location: Boston", limit:3) hides its delimiters instead of splitting the value into garbage keys; apostrophes keep the value-start-only rule so contractions stay prose. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: strict GLM refuses in-quote close fallback, Gemma guard covers preambles - _glm_value_close gains a strict flag: a truncated value whose only close candidates sit inside a string literal rejects the call in strict mode (Auto-Heal keeps the lenient partial), restoring the strict contract the quote-aware fallback had weakened. - The leading wrapper-less Gemma guard no longer requires the call to open the response: a visible preamble before call:NAME{...} is the normal shape, and the quoted foreign literal inside the argument was promoted again in that shape. An enabled balanced call beginning before the first foreign signal owns it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: contextual GLM quote openers, disabled Gemma examples stay prose, JSON array answers - The GLM value-close quote tracker uses the same contextual openers as the Gemma scanners (single quote after punctuation context, double quote also at word start), so strict mode accepts a normal apostrophe value again while still rejecting a truncated value whose only close candidates sit inside a string literal. - A disabled wrapper-less Gemma call is prose by design, so a tool literal quoted inside it no longer promotes: the span is dropped for parsing and the tail parsed, mirroring the nameless-JSON guard. - Leading JSON ARRAY answers join the leading-JSON envelope guard, so a marker quoted inside a structured array response stays data. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Align closed-envelope regression test with the document-order contract The test asserted the pre-round-13 behavior (trailing DeepSeek/Kimi block wins over a leading closed envelope) while the shipped rule is document order: the leading closed call owns the turn. Rename the test and assert the leading call so the suite matches the contract exercised by test_leading_xml_call_wins_over_trailing_kimi_example. * Parse a leading Llama-3.2 bare-JSON call before the markerless Gemma scan The bare-JSON form only ever matches a leading call object, and document order says that call owns the turn. Running the Gemma wrapper-less scan first let an enabled call:NAME{...} snippet quoted inside the leading call's string arguments steal the turn when the JSON was not the whole content (trailing prose or a second ;-separated call), executing the quoted tool instead of the real one. Reordering cannot take a leading Gemma call's turn since that content never starts with an object brace. * Leading-call ownership: Mistral trigger in Gemma guards, closed bare JSON before markers, depth-aware nested Gemma values Three parser gaps against the document-order contract: The wrapperless Gemma leading guards did not count [TOOL_CALLS] as a foreign signal, so a leading Gemma call quoting a Mistral snippet in its argument lost the turn to the quoted literal. Both the enabled-call and disabled-example guards now include the trigger, matching the bare-JSON guard's local inclusion. _marker_inside_leading_envelope required the DeepSeek/Kimi marker to sit inside the first closed bare-JSON or Mistral call. A marker after that closed call (a trailing example or data in a later ;-chained call's strings) now also defers to the leading call, the same inside-or-after rule the closed XML envelope patterns already applied. The nested Gemma primitive value scan split on every comma, corrupting arguments like opts:{code:print(1,2),lang:py}. It now applies the same paren/brace depth, contextual quote openers, and comma-only-before-a-key mapping rule as the top-level scan. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gemma leading guard: a closed enabled call preceding the signal owns the turn The wrapperless Gemma guard only claimed the turn when the first foreign signal sat inside the first enabled balanced call. When that call closed before the signal (a second call quoting a Mistral or Kimi literal, or a trailing prose example), the guard forfeited the turn and the foreign parser promoted the quoted literal, dropping the real Gemma calls. Apply the same inside-or-after ownership rule as the closed bare-JSON and Mistral envelopes, gated on an enabled name so the name-agnostic legacy path is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Marker guard: only an executable leading bare-JSON call owns the turn The bare-JSON branch of the leading-envelope marker guard claimed the turn for any NAMED leading object. A disabled-name object is prose by design (the bare-JSON parser will not execute it), so deferring the DeepSeek/Kimi pre-pass to it lost the real later call entirely. Gate the ownership claim on the enabled set (or the name-agnostic None path). A marker inside the disabled object's own strings stays data, matching the tail-exclusion contract; a marker after it now falls through so the pre-pass parses the real call. The Mistral branch stays ungated since [TOOL_CALLS] parsing is never name-gated. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gemma scan skips leading JSON answers; GLM heal bounds values at structural tags Two fixes to the document-order data contracts: The markerless Gemma scan only exempted whole-content JSON, so a leading JSON answer followed by prose had an enabled call:NAME{...} snippet inside its strings promoted to a real executed call and stripped from the displayed answer. Both the parse and strip scans now start after a balanced json-valid leading value span, keeping parse and strip mirrored. Real calls after the answer still parse; mid-prose JSON gets no exemption. The GLM heal fallback for a missing closing arg_value tag took the entire remainder as the value, executing markup-contaminated arguments like city="NYC" and swallowing trailing prose. The healed value now stops at the next arg_key or tool_call close and the pair walk resumes there. EOF-truncated values keep the partial heal, strict mode still rejects, and closed values holding a literal close tag in quotes are untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Compress docstrings in the multi-format tool parser to their contract essence * Condense parser guard comments and test narration to contract essentials * verify_import_hoist: exempt __future__ imports and same-diff relocations Two false positives fired on this PR's refactor. A from __future__ import is a compiler directive whose name never appears as a runtime load, so HOISTED-IMPORT-UNUSED can never see it used, yet the file requires it for PEP 604 annotations on Python 3.9. TARGET-CHANGED flagged the deliberate move of the strip-pattern constants into core.inference.tool_call_parser as a silent re-point even though the old module-level target was removed and the new one added in the same diff. Both get narrow exemptions; a re-point to a pre-existing target is still caught, and the self-test negative controls all pass unchanged. * Leading bare-JSON calls own the turn; function calls end at the first balanced close The XML-signal guard for a leading bare-JSON call required the signal strictly inside the object, so a trailing XML example stole the turn from the leading call; it now applies the same inside-or-after rule as the Mistral guard. Function-XML calls also ended at the LAST close tag, which let prose after a closed call that mentions a literal close tag get swallowed into the final parameter value; calls now end at the first close tag that is not inside an open parameter, and the strip mirrors the same rule so parse and strip agree. * [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 * Attribute-form calls end at the first balanced close; bare-JSON strip requires the call shape The attribute form parser still kept the last close tag in the call window, folding prose after a closed call into the final parameter value. It now takes the first close not inside an open parameter, the same rule the equals form and the strip already use. The leading bare-JSON strip deleted any closed object whose top-level name matched an enabled tool, including plain JSON answers the parser correctly rejects as non-calls. The strip (and the drain gate that delegates to it) now requires the parser's exact call shape, so answers like {"name":"web_search","result":...} stream and display intact. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * False-alarm markers keep the answer; the bare-JSON strip consumes the whole chain The trailing strip arms dropped everything from a bare marker to EOF, so a normal answer that mentions [TOOL_CALLS] or another marker literally was truncated (or fully swallowed when it started with the literal) after the no-call drain fallback. Those arms now require a call-shaped lookahead or marker-at-EOF before dropping; truncated real calls still strip. Chained bare-JSON turns executed both calls but stripped only the first object, so the second call's raw JSON replayed into the next assistant history message alongside the structured tool_calls. The strip now consumes the entire chained run of call-shaped enabled objects while non-call answers, disabled names, and trailing prose stay intact. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * DeepSeek and Kimi trailing strip arms require a call-shaped lookahead Same false-alarm rule as the bare-word markers: a prose answer that mentions a DeepSeek or Kimi marker literally keeps its tail, while truncated real envelopes and bare end-of-text fragments still drop. * Attribute-form containment, parameter-close-decides rule, preamble-tolerant Mistral guard, strict strip shape Four document-order and containment fixes. A leading attribute-form call now parses before the shared XML pass, so markup quoted in its parameter stays data. The open-parameter scan lets the parameter's own close tag decide, so any number of literal function closes inside one value stay data, restoring the pre-close-scan behavior for multi-close arguments. The leading-Mistral guard tolerates a visible preamble, with the leading-bare-JSON guard running first so a trigger quoted inside a leading JSON object stays data. The bare-JSON strip requires the parser's top-level name in every mode, so nested-name JSON answers survive name-agnostic stripping. * Keep buffering long wrapper-less Gemma tool names instead of leaking the prefix The streaming buffer stopped holding a call:NAME prefix at a fixed 32-char cap, so a Gemma wrapper-less call to a tool whose name exceeds that (OpenAI allows 64 chars, MCP names run longer) streamed its raw call:longname text as visible content before the end-of-turn parser executed it. Hold the variable-length prefix while it still matches the call: shape, bounded like the bare-JSON path and self-terminating into prose, draining once the opening brace arrives. * Keep prose that only mentions DeepSeek/Kimi markers in the route display strip The route-level _TOOL_XML_RE DeepSeek/Kimi arms consumed from an opener up to the end of text whenever the marker appeared, so an answer that merely refers to a marker (for example "See <|tool_call_begin|> in the docs") had the rest of the reply truncated. The parser-level _TOOL_ALL_PATS already gates these arms with a call-shaped lookahead. Mirror it here so a marker is only stripped when a real call follows it or it is a bare fragment at end of text. * Tighten tool-calling parser and backend comments * Pass trust_remote_code when reloading native tokenizers The native-template fallback re-fetches a model's native chat template from its repo when an Unsloth override template drops the tools schema. The secondary AutoTokenizer.from_pretrained threaded hf_token but not trust_remote_code, so for a model loaded with trust_remote_code=True whose tokenizer repo carries custom code the reload raised, was swallowed, and the request silently kept the tool-dropping prompt for a model that supports tools. Store the loaded trust_remote_code on each backend's per-model info dict and source it in render_native_template, so the reload re-uses exactly the consent granted at load. For a LoRA adapter the reload targets the base model, whose remote code was gated and loaded under the same stored flag, so re-passing it executes no unconsented code. Falsy stored flag preserves the prior behaviour. Adds a regression test that fails without the flag (custom-code reload raises, returns None) and passes with it (tools-advertising native prompt returned). * Treat <|python_tag|> as an outer marker envelope A Llama-3 <|python_tag|> tool call (built-in NAME.call(...) or custom {json} form) whose argument quotes a complete DeepSeek/Kimi example was hijacked by the DeepSeek/Kimi marker pre-pass: the embedded example (for example delete_all) executed instead of the real outer call. python_tag is Llama-3's tool-call envelope, so a marker quoted inside its arguments is data, the same as for , , bare JSON, Mistral and wrapper-less Gemma, which the guard already covers. Add <|python_tag|> to _OUTER_ENVELOPE_OPEN_RE with a call-shaped lookahead (mirroring the _TOOL_ALL_PATS python_tag arm) so the marker pre-pass is suppressed when a python_tag call opens before the first marker, while a bare prose <|python_tag|> mention is left untouched. * Tighten tool-call parser comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: Daniel Han Co-authored-by: danielhanchen --- .gitignore | 2 + scripts/verify_import_hoist.py | 15 +- .../core/inference/chat_template_helpers.py | 184 +- studio/backend/core/inference/inference.py | 34 + studio/backend/core/inference/llama_cpp.py | 124 +- .../backend/core/inference/mlx_inference.py | 32 + .../core/inference/passthrough_healing.py | 9 +- .../core/inference/safetensors_agentic.py | 138 +- .../core/inference/tool_call_parser.py | 1546 +++++++++++++++-- studio/backend/routes/inference.py | 168 +- .../tests/test_gemma_tool_parse_edge_cases.py | 61 +- .../backend/tests/test_llama_cpp_tool_loop.py | 136 +- studio/backend/tests/test_mcp_servers.py | 4 +- .../tests/test_mlx_inference_backend.py | 51 +- .../test_native_template_trust_remote_code.py | 176 ++ .../backend/tests/test_pr5624_regressions.py | 1011 +++++++++++ .../tests/test_responses_tool_passthrough.py | 14 +- .../test_safetensors_capability_advertise.py | 133 +- .../test_safetensors_reasoning_stream.py | 12 +- .../tests/test_safetensors_tool_loop.py | 1247 ++++++++++++- .../tests/test_tool_call_parser_strict.py | 632 ++++++- studio/backend/tests/test_tool_xml_strip.py | 194 ++- 22 files changed, 5472 insertions(+), 451 deletions(-) create mode 100644 studio/backend/tests/test_native_template_trust_remote_code.py create mode 100644 studio/backend/tests/test_pr5624_regressions.py diff --git a/.gitignore b/.gitignore index 9f7d4b8c60..39ca2226ca 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,8 @@ outputs/ exports/ /datasets/ studio/backend/assets/datasets/ +# Generated async worker / reviewer transcripts (never part of the product). +studio/backend/async_task_outputs/ unsloth_training_checkpoints/ *.gguf *.safetensors diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 2d30265abe..22a21a2ebc 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -564,7 +564,10 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] for n, tids in b["module_import_targets"].items(): if tids & after_used: continue # resolved -> fine - # `from __future__ import ...` is a compiler directive whose name is never loaded; skip it. + # `from __future__ import ...` is a compiler directive, not a runtime + # binding: the name (`annotations`, ...) is never loaded, so it can never + # "resolve" to a use. Skip it so a legitimately-added future import + # (e.g. `annotations` for lazy PEP 604 `X | None` on py3.9) is not flagged. if all(t.startswith("from:__future__:") for t in tids): continue newly_added = bool(tids - before_module_targets) @@ -592,9 +595,13 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] # `import urllib.error` next to `import urllib.request`). Nothing the name # resolved to before is lost, so no reference is re-pointed -- skip it. # - # A deliberate *relocation* is also benign: a name's import source moves A -> B in - # THIS diff (old `from A import x` removed, new `from B import x` added). Mirrors the - # TARGET-MISSING tolerance. Re-pointing to a pre-existing target (clash) is NOT exempted. + # A deliberate *relocation* is also benign and must not block: when a name + # keeps its spelling but its import source is moved A -> B in THIS diff (the + # old `from A import x` is removed at module level and a new `from B import x` + # is added), the swap is intentional, not a silent re-point to a pre-existing + # different object. This mirrors the relocation tolerance already applied to + # TARGET-MISSING. The dangerous case -- the name now resolving to a target + # that already existed before (shadow/clash) -- is NOT exempted. removed_module_targets = before_module_targets - after_module_targets for key, tafter in b["target_by_use"].items(): tbefore = a["target_by_use"].get(key) diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index f58c93b7fe..dfd4c1c0bc 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -3,13 +3,19 @@ """ Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg -fallback for templates that reject reasoning/tools args. +fallback for templates that reject reasoning/tools args, plus the shared +native-chat-template fallback used by the transformers and MLX backends. """ +import copy import json +import logging from typing import Optional +logger = logging.getLogger(__name__) + + def _normalize_tool_call_arguments(messages: list) -> list: """Coerce each assistant ``tool_calls[].function.arguments`` from a JSON string to a dict. @@ -110,3 +116,179 @@ def apply_chat_template_for_generation( if normalized is messages: raise return _render(normalized) + + +def render_native_template( + *, + model_info: dict, + active_model_name: Optional[str], + messages: list, + tools: list, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + apply_fn = None, + hf_token: Optional[str] = None, +) -> Optional[str]: + """Render ``messages`` + ``tools`` with the model's NATIVE chat template. + + Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit + the ``tools`` schema, so a tool-calling turn silently stops advertising tools. + The native template ships in the model repo and carries the family's + tool-calling syntax. It is loaded straight from the repo (bypassing any + override on the live tokenizer) and cached on ``model_info``. Returns the + rendered prompt only if the native template actually emits the tools (render + differs with vs without tools); otherwise ``None``. + + ``hf_token`` is the token the model was loaded with -- passed to the repo load + so a gated/private model's native template can still be fetched (otherwise the + fallback fails silently and keeps the override prompt that dropped tools). + + ``trust_remote_code`` is sourced from ``model_info`` (the value the model was + actually loaded with) rather than a call-site argument, so the native-template + reload uses exactly the consent already granted at load. A custom-code tokenizer + repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is + passed, so without this the fallback fails silently and keeps the tool-dropping + prompt for a model the user already consented to run remote code for. For a LoRA + adapter the reload targets the base model, whose remote code was gated and loaded + under the same stored flag, so re-passing it executes no unconsented code. + """ + # ``apply_fn`` lets a backend inject its own render; defaults to the module helper. + if apply_fn is None: + apply_fn = apply_chat_template_for_generation + native_tpl = model_info.get("native_chat_template") + if native_tpl is None: + # A LoRA adapter's native template lives on the base model, not the adapter id. + template_source = model_info.get("base_model") or active_model_name + # Re-use the load-time trust_remote_code so a custom-code tokenizer repo can + # instantiate its class (the stored flag already covers template_source). + trust_remote_code = bool(model_info.get("trust_remote_code", False)) + try: + from transformers import AutoTokenizer + nt = AutoTokenizer.from_pretrained( + template_source, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, + ) + native_tpl = nt.chat_template or False + except Exception as exc: + logger.warning( + "Could not load native chat template for '%s': %s", + template_source, + exc, + ) + # A failed fetch is not "no template": leave the sentinel unset so the next + # call retries (caching False would pin the tool-dropping override). + return None + model_info["native_chat_template"] = native_tpl + if not native_tpl: + return None + + tokenizer = model_info.get("tokenizer") or model_info.get("processor") + if tokenizer is None: + return None + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) + # Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the + # generation lock) races concurrent requests. + try: + render_tokenizer = copy.copy(tokenizer) + render_tokenizer.chat_template = native_tpl + except Exception as exc: + logger.warning( + "Could not clone tokenizer for native-template render of '%s': %s", + active_model_name, + exc, + ) + return None + try: + with_tools = apply_fn( + render_tokenizer, + messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + no_tools = apply_fn( + render_tokenizer, + messages, + tools = None, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + logger.warning( + "Native-template tool render failed for '%s': %s", + active_model_name, + exc, + ) + return None + return with_tools if with_tools != no_tools else None + + +def render_with_native_template_fallback( + *, + formatted_prompt: str, + tokenizer, + model_info: dict, + active_model_name: Optional[str], + messages: list, + tools: Optional[list], + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + apply_fn = None, + hf_token: Optional[str] = None, +) -> str: + """Return ``formatted_prompt``, swapping in a native-template render when an + override template dropped the ``tools`` schema. + + If ``tools`` were requested but the live render is identical with and without + them (detected by comparison, robust against tool names in the system prompt), + re-render with the model's native template. Shared by the transformers and MLX + backends so both advertise tools consistently. ``hf_token`` is forwarded so a + gated/private model's native template can still be fetched.""" + if not tools: + return formatted_prompt + if apply_fn is None: + apply_fn = apply_chat_template_for_generation + # Probe whether the live template dropped the schema. A tools-requiring template + # can raise here; on any error keep the valid tools prompt rather than lose it. + try: + probe_no_tools = apply_fn( + tokenizer, + messages, + tools = None, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + logger.warning( + "No-tools probe failed for '%s'; keeping the existing tools prompt: %s", + active_model_name, + exc, + ) + return formatted_prompt + if formatted_prompt != probe_no_tools: + return formatted_prompt # template already emits the tools schema + native_prompt = render_native_template( + model_info = model_info, + active_model_name = active_model_name, + messages = messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + apply_fn = apply_fn, + hf_token = hf_token, + ) + if native_prompt: + logger.info( + "Override template for '%s' dropped tool schemas; using the model's " + "native template for this tool-calling turn.", + active_model_name, + ) + return native_prompt + return formatted_prompt diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index eaee5a213a..164f202681 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -269,6 +269,9 @@ class InferenceBackend: gpu_ids: Optional[list[int]] = None, ) -> bool: """Load any model: base, LoRA adapter, text, or vision.""" + # Keep the token so the native-template fallback can fetch a + # gated model's repo template later during generation. + self._hf_token = hf_token # GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it. if max_seq_length <= 0: max_seq_length = 2048 @@ -279,6 +282,8 @@ class InferenceBackend: # Already loaded? if model_name in self.models and self.models[model_name].get("model"): logger.info(f"Model {model_name} already loaded") + if hf_token: + self.models[model_name]["hf_token"] = hf_token self.active_model_name = model_name return True @@ -294,6 +299,14 @@ class InferenceBackend: ) self.models[model_name] = { + # Per-model token: the native-template fallback must use the + # token this model was loaded with, not whichever loaded last. + "hf_token": hf_token, + # Per-model consent: the native-template reload must re-use the + # exact trust_remote_code this model (and a LoRA's base) was loaded + # with, so a custom-code tokenizer repo can be re-fetched without + # executing any code the user did not already consent to. + "trust_remote_code": trust_remote_code, "is_vision": config.is_vision, "is_lora": config.is_lora, "is_audio": config.is_audio, @@ -1040,6 +1053,27 @@ class InferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, ) + + # If tools were requested but the (possibly overridden) template ignored + # them, fall back to the model's native template (shared with MLX). + from core.inference.chat_template_helpers import ( + render_with_native_template_fallback, + ) + + formatted_prompt = render_with_native_template_fallback( + formatted_prompt = formatted_prompt, + tokenizer = tokenizer, + model_info = model_info, + active_model_name = self.active_model_name, + messages = template_messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + apply_fn = self._apply_chat_template_for_generation, + hf_token = model_info.get("hf_token"), + ) + logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: logger.error(f"Error applying chat template: {e}") diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5e67f6b484..455d1d084c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -40,11 +40,15 @@ from core.inference.llama_server_args import ( ) # Share strip / signal constants with the multi-format parser so BUFFERING also -# catches Llama-3 / Mistral / Gemma 4. +# catches Llama-3 / Mistral / Gemma 4 (legacy helper only knew / str: if not (auto_heal_tool_calls or force): return text - return _shared_strip_tool_markup(text, final = final) + return _shared_strip_tool_markup( + text, final = final, enabled_tool_names = _enabled_tool_names + ) def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text - # Shared patterns so a textual Mistral/Llama call entering DRAINING is stripped, not - # leaked. Mistral first; no final trim so incremental length comparisons hold. + # Shared parser patterns (not the legacy tool_healing set) so textual + # Mistral/python_tag calls entering DRAINING never leak. Balanced strips + # first (nested JSON removed whole); no final trim so length compares hold. text = _strip_mistral_closed_calls(text) - # Parser-accurate function-XML scan before the regex arms so a literal ```` - # in a value doesn't make the tail eat trailing prose after the real ````. + text = _strip_gemma_wrapperless_calls(text, _enabled_tool_names) + # Parser-accurate scans close at each call's REAL terminator before + # the regex arms: literal markup inside a value is data. text = _strip_function_xml_calls(text, final = True) + text = _strip_glm_calls(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text @@ -8507,8 +8531,8 @@ class LlamaCppBackend: # "Hello!" won't match. Pattern compiled at module level # (_INTENT_SIGNAL). _reprompt_count = 0 - # Gates ``max_tool_iterations`` on real tool turns so reserved re-prompt slots don't - # extend the budget. Mirrors the safetensors guard. + # Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved + # re-prompt slots don't extend the budget. Mirrors the safetensors guard. _tool_iters_done = 0 _forced_tool_call_pending = False @@ -8525,13 +8549,13 @@ class LlamaCppBackend: if not active_tools: _append_budget_exhausted_nudge = False break - # Gate the markerless bare-JSON form on enabled names so a JSON answer isn't misread as a call. + # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call. _enabled_tool_names = { (tool.get("function") or {}).get("name") for tool in active_tools if (tool.get("function") or {}).get("name") } - # Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows. + # Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows (like safetensors). _tool_xml_signals = _SHARED_TOOL_XML_SIGNALS # Build payload -- stream: True so we detect tool signals @@ -8815,8 +8839,9 @@ class LlamaCppBackend: is_prefix = True break - # Bare Llama-3.2 {"name":..} has no XML signal: hold an - # incomplete object, drain a complete one (mirrors safetensors). + # Signal-less call shapes (mirror the safetensors + # loop): Llama-3.2 bare {"name":..} and Gemma + # call:NAME{...} would otherwise stream raw. _hold_buffer = False # Whole buffer is the call (no visible prefix) -- drain silently. _drain_silently = False @@ -8829,9 +8854,9 @@ class LlamaCppBackend: elif _looks_like_enabled_bare_json( _bare, _enabled_tool_names ): - # Oversized still-open ENABLED-tool call: stop - # holding (memory bound) but DRAIN, not leak; - # a giant ordinary JSON answer still streams. + # Oversized still-open enabled call: drain + # rather than leak; a giant ordinary JSON + # answer still streams. _drain_silently = True elif self._parse_tool_calls_from_text( content_buffer, @@ -8839,6 +8864,17 @@ class LlamaCppBackend: enabled_tool_names = _enabled_tool_names, ): _drain_silently = True + elif ( + "call:".startswith(stripped_buf) + or _GEMMA_BARE_TC_PREFIX_RE.match(stripped_buf) + is not None + or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None + ): + # Whitespace-tolerant like the parser. + if _GEMMA_BARE_TC_RE.match(stripped_buf): + _drain_silently = True + elif len(stripped_buf) < _MAX_BUFFER_CHARS: + _hold_buffer = True if _drain_silently: # No visible prefix -- the buffered text IS @@ -8890,9 +8926,10 @@ class LlamaCppBackend: # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: stripped_buf = content_buffer.lstrip() - # A held bare-JSON fragment has no XML signal; route it to DRAINING. + # A held bare-JSON fragment has no XML signal; route it to DRAINING (the signal-only + # gate below would flush the raw JSON to the user). _bare_eos = strip_llama3_leading_sentinels(stripped_buf) - # Gate on enabled names so a JSON answer isn't routed to DRAINING and dropped. + # Gate on enabled names so an ordinary JSON answer isn't routed to DRAINING and dropped. _is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json( _bare_eos, _enabled_tool_names ) @@ -8925,8 +8962,8 @@ class LlamaCppBackend: "text": cumulative_display, } else: - # No tool signal and no enabled bare-JSON call: a leading ``{`` is an ordinary - # JSON answer and must be shown; any other partial-markup prefix is dropped. + # Held buffer was no tool signal and no enabled bare-JSON call: a leading ``{`` is an + # ordinary JSON answer and must be shown; any other partial-markup prefix is dropped. _held = strip_llama3_leading_sentinels(content_buffer.lstrip()) if _held.startswith("{") and not _suppress_visible_output: yield {"type": "content", "text": _held} @@ -8934,10 +8971,12 @@ class LlamaCppBackend: # ── STREAMING path: no tool call ── if detect_state == _S_STREAMING: - # Safety net: re-parse the full content for tool calls. The route layer resets - # prev_text on tool_start, so post-tool synthesis streams correctly even if - # content was emitted before the tool XML. Unconditional (not gated on - # _tool_xml_signals): bare-JSON and Gemma wrapper-less calls carry no signal. + # Safety net: re-parse the full content for tool calls. The + # route layer resets prev_text on tool_start, so post-tool + # synthesis streams correctly even if content was emitted + # before the tool XML. + # Unconditional (not gated on _tool_xml_signals): bare-JSON and Gemma wrapper-less + # calls carry no XML signal, so a signal gate would let them slip past. _safety_tc = self._parse_tool_calls_from_text( content_accum, allow_incomplete = auto_heal_tool_calls, @@ -9060,8 +9099,8 @@ class LlamaCppBackend: if (tool_calls_acc[i].get("function", {}).get("name", "").strip()) ] or None if not tool_calls: - # Unconditional re-parse: DRAINING means the buffer looked like a call, and - # bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on. + # Unconditional re-parse: we only reach DRAINING when the buffer looked like a + # call, and bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on. tool_calls = self._parse_tool_calls_from_text( content_accum, allow_incomplete = auto_heal_tool_calls, @@ -9073,8 +9112,8 @@ class LlamaCppBackend: final = True, force = True, ) - # ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call - # so the executed call isn't replayed as text or next-turn history. + # ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call so the + # executed call isn't replayed as text or next-turn history. content_text = strip_leading_bare_json_call( content_text, _enabled_tool_names ) @@ -9091,8 +9130,8 @@ class LlamaCppBackend: if content_accum: # Strip leaked tool-call XML before yielding. content_accum = _strip_tool_markup(content_accum, final = True) - # A truncated bare-JSON call has no XML to strip and didn't parse. With - # Auto-Heal on drop a leading ENABLED-tool fragment (plain JSON untouched); + # A truncated bare-JSON call has no XML markup to strip and didn't parse. With + # Auto-Heal on, drop a leading ENABLED-tool fragment (ordinary JSON answers untouched); # off keeps it visible per the strict contract. if content_accum and active_tools and auto_heal_tool_calls: content_accum = strip_leading_bare_json_call( @@ -9115,6 +9154,29 @@ class LlamaCppBackend: _accumulated_predicted_ms += _it.get("predicted_ms", 0) _accumulated_predicted_n += _it.get("predicted_n", 0) + # Collapse exact-duplicate calls and cap the count for the TEXTUAL + # fallback (mirrors the safetensors loop; see _MAX_TOOL_CALLS_PER_TURN). + if tool_calls and not has_structured_tc and len(tool_calls) > 1: + _seen_keys: set = set() + _deduped: list = [] + for _tc in tool_calls: + _fn = _tc.get("function", {}) or {} + _key = (_fn.get("name", ""), str(_fn.get("arguments", ""))) + if _key in _seen_keys: + continue + _seen_keys.add(_key) + _deduped.append(_tc) + if len(_deduped) >= _MAX_TOOL_CALLS_PER_TURN: + break + if len(_deduped) != len(tool_calls): + logger.info( + "GGUF textual fallback: collapsed %d repeated tool call(s) " + "in one turn to %d", + len(tool_calls), + len(_deduped), + ) + tool_calls = _deduped + # disable_parallel_tool_use: execute only the first tool call # this turn. Truncate before building assistant_msg so the # conversation stays consistent and extra calls are never executed. @@ -9265,8 +9327,8 @@ class LlamaCppBackend: if tool_controller.force_final_answer or not tool_controller.active_tools(): _append_budget_exhausted_nudge = False break - # Count only real tool turns against the cap so reserved re-prompt slots can't - # become extra tool rounds; a no-op turn doesn't consume budget (GGUF parity). + # Count only real tool turns against the cap so reserved re-prompt slots can't become + # extra tool rounds; a no-op correction turn doesn't consume budget (GGUF parity). if _turn_executed_real_tool: _tool_iters_done += 1 if _tool_iters_done >= max_tool_iterations: diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 5c7799152f..45f46fef2f 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -104,6 +104,9 @@ class MLXInferenceBackend: ) -> bool: import mlx.core as mx + # Keep the token so the native-template fallback can fetch a + # gated model's repo template later during generation. + self._hf_token = hf_token model_name = config.identifier if hasattr(config, "identifier") else str(config) is_vision = getattr(config, "is_vision", False) @@ -168,11 +171,20 @@ class MLXInferenceBackend: self.active_model_name = model_name self.models[model_name] = { + # Per-model token for the native-template fallback (matches transformers). + "hf_token": hf_token, + # Per-model consent for the native-template reload: re-use the exact + # trust_remote_code this model was loaded with (matches transformers). + "trust_remote_code": trust_remote_code, "model": self._model, "tokenizer": self._tokenizer, "processor": self._processor, "is_vision": is_vision, "is_lora": getattr(config, "is_lora", False), + # For a LoRA adapter the native chat template lives on the base model. + "base_model": getattr(config, "base_model", None) + if getattr(config, "is_lora", False) + else None, "is_audio": False, "audio_type": None, "has_audio_input": False, @@ -355,6 +367,7 @@ class MLXInferenceBackend: from core.inference.chat_template_helpers import ( apply_chat_template_for_generation, + render_with_native_template_fallback, ) prompt = apply_chat_template_for_generation( @@ -368,6 +381,25 @@ class MLXInferenceBackend: if prompt is None: raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible") + # Same parity fix as the transformers backend: if the template dropped the + # requested tools, fall back to the native template so MLX text models keep + # advertising them. ``self._tokenizer`` is this entry's model_info tokenizer, + # so probe and native render share a renderer. (The VLM path renders via the + # processor for image tokens and is intentionally not wired here.) + model_info = self.models.get(self.active_model_name, {}) + prompt = render_with_native_template_fallback( + formatted_prompt = prompt, + tokenizer = self._tokenizer, + model_info = model_info, + active_model_name = self.active_model_name, + messages = messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + hf_token = model_info.get("hf_token"), + ) + sampler = make_sampler( temp = temperature, top_p = top_p, diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index 35855cc34d..fe1aca0e4a 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -32,9 +32,12 @@ from typing import Any, Optional from core.inference.tool_loop_controller import coerce_tool_arguments from core.tool_healing import parse_tool_calls_from_text -# Only the formats this healer can promote. The parser's broader list adds Llama -# <|python_tag|> / Mistral [TOOL_CALLS], but buffering those here would flush a -# streamed call as prose, so keep a healer-aligned list. +# Signals limited to the formats parse_tool_calls_from_text (core.tool_healing) +# actually promotes. The parser module's broader signal list also covers Llama +# <|python_tag|> and Mistral [TOOL_CALLS] for the streaming DRAIN buffers whose +# full parser handles them; buffering those here would hold a streamed +# client-tool call until finalization and then flush it as prose (this healer +# cannot promote them), so the passthrough keeps its own aligned list. _HEAL_SIGNALS = ( "", "<|tool_call>", diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index b67c6cf7e7..8e86d09754 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -21,9 +21,13 @@ from typing import Callable, Generator, Optional from loggers import get_logger from core.inference.tool_call_parser import ( + _GEMMA_BARE_TC_PREFIX_RE, + _GEMMA_BARE_TC_RE, _TOOL_ALL_PATS, _balanced_brace_end, _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, _strip_mistral_closed_calls, _strip_mistral_reasoning, BUDGET_EXHAUSTED_NUDGE, @@ -59,8 +63,8 @@ _MAX_BUFFER_CHARS = 32 # Memory bound for holding a leading bare-JSON object whose top-level "{" never balances. _MAX_BARE_JSON_BUFFER = 16384 -# Forward-looking intent ("I'll", "First,", "Step 1:") = planning; nudge a call. Negative -# lookahead drops negated forms ("I will not"). Mirrors GGUF. +# Forward-looking intent ("I'll", "First,", "Step 1:") = planning, not answering; nudge a call. +# Negative lookahead drops negated forms ("I will not") so a refusal doesn't trigger it. Mirrors GGUF. _INTENT_SIGNAL = re.compile( r"(?i)(" r"\b(i['’](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" @@ -70,11 +74,15 @@ _INTENT_SIGNAL = re.compile( ) _MAX_REPROMPTS = 3 _REPROMPT_MAX_CHARS = 2000 -# Templated so the nudge names the caller's enabled tools. Mirrors GGUF tool_hint. +# Templated so the nudge names the caller's enabled tools, not a hardcoded set. Mirrors GGUF tool_hint. _REPROMPT_INSTRUCTION_TEMPLATE = ( "STOP. Do NOT write code or explain. You MUST call a tool NOW. Call {tool_hint} immediately." ) +# No grammar constraint here (unlike llama-server's lazy grammar): collapse +# exact-duplicate calls and cap the count so a runaway turn cannot fan out. +_MAX_TOOL_CALLS_PER_TURN = 8 + def _active_tool_names(active_tools: list[dict]) -> list[str]: names = [ @@ -90,16 +98,25 @@ def strip_tool_markup_streaming( *, auto_heal_tool_calls: bool = True, tool_protocol_active: bool = False, + enabled_tool_names: Optional[set] = None, ) -> str: - """Strip open-ended tool XML from display text without trimming whitespace.""" + """Strip open-ended tool XML from display text without trimming whitespace. + ``enabled_tool_names`` gates the markerless Gemma ``call:NAME{...}`` strip so a + disabled/example name in prose is kept (mirrors the parser gate).""" if not (auto_heal_tool_calls or tool_protocol_active): return text - # Mirror the final strip (no final trim): drop a leading Magistral ``[THINK]...[/THINK]`` - # block, then Mistral calls, then a parser-accurate function-XML scan before the regex - # arms. An unclosed ``[THINK]`` holds until ``[/THINK]`` so text stays monotonic. + # Mirror the final strip's scan order so streaming and final display agree: + # balanced strips first (nested JSON removed whole), then the guarded + # function-XML/GLM scans that close at each call's REAL terminator, so literal + # markup inside argument values is data and trailing prose survives. No final + # trim so streaming length comparisons hold. Leading Magistral [THINK]...[/THINK] + # is dropped (bracket form, not the reasoning channel's ); an unclosed + # [THINK] holds until [/THINK] so the cleaned text stays monotonic. text = _strip_mistral_reasoning(text) text = _strip_mistral_closed_calls(text) + text = _strip_gemma_wrapperless_calls(text, enabled_tool_names) text = _strip_function_xml_calls(text, final = True) + text = _strip_glm_calls(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text @@ -110,10 +127,11 @@ def _strip_tool_markup_final( *, auto_heal_tool_calls: bool, tool_protocol_active: bool = False, + enabled_tool_names: Optional[set] = None, ) -> str: if not (auto_heal_tool_calls or tool_protocol_active): return text - return strip_tool_markup(text, final = True) + return strip_tool_markup(text, final = True, enabled_tool_names = enabled_tool_names) def _status_for_tool(tool_name: str, arguments: dict) -> str: @@ -247,8 +265,9 @@ def run_safetensors_tool_loop( final_attempt_done = False next_call_id = 0 reprompt_count = 0 - # Only turns that executed a tool count against ``max_tool_iterations``; a no-op or - # re-prompt turn must not consume budget (GGUF parity). + # Real tool-call turns completed. Only turns that actually executed a tool count + # against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a + # plan-without-action re-prompt) must not consume budget, matching the GGUF loop. _executed_tool_iters = 0 def _tool_succeeded(tool_name: str) -> bool: @@ -285,7 +304,7 @@ def run_safetensors_tool_loop( tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools)) tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else () - # Gate the markerless bare-JSON form on enabled names so a JSON answer isn't misread as a call. + # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call. _enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools)) detect_state = _state_buffering @@ -373,6 +392,7 @@ def run_safetensors_tool_loop( before_tool, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_tool_names, ) if len(cleaned_before) > len(last_emitted): last_emitted = cleaned_before @@ -403,6 +423,7 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_tool_names, ) if len(cleaned) > len(last_emitted): last_emitted = cleaned @@ -425,8 +446,9 @@ def run_safetensors_tool_loop( is_prefix = True break - # Bare Llama-3.2 ``{"name":..,"parameters":..}`` carries no XML signal. Hold a leading - # ``{`` (after any sentinel) until it closes: drain if it parses as a call, else stream. + # Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML + # signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses + # as a call, else stream as content. Non-call text is always recovered downstream. bare_probe = strip_llama3_leading_sentinels(stripped) if ( not is_match @@ -439,7 +461,7 @@ def run_safetensors_tool_loop( continue # object still open -- keep buffering elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names): # Oversized still-open ENABLED-tool call: stop holding (memory bound) but - # DRAIN, not leak; a giant ordinary JSON answer still streams. + # DRAIN instead of leaking the raw prefix; a giant ordinary JSON answer still streams. detect_state = _state_draining continue elif parse_tool_calls_from_text( @@ -453,6 +475,35 @@ def run_safetensors_tool_loop( continue # Closed non-call object (or oversized non-call) -- stream as text. + # Gemma wrapper-less ``call:NAME{...}`` has no tool_xml_signals entry: + # buffer it here or it streams raw until the end-of-turn safety net. + # ``(? len(last_emitted): last_emitted = cleaned @@ -493,6 +545,7 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_tool_names, ) if len(cleaned) > len(last_emitted): last_emitted = cleaned @@ -515,23 +568,25 @@ def run_safetensors_tool_loop( elif tool_protocol_active and _looks_like_enabled_bare_json( _bare_eos, _enabled_tool_names ): - # Held ENABLED-tool bare-JSON fragment has no XML signal; DRAIN it (a JSON answer - # falls through to the else and streams, GGUF parity). + # A held bare-JSON ENABLED-tool fragment has no XML signal; DRAIN it (an ordinary + # JSON answer falls through to the else and streams as content, GGUF parity). detect_state = _state_draining else: # Drain and fall through to STREAMING so the intent re-prompt + safety-net parser # still fire on short emissions like "Let me search." that never exit BUFFERING. if content_buffer: cumulative_display += content_buffer - cleaned = strip_tool_markup(cumulative_display, final = True) + cleaned = strip_tool_markup( + cumulative_display, final = True, enabled_tool_names = _enabled_tool_names + ) if len(cleaned) > len(last_emitted): last_emitted = cleaned yield {"type": "content", "text": cleaned} detect_state = _state_streaming if detect_state == _state_streaming: - # Run the parser even with no XML signal (bare-JSON carries none); it's strict so - # plain answers stay untouched. Mirrors GGUF. + # Run the parser even with no XML signal (the Llama-3.2 bare-JSON form carries none); it's + # strict so plain answers stay untouched. Mirrors GGUF. safety_tc = parse_tool_calls_from_text( content_accum, id_offset = next_call_id, @@ -539,8 +594,8 @@ def run_safetensors_tool_loop( enabled_tool_names = _enabled_tool_names, ) if not safety_tc: - # Re-prompt only when the model planned without acting (intent signal); - # "4" / "Hello!" never trigger. Mirrors GGUF. + # Re-prompt only when the model planned without acting (intent + # signal); "4" / "Hello!" never trigger. Mirrors GGUF. _stripped = content_accum.strip() if ( tools @@ -569,9 +624,9 @@ def run_safetensors_tool_loop( yield {"type": "status", "text": ""} continue - # Final answer. If a literal tool marker in prose was buffered but never - # parsed as a call, restore the raw text so the prose surfaces; route - # cleanup still applies the Auto-Heal policy. + # Final answer. If a literal tool marker in prose was buffered but + # never parsed as a call, restore the raw text so the prose surfaces + # in full; route-level cleanup still applies the Auto-Heal policy. if content_accum and any(sig in content_accum for sig in tool_xml_signals): yield {"type": "content", "text": content_accum} yield {"type": "status", "text": ""} @@ -581,6 +636,7 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, + enabled_tool_names = _enabled_tool_names, ) logger.info( "Safetensors safety net: parsed %d tool call(s) from streamed content", @@ -603,9 +659,10 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = False, + enabled_tool_names = _enabled_tool_names, ) - # Drained bare-JSON call that didn't parse: with Auto-Heal on drop the fragment - # (plain JSON untouched); off keeps it visible per the strict contract. + # Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment + # (plain JSON answers are left untouched); off keeps it visible per the strict contract. if tool_protocol_active and auto_heal_tool_calls: _drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names) if _drain_text: @@ -625,12 +682,13 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, + enabled_tool_names = _enabled_tool_names, ) if tool_calls: next_call_id += len(tool_calls) - # Strip a leading bare-JSON call so it isn't replayed as text or next-turn history - # (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers. + # Strip a leading bare-JSON call from the kept content so it isn't replayed as text or + # next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers. content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names) if final_attempt_done: @@ -640,6 +698,27 @@ def run_safetensors_tool_loop( yield {"type": "status", "text": ""} return + # Collapse exact-duplicate calls and cap the count (runaway-turn guard). + if tool_calls: + seen_keys: set = set() + deduped: list = [] + for _tc in tool_calls: + _fn = _tc.get("function", {}) or {} + _key = (_fn.get("name", ""), str(_fn.get("arguments", ""))) + if _key in seen_keys: + continue + seen_keys.add(_key) + deduped.append(_tc) + if len(deduped) >= _MAX_TOOL_CALLS_PER_TURN: + break + if len(deduped) != len(tool_calls): + logger.info( + "Safetensors: collapsed %d repeated tool call(s) in one turn to %d", + len(tool_calls), + len(deduped), + ) + tool_calls = deduped + assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False @@ -771,7 +850,8 @@ def run_safetensors_tool_loop( if not unrestricted_tools and not tool_controller.active_tools(): final_attempt_done = True continue - # Count only real tool turns against the cap so a no-op turn doesn't consume budget (GGUF parity). + # Count only turns that executed a tool against the cap; a no-op correction turn doesn't + # consume budget so the model gets its nudge and another tool-enabled turn (GGUF parity). if _turn_executed_real_tool: _executed_tool_iters += 1 if _executed_tool_iters >= max_tool_iterations and not final_attempt_done: diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 9e82e40de2..08a6bf418a 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -14,18 +14,22 @@ safetensors + MLX agentic loop sees the same call shape llama-server gives GGUF: - ``[TOOL_CALLS]name{json}`` (Mistral v11+ / Magistral) - ``[TOOL_CALLS]name[ARGS]{json}`` (Ministral / Mistral Large 3) - ``<|tool_call>call:NAME{k:<|"|>v<|"|>}`` (Gemma 4) + - ``<｜tool▁calls▁begin｜>...function<｜tool▁sep｜>NAME\\n``\\`\\`\\`json\\n{...}\\n\\`\\`\\`...`` (DeepSeek R1) + - ``<｜tool▁calls▁begin｜>...<｜tool▁call▁begin｜>NAME<｜tool▁sep｜>{json}<｜tool▁call▁end｜>...`` (DeepSeek V3 / V3.1) + - ``NAME\\nk\\nv...`` (GLM 4.5 / 4.6 / 4.7) + - ``<|tool_calls_section_begin|>...<|tool_call_begin|>functions.NAME:IDX<|tool_call_argument_begin|>{json}<|tool_call_end|>...`` (Kimi K2) Missing closing tags / brackets are tolerated: models often truncate mid-stream. """ -# Keeps PEP 604 `X | None` lazy for python 3.9 (imported standalone by external servers). +# Lazy annotations keep the standalone python 3.9 import working. from __future__ import annotations import json import re from typing import Any, Optional -# Shared parser handles Qwen/Hermes, Qwen3.5 XML, Gemma 4; this module adds Llama-3, Mistral, bare JSON. +# Qwen/Hermes, Qwen3.5 XML and Gemma 4 live in core.tool_healing; this module adds the rest. from core import tool_healing as _tool_healing @@ -37,14 +41,31 @@ TOOL_XML_SIGNALS = ( "<|python_tag|>", "[TOOL_CALLS]", "<|tool_call>", + # DeepSeek R1 / V3 / V3.1 -- 5 opener variants llama.cpp keeps. + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>", + "<｜tool_calls_begin｜>", + "<｜tool▁calls｜>", + "<｜tool calls begin｜>", + "<｜tool\\_calls\\_begin｜>", + # Kimi K2 / Moonshot. + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>", ) -# Closed pairs only (mid-stream); _TOOL_ALL_PATS eats unclosed tails at end-of-turn. +# DeepSeek opener variants; shared by parse and strip so a parsed signal is always stripped. +_DEEPSEEK_OPEN_ALT = ( + r"tool▁calls▁begin|tool_calls_begin|tool calls begin|tool\\_calls\\_begin|tool▁calls" +) +_DEEPSEEK_OPEN_RE_SRC = r"<｜(?:" + _DEEPSEEK_OPEN_ALT + r")｜>" + +# Closed pairs only (mid-stream); _TOOL_ALL_PATS also eats unclosed tails at +# end-of-turn. ``[\w-]+`` on ```` tracks OpenAI's +# ``^[a-zA-Z0-9_-]{1,64}$`` so hyphenated MCP names parse like built-ins. _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), - # Match to the real ```` (lookahead, not greedy ``.*``) so a literal - # ```` in a value doesn't truncate and each call stays separate. + # Span to the real ```` so a literal one inside a value can't truncate the strip. re.compile( r'' r'(?:(?!).)*' @@ -52,12 +73,21 @@ _TOOL_CLOSED_PATS = [ re.DOTALL, ), re.compile(r"<\|tool_call>.*?", re.DOTALL), + re.compile(r"\[TOOL_CALLS\]\s*\[.*?\](?:\s*)?", re.DOTALL), + # Mistral v11+ ``[TOOL_CALLS]name{json}`` (may chain), close at ``}``. + re.compile(r"\[TOOL_CALLS\]\s*[\w\.\-]+\s*(?:\[ARGS\])?\s*\{.*?\}", re.DOTALL), + # DeepSeek R1 / V3 / V3.1: full envelope (any opener variant) ... end. + re.compile(_DEEPSEEK_OPEN_RE_SRC + r".*?<｜tool▁calls▁end｜>", re.DOTALL), + # Kimi K2: ``<|tool_calls_section_begin|>...<|tool_calls_section_end|>``. + re.compile(r"<\|tool_calls_section_begin\|>.*?<\|tool_calls_section_end\|>", re.DOTALL), + # Kimi K2 section-less closed call; else the catch-all below eats trailing prose to EOS. + re.compile(r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>", re.DOTALL), ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ re.compile(r".*$", re.DOTALL), re.compile(r'.*$', re.DOTALL), - # Bare-word markers drop a trailing truncated call only when the next chars look like - # a call start, so prose mentioning the marker is kept; a marker at end-of-text drops. + # Bare-word markers drop a trailing truncated call only when a call-shaped start + # follows; a prose mention (``See [TOOL_CALLS] docs...``) keeps its tail. Bare marker at EOF drops. re.compile(r"<\|tool_call>(?=\s*call\s*:|\s*$).*$", re.DOTALL), re.compile( r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*[\[{])|\s*$).*$", @@ -67,6 +97,22 @@ _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ r"<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\()|\s*$).*$", re.DOTALL, ), + # DeepSeek envelopes truncated mid-stream (any opener); same call-shaped lookahead as above. + re.compile( + _DEEPSEEK_OPEN_RE_SRC + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*$", + re.DOTALL, + ), + re.compile(r"<｜tool▁call▁begin｜>(?=\s*function|\s*$).*$", re.DOTALL), + # Kimi K2 envelope truncated. + re.compile( + r"<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*$", + re.DOTALL, + ), + re.compile( + r"<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*$", + re.DOTALL, + ), + # Gemma wrapper-less ``call:NAME{...}`` is handled by ``_strip_gemma_wrapperless_calls`` (enabled-name gate). ] @@ -105,8 +151,7 @@ BUDGET_EXHAUSTED_NUDGE = ( "any more tools." ) -# The exact-args dup guard misses paraphrased re-searches, so also cap executed -# KB searches per turn, then nudge. +# The exact-args dup guard misses paraphrased re-searches, so also cap KB searches per turn. RAG_MAX_SEARCHES_PER_TURN = 3 RAG_SEARCH_CAP_NUDGE = ( "You have already searched the knowledge base several times this turn. " @@ -117,14 +162,16 @@ RAG_SEARCH_CAP_NUDGE = ( # Qwen / Hermes ``{json}``. _TC_JSON_START_RE = re.compile(r"\s*\{") -# Qwen3.5 ```` plus attribute form ```` (MiniCPM-5, -# MiniMax-M2); name in group(1) or group(2). +# Qwen3.5 ```` and the attribute form ```` +# (MiniCPM-5, MiniMax-M2); name class ``[\w.\-]+`` lands in group(1) or group(2). _TC_FUNC_START_RE = re.compile(r'\s*') -# Body ends at ```` or ```` so trailing prose stays out of args. +# Body ends at ```` (Hermes) or ```` (Qwen3.5 / MiniCPM-5) +# so it stops at the close even when prose follows (else prose leaked into args). _TC_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -# Horizontal whitespace only so the wrapping newline + indent survive (``_trim_param_value`` -# trims one newline), preserving code indent. +# Horizontal whitespace only (``[^\S\n]*``, not ``\s*``) so the wrapping newline + +# first-line indentation survive; ``_trim_param_value`` trims one newline, preserving +# code indentation (SGLang qwen3_coder). _TC_PARAM_START_RE = re.compile( r'<(?:parameter|param)(?:=([\w\.\-]+)|\s+name="([\w\.\-]+)")>[^\S\n]*' ) @@ -135,35 +182,81 @@ _LLAMA3_PYTHON_TAG = "<|python_tag|>" _LLAMA3_PY_CALL_RE = re.compile( r"<\|python_tag\|>\s*([\w\.\-]+)\s*\.\s*call\s*\(", ) -# Anchored at the char after ``<|python_tag|>`` plus the ``; NAME.call(`` chain sep, so -# a ``.call(`` inside JSON args is ignored. +# Anchored at a fixed offset (char after ``<|python_tag|>``) plus the ``; NAME.call(`` +# chain separator; fixed-offset (not a free scan) ignores ``.call(`` inside JSON args. _LLAMA3_PY_CALL_HEAD_RE = re.compile(r"\s*([\w\.\-]+)\s*\.\s*call\s*\(") _LLAMA3_CALL_CHAIN_RE = re.compile(r"\s*;\s*([\w\.\-]+)\s*\.\s*call\s*\(") -# ``.call(k=v)`` kwarg tokens, hand-scanned below (not finditer) to stay linear on a -# truncated body (ReDoS). +# Llama-3 ``.call(k=v)`` kwarg tokens, hand-scanned below (not finditer) to stay +# linear on a truncated body; finditer retries every offset of a long run (ReDoS). _LLAMA3_KEY_RE = re.compile(r"\w+") _LLAMA3_WS_RE = re.compile(r"\s*") -# ints, decimals, sci notation; trailing ``(?![\w.])`` stops ``1.2.3`` truncating to ``1.2``. +# ints, decimals (1.5, 1., .5) and sci notation; trailing ``(?![\w.])`` stops a token +# like ``1.2.3`` being truncated to ``1.2`` (which would mis-parse the remainder). _LLAMA3_NUM_RE = re.compile(r"-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?(?![\w.])") _LLAMA3_LIT_RE = re.compile(r"true|false|null") -# Mistral ``[TOOL_CALLS]`` trigger. v11+ chains ``name{json}`` (Magistral) or -# ``name[ARGS]{json}`` (Ministral / Large 3). +# Mistral ``[TOOL_CALLS]`` trigger. v11+ chains them, each followed by a bare name +# plus ``{json}`` (Magistral) or ``[ARGS]{json}`` (Ministral / Large 3). _MISTRAL_TRIGGER = "[TOOL_CALLS]" _MISTRAL_ARGS_MARKER = "[ARGS]" -# Mistral Small 3.2 emits ``name[CALL_ID][ARGS]{json}`` (absent on Ministral / Magistral). +# Mistral Small 3.2 emits ``name[CALL_ID][ARGS]{json}`` (absent on Ministral / +# Magistral); llama.cpp distinguishes the two on ``[CALL_ID]`` (common/chat.cpp). _MISTRAL_CALL_ID_MARKER = "[CALL_ID]" -# Magistral wraps reasoning in ``[THINK]...[/THINK]``; a ``[TOOL_CALLS]`` inside is not a real call. +# Magistral wraps reasoning in ``[THINK]...[/THINK]``; a ``[TOOL_CALLS]`` inside +# that block is chain-of-thought, not a real call. _MISTRAL_THINK_OPEN = "[THINK]" _MISTRAL_THINK_CLOSE = "[/THINK]" _MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") +# DeepSeek markers (full-width pipe U+FF5C, block U+2581); five outer-open variants like llama.cpp. +_DEEPSEEK_BEGIN_RE = re.compile(_DEEPSEEK_OPEN_RE_SRC) +_DEEPSEEK_END = "<｜tool▁calls▁end｜>" +_DEEPSEEK_CALL_BEGIN = "<｜tool▁call▁begin｜>" +_DEEPSEEK_SEP = "<｜tool▁sep｜>" +_DEEPSEEK_CALL_END = "<｜tool▁call▁end｜>" +# R1 wraps args in a ```json fence with a ``function`` prefix; V3/V3.1 do not. +# Scanned with ``str.find`` -- the regex forms are O(N^2) on truncated bodies. +_DEEPSEEK_R1_FUNC_MARKER = "function" + _DEEPSEEK_SEP +_DEEPSEEK_R1_FENCE = "\n```json\n" +_DEEPSEEK_R1_CLOSE_RE = re.compile(r"```[\s\r\n]*" + re.escape(_DEEPSEEK_CALL_END)) + +# GLM 4.5-4.7: ``NAME[\n]K...``; the lookahead also allows a +# direct ````/```` (4.7 drops the newline, zero-arg calls close at once). +# Name class ``[\w.\-]+`` keeps prose like ``not a call`` unparsed; +# ``{`` stays with the Qwen JSON parser. +_GLM_TC_OPEN_RE = re.compile(r"\s*([\w.\-]+)\s*(?=\n||)") +_GLM_TC_CLOSE = "" +_GLM_ARG_KEY_OPEN = "" +_GLM_ARG_KEY_CLOSE = "" +_GLM_ARG_VAL_OPEN = "" +_GLM_ARG_VAL_CLOSE = "" +# Strings arrive raw, non-strings via tojson; only unambiguous JSON literals decode +# (bare ``42``/``true``/``null`` stay strings). +_GLM_JSON_NUMERIC_RE = re.compile(r"-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?") + +# Kimi K2 / Moonshot (ASCII pipes). Id ``functions.NAME:IDX`` -- strip ``functions.``/``:N`` for the name. +_KIMI_SECTION_BEGIN = "<|tool_calls_section_begin|>" +_KIMI_SECTION_END = "<|tool_calls_section_end|>" +_KIMI_CALL_BEGIN = "<|tool_call_begin|>" +_KIMI_ARG_BEGIN = "<|tool_call_argument_begin|>" +_KIMI_CALL_END = "<|tool_call_end|>" +_KIMI_ID_RE = re.compile(r"^(?:functions\.)?([\w\.\-]+)(?::(\d+))?$") + # Gemma 4: ``<|tool_call>call:NAME{...}``, ``<|"|>`` wraps strings. _GEMMA_TC_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w\.\-]+)\s*\{") _GEMMA_STR_BEGIN = '<|"|>' _GEMMA_STR_END = '<|"|>' _GEMMA_TC_END = "" +# skip_special_tokens strips the wrapper and ``<|"|>`` markers, so streamed Gemma calls +# arrive as bare ``call:NAME{k:v, ...}``; ``(? int | None: """Index of the ``]`` matching ``[`` at ``text[start]`` (ignores brackets in JSON strings).""" @@ -215,7 +308,8 @@ def _skip_mistral_call_id(text: str, pos: int) -> int: def _strip_mistral_reasoning(content: str) -> str: - """Drop a leading Magistral ``[THINK]`` block so rehearsed calls inside reasoning are not promoted; unclosed drops to EOF.""" + """Drop a leading Magistral ``[THINK]...[/THINK]`` so a ``[TOOL_CALLS]`` inside + reasoning is not taken as a real call; an unclosed ``[THINK]`` drops from it on.""" i = 0 n = len(content) while i < n and content[i] in " \t\n\r": @@ -229,7 +323,10 @@ def _strip_mistral_reasoning(content: str) -> str: def _strip_mistral_closed_calls(text: str) -> str: - """Strip cleanly-closed ``[TOOL_CALLS]`` blocks via balanced scanning (a non-greedy regex would truncate nested JSON); unclosed runs wait for ``final=True``.""" + """Strip cleanly-closed ``[TOOL_CALLS]`` blocks (array, ``name{json}``, + ``name[ARGS]{json}``) via balanced scanning -- a non-greedy ``\\{.*?\\}`` would + truncate at the first ``}`` and lose nested JSON. Unclosed runs are left for + ``final=True`` cleanup.""" n = len(text) out = [] cursor = 0 @@ -254,7 +351,8 @@ def _strip_mistral_closed_calls(text: str) -> str: if text.startswith("", cursor): cursor += len("") continue - # Single-object shape ``[TOOL_CALLS] { json }``: the parser accepts it, so strip it too. + # Single-object shape ``[TOOL_CALLS] { json }`` (no name/array): the parser + # accepts it, so the display strip must remove it too (else it leaks). if i < n and text[i] == "{": end = _balanced_brace_end(text, i) if end is None: @@ -287,12 +385,50 @@ def _strip_mistral_closed_calls(text: str) -> str: out.append(text[idx:]) break cursor = end + 1 - # Consume the optional EOS marker so ``...{json}`` doesn't leave ```` as content. + # Consume the optional EOS marker too, mirroring the array shape, so a + # ``[TOOL_CALLS]name{json}`` tail doesn't leave ```` as content. if text.startswith("", cursor): cursor += len("") return "".join(out) +def _strip_gemma_wrapperless_calls(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Strip closed wrapper-less Gemma ``call:NAME{...}`` calls with balanced brace + scanning (nested arguments are removed whole). ``enabled_tool_names`` gates the + strip like the parser gate: a disabled/example name stays visible; ``None`` + strips every closed call.""" + if _whole_content_is_json_value(text): + return text + n = len(text) + out = [] + # Mirror the parse scan: a leading JSON answer's span is data, kept visible. + cursor = _leading_json_value_end(text) or 0 + if cursor: + out.append(text[:cursor]) + while cursor < n: + m = _GEMMA_BARE_TC_RE.search(text, cursor) + if not m: + out.append(text[cursor:]) + break + disabled = enabled_tool_names is not None and m.group(1) not in enabled_tool_names + brace = m.end() - 1 # _GEMMA_BARE_TC_RE consumes through the opening ``{`` + # Same boundary scanner as the parser: strip exactly what it consumed. + end = _gemma_body_brace_end(text, brace) + closed = end is not None + next_index = (end + 1) if closed else len(text) + if not closed: + # Unclosed call: drop an enabled call to EOS; keep a disabled/example name as prose. + out.append(text[cursor:] if disabled else text[cursor : m.start()]) + break + if disabled: + # Disabled/example name is prose: keep it whole. + out.append(text[cursor:next_index]) + else: + out.append(text[cursor : m.start()]) + cursor = next_index # already past the matching ``}`` + return "".join(out) + + _FUNC_CLOSE_TAG_RE = re.compile(r"") @@ -327,16 +463,131 @@ def _strip_function_xml_calls(text: str, *, final: bool) -> str: return "".join(out) -def strip_tool_markup(text: str, *, final: bool = False) -> str: - """Strip tool-call markup; ``final=True`` also drops trailing unclosed runs and trims.""" +def _glm_value_close( + text: str, + vs: int, + *, + strict: bool = False, +) -> int: + """Index of the ```` that really ends the GLM value at ``vs``: the + first one whose next non-space token is ````, ```` or + end-of-text AND that sits at balanced quote state (an embedded literal pair + like ``print("")`` lives inside a still-open string). + Quote openers are contextual (single quote only after punctuation, so + apostrophes are prose; double quote also at word start), mirroring the Gemma + scanners. If no candidate balances, the first token-valid one wins -- except + in ``strict`` mode (Auto-Heal off), which refuses the in-quote fallback rather + than execute truncated arguments. Returns -1 if unclosed.""" + n = len(text) + search = vs + first_candidate = -1 + quote = "" + prev = ":" + prev_raw = ":" + qpos = vs # quote-state cursor; advanced incrementally to each candidate + while True: + ve = text.find(_GLM_ARG_VAL_CLOSE, search) + if ve < 0: + return -1 if strict else first_candidate + j = ve + len(_GLM_ARG_VAL_CLOSE) + while j < n and text[j] in " \t\r\n": + j += 1 + if j >= n or text.startswith(_GLM_ARG_KEY_OPEN, j) or text.startswith(_GLM_TC_CLOSE, j): + while qpos < ve: + ch = text[qpos] + if quote: + if ch == "\\" and qpos + 1 < ve: + qpos += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + if not ch.isspace(): + prev = ch + prev_raw = ch + qpos += 1 + if not quote: + return ve + if first_candidate < 0: + first_candidate = ve + search = ve + len(_GLM_ARG_VAL_CLOSE) + + +def _strip_glm_calls(text: str, *, final: bool) -> str: + """Strip GLM 4.x calls by scanning to each call's REAL ```` (the one + after the last consumed ````, mirroring ``_parse_glm_tool_calls``), so + a literal ```` inside a value is data. Qwen ``{json}`` has + no NAME token and is left to the regex arms. ``final`` drops a truncated call to + EOS; otherwise it stays buffered.""" + out: list[str] = [] + cursor = 0 + n = len(text) + while True: + m = _GLM_TC_OPEN_RE.search(text, cursor) + if not m: + break + apos = m.end() + close = -1 + while True: + ks = text.find(_GLM_ARG_KEY_OPEN, apos) + tc = text.find(_GLM_TC_CLOSE, apos) + if tc >= 0 and (ks < 0 or tc < ks): + close = tc + break + if ks < 0: + break # no close and no more keys -- truncated body + ke = text.find(_GLM_ARG_KEY_CLOSE, ks + len(_GLM_ARG_KEY_OPEN)) + if ke < 0: + break + vstart = ke + len(_GLM_ARG_KEY_CLOSE) + while vstart < n and text[vstart] in " \t\r\n": + vstart += 1 + if not text.startswith(_GLM_ARG_VAL_OPEN, vstart): + apos = ke + len(_GLM_ARG_KEY_CLOSE) + continue + vs = vstart + len(_GLM_ARG_VAL_OPEN) + ve = _glm_value_close(text, vs) + if ve < 0: + break # unclosed -- truncated + apos = ve + len(_GLM_ARG_VAL_CLOSE) + if close >= 0: + out.append(text[cursor : m.start()]) + cursor = close + len(_GLM_TC_CLOSE) + continue + # Truncated GLM call (no real close yet). + if final: + out.append(text[cursor : m.start()]) + cursor = n + # Non-final: leave the unclosed call (and any tail) buffered as-is. + break + out.append(text[cursor:]) + return "".join(out) + + +def strip_tool_markup( + text: str, + *, + final: bool = False, + enabled_tool_names: Optional[set] = None, +) -> str: + """Strip tool-call markup. ``final=False`` keeps in-progress markup buffered; + ``final=True`` also drops trailing unclosed runs and trims. ``enabled_tool_names`` + gates the markerless Gemma ``call:NAME{...}`` strip so a disabled/example name in + prose is kept (mirrors the parser gate); ``None`` strips every closed call.""" if final: - # End-of-turn only: drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket form, - # not the ```` reasoning channel) so raw reasoning doesn't leak into display/history. + # Drop a leading Magistral ``[THINK]...[/THINK]`` at end-of-turn; its bracket + # form is not the ```` the reasoning channel renders. text = _strip_mistral_reasoning(text) text = _strip_mistral_closed_calls(text) - # Scan-strip the function-XML form first (parser-accurate: a literal ```` in - # a value is data, not a call); the regex arms below cover the other formats. + if final: + text = _strip_gemma_wrapperless_calls(text, enabled_tool_names) + # Scan-strip the function-XML form (a literal ```` inside a value is + # data). The regex arms below cover the other formats but no-op on function calls here. text = _strip_function_xml_calls(text, final = final) + # GLM 4.x: scan to the call's real so a literal one inside a value is data, + # not a leak. Qwen {json} is left to the regex arms. + text = _strip_glm_calls(text, final = final) pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS for pat in pats: text = pat.sub("", text) @@ -347,8 +598,77 @@ def has_tool_signal(text: str) -> bool: return any(s in text for s in TOOL_XML_SIGNALS) +# A Qwen/Hermes ````/```` envelope whose arguments carry literal +# DeepSeek/Kimi markers must parse as the OUTER call. Detect it opening before the first +# marker so the pre-pass skips it. +_EMBEDDED_MARKER_RE = re.compile( + _DEEPSEEK_OPEN_RE_SRC + "|" + re.escape(_KIMI_SECTION_BEGIN) + "|" + re.escape(_KIMI_CALL_BEGIN) +) +# Covers ```` and the attribute form. ``<|python_tag|>`` is Llama-3's +# envelope too (built-in ``NAME.call(`` and custom ``{json}``), so a quoted DeepSeek/Kimi +# example is data; the call-shaped lookahead mirrors the ``_TOOL_ALL_PATS`` python_tag arm +# so a bare prose ``<|python_tag|>`` mention isn't treated as one. +_OUTER_ENVELOPE_OPEN_RE = re.compile( + r'|' + r"|<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\())" +) +# CLOSED outer envelopes, each spanning to its REAL final close so a literal +# ````/```` inside a value is data. Wrapped Gemma counts too. +_OUTER_ENVELOPE_CLOSED_PATS = ( + re.compile(r"(?:(?!).)*", re.DOTALL), + _TOOL_CLOSED_PATS[1], + re.compile(r"<\|tool_call>.*?", re.DOTALL), +) + + +def _marker_inside_leading_envelope(content: str, enabled_tool_names: Optional[set] = None) -> bool: + first_marker = _EMBEDDED_MARKER_RE.search(content) + if first_marker is None: + return False + # A leading bare-JSON or Mistral [TOOL_CALLS] call is an outer envelope too: + # a DS/Kimi marker in its argument strings is data. + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if content.startswith("{", i): + end = _balanced_brace_end(content, i) + if end is not None and i < first_marker.start(): + name = _top_level_bare_json_name(content[i : end + 1]) + if name is not None and (enabled_tool_names is None or name in enabled_tool_names): + # The closed leading call owns the turn: a marker inside it is argument + # data, one after it a trailing example (same rule as the XML envelopes below). + return True + if name is not None and first_marker.start() <= end: + # A disabled-name leading object is prose (can't own the turn), but a marker + # inside its own strings stays data. A marker AFTER it falls through to the pre-pass. + return True + elif content.startswith(_MISTRAL_TRIGGER, i): + end = _mistral_region_end(content, i) + if end is not None and i < first_marker.start(): + return True + # A closed outer call PRECEDING the first marker owns the turn; the pre-pass must + # not steal a trailing example or argument data. + for _pat in _OUTER_ENVELOPE_CLOSED_PATS: + m = _pat.search(content) + if m is not None and m.start() < first_marker.start(): + return True + residue = content + for _pat in _OUTER_ENVELOPE_CLOSED_PATS: + residue = _pat.sub("", residue) + marker = _EMBEDDED_MARKER_RE.search(residue) + if marker is None: + return True + # A marker still stands; any opener left in the residue is UNCLOSED. One before the + # marker is a truncated outer call holding the marker as data: skip the pre-pass. + opener = _OUTER_ENVELOPE_OPEN_RE.search(residue) + return opener is not None and opener.start() < marker.start() + + def _mistral_region_end(text: str, idx: int) -> int | None: - """Exclusive end of the balanced ``[TOOL_CALLS]`` call at ``idx``, or ``None`` when truncated (array, object, and named forms).""" + """Exclusive end of the balanced ``[TOOL_CALLS]`` call starting at ``idx``, + or ``None`` when truncated/unrecognised (same shapes as the strip scan: + array, single-object, and named ``name [CALL_ID]? [ARGS]? {json}``).""" n = len(text) i = idx + len(_MISTRAL_TRIGGER) while i < n and text[i] in " \t\n\r": @@ -384,8 +704,10 @@ def _xml_signal_inside_leading_mistral(content: str) -> bool: first_xml = _first_foreign_tool_signal(content) if first_xml is not None and first_xml < trig: return False - # Only plain prose precedes the trigger (preamble-tolerant); prose merely mentioning - # the marker has no parseable region and keeps the normal order. + # Only plain prose precedes the trigger: a visible preface must not hand + # the turn to a later XML literal (preamble-tolerant, like the + # wrapperless-Gemma guard). Prose that merely mentions the marker has no + # parseable region and keeps the normal order. return _mistral_region_end(content, trig) is not None @@ -393,7 +715,8 @@ _ATTR_FUNC_OPEN_RE = re.compile(r' int | None: - """Offset of the first signal a non-envelope parser would fire on (XML forms plus the Llama-3 ``<|python_tag|>`` marker).""" + """Offset of the first tool signal a non-envelope parser would fire on + (XML forms plus ``<|python_tag|>``, which also runs before the Mistral parser).""" first = None for sig in ("", "<|tool_call>", ""): p = content.find(sig) @@ -402,37 +725,126 @@ def _first_foreign_tool_signal(content: str) -> int | None: attr = _ATTR_FUNC_OPEN_RE.search(content) if attr is not None and (first is None or attr.start() < first): first = attr.start() + # DeepSeek/Kimi markers are foreign to a JSON envelope too: a marker inside a leading + # object routes through the same guard (and, if disabled, the drop-and-parse-the-tail + # recursion, so a real call after the object is still reached). + marker = _EMBEDDED_MARKER_RE.search(content) + if marker is not None and (first is None or marker.start() < first): + first = marker.start() return first def _xml_signal_inside_leading_bare_json(content: str) -> bool: - """True when the first foreign signal sits inside a LEADING bare-JSON call's balanced body: quoted argument data, so the bare-JSON parser takes the outer call first.""" + """True when the first foreign tool signal is a quoted literal inside a + LEADING bare-JSON call object or JSON answer -- data, not a real call + (sibling of ``_xml_signal_inside_leading_mistral``).""" i = 0 n = len(content) while i < n and content[i] in " \t\n\r": i += 1 - if i >= n or content[i] != "{": + if i >= n or content[i] not in "{[": return False + if content[i] == "[": + # A leading array is only ever a structured answer; its literals are data. + end = _balanced_bracket_end(content, i) + if end is None: + return False + try: + json.loads(content[i : end + 1]) + except ValueError: + return False + first_xml = _first_foreign_tool_signal(content) + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first_xml is None or trig < first_xml): + first_xml = trig + return first_xml is not None and i < first_xml < end end = _balanced_brace_end(content, i) if end is None: return False if _top_level_bare_json_name(content[i : end + 1]) is None: - # Not a call object, but a nameless object that parses as real JSON is an envelope - # too (markup in its strings is data); non-JSON braced prose keeps the old behaviour. + # A NAMELESS object that parses as real JSON is a structured answer / envelope too: + # quoted markup is data, and the decline path drops it and parses the tail. + # Non-JSON braced prose keeps the old behaviour. try: json.loads(content[i : end + 1]) except ValueError: return False first_xml = _first_foreign_tool_signal(content) - # The Mistral trigger is foreign to a JSON envelope too, so fold it into first_xml. + # The Mistral trigger is foreign to a JSON envelope too (its parser runs first). trig = content.find(_MISTRAL_TRIGGER) if trig >= 0 and (first_xml is None or trig < first_xml): first_xml = trig - # Inside the balanced body the signal is quoted argument data, so the leading call owns - # the turn; a non-call object takes the decline path (dropped, only the tail parsed). + # Inside the balanced body the signal is quoted data; after the closed object the + # leading call still owns the turn (mirrors the leading-Mistral rule). return first_xml is not None and i < first_xml +def _signal_inside_leading_wrapperless_gemma( + content: str, enabled_tool_names: Optional[set] +) -> bool: + """True when the first foreign tool signal is a quoted literal inside (or + after) a LEADING enabled wrapper-less Gemma call (sibling of the + Mistral/bare-JSON leading guards). Markerless form, so gated on an enabled + name (``None`` keeps the name-agnostic behaviour).""" + first = _first_foreign_tool_signal(content) + # The Mistral trigger is foreign to a Gemma call too (its parser runs first). + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first is None or trig < first): + first = trig + if first is None: + return False + # A preamble before ``call:NAME{...}`` is normal; what matters is an ENABLED balanced + # call beginning before the first foreign signal. + cursor = 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None or m.start() > first: + return False + if enabled_tool_names is not None and m.group(1) not in enabled_tool_names: + cursor = m.end() + continue + end = _gemma_body_brace_end(content, m.end() - 1) + if end is None: + return False + if m.end() - 1 < first <= end: + return True + # An enabled call that CLOSES before the signal still owns the turn (inside-or-after + # rule, as for closed bare-JSON/Mistral envelopes), gated on an enabled name. + return enabled_tool_names is not None and end < first + + +def _disabled_gemma_call_end_containing_signal( + content: str, enabled_tool_names: Optional[set] +) -> int | None: + """End offset (exclusive) of the earliest DISABLED wrapper-less Gemma call + whose balanced body contains the first foreign signal, else None. A disabled + name is prose, so the quoted literal is data: the caller drops the span and + recurses on the tail. An ENABLED call defers to the enabled-call guard.""" + if enabled_tool_names is None: + return None + first = _first_foreign_tool_signal(content) + # Mirror the enabled-call guard: the Mistral trigger is foreign here too. + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first is None or trig < first): + first = trig + if first is None: + return None + cursor = 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None or m.start() > first: + return None + if m.group(1) in enabled_tool_names: + return None + end = _gemma_body_brace_end(content, m.end() - 1) + if end is None: + cursor = m.end() + continue + if m.end() - 1 < first <= end: + return end + 1 + cursor = end + 1 + + def parse_tool_calls_from_text( content: str, *, @@ -440,26 +852,35 @@ def parse_tool_calls_from_text( allow_incomplete: bool = True, enabled_tool_names: Optional[set] = None, ) -> list[dict]: - """Return OpenAI-format tool calls, first-match wins. ``allow_incomplete`` heals truncated calls (``False`` = strict closed-only); ``enabled_tool_names`` gates the markerless bare-JSON form.""" - # Drop Magistral reasoning before any dispatch so a rehearsed call inside - # [THINK]...[/THINK] is not promoted; keeps the parse path aligned with the display strip. + """Return OpenAI-format tool calls, first-match wins so calls are never double-counted. + + ``allow_incomplete=True`` (default) heals truncated calls (missing close tag / + unclosed parameter); ``False`` accepts only well-formed closed calls (trailing + prose tolerated), matching llama-server's strict path when Auto-Heal is off. + + ``enabled_tool_names`` gates only the markerless Llama-3.2 bare-JSON form (the + marker-based forms carry an explicit signal, so a disabled-tool name there is a + real call attempt). ``None`` keeps the name-agnostic behaviour.""" + # Drop Magistral [THINK]...[/THINK] BEFORE dispatch: a rehearsed call inside it must + # never be promoted, and the parse path must agree with the display strip. content = _strip_mistral_reasoning(content) - # A leading bare-JSON value is decided FIRST so markup quoted in its arguments stays - # data. Must precede the Mistral guard, whose preamble tolerance would else claim a - # trigger quoted inside the leading object. + # A leading bare-JSON value is decided FIRST: a string argument quoting tool markup + # (XML or a Mistral trigger) must stay data, so the bare-JSON parser takes the outer + # call before any other pass. Precedes the Mistral guard, whose preamble tolerance + # would otherwise claim a trigger quoted inside the leading object. if _xml_signal_inside_leading_bare_json(content): calls = _parse_llama3_bare_json( content, id_offset = id_offset, enabled_tool_names = enabled_tool_names ) if calls: return calls - # Disabled/example name: the leading object is ordinary content. Drop it and parse - # only the tail -- a real call after it still parses, nothing inside it is promoted. + # Disabled/example name: the leading object is content. Drop it and parse the tail. i = 0 while i < len(content) and content[i] in " \t\n\r": i += 1 - end = _balanced_brace_end(content, i) # guard guarantees a balanced object + # The guard guarantees a balanced leading value (object or array). + end = (_balanced_brace_end if content[i] == "{" else _balanced_bracket_end)(content, i) return parse_tool_calls_from_text( content[end + 1 :], id_offset = id_offset, @@ -467,8 +888,32 @@ def parse_tool_calls_from_text( enabled_tool_names = enabled_tool_names, ) + # A leading enabled wrapper-less Gemma call is decided BEFORE the Mistral guard: its + # body reads as prose to the preamble tolerance below, so a quoted [TOOL_CALLS] would + # otherwise steal the turn. + if _signal_inside_leading_wrapperless_gemma(content, enabled_tool_names): + calls = _parse_gemma_tool_calls( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + if calls: + return calls + + # A DISABLED wrapper-less Gemma call is prose: drop the span and parse the tail BEFORE + # the Mistral guard, whose preamble tolerance would otherwise parse a quoted trigger. + _prose_end = _disabled_gemma_call_end_containing_signal(content, enabled_tool_names) + if _prose_end is not None: + return parse_tool_calls_from_text( + content[_prose_end:], + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + # A [TOOL_CALLS] call that is the first tool emission owns the turn: XML quoted in its - # arguments or in trailing prose is not promoted, and a plain-prose preface keeps it. + # arguments or trailing prose is not promoted over it, nor does a prose preface forfeit it. if _xml_signal_inside_leading_mistral(content): calls = _parse_mistral_tool_calls( content, id_offset = id_offset, allow_incomplete = allow_incomplete @@ -476,8 +921,29 @@ def parse_tool_calls_from_text( if calls: return calls - # A leading MiniCPM/MiniMax ```` call owns the turn: tool_healing - # does not know the wrapper, so gate it here. A signal before the opener keeps normal order. + # DeepSeek/Kimi markers are unique, so try them first -- unless an outer envelope + # opens before the first marker (then the marker is argument data). + if not _marker_inside_leading_envelope(content, enabled_tool_names): + # Dispatch by earliest opener so a quoted DS example inside a Kimi call (or vice + # versa) can't hijack the turn via fixed parser order. + _ds = _DEEPSEEK_BEGIN_RE.search(content) + _ds_pos = _ds.start() if _ds else len(content) + _km_section = content.find(_KIMI_SECTION_BEGIN) + _km_bare = content.find(_KIMI_CALL_BEGIN) + _km_pos = min(p for p in (_km_section, _km_bare, len(content)) if p >= 0) + pre_pass = [ + (_ds_pos, _parse_deepseek_tool_calls), + (_km_pos, _parse_kimi_tool_calls), + ] + pre_pass.sort(key = lambda pair: pair[0]) + for _pos, parser in pre_pass: + calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # A leading MiniCPM/MiniMax attribute-form call owns the turn: tool_healing doesn't know + # the wrapper, so a quoted in its parameter would beat + # the outer call. Any earlier signal keeps normal order. attr = _ATTR_FUNC_OPEN_RE.search(content) if attr is not None: first_other = None @@ -518,8 +984,9 @@ def parse_tool_calls_from_text( if calls: return calls - # Qwen/Hermes, Qwen3.5 XML, and Gemma 4 use the shared tool_healing parser (the - # strict/Auto-Heal + nested-marker + ``<|"|>`` handling GGUF relies on). + # Qwen/Hermes, Qwen3.5 XML, and Gemma 4 go through the shared tool_healing + # parser (strict/Auto-Heal contract + nested-marker, trailing-prose, and + # ``<|"|>`` quoted-string handling the GGUF path relies on). calls = _tool_healing.parse_tool_calls_from_text( content, id_offset = id_offset, @@ -528,11 +995,11 @@ def parse_tool_calls_from_text( if calls: return calls - # Formats tool_healing does not cover: ```` (MiniCPM-5 / MiniMax-M2), - # Llama-3 and Mistral. Run only after tool_healing found nothing, so a strict-rejected - # call is never re-healed here. Blank any JSON/Gemma marker coverage first: markup inside - # a marker's span (even one that failed to parse) is that call's data, not a sibling, so - # a nested ```` / ``<|python_tag|>`` / ``[TOOL_CALLS]`` must not be promoted. + # Formats tool_healing does not cover; these run only after it finds + # nothing, so a strict-rejected call is never re-healed here. Blank any + # JSON/Gemma marker coverage first: markup inside a marker's span (even one + # that failed to parse) is that call's data, not a sibling, so a nested + # ```` / ``<|python_tag|>`` / ``[TOOL_CALLS]`` must not be promoted. fallback_content = content coverage = _tool_healing.marker_coverage(content) if coverage: @@ -542,6 +1009,7 @@ def parse_tool_calls_from_text( chars[i] = " " fallback_content = "".join(chars) for parser in ( + _parse_glm_tool_calls, # GLM 4.x name _parse_function_xml, # attribute form _parse_llama3_python_tag, # Llama-3 <|python_tag|> _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] @@ -550,11 +1018,22 @@ def parse_tool_calls_from_text( if calls: return calls - # Llama-3.2 bare ``{"name":..., "parameters":...}``. Strict (starts with ``{`` - # and parses to the right shape) so plain prose stays untouched. - return _parse_llama3_bare_json( + # Llama-3.2 bare ``{"name":..., "parameters":...}`` (strict shape). Only a LEADING call + # object matches and owns the turn, so an enabled ``call:NAME{...}`` in its arguments + # stays data (Gemma never starts ``{``). + calls = _parse_llama3_bare_json( content, id_offset = id_offset, enabled_tool_names = enabled_tool_names ) + if calls: + return calls + + # Gemma wrapper-less ``call:NAME{...}``: markerless, so the same enabled-name gate applies. + return _parse_gemma_tool_calls( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) def _parse_tool_call_json( @@ -569,8 +1048,9 @@ def _parse_tool_call_json( end = _balanced_brace_end(content, brace_start) if end is None: continue - # Strict mode: a balanced body that never closed its ```` is truncated - # (trailing prose after the close is still tolerated). + # Strict mode: a balanced JSON body that never closed its ```` + # is a truncated call, not a finished one. Trailing prose after the close + # is still tolerated (matches the GGUF strict path). if not allow_incomplete and not content[end + 1 :].lstrip().startswith(""): continue try: @@ -578,7 +1058,7 @@ def _parse_tool_call_json( except (json.JSONDecodeError, ValueError): continue name = obj.get("name", "") - # Accept both ``arguments`` (Hermes/Qwen) and ``parameters`` (Llama-3 drift). + # Accept ``arguments`` (Hermes/Qwen) and ``parameters`` (Llama-3 drift). args = obj.get("arguments") if args is None: args = obj.get("parameters", {}) @@ -601,7 +1081,10 @@ def _parse_tool_call_json( def _trim_param_value(val: str) -> str: - """Trim only the template's wrapping newline around an XML parameter value; ``str.strip()`` destroyed code/diff indentation.""" + """Trim one wrapping newline the template adds around an XML parameter value + (``\nVALUE\n``), preserving inner indentation. + ``str.strip()`` destroyed code/diff indentation; SGLang's qwen3_coder trims only + the wrapping newline.""" if val.startswith("\n"): val = val[1:] if val.endswith("\n"): @@ -610,14 +1093,19 @@ def _trim_param_value(val: str) -> str: def _inside_open_parameter(text: str, pos: int) -> bool: - """True if ``pos`` is inside an unclosed ```` block, i.e. the opener at ``pos`` is literal argument data, not a nested call.""" + """True if ``pos`` sits inside an unclosed ````/```` block -- + i.e. a ```` / ```` opener at ``pos`` is a literal inside an + argument value (e.g. code that prints tool-call XML), not a real nested call. + Compares the last parameter opener before ``pos`` against the last + parameter/function close before it.""" last_param_open = -1 for m in _TC_PARAM_START_RE.finditer(text, 0, pos): last_param_open = m.start() if last_param_open < 0: return False - # The parameter's OWN close tag decides: if it closes after ``pos`` the position is - # argument data (even across literal ````); an unclosed one falls back to func close. + # The parameter's OWN close tag decides: while it closes after ``pos`` the position is + # argument data, even across several literal function closes. Only an unclosed + # parameter (heal mode) falls back to the first function close. own_closes = [ c for c in ( @@ -647,7 +1135,7 @@ def _parse_function_xml( ) -> list[dict]: out: list[dict] = [] # Skip ```` openers that are literals inside an open parameter value, - # else the nested marker becomes a second call and truncates the real argument. + # else the nested marker is promoted to a second call and truncates the real argument. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) @@ -658,9 +1146,10 @@ def _parse_function_xml( func_name = fm.group(1) or fm.group(2) body_start = fm.end() next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - # The call ends at the FIRST / not inside an open parameter: - # a literal close in an argument is skipped as data, prose after the real close is not - # folded in (mirrors _strip_function_xml_calls). + # The call ends at the FIRST / not inside an open + # parameter: a literal close in a code/search argument is skipped as data, and + # prose after the real close isn't folded into the last argument (mirrors + # _strip_function_xml_calls and tool_healing._func_close_index). close_match = None for cm in _TC_END_TAG_RE.finditer(content, body_start, next_func): if not _inside_open_parameter(content, cm.start()): @@ -671,14 +1160,14 @@ def _parse_function_xml( body_end = close_match.start() else: body_end = min(len(content), next_func) - # Strict mode: a call that never reached its close is truncated; do not heal it. + # Strict mode: an unclosed function call is truncated -- do not heal it. if not allow_incomplete and not has_close: continue body = _TC_FUNC_CLOSE_RE.sub("", content[body_start:body_end]) args: dict = {} param_unclosed = False - # Same nested-literal guard: a ```` opener inside an open value is literal text. + # A ```` opener inside an open parameter value is literal text. param_starts = [ pm for pm in _TC_PARAM_START_RE.finditer(body) @@ -703,8 +1192,8 @@ def _parse_function_xml( val = _TC_PARAM_CLOSE_RE.sub("", raw_val) args[pm.group(1) or pm.group(2)] = _trim_param_value(val) - # Strict mode: every parameter must close; a dangling one means the call was cut off. - # A closed call with no parameters is a valid zero-argument call, so keep it. + # Strict mode: a dangling parameter means the call was cut off; a closed + # zero-parameter call stays valid. if not allow_incomplete and param_unclosed: continue @@ -719,7 +1208,8 @@ def _parse_function_xml( def _llama3_kv_value(body: str, p: int, n: int) -> tuple[Any, int | None]: - """One ``.call`` value at ``body[p:]``; returns ``(value, len)`` or ``(None, None)``.""" + """One ``.call`` value (string/number/true/false/null) at ``body[p:]``. + Returns ``(value, consumed_len)`` or ``(None, None)`` if none matches.""" if p >= n: return None, None if body[p] == '"': @@ -745,7 +1235,8 @@ def _llama3_kv_value(body: str, p: int, n: int) -> tuple[Any, int | None]: nm = _LLAMA3_NUM_RE.match(body, p) if nm: v = nm.group(0) - # Sci notation and decimals decode as float; a bare integer stays int. + # Scientific notation (1e-3, -2E+4, 0.5e2) and decimals decode as float; a bare + # integer stays int. ``"." in v`` alone missed the exponent forms (1e-3 -> 1). return (float(v) if any(c in v for c in ".eE") else int(v)), nm.end() - p lm = _LLAMA3_LIT_RE.match(body, p) if lm: @@ -754,7 +1245,8 @@ def _llama3_kv_value(body: str, p: int, n: int) -> tuple[Any, int | None]: def _parse_llama3_kv_args(body: str) -> dict[str, Any]: - """Left-to-right ``k=v`` kwargs from a ``.call(...)`` body (linear scan; later keys win).""" + """``k=v, ...`` kwargs from a ``.call(...)`` body, left to right (later keys win). + Linear hand-scan replacing the quadratic ``_LLAMA3_KV_RE.finditer`` walk.""" args: dict[str, Any] = {} n = len(body) i = 0 @@ -783,13 +1275,17 @@ def _parse_llama3_python_tag( id_offset: int, allow_incomplete: bool = True, ) -> list[dict]: - """Parse Llama-3 ``<|python_tag|>`` emissions: ``NAME.call(...)``, bare JSON, ``; `` multi-call, ``parameters``/``arguments`` keys.""" + """Parse the Llama-3 emissions: ``<|python_tag|>NAME.call(...)`` (built-in), + ``<|python_tag|>{"name":..., "parameters":...}`` (custom), multi-call via + ``; ``, ``parameters`` or ``arguments`` key.""" out: list[dict] = [] if _LLAMA3_PYTHON_TAG not in content: return out - # 1. ``NAME.call(...)`` built-in form, anchored to ``<|python_tag|>`` (optionally - # ``; ``-chained) so a ``.call(...)`` inside a JSON string argument isn't mistaken for one. + # 1. ``NAME.call(...)`` built-in form, anchored to ``<|python_tag|>`` and optionally + # ``; ``-chained within one emission. Anchoring to the tag boundary (not a free scan) + # keeps a literal ``<|python_tag|>x.call(...)`` quoted in a custom-form JSON argument + # from being mistaken for a real built-in call. pos = content.find(_LLAMA3_PYTHON_TAG) truncated = False while pos >= 0 and not truncated: @@ -824,7 +1320,8 @@ def _parse_llama3_python_tag( if depth == 0: break i += 1 - # Truncated ``.call(...)`` (no closing paren): reject in strict mode. + # Truncated ``.call(...)`` with no closing paren: reject in strict mode + # instead of executing a partial. if not allow_incomplete and depth > 0: truncated = True break @@ -848,7 +1345,8 @@ def _parse_llama3_python_tag( # Past the consumed region: a second ``<|python_tag|>`` may carry more calls. pos = content.find(_LLAMA3_PYTHON_TAG, i + 1) - # 2. ``<|python_tag|>{"name":.., "parameters":..}``; raw_decode peels ``; ``-separated objects. + # 2. ``<|python_tag|>{"name":..., "parameters":...}``. ``raw_decode`` peels multiple + # ``; ``-separated objects from one emission. if not out: decoder = json.JSONDecoder() idx = content.find(_LLAMA3_PYTHON_TAG) @@ -873,12 +1371,14 @@ def _parse_llama3_python_tag( continue name = obj.get("name") or obj.get("function") or "" args = obj.get("parameters") if "parameters" in obj else obj.get("arguments", {}) + # Skip rather than fabricate ``{"value": args}`` for a non-dict/non-string value. if isinstance(args, dict): args_str = json.dumps(args) elif isinstance(args, str): args_str = args else: - args_str = json.dumps({"value": args}) + cursor = brace + end_offset + continue if name: out.append( { @@ -892,7 +1392,8 @@ def _parse_llama3_python_tag( return out -# Llama-3 special-token sentinels (chainable, any order) plus the header role label. +# Llama-3 special-token sentinels (chainable, any order) plus the role label the +# template inserts between ``<|start_header_id|>`` and ``<|end_header_id|>``. _LLAMA3_BARE_JSON_SENTINELS = ( "<|begin_of_text|>", "<|eot_id|>", @@ -904,7 +1405,10 @@ _LLAMA3_HEADER_ROLES = ("assistant", "user", "system", "tool", "ipython") def strip_llama3_leading_sentinels(content: str) -> str: - """Strip leading Llama-3 sentinels leaked from a prior turn; shared by the parser and the streaming guards.""" + """Strip leading Llama-3 special-token sentinels (and the role label after + ``<|start_header_id|>``) that can leak from a prior turn before a bare-JSON tool + call. Shared by the parser and the streaming buffering guards so a + sentinel-prefixed ``{"name":...}`` is recognised the same everywhere.""" stripped = content.lstrip() while True: stripped = stripped.lstrip() @@ -930,7 +1434,9 @@ def _parse_llama3_bare_json( allow_incomplete: bool = True, enabled_tool_names: Optional[set] = None, ) -> list[dict]: - """Llama-3.2 bare ``{"name":.., "parameters":..}`` (strict). ``enabled_tool_names`` keeps ordinary JSON answers from being misread; ``None`` is name-agnostic.""" + """Llama-3.2 ``custom_tools`` bare ``{"name":.., "parameters":{..}}`` (no ``<|python_tag|>``), + strict so prose/echoes don't fire. ``enabled_tool_names`` gates on the parsed name so an + ordinary JSON answer isn't misread as a call to a disabled tool; ``None`` is name-agnostic.""" out: list[dict] = [] stripped = strip_llama3_leading_sentinels(content) if not stripped.startswith("{"): @@ -954,11 +1460,12 @@ def _parse_llama3_bare_json( name = obj.get("name") or obj.get("function") or "" if not isinstance(name, str) or not name: break - # Markerless JSON is ambiguous: only a call when the name is an enabled tool. + # Markerless JSON is ambiguous: treat it as a call only when the name is an enabled + # tool, else it is an ordinary JSON answer. if enabled_tool_names is not None and name not in enabled_tool_names: break - # ``parameters`` must be a dict (Llama-3 spec); ``arguments`` may be a dict or a - # JSON-string of one (OpenAI). + # ``parameters`` must be a dict (Llama-3 spec); ``arguments`` may be a dict or + # JSON-string of one (OpenAI). Looser would fire on ``{"name":"x","parameters":"sentence"}``. if "parameters" in obj: args = obj.get("parameters") if not isinstance(args, dict): @@ -997,14 +1504,15 @@ def _parse_mistral_tool_calls( id_offset: int, allow_incomplete: bool = True, ) -> list[dict]: - """Parse Mistral ``[TOOL_CALLS]`` emissions: pre-v11 array/object and v11+ named forms.""" + """Parse all Mistral emissions: pre-v11 ``[TOOL_CALLS][...]`` / ``[TOOL_CALLS]{...}`` + and v11+ ``[TOOL_CALLS]name{json}`` / ``[TOOL_CALLS]name[ARGS]{json}``.""" out: list[dict] = [] content = _strip_mistral_reasoning(content) idx = content.find(_MISTRAL_TRIGGER) if idx < 0: return out - # Disambiguate the first occurrence: array / single object (pre-v11) or bare-name (v11+). + # Disambiguate the first occurrence: array / single object (pre-v11), or bare-name (v11+). j = idx + len(_MISTRAL_TRIGGER) k = j while k < len(content) and content[k] in " \t\n\r": @@ -1016,7 +1524,7 @@ def _parse_mistral_tool_calls( return _parse_mistral_array(content, k, id_offset, allow_incomplete = allow_incomplete) if content[k] == "{": - # Pre-v11 single ``{"name":...}``; fall through to v11+ if it carries no ``name``. + # Pre-v11 single ``{"name":...}``; fall through without a ``name`` so v11+ still runs. end = _balanced_brace_end(content, k) if end is not None: try: @@ -1027,7 +1535,8 @@ def _parse_mistral_tool_calls( except (json.JSONDecodeError, ValueError): pass - # v11+: walk every ``[TOOL_CALLS]``, parsing ``name{json}`` or ``name[ARGS]{json}``. + # v11+: walk every ``[TOOL_CALLS]``, parsing ``name{json}`` or + # ``name[ARGS]{json}`` after each trigger. pos = idx while pos >= 0: cur = pos + len(_MISTRAL_TRIGGER) @@ -1101,7 +1610,8 @@ def _parse_mistral_array( if depth == 0: break j += 1 - # An unclosed array (no matching ]) is truncated; reject in strict mode. + # An unclosed array (no matching ]) is a truncated call. In strict mode reject it + # instead of recovering objects by hand below. if not allow_incomplete and depth != 0: return out body = content[start : j + 1] if depth == 0 else content[start:] @@ -1117,8 +1627,8 @@ def _parse_mistral_array( if not allow_incomplete: return out - # Healing path for unclosed arrays: walk top-level objects, advancing past each - # balanced ``{...}`` (re-scanning from every ``{`` would be quadratic ReDoS). + # Healing path for unclosed arrays: walk top-level objects, advancing past each balanced + # ``{...}`` instead of re-scanning from every ``{`` (quadratic ReDoS). pos = 0 blen = len(body) while pos < blen: @@ -1141,7 +1651,8 @@ def _consume_mistral_call(obj_text: str, out: list[dict], id_offset: int) -> Non if not isinstance(obj, dict): return name = obj.get("name") or "" - # Mistral uses ``arguments``; accept the ``parameters`` alias too. + # Mistral uses ``arguments``; accept the ``parameters`` alias too (sibling paths and + # SGLang's base detector alias it) so an array object keyed on it keeps args. args = obj.get("arguments") if args is None: args = obj.get("parameters", {}) @@ -1161,28 +1672,86 @@ def _consume_mistral_call(obj_text: str, out: list[dict], id_offset: int) -> Non ) +def _whole_content_is_json_value(text: str) -> bool: + """True when the entire content is one valid JSON value (a structured + answer, e.g. a response_format turn). Markerless scans must treat text + inside it as data: an answer documenting an enabled tool's syntax must + not execute that tool or have the example stripped from display.""" + t = text.strip() + if t[:1] not in "{[": + return False + try: + json.loads(t) + except ValueError: + return False + return True + + +def _leading_json_value_end(text: str) -> int | None: + """End index (exclusive) of a balanced LEADING JSON value that parses as + JSON: a structured answer possibly followed by prose. Markerless scans treat + its contents as data (extends ``_whole_content_is_json_value``); leading-keyed, + so a JSON blob mid-prose is not an answer span.""" + i = 0 + n = len(text) + while i < n and text[i].isspace(): + i += 1 + if i >= n or text[i] not in "{[": + return None + end = (_balanced_brace_end if text[i] == "{" else _balanced_bracket_end)(text, i) + if end is None: + return None + try: + json.loads(text[i : end + 1]) + except ValueError: + return None + return end + 1 + + def _parse_gemma_tool_calls( content: str, *, id_offset: int, allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, ) -> list[dict]: - """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``.""" + """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``, plus the + ``skip_special_tokens`` stream where the wrapper and string markers were + stripped (bare ``call:NAME{k:v, ...}``). + + ``enabled_tool_names`` gates on the parsed name: the wrapper-less shape is + indistinguishable from prose documenting the syntax, so a disabled/example + name must not be stolen as a call. ``None`` keeps the name-agnostic behaviour.""" out: list[dict] = [] - for m in _GEMMA_TC_RE.finditer(content): + # The WRAPPED form (strict + nested-marker handling) is tool_healing's, which runs + # first: defer content with a wrapped opener. A marker literal alone is not enough -- + # a wrapper-less call mentioning ``<|tool_call>`` would be lost if deferred. + if _GEMMA_TC_RE.search(content): + return out + # A whole-content JSON value is a structured answer: quoted examples must not become calls. + if _whole_content_is_json_value(content): + return out + # Manual cursor: resume AFTER each consumed balanced body so a nested ``call:OTHER{...}`` + # in an argument is never re-matched. A leading JSON answer's span is data -- scan after it. + cursor = _leading_json_value_end(content) or 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None: + break name = m.group(1) body_start = m.end() - 1 - end_marker = content.find(_GEMMA_TC_END, body_start) - # No closing tag: truncated call, reject in strict mode. - if not allow_incomplete and end_marker < 0: - continue - scan_end = end_marker if end_marker >= 0 else len(content) - end = _gemma_balanced_brace_end(content, body_start, scan_end) + end = _gemma_body_brace_end(content, body_start) if end is None: + # Unclosed call: nothing parseable follows (mirrors the strip contract); + # scanning on would promote quoted argument text. + break + cursor = end + 1 + # Markerless: a disabled/example name is prose, not a call. + if enabled_tool_names is not None and name not in enabled_tool_names: continue body = content[body_start + 1 : end] try: - args = _gemma_parse_mapping_body(body) + args = _gemma_parse_stripped_body(body) except Exception: args = {} out.append( @@ -1225,11 +1794,54 @@ def _balanced_brace_end(text: str, brace_pos: int) -> int | None: return None +def _gemma_body_brace_end(text: str, brace_pos: int) -> int | None: + """Index of the ``}`` closing the wrapper-less Gemma body at ``brace_pos``. + + Values are raw after ``skip_special_tokens``, so quoted strings (single or + double) hide braces; the quote rules mirror ``_gemma_parse_stripped_body`` so + the boundary always agrees with the body parser. Contextual openers: a single + quote opens only at value-start context (after ``:{[(,=`` -- apostrophes in + ``what's the weather`` are prose), a double quote also at word start (so + ``query:find "a, b"`` hides its delimiters).""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + quote = "" + prev = "" + prev_raw = "" + i = brace_pos + n = len(text) + while i < n: + ch = text[i] + if quote: + if ch == "\\" and i + 1 < n: + i += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + if not ch.isspace(): + prev = ch + prev_raw = ch + i += 1 + return None + + _BARE_JSON_NAME_RE = re.compile(r'"name"\s*:\s*"([^"]+)"') def _top_level_bare_json_name(probe: str) -> Optional[str]: - """Top-level ``"name"`` (or ``"function"`` alias) of a bare-JSON object, else None; nested objects are skipped and truncated tails return None.""" + """TOP-LEVEL ``"name"`` (or ``"function"`` alias, name wins) of a bare-JSON object, else None. + + Skips nested objects/arrays so a nested ``"name"`` isn't mistaken for the call name; a + truncated tail returns None so the caller keeps the text.""" if not probe.startswith("{"): return None decoder = json.JSONDecoder() @@ -1240,7 +1852,7 @@ def _top_level_bare_json_name(probe: str) -> Optional[str]: while i < n and probe[i] in " \t\r\n,": i += 1 if i >= n or probe[i] == "}": - # End of object, no top-level ``"name"``: fall back to the ``"function"`` alias. + # End of the object with no top-level ``"name"``: fall back to a recorded ``"function"`` alias. return function_value if probe[i] != '"': return None @@ -1267,7 +1879,8 @@ def _top_level_bare_json_name(probe: str) -> Optional[str]: return value if isinstance(value, str) else None return None if key == "function" and function_value is None and i < n and probe[i] == '"': - # ``"function"`` is an alias; record it but keep scanning (``"name"`` wins). + # ``"function"`` aliases the call name. Record it but keep scanning: a top-level + # ``"name"`` still wins. try: value, consumed = decoder.raw_decode(probe[i:]) except (json.JSONDecodeError, ValueError): @@ -1276,7 +1889,8 @@ def _top_level_bare_json_name(probe: str) -> Optional[str]: function_value = value i += consumed continue - # Skip a non-name top-level value; a truncated one returns None (keep the text). + # Skip a non-name top-level value; a truncated one can't prove a top-level name + # exists, so return None (keep the text). if i < n and probe[i] == "{": end = _balanced_brace_end(probe, i) if end is None: @@ -1314,16 +1928,18 @@ def strip_leading_bare_json_call(text: str, enabled_tool_names: Optional[set] = if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): return probe.lstrip() if stripped_any else text if enabled_tool_names is not None: - # Only suppress when the leading object's TOP-LEVEL name is an enabled tool - # (a nested ``"name"`` is data); an unknown name is kept. + # Only suppress when the leading object's TOP-LEVEL name is an enabled tool. A + # nested ``"name"`` (e.g. {"result":{"name":"web_search",...}}) is data, not the + # call name, so it must not gate the strip. An un-extractable name is kept. name = _top_level_bare_json_name(probe) if name not in enabled_tool_names: return probe.lstrip() if stripped_any else text end = _balanced_brace_end(probe, 0) if end is None: return "" # truncated bare-JSON call -- nothing recoverable - # A closed object must have the CALL SHAPE the parser accepts; an ordinary JSON - # answer it rejects is content, so keep it visible. + # A closed object must have the CALL SHAPE the parser accepts (dict ``parameters``, + # or dict / JSON-string ``arguments``). An ordinary JSON answer like + # {"name":"web_search","result":"no call"} is content, so the strip keeps it visible. try: obj = json.loads(probe[: end + 1]) except (json.JSONDecodeError, ValueError): @@ -1338,7 +1954,8 @@ def _bare_json_call_shaped(obj) -> bool: """The shape gate ``_parse_llama3_bare_json`` applies to a decoded object.""" if not isinstance(obj, dict): return False - # The parser requires a TOP-LEVEL name; a nested one is data, not the call name. + # The parser requires a TOP-LEVEL name; a nested one (e.g. in a "result" value of an + # ordinary JSON answer) is data, and stripping it name-agnostically would delete content. name = obj.get("name") or obj.get("function") or "" if not isinstance(name, str) or not name: return False @@ -1379,101 +1996,670 @@ def _gemma_balanced_brace_end(text: str, brace_pos: int, hard_stop: int) -> int return None -def _gemma_parse_value(text: str, i: int): - """Parse one Gemma arg value at ``i``; returns ``(value, next_index)``.""" +def _gemma_parse_value( + text: str, + i: int, + *, + in_mapping: bool = False, +): + """Parse one Gemma arg value at ``i`` in a single O(n) forward pass; returns + ``(value, next_index, closed)``. ``closed`` is False when a string/object/array + runs off the end without its terminator, so the caller can fall back to raw. + ``in_mapping`` applies the top-level rule that a comma only ends the value + when a ``key:`` follows (array elements split on every top-level comma).""" if text.startswith(_GEMMA_STR_BEGIN, i): close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) if close < 0: - return text[i + len(_GEMMA_STR_BEGIN) :], len(text) - return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END) + return text[i + len(_GEMMA_STR_BEGIN) :], len(text), False + return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END), True if text[i] == "{": - end = _gemma_balanced_brace_end(text, i, len(text)) - if end is None: - return {}, len(text) - return _gemma_parse_mapping_body(text[i + 1 : end]), end + 1 + return _gemma_parse_mapping(text, i) if text[i] == "[": - j, depth = i, 0 - while j < len(text): - if text.startswith(_GEMMA_STR_BEGIN, j): - k = text.find(_GEMMA_STR_END, j + len(_GEMMA_STR_BEGIN)) - if k < 0: - j = len(text) - break - j = k + len(_GEMMA_STR_END) + return _gemma_parse_array(text, i) + if text[i] in "\"'": + # Raw-quoted string: delimiters inside are data (``{city:"New, York"}`` is one + # value); returned unquoted like the top-level scalar coercion. + quote = text[i] + j = i + 1 + n = len(text) + while j < n: + if text[j] == "\\" and j + 1 < n: + j += 2 continue - ch = text[j] - if ch == "[": - depth += 1 - elif ch == "]": - depth -= 1 - if depth == 0: - break + if text[j] == quote: + return text[i + 1 : j], j + 1, True j += 1 - body = text[i + 1 : j] - items: list[Any] = [] - k = 0 - while k < len(body): - if body[k] in " \t\n\r,": - k += 1 - continue - v, k = _gemma_parse_value(body, k) - items.append(v) - return items, j + 1 - # Primitive: number / true/false/null / bare identifier. + return text[i + 1 :], n, False + # Primitive / unquoted code: same delimiter rules as the top-level scan (bracket depth + # + contextual quote openers hide commas and closers). end = i - while end < len(text) and text[end] not in ",}]" and not text.startswith(_GEMMA_STR_BEGIN, end): + n = len(text) + depth = 0 + quote = "" + prev = ":" + prev_raw = ":" + while end < n and not text.startswith(_GEMMA_STR_BEGIN, end): + ch = text[end] + if quote: + if ch == "\\" and end + 1 < n: + end += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch in "{[(": + depth += 1 + elif ch in "}])": + if depth == 0: + break + depth -= 1 + elif ch == "," and depth == 0: + if not in_mapping or _GEMMA_KEY_RE.match(text, end + 1): + break + if not ch.isspace(): + prev = ch + prev_raw = ch end += 1 if end == i: - # Stray delimiter, nothing consumed: advance past it so callers can't spin forever. - return "", i + 1 + # Stray delimiter where a value was expected: consume one char so callers always + # advance (no infinite loop on malformed input). + return "", i + 1, True raw = text[i:end].strip() if raw == "true": - return True, end + return True, end, True if raw == "false": - return False, end + return False, end, True if raw == "null": - return None, end + return None, end, True try: - return int(raw), end + return int(raw), end, True except ValueError: pass try: - return float(raw), end + return float(raw), end, True except ValueError: pass - return raw, end + return raw, end, True -def _gemma_parse_mapping_body(body: str) -> dict[str, Any]: - """Parse a Gemma argument mapping (content between `{` and `}`).""" - out: dict[str, Any] = {} - i = 0 - n = len(body) +def _gemma_parse_array(text: str, start: int): + """Parse a Gemma ``[...]`` array at ``text[start] == '['`` in one forward + pass; returns ``(list, next_index, closed)``.""" + items: list[Any] = [] + i, n = start + 1, len(text) while i < n: - while i < n and body[i] in " \t\n\r,": + while i < n and text[i] in " \t\n\r,": i += 1 + if i < n and text[i] == "]": + return items, i + 1, True if i >= n: break - if body.startswith(_GEMMA_STR_BEGIN, i): - close = body.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + v, i, _closed = _gemma_parse_value(text, i) + items.append(v) + return items, i, False + + +def _gemma_coerce_scalar(raw: str) -> Any: + """Coerce an unquoted Gemma value to bool/int/float/None, else keep str + (quotes stripped first so quoted/unquoted variants compare identical).""" + raw = raw.strip() + if len(raw) >= 2 and raw[0] == raw[-1] and raw[0] in "\"'": + return raw[1:-1] + if raw == "true": + return True + if raw == "false": + return False + if raw == "null": + return None + try: + return int(raw) + except ValueError: + pass + try: + return float(raw) + except ValueError: + pass + return raw + + +def _gemma_strip_quoted_leaves(value: Any) -> Any: + """Recursively unquote quoted string leaves of a nested stripped-stream value, + so nested ``city:"New York"`` matches the top-level coercion (no stray quotes).""" + if isinstance(value, str): + v = value.strip() + if len(v) >= 2 and v[0] == v[-1] and v[0] in "\"'": + return v[1:-1] + return value + if isinstance(value, dict): + return {k: _gemma_strip_quoted_leaves(v) for k, v in value.items()} + if isinstance(value, list): + return [_gemma_strip_quoted_leaves(v) for v in value] + return value + + +def _gemma_parse_stripped_body(body: str) -> dict[str, Any]: + """Parse a quote-less Gemma arg body ``key:value, key2:value2`` (the + ``skip_special_tokens`` stream with ``<|"|>`` markers removed). Each value runs + to the next top-level ``, key:`` boundary, tracking ``{}``/``[]``/``()`` depth so + commas/braces inside a ``code`` / ``command`` value aren't truncated.""" + out: dict[str, Any] = {} + i, n = 0, len(body) + while i < n: + m = _GEMMA_KEY_RE.match(body, i) + if not m: + break + key = m.group(1) + i = m.end() + vstart = i + depth = 0 + quote = "" + # Contextual quote openers mirror _gemma_body_brace_end. + prev = ":" + prev_raw = ":" + while i < n: + ch = body[i] + if quote: + # A ``, key:`` shape inside the quoted string is not a boundary. + if ch == "\\" and i + 1 < n: + i += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch in "{[(": + depth += 1 + elif ch in "}])": + if depth > 0: + depth -= 1 + elif ch == "," and depth == 0 and _GEMMA_KEY_RE.match(body, i + 1): + break + if not ch.isspace(): + prev = ch + prev_raw = ch + i += 1 + raw_val = body[vstart:i].strip() + if raw_val[:1] in "{[": + # Nested object/array: accept only a fully consumed, closed parse; a + # truncated/malformed value falls back to the raw string. + parsed, end, closed = _gemma_parse_value(raw_val, 0) + out[key] = ( + _gemma_strip_quoted_leaves(parsed) + if (closed and end == len(raw_val)) + else _gemma_coerce_scalar(raw_val) + ) + else: + out[key] = _gemma_coerce_scalar(raw_val) + if i < n and body[i] == ",": + i += 1 + return out + + +def _gemma_parse_mapping(text: str, start: int): + """Parse a Gemma ``{key:value, ...}`` mapping at ``text[start] == '{'`` in one + forward pass; returns ``(dict, next_index, closed)`` (``closed`` True iff the + matching ``}`` was reached).""" + out: dict[str, Any] = {} + i, n = start + 1, len(text) + while i < n: + while i < n and text[i] in " \t\n\r,": + i += 1 + if i < n and text[i] == "}": + return out, i + 1, True + if i >= n: + break + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) if close < 0: break - key = body[i + len(_GEMMA_STR_BEGIN) : close] + key = text[i + len(_GEMMA_STR_BEGIN) : close] i = close + len(_GEMMA_STR_END) else: kstart = i - while i < n and body[i] != ":": + while i < n and text[i] not in ":}": i += 1 - key = body[kstart:i].strip() - while i < n and body[i] in " \t\n\r": + key = text[kstart:i].strip() + while i < n and text[i] in " \t\n\r": i += 1 - if i < n and body[i] == ":": + if i < n and text[i] == ":": i += 1 - while i < n and body[i] in " \t\n\r": + while i < n and text[i] in " \t\n\r": i += 1 if i >= n: out[key] = None break - v, i = _gemma_parse_value(body, i) + if text[i] == "}": + out[key] = None + return out, i + 1, True + v, i, _closed = _gemma_parse_value(text, i, in_mapping = True) out[key] = v + return out, i, False + + +# ── DeepSeek R1 / V3 / V3.1 ───────────────────────────────────────── + + +def _find_outside_json_strings(text: str, needle: str, start: int) -> int: + """Index of ``needle`` at/after ``start`` OUTSIDE any JSON string, or -1: a + marker inside an argument string must not be taken as the structural terminator.""" + i = start + n = len(text) + in_string = False + esc = False + while i < n: + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + i += 1 + continue + if ch == '"': + in_string = True + i += 1 + continue + if text.startswith(needle, i): + return i + i += 1 + return -1 + + +def _parse_deepseek_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """DeepSeek R1 / V3 / V3.1. + + R1: ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>NAME\\n``\\`\\`\\`json\\n{...}\\n\\`\\`\\`<｜tool▁call▁end｜>...`` + V3.x: ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>NAME<｜tool▁sep｜>{json}<｜tool▁call▁end｜>...`` + + Mirrors llama.cpp's pre-autoparser ``common_chat_parse_deepseek_r1`` / + ``_v3_1`` handling; tolerates the 5 opener variants llama.cpp keeps. + """ + out: list[dict] = [] + begin = _DEEPSEEK_BEGIN_RE.search(content) + if not begin: + return out + scan_start = begin.end() + # Envelope end OUTSIDE JSON strings: an argument may contain the literal end token, + # and a raw find would truncate the call. + end_pos = _find_outside_json_strings(content, _DEEPSEEK_END, scan_start) + # Strict mode: an unclosed envelope is truncated; reject, don't heal to EOF. + if not allow_incomplete and end_pos < 0: + return out + scan_end = end_pos if end_pos >= 0 else len(content) + body = content[scan_start:scan_end] + + # R1 path first: ``function<｜tool▁sep｜>NAME\n```json\n{...}\n```<｜tool▁call▁end｜>``. + pos = 0 + while pos < len(body): + fpos = body.find(_DEEPSEEK_R1_FUNC_MARKER, pos) + if fpos < 0: + break + name_start = fpos + len(_DEEPSEEK_R1_FUNC_MARKER) + nl = body.find("\n", name_start) + if nl < 0: + break + if not body.startswith(_DEEPSEEK_R1_FENCE, nl): + pos = name_start + continue + name = body[name_start:nl].strip() + json_start = nl + len(_DEEPSEEK_R1_FENCE) + # Walk a balanced ``{`` even if the trailing fence is truncated. + if json_start >= len(body) or body[json_start] != "{": + pos = json_start + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + break + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + # The closing fence + <｜tool▁call▁end｜> must IMMEDIATELY follow the JSON, else an + # unbounded search lands on a LATER call's terminator. Absent close: heal past the + # JSON (strict rejects); later well-formed calls are still kept. + after = brace_end + 1 + while after < len(body) and body[after] in " \t\r\n": + after += 1 + close_m = _DEEPSEEK_R1_CLOSE_RE.match(body, after) + if not allow_incomplete and close_m is None: + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = close_m.end() if close_m else brace_end + 1 + if out: + return out + + # V3 / V3.1: name then bare JSON. Use ``str.find`` for the sep marker and walk + # back for the name (a ``[^\n<]+`` regex search is O(N^2) on truncated bodies). + pos = 0 + while pos < len(body): + sep_pos = body.find(_DEEPSEEK_SEP, pos) + if sep_pos < 0: + break + # Walk left from sep_pos to the name start; stop at ``\n`` (turn boundary), ``<`` + # (tag start), or ``>`` (end of an optional ``<｜tool▁call▁begin｜>``). + name_start = sep_pos + while name_start > pos and body[name_start - 1] not in "\n<>": + name_start -= 1 + name = body[name_start:sep_pos].strip() + json_start = sep_pos + len(_DEEPSEEK_SEP) + while json_start < len(body) and body[json_start] in " \t\n\r": + json_start += 1 + if json_start >= len(body) or body[json_start] != "{": + pos = sep_pos + len(_DEEPSEEK_SEP) + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + break + # Strict mode: a real V3 call closes with the per-call <｜tool▁call▁end｜>; without + # it the call is truncated/merged, so skip it but keep scanning for a later + # well-formed call (matches Kimi strict). + if not allow_incomplete: + after = brace_end + 1 + while after < len(body) and body[after] in " \t\r\n": + after += 1 + if not body.startswith(_DEEPSEEK_CALL_END, after): + pos = brace_end + 1 + continue + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + # Advance just past the JSON; seeking the optional <｜tool▁call▁end｜> could land on + # a LATER call's end marker and skip the call between. + pos = brace_end + 1 + return out + + +# ── GLM 4.5 / 4.6 / 4.7 ───────────────────────────────────────────── + + +def _parse_glm_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """GLM 4.5 / 4.6 / 4.7. + + ``NAME[\\n]K[\\n]V + ...``. Multi-call is back-to-back blocks, no envelope. + Mirrors llama.cpp's GLM 4.x tool-call handling (``common_chat_params_init_glm_4_5`` + plus its generalized XML-style parser, llama.cpp PRs #15904 / #16932). + """ + out: list[dict] = [] + pos = 0 + while pos < len(content): + m = _GLM_TC_OPEN_RE.search(content, pos) + if not m: + break + name = m.group(1).strip() + apos = m.end() # absolute position in ``content``; advances past each pair + + args: dict[str, Any] = {} + valid = True + close = -1 + # Walk arg pairs directly against ``content``: a value may contain a literal + # , so the real close is the before the next . + # ``str.find`` keeps this linear. + while True: + ks = content.find(_GLM_ARG_KEY_OPEN, apos) + tc = content.find(_GLM_TC_CLOSE, apos) + if tc >= 0 and (ks < 0 or tc < ks): + close = tc + break + if ks < 0: + break # no close and no more keys -- truncated body + ke = content.find(_GLM_ARG_KEY_CLOSE, ks + len(_GLM_ARG_KEY_OPEN)) + if ke < 0: + break + vstart = ke + len(_GLM_ARG_KEY_CLOSE) + while vstart < len(content) and content[vstart] in " \t\r\n": + vstart += 1 + if not content.startswith(_GLM_ARG_VAL_OPEN, vstart): + # Key without : strict rejects the call; Auto-Heal skips it. + if not allow_incomplete: + valid = False + apos = ke + len(_GLM_ARG_KEY_CLOSE) + continue + vs = vstart + len(_GLM_ARG_VAL_OPEN) + # A first-match find on would truncate values containing literal + # close tags and execute corrupted arguments. + ve = _glm_value_close(content, vs, strict = not allow_incomplete) + key = content[ks + len(_GLM_ARG_KEY_OPEN) : ke].strip() + if ve < 0: + # Unclosed : strict rejects the whole call; Auto-Heal keeps the + # partial value (a truncated query is not a no-arg call). + if not allow_incomplete: + valid = False + break + # Bound the healed value at the next structural tag, not EOF, so a value + # missing only its can't swallow the markup after it. + nk = content.find(_GLM_ARG_KEY_OPEN, vs) + tc = content.find(_GLM_TC_CLOSE, vs) + bounds = [b for b in (nk, tc) if b >= 0] + if not bounds: + args[key] = content[vs:].rstrip() + break + bound = min(bounds) + args[key] = content[vs:bound].rstrip() + apos = bound + continue + raw_val = content[vs:ve] + apos = ve + len(_GLM_ARG_VAL_CLOSE) + # Decode only unambiguous JSON literals; else keep the value RAW so whitespace + # in string args survives (matches vLLM glm4_moe). ``"`` is left out of the + # probe: a verbatim string's quotes are meaningful. + probe = raw_val.strip() + if ( + probe[:1] in "{[" + or probe in ("true", "false", "null") + or _GLM_JSON_NUMERIC_RE.fullmatch(probe) + ): + try: + args[key] = json.loads(probe) + continue + except (json.JSONDecodeError, ValueError): + pass + args[key] = raw_val + + # Strict mode: a block with no is truncated; reject it. + if not allow_incomplete and close < 0: + valid = False + + if name and valid: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = close + len(_GLM_TC_CLOSE) if close >= 0 else len(content) + return out + + +# ── Kimi K2 / Moonshot ────────────────────────────────────────────── + + +def _parse_kimi_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Kimi K2. + + ``<|tool_calls_section_begin|><|tool_call_begin|>functions.NAME:IDX + <|tool_call_argument_begin|>{json}<|tool_call_end|>... + <|tool_calls_section_end|>``. Full id is preserved on ``tool_calls + [i].id`` for round-trip through the chat template. Outer loop walks + every section in the stream (vLLM / SGLang parity); mirrors llama.cpp's + Kimi K2 handling via its generalized XML-style parser (llama.cpp PR #16932). + """ + out: list[dict] = [] + outer_pos = 0 + while True: + section_start = content.find(_KIMI_SECTION_BEGIN, outer_pos) + if section_start < 0: + break + scan_start = section_start + len(_KIMI_SECTION_BEGIN) + # Section end OUTSIDE JSON strings: an argument may contain the literal end token, + # and a raw find would drop the later valid call. + section_end = _find_outside_json_strings(content, _KIMI_SECTION_END, scan_start) + scan_end = section_end if section_end >= 0 else len(content) + body = content[scan_start:scan_end] + # Truncated tail: parse what we have, then exit. In strict mode a section with no + # <|tool_calls_section_end|> is truncated; reject it instead. + if section_end < 0: + if allow_incomplete: + out.extend( + _parse_kimi_section_body( + body, id_offset = id_offset + len(out), allow_incomplete = True + ) + ) + return out + outer_pos = section_end + len(_KIMI_SECTION_END) + out.extend( + _parse_kimi_section_body( + body, id_offset = id_offset + len(out), allow_incomplete = allow_incomplete + ) + ) + + # The section wrapper is optional (llama.cpp): a bare <|tool_call_begin|> call parses + # as one section when the loop matched nothing. + if not out and _KIMI_CALL_BEGIN in content: + out.extend( + _parse_kimi_section_body( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + ) + return out + + +def _parse_kimi_section_body( + body: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse one Kimi K2 section body (between begin / end markers).""" + out: list[dict] = [] + pos = 0 + while pos < len(body): + call_start = body.find(_KIMI_CALL_BEGIN, pos) + if call_start < 0: + break + id_start = call_start + len(_KIMI_CALL_BEGIN) + arg_begin = body.find(_KIMI_ARG_BEGIN, id_start) + if arg_begin < 0: + break + full_id = body[id_start:arg_begin].strip() + m = _KIMI_ID_RE.match(full_id) + if m: + # group(1) is the whole name; do NOT split on ``.`` -- a dotted MCP name stays intact. + name = m.group(1) + else: + base = full_id.split(":")[0] + name = base[len("functions.") :] if base.startswith("functions.") else base + # Drop bare-counter ids (``3``, ``42``) -- matches vLLM; SGLang infers the name + # from the tool schema, which we don't have here. + if name.isdigit(): + json_start = arg_begin + len(_KIMI_ARG_BEGIN) + brace_end = ( + _balanced_brace_end(body, json_start) + if (json_start < len(body) and body[json_start] == "{") + else None + ) + if brace_end is None: + pos = arg_begin + len(_KIMI_ARG_BEGIN) + else: + pos = brace_end + 1 + continue + json_start = arg_begin + len(_KIMI_ARG_BEGIN) + # Balanced brace lets a truncated trailing end marker still surface a call. + while json_start < len(body) and body[json_start] in " \t\n\r": + json_start += 1 + if json_start >= len(body) or body[json_start] != "{": + pos = arg_begin + len(_KIMI_ARG_BEGIN) + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + # Malformed / truncated JSON: skip this call but keep parsing later ones + # instead of dropping the rest of the section (vLLM recovers them). + nxt = body.find(_KIMI_CALL_BEGIN, json_start) + if nxt < 0: + break + pos = nxt + continue + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + if not allow_incomplete: + # Strict mode: this call must close with <|tool_call_end|> before the next + # <|tool_call_begin|>; otherwise it is truncated, so reject it. + end_marker = body.find(_KIMI_CALL_END, brace_end + 1) + next_call = body.find(_KIMI_CALL_BEGIN, brace_end + 1) + if end_marker < 0 or (next_call >= 0 and end_marker > next_call): + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": full_id or f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + # Advance past the JSON; seeking <|tool_call_end|> could skip a following call + # when this one's end marker is missing. + pos = brace_end + 1 return out diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1a1a934009..3341a9c628 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1150,7 +1150,13 @@ from core.inference.key_exchange import decrypt_api_key from core.inference.model_ids import public_model_id from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client -from core.inference.tool_call_parser import _strip_function_xml_calls, _strip_mistral_closed_calls +from core.inference.tool_call_parser import ( + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) +from core.inference.tool_call_parser import TOOL_XML_SIGNALS as _PARSER_TOOL_SIGNALS from core.inference.passthrough_healing import ( StreamToolCallHealer, heal_gate, @@ -1309,8 +1315,8 @@ async def artifact_preview_frame(allow_network: bool = False): ) -# Whitespace/escape-tolerant bare-JSON tool-template detector: matches pretty-printed and -# JSON-escaped ``{"name":`` plus the ``"function"`` alias. +# Whitespace/escape-tolerant bare-JSON tool-template detector (matches pretty-printed and +# JSON-escaped ``{"name":`` plus the ``"function"`` alias), mirroring the parser's tolerance. _BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:') @@ -1324,15 +1330,16 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: model_identifier = model_id, log_source = "safetensors", ) - # Markers the parser recognises; drop the pill if a template advertises tools but uses none. - # The bare-JSON ``{"name":`` form is matched whitespace-tolerantly below. + # Markers any supported parser recognises (template advertises tools but + # uses none -> drop the pill). Reuse the parser's own signal list so this + # gate never drifts (a hand-maintained copy lost the DeepSeek variants); + # ```` is GLM's unique signal, absent from the shared set. The + # bare-JSON ``{"name":`` form is matched below with the whitespace/escape- + # tolerant ``_BARE_JSON_NAME_MARKER_RE`` so pretty-printed or escaped + # templates are not mis-classified as tool-less. _PARSER_MARKERS = ( - "", - "", - "[TOOL_CALLS]", - "<|tool_call>", + *_PARSER_TOOL_SIGNALS, + "", ) if ( flags.get("supports_tools") @@ -1365,7 +1372,12 @@ def _sf_reasoning_prefill_mode( template: Optional[str] = None, reasoning_effort: Optional[str] = None, ) -> bool: - """Whether this request begins inside an unclosed ```` (Qwen3/GLM prefill it). Gated on the standard markers; bespoke channels, gpt-oss, and thinking-disabled requests are excluded. ``enable_thinking=None`` defaults ON.""" + """Whether this request begins INSIDE an unclosed ```` (Qwen3/Qwen3.5/GLM prefill it). + + Gated on the STANDARD ````/```` markers: a bespoke reasoning channel (e.g. gemma) + never emits ````, so prefilled mode would swallow the whole answer -- excluded, as are + gpt-oss and thinking-disabled requests. ``enable_thinking=None`` defaults ON, so plain requests prefill. + """ if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"): return False tpl = template or "" @@ -1377,8 +1389,11 @@ def _sf_reasoning_prefill_mode( return False if enable_thinking is False: return False - # reasoning_effort="none" disables thinking on enable_thinking_effort (GLM-5.2) models like - # enable_thinking=False; without this the answer is swallowed into empty reasoning_content. + # A reasoning_effort="none" request disables thinking for enable_thinking_effort + # (GLM-5.2) models the same way enable_thinking=False does (see + # ``_request_reasoning_kwargs``). Without this, the model emits no ```` and + # a plain answer is swallowed whole into reasoning_content, leaving the visible + # response empty. if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none": return False return True @@ -1654,41 +1669,83 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str: return nudge + " " + _RAG_GROUNDING_NUDGE -# Strip leaked tool-call markup: every shared-parser format plus the leak shapes -# ``llama_cpp.py``'s speculative buffer splits across the visible/DRAIN boundary. Mistral -# ``[TOOL_CALLS]`` uses the parser's balanced-brace helper (``\{.*?\}`` would truncate nested JSON). +# Strip leaked tool-call markup: every shared-parser format plus the four leak +# shapes llama_cpp.py's speculative buffer splits across the visible/DRAIN +# boundary. Mistral [TOOL_CALLS] uses the parser's balanced-brace helper (a +# non-greedy regex would truncate nested JSON); the DeepSeek opener alternation +# is the parser's own, so a signal we parse is never left un-stripped. +from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC + _TOOL_XML_RE = _re.compile( - # Hyphen in the name char-class matches MCP tool names with dashes - # (mcp__srv__list-issues) that would otherwise leak past this strip. - # The ``<|python_tag|>`` arm runs to the next REAL Llama sentinel or EOF, so a literal - # ``<|...|>`` token in an argument (e.g. ``<|cite|>``) doesn't truncate the strip. - # ```` plus the ```` attribute form; name class mirrors the parser. - # A CLOSED ``...`` extends to the last ```` before the next - # opener (so a literal ```` in a value can't truncate); this arm runs first. + # Arm order/notes: the closed ```` arm runs first and extends + # to the call's REAL close so a literal ```` in a value does not + # leak the tail; the combined arm still catches ```` and orphan + # tails. The python_tag arm bounds only on REAL Llama control sentinels + # (stopping at any ``<|`` truncated on literal ``<|x|>`` tokens in values). + # The last arms cover DeepSeek envelopes (all opener variants), Kimi section + # blocks, and bare Kimi calls. Name class ``[\w.\-]`` mirrors the parser. + # Those three arms carry a call-shaped lookahead (matching the parser's + # ``_TOOL_ALL_PATS``): a prose answer that merely mentions a marker + # (``See <|tool_call_begin|> in the docs``) is only stripped when a real + # call actually follows the marker, or the marker is a bare fragment at EOF. r'(?:(?!).)*' r'|<(?:tool_call|function(?:=[\w.\-]+|\s+name="[\w.\-]+"))>.*?(?:|\Z)' r"|<\|tool_call>.*?(?:|\Z)" r"|" r"|" r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*" - # ```` is the attribute-form alias of ````; strip a tail-only orphan. + r"|" + + _DS_OPEN_SRC + + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*?(?:<｜tool▁calls▁end｜>|\Z)" + r"|<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*?(?:<\|tool_calls_section_end\|>|\Z)" + r"|<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*?(?:<\|tool_call_end\|>|\Z)" + # ```` is the attribute-form alias of ```` (the parser accepts + # both); strip a tail-only orphan close of either spelling. r"|\s*\Z", _re.DOTALL, ) -def _strip_tool_xml(text: str) -> str: - """Mistral balanced-brace helper + guarded function-XML scan + ``_TOOL_XML_RE`` (skips openers inside an open ````).""" - return _TOOL_XML_RE.sub( - "", _strip_function_xml_calls(_strip_mistral_closed_calls(text), final = True) +def _gemma_strip_gate(tools) -> set: + """Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the + parser/loop gate: only an enabled ``call:foo{...}`` is a call). With NO tools + enabled this returns an EMPTY set, not ``None``: every ``call:NAME{...}`` is + then prose, and ``None`` would strip-all and delete a legitimate answer.""" + names = { + (t.get("function") or {}).get("name") + for t in (tools or []) + if isinstance(t, dict) and isinstance(t.get("function"), dict) + } + names.discard(None) + return names + + +def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Combine the parser's scan-based strips (Mistral balanced-brace, gated + Gemma wrapper-less, GLM real-close, guarded function-XML) with + ``_TOOL_XML_RE`` -- the scan strips close at each call's REAL terminator so + literal markup inside argument values is data, not a leaked tail. + ``enabled_tool_names`` gates the Gemma strip; ``None`` strips every closed call.""" + cleaned = _strip_glm_calls( + _strip_gemma_wrapperless_calls(_strip_mistral_closed_calls(text), enabled_tool_names), + final = True, ) + cleaned = _strip_function_xml_calls(cleaned, final = True) + return _TOOL_XML_RE.sub("", cleaned) -def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str: - """Route-level tool-call leak cleanup (Auto-Heal only) via ``_strip_tool_xml``.""" +def _strip_tool_xml_for_display( + text: str, + *, + auto_heal_tool_calls: bool, + enabled_tool_names: Optional[set] = None, +) -> str: + """Route-level leak cleanup (Auto-Heal only). Delegates to ``_strip_tool_xml`` + so the Mistral balanced-brace pass runs too (``_TOOL_XML_RE`` alone has no + ``[TOOL_CALLS]`` arm). ``enabled_tool_names`` gates the Gemma strip.""" if not auto_heal_tool_calls: return text - return _strip_tool_xml(text) + return _strip_tool_xml(text, enabled_tool_names) logger = get_logger(__name__) @@ -5960,6 +6017,7 @@ async def openai_chat_completions( _msg["content"] = _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(tools_to_use), ).strip() def gguf_generate_with_tools(): @@ -6093,6 +6151,7 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(tools_to_use), ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative @@ -6199,6 +6258,7 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(tools_to_use), ) return full_text, usage, finish finally: @@ -6572,7 +6632,7 @@ async def openai_chat_completions( _sf_features = _detect_safetensors_features(backend, _sf_tpl) # Split prefilled-```` output into reasoning_content deltas (GGUF parity) so the UI - # renders the thinking block for safetensors and MLX. + # renders the thinking block for safetensors AND MLX. _sf_parse_think = bool( _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on") ) @@ -6673,6 +6733,7 @@ async def openai_chat_completions( "content": _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), ).strip(), } ) @@ -6731,7 +6792,7 @@ async def openai_chat_completions( reasoning_extractor = _new_sf_reasoning_extractor() def _sf_flush_reasoning(): - # Drain the extractor at a turn boundary / stream end; only visible text reaches the monitor. + # Drain the extractor at a turn boundary / stream end (GGUF parity); only visible text reaches the monitor. fr, fv = reasoning_extractor.finish() out = [] if fr: @@ -6757,7 +6818,7 @@ async def openai_chat_completions( if event["type"] == "status": if not event["text"]: - # Turn boundary: flush reasoning, then start a fresh extractor. + # Iteration boundary: flush reasoning, then start a fresh extractor for the next turn. for _c in _sf_flush_reasoning(): yield _c prev_text = "" @@ -6773,7 +6834,7 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": - # Flush reasoning before tool_start so the thinking block closes ahead of the tool card. + # Flush reasoning before the tool_start line so the thinking block closes ahead of the tool card. for _c in _sf_flush_reasoning(): yield _c prev_text = "" @@ -6786,6 +6847,7 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative @@ -6877,11 +6939,12 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), ) return full_text content_text = await asyncio.to_thread(_drain_to_text) - # Split prefilled reasoning from the visible answer; monitor gets visible text only. + # Split prefilled reasoning out of the visible answer (GGUF parity); monitor gets visible text only. _reasoning_text, _visible_text = _extract_responses_reasoning( content_text, parse_think_markers = _sf_parse_think, @@ -6980,7 +7043,7 @@ async def openai_chat_completions( yield _chat_role_chunk(completion_id, created, model_name) prev_text = "" - # Split prefilled into reasoning_content deltas. Single turn (no per-turn reset); also MLX. + # Split prefilled into reasoning_content deltas (GGUF parity). Single turn (no per-turn reset); also serves MLX. reasoning_extractor = _new_sf_reasoning_extractor() # Run the sync generator in a thread pool to avoid blocking the # event loop. Critical for compare mode: two SSE requests arrive @@ -7087,7 +7150,7 @@ async def openai_chat_completions( for token in generate(): full_text = token - # Split prefilled reasoning from the visible answer; also covers MLX. + # Split prefilled reasoning from the visible answer (GGUF parity); also covers MLX. _reasoning_text, _visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _sf_parse_think, @@ -7937,8 +8000,8 @@ class _ResponsesReasoningExtractor: reasoning_prefilled: bool = False, ) -> None: self._buffer = "" - # ``reasoning_prefilled``: output begins inside an unclosed ```` (Qwen3/GLM prefill), - # so start in reasoning to capture leading text until the first ````. + # ``reasoning_prefilled``: output begins INSIDE an unclosed ```` (Qwen3/GLM prefill), + # so start in reasoning to capture leading text until the first ````. Callers default False. self._in_reasoning = reasoning_prefilled # Splitting requires marker parsing; a prefilled open implies it. self._parse_think_markers = parse_think_markers or reasoning_prefilled @@ -7970,7 +8033,7 @@ class _ResponsesReasoningExtractor: self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] self._in_reasoning = False continue - # Hold back a trailing partial of either marker: the close (clean chunk-boundary split) + # Hold back a trailing partial of EITHER marker: the close (clean chunk-boundary split) # and a stray open (so a re-emitted ```` isn't leaked into the reasoning drawer). keep = _responses_marker_holdback( self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN) @@ -9859,7 +9922,9 @@ async def anthropic_messages( # Strip stale tool-call XML from conversation for _msg in openai_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): - _msg["content"] = _strip_tool_xml(_msg["content"]).strip() + _msg["content"] = _strip_tool_xml( + _msg["content"], _gemma_strip_gate(openai_tools) + ).strip() def _run_tool_gen(): return llama_backend.generate_chat_completion_with_tools( @@ -9904,6 +9969,7 @@ async def anthropic_messages( message_id, model_name, disable_parallel_tool_use = _disable_parallel, + openai_tools = openai_tools, ) ) @@ -10010,7 +10076,7 @@ async def _anthropic_tool_stream( # content event that was purely tool XML doesn't count as text. if etype == "content": event = dict(event) - event["text"] = _strip_tool_xml(event["text"]) + event["text"] = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools)) # disable_parallel_tool_use: keep only the first tool_use block, # dropping every later tool_start and its paired tool_end (robust # to empty tool-call ids — tracked by state, not id matching). @@ -10165,6 +10231,7 @@ async def _anthropic_tool_non_streaming( message_id, model_name, disable_parallel_tool_use = False, + openai_tools = None, ): """Non-streaming response for the tool-calling path. @@ -10193,7 +10260,7 @@ async def _anthropic_tool_non_streaming( etype = event.get("type", "") if etype == "content": # Strip leaked tool-call XML - clean = _strip_tool_xml(event["text"]) + clean = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools)) new = clean[len(prev_text) :] prev_text = clean if new: @@ -10662,11 +10729,14 @@ async def _anthropic_passthrough_non_streaming( else: text = message.get("content") or "" if text: - # Keep unpromoted bytes when healing is active; legacy stripping is only for opted-out - # or no-client-tool requests. _strip_tool_xml also cleans Mistral [TOOL_CALLS] and - # guarded function-XML, not just _TOOL_XML_RE. + # Keep unpromoted bytes when healing is active; legacy stripping is + # only for opted-out or no-client-tool requests. Use the full + # _strip_tool_xml pass so Mistral [TOOL_CALLS] and guarded + # function-XML leaks are cleaned too, not just _TOOL_XML_RE forms, + # with the Gemma display gate so a disabled/example call:NAME{...} + # in prose survives. if not healing_active: - text = _strip_tool_xml(text) + text = _strip_tool_xml(text, _gemma_strip_gate(openai_tools)) text = text.strip() if text: content_blocks.append(AnthropicResponseTextBlock(text = text)) diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index fff6b240c5..7b653f47aa 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -41,7 +41,7 @@ def test_normal_multi_key_arguments_still_split(): def test_empty_bare_value_becomes_empty_string_not_dropped(): - # An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON). + # An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON and dropped the call). calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}") assert len(calls) == 1, calls assert _args(calls[0]) == {"query": "", "unit": "celsius"} @@ -60,6 +60,15 @@ def test_bare_value_with_timestamps_after_comma_is_kept(): assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"} +def test_wrapperless_bare_value_with_timestamps_after_comma_is_kept(): + # The wrapper-less Gemma form (no <|tool_call> markers) goes through the + # _gemma_parse_stripped_body scanner and its _GEMMA_KEY_RE. + calls = parse_tool_calls_from_text("call:web_search{query:meet at 10:00, 11:00 tomorrow}") + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "web_search" + assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow"} + + def test_marker_inside_json_argument_is_not_a_second_call(): content = ( '{"name":"python","arguments":{"code":' @@ -97,8 +106,8 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call(): def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call(): - # The outer object fails to normalize, but the nested marker is covered by - # its span; safe outcome is no executed call at all. + # An UNQUOTED Gemma value containing a literal marker: the marker is nested in the outer + # candidate span, so it must not be promoted to a standalone `terminal` call (no tool call). content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}}" calls = parse_tool_calls_from_text(content) assert "terminal" not in [c["function"]["name"] for c in calls], calls @@ -151,6 +160,43 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): assert [c["function"]["name"] for c in calls] == ["python"], calls +def test_wrapperless_nested_object_argument_is_parsed(): + # skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare. + calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}") + assert len(calls) == 1 + assert _args(calls[0]) == {"loc": {"city": "NYC"}, "n": 3} + + +def test_wrapperless_array_argument_is_parsed(): + calls = parse_tool_calls_from_text("call:label{labels:[bug,ui],n:2}") + assert len(calls) == 1 + assert _args(calls[0]) == {"labels": ["bug", "ui"], "n": 2} + + +def test_wrapperless_deeply_nested_object_and_array_are_preserved(): + # The single-pass parser must keep multi-level nesting (objects inside + # objects, arrays inside arrays) intact, not flatten or drop it. + calls = parse_tool_calls_from_text( + "call:f{loc:{city:NYC,geo:{lat:1,lng:2}},tags:[a,b,[c,d]],n:3}" + ) + assert len(calls) == 1 + assert _args(calls[0]) == { + "loc": {"city": "NYC", "geo": {"lat": 1, "lng": 2}}, + "tags": ["a", "b", ["c", "d"]], + "n": 3, + } + + +def test_gemma_parse_array_advances_on_stray_brace(): + # Regression: a stray '}' / ']' / ',' where an array element is expected must + # not stall _gemma_parse_value at the same index (it looped forever before). + from core.inference.tool_call_parser import _gemma_parse_array + + items, end, closed = _gemma_parse_array("[a,}]", 0) + assert end == 5 and closed is True # consumed through the closing ']' + assert items[0] == "a" + + def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): # Parse keeps the quoted close marker as data; strip removes the whole span. text = '<|tool_call>call:python{code:<|"|>print("")<|"|>}' @@ -312,16 +358,17 @@ def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered def test_gemma_parse_value_always_advances_on_stray_delimiter(): # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the - # parser, or a looping caller spins forever (DoS). + # index by at least one, or a caller looping on it spins forever at 100% CPU (DoS). for delim in (",", "}", "]"): text = delim + "rest" - value, nxt = _gemma_parse_value(text, 0) + value, nxt, _explicit = _gemma_parse_value(text, 0) assert nxt > 0, (delim, value, nxt) def test_malformed_gemma_array_does_not_hang(): - # ``[},]`` (stray ``}`` in a list body) hung the buggy parser; the timeout fails - # the regression loudly instead of blocking CI forever. + # ``[},]`` puts a stray ``}`` at the primitive position inside a list body. + # On the buggy parser this hangs the server; guard with a wall-clock timeout + # so the regression fails loudly instead of blocking CI forever. import threading result: dict = {} diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 8977d6e92a..dcc759a210 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -1040,7 +1040,7 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch): def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): """No-tool re-prompt attempts should not concatenate into the UI.""" - # One initial response plus one stream per re-prompt (count from the shared cap). + # One initial response plus one stream per re-prompt; derive the count from the shared cap. streams = [[_sse({"content": "I will use render_html now."}), _done()]] streams += [ [_sse({"content": "Understood. I will use render_html now."}), _done()] @@ -1207,8 +1207,8 @@ def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatc def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch): - # Inline Mistral ``[TOOL_CALLS]`` after a visible preface: the DRAINING flush must use the - # shared parser patterns (the legacy set leaked the marker to clients). + # Textual Mistral ``[TOOL_CALLS]`` inline with visible preface: the DRAINING flush must use the + # shared parser patterns (which know ``[TOOL_CALLS]``); the legacy set leaked the marker to clients. streams = [ [_sse({"content": 'Let me search. [TOOL_CALLS]web_search{"query":"cats"}'}), _done()], [_sse({"content": "done"}), _done()], @@ -1836,6 +1836,7 @@ def test_bare_json_tool_call_streamed_is_not_leaked_and_executes(monkeypatch): ) ) + # The tool ran with the parsed arguments. assert calls == [("web_search", {"query": "weather in Sydney"})] assert any( event.get("type") == "tool_end" and event.get("tool_name") == "web_search" @@ -1903,6 +1904,37 @@ def test_incomplete_bare_json_truncation_is_not_leaked(monkeypatch): assert all('{"name"' not in t for t in content_texts), content_texts +def test_gguf_truncated_ordinary_json_with_name_key_is_shown_not_suppressed(monkeypatch): + """A truncated markerless object whose "name" is NOT an enabled tool (a person + record cut off mid-stream, ``{"name":"Alice","age":``) must still be shown. The + end-of-stream ``_is_bare_tc`` heuristic routed any ``{...,"name",...}`` fragment + to DRAINING (dropped); it is now gated on the enabled tool names so only a real + truncated tool call is suppressed, ordinary JSON streams through.""" + + truncated = '{"name": "Alice", "age": 30, "bio": "loves ' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "start a person record"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + def test_gguf_truncated_disabled_name_json_is_preserved_when_tools_active(monkeypatch): """A truncated JSON answer with a non-enabled name must still be shown (resolvers are gated on enabled names).""" @@ -1987,6 +2019,40 @@ def test_gguf_oversized_disabled_name_json_is_preserved(monkeypatch): assert any("Alice" in t for t in content_texts), content_texts[:1] +def test_gemma_wrapperless_call_streamed_is_not_leaked_and_executes(monkeypatch): + """Gemma 4 GGUF (skip_special_tokens) streams a wrapper-less ``call:NAME{..}`` + with no XML signal. Like bare JSON, the BUFFERING scan must recognise it via + _GEMMA_BARE_TC_RE, drain it silently, and execute the tool -- never leaking + the ``call:`` markup to the user-visible stream.""" + + gemma_call = 'call:web_search{query:"weather in Sydney"}' + first_stream = _streamed_content(gemma_call) + final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Weather: sunny, 22C." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather in Sydney?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "weather in Sydney"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("call:" not in t for t in content_texts), content_texts + assert any("sunny in Sydney" in t for t in content_texts), content_texts + + def _usage_done(usage: dict, finish_reason: str = "stop") -> str: """A terminal SSE chunk carrying llama-server's ``usage`` block, the way the real server reports it on the final chunk of a completion.""" @@ -2124,6 +2190,66 @@ def test_gguf_bare_json_call_not_replayed_in_next_turn_content(monkeypatch): assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst +def test_gguf_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch): + """A single textual-fallback turn that parses many DISTINCT tool calls must be + capped at _MAX_TOOL_CALLS_PER_TURN (structured delta.tool_calls are grammar + bounded by llama-server; text parsed from content is not). Mirrors the + safetensors loop so one runaway turn cannot fan out into dozens of executions.""" + from core.inference.llama_cpp import _MAX_TOOL_CALLS_PER_TURN + + n = _MAX_TOOL_CALLS_PER_TURN + 4 + blocks = "".join( + '{"name":"t%d","arguments":{"i":%d}}' % (i, i) for i in range(n) + ) + first_stream = [_sse({"content": blocks}), _done()] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": f"t{i}"}} for i in range(n)], + max_tool_iterations = 1, + ) + ) + + assert len(calls) == _MAX_TOOL_CALLS_PER_TURN, [c[0] for c in calls] + # The cap keeps the first calls in order (no reordering / drop of leading ones). + assert [c[0] for c in calls] == [f"t{i}" for i in range(_MAX_TOOL_CALLS_PER_TURN)] + + +def test_gguf_textual_fallback_collapses_duplicate_tool_calls(monkeypatch): + """Exact-duplicate textual calls in one turn collapse to a single execution.""" + blocks = '{"name":"web_search","arguments":{"query":"cats"}}' * 5 + first_stream = [_sse({"content": blocks}), _done()] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert len(calls) == 1, [c[0] for c in calls] + + def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(monkeypatch): """Auto-Heal OFF keeps a truncated enabled-name fragment visible; ON suppresses it (strip gated on auto_heal_tool_calls).""" @@ -2159,8 +2285,8 @@ def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disable def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch): """Re-prompt slots must not extend the tool budget: stop after ``max_tool_iterations`` executed rounds.""" - # More tool-call streams than the budget: leaked re-prompt slots would run 2+3=5 rounds; - # honouring the budget stops after 2, then a tool-less final-answer pass. + # More tool-call streams than the budget: if re-prompt slots leaked into the budget (the bug) the + # loop would run 2+3=5 rounds; honouring it stops after 2, then a tool-less final-answer pass. streams = [ _structured_tool_call("web_search", {"query": f"q{i}"}, f"call_{i}") for i in range(6) ] diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 12239e7113..6d26d075cf 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -587,10 +587,12 @@ def test_tool_xml_strip_handles_hyphenated_function_names(): import re as _re from pathlib import Path + from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC + src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text() m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL) assert m, "could not extract _TOOL_XML_RE" - ns: dict = {"_re": _re} + ns: dict = {"_re": _re, "_DS_OPEN_SRC": _DS_OPEN_SRC} exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns) rx = ns["_TOOL_XML_RE"] stripped = rx.sub( diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 9871965ce8..ac4088fb25 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -100,6 +100,32 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): ] assert backend._is_vlm is False assert isinstance(backend._tokenizer, _DummyTokenizer) + # Non-LoRA text model: no base_model on the record. + assert backend.models["fake/text"]["base_model"] is None + + +def test_mlx_text_lora_record_keeps_base_model_for_native_template(monkeypatch): + # A LoRA adapter's own tokenizer often ships no chat template; the native tool-calling template + # lives on the base model. + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + config = SimpleNamespace( + identifier = "fake/text-adapter", + is_vision = False, + is_lora = True, + base_model = "fake/text-base", + ) + + assert backend.load_model(config, max_seq_length = 4096, hf_token = "hf-token") + + record = backend.models["fake/text-adapter"] + assert record["is_lora"] is True + assert record["base_model"] == "fake/text-base" def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite( @@ -188,12 +214,12 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): _install_fake_mlx(monkeypatch) from core.inference.mlx_inference import MLXInferenceBackend - captured = {} + # The text path renders once with tools, then the native-template fallback makes a second no- + # tools probe call (tools=None) to detect whether the template dropped the schema. + captured_calls = [] def _fake_apply(tokenizer, messages, **kwargs): - captured["tokenizer"] = tokenizer - captured["messages"] = messages - captured["kwargs"] = kwargs + captured_calls.append({"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs}) return "" monkeypatch.setattr( @@ -248,8 +274,15 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): ) ) assert out == ["hi"] - # The toggled kwargs must reach the chat-template helper. - assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}] - assert captured["kwargs"]["enable_thinking"] is True - assert captured["kwargs"]["reasoning_effort"] == "medium" - assert captured["kwargs"]["preserve_thinking"] is True + # The toggled kwargs must reach the chat-template helper on the real render + # (one of the calls carries the tools; the fallback probe passes tools=None). + tool_renders = [ + c + for c in captured_calls + if c["kwargs"].get("tools") == [{"function": {"name": "web_search"}}] + ] + assert tool_renders, captured_calls + render = tool_renders[0] + assert render["kwargs"]["enable_thinking"] is True + assert render["kwargs"]["reasoning_effort"] == "medium" + assert render["kwargs"]["preserve_thinking"] is True diff --git a/studio/backend/tests/test_native_template_trust_remote_code.py b/studio/backend/tests/test_native_template_trust_remote_code.py new file mode 100644 index 0000000000..60dc80f64c --- /dev/null +++ b/studio/backend/tests/test_native_template_trust_remote_code.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for trust_remote_code in the native-template fallback. + +``render_native_template`` re-fetches a model's native chat template from its +repo when an Unsloth override template (mistral, gemma-4) dropped the tools +schema. For a model loaded with ``trust_remote_code=True`` whose tokenizer repo +carries custom code, the secondary ``AutoTokenizer.from_pretrained`` must re-use +that same consent or transformers raises (it requires ``trust_remote_code`` to +instantiate a custom tokenizer class), the ``except`` swallows it, and the +request silently keeps the tool-dropping prompt even though the user already +consented to remote code for the model load. + +These tests pin that the stored ``trust_remote_code`` is threaded to the reload, +that the reload is skipped (returns ``None`` without executing code) when no +consent is stored, and that both backend ``model_info`` dicts persist the flag at +load time so the read lands on a value ``load_model`` actually set. +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# ``chat_template_helpers`` is dependency-light (copy / logging / typing, with the +# transformers import deferred inside the function). Load it directly so the test +# runs without importing the heavy ``core.inference`` package (unsloth / torch). +_HELPERS_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_template_helpers.py" +_spec = importlib.util.spec_from_file_location("_native_tpl_trc_test", _HELPERS_PATH) +chat_template_helpers = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(chat_template_helpers) + +render_native_template = chat_template_helpers.render_native_template + + +# A native template that emits a tools section only when tools are provided, so the +# with-tools vs no-tools render differs and ``render_native_template`` accepts it. +_NATIVE_TEMPLATE = ( + "{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n{% endfor %}" + "{% if tools %}[AVAILABLE_TOOLS]{{ tools }}[/AVAILABLE_TOOLS]\n{% endif %}" + "{% if add_generation_prompt %}assistant:{% endif %}" +) + +_MESSAGES = [{"role": "user", "content": "what is the weather"}] +_TOOLS = [{"type": "function", "function": {"name": "get_weather"}}] + + +class _JinjaTokenizer: + """Minimal tokenizer whose ``apply_chat_template`` renders ``self.chat_template``. + + Stands in for the live model tokenizer that ``render_native_template`` shallow- + copies and re-points at the native template before rendering. + """ + + def __init__(self, chat_template): + self.chat_template = chat_template + + def apply_chat_template( + self, + messages, + tokenize = False, + add_generation_prompt = True, + tools = None, + **kwargs, + ): + from jinja2 import BaseLoader, Environment + env = Environment(loader = BaseLoader()) + return env.from_string(self.chat_template).render( + messages = messages, + tools = tools, + add_generation_prompt = add_generation_prompt, + ) + + +def _install_custom_code_tokenizer(monkeypatch): + """Patch ``AutoTokenizer.from_pretrained`` to mimic a custom-code repo: raise + unless ``trust_remote_code`` is truthy, else return a tokenizer carrying the + native template. Records the ``trust_remote_code`` it was called with.""" + pytest.importorskip("jinja2") + from transformers import AutoTokenizer + + calls = {} + + def fake_from_pretrained( + model_id, + *args, + trust_remote_code = False, + token = None, + **kwargs, + ): + calls["trust_remote_code"] = trust_remote_code + calls["model_id"] = model_id + calls["token"] = token + if not trust_remote_code: + # Mirrors transformers.dynamic_module_utils.resolve_trust_remote_code: + # has_remote_code and not has_local_code and not trust_remote_code -> ValueError. + raise ValueError( + f"The repository {model_id} contains custom code which must be executed " + "to correctly load the model. Please pass the argument " + "`trust_remote_code=True` to allow custom code to be run." + ) + return _JinjaTokenizer(_NATIVE_TEMPLATE) + + monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained)) + return calls + + +def _model_info(trust_remote_code): + return { + "native_chat_template": None, # force the repo reload path + "base_model": None, # non-LoRA: template_source == active_model_name + "trust_remote_code": trust_remote_code, + # Live tokenizer that gets shallow-copied + re-pointed at the native template. + "tokenizer": _JinjaTokenizer("OVERRIDE-THAT-DROPS-TOOLS"), + } + + +def test_native_reload_passes_stored_trust_remote_code(monkeypatch): + """With ``trust_remote_code`` stored on ``model_info`` the custom-code reload + succeeds and the tools-advertising native prompt is returned. This FAILS before + the fix (reload omits the flag, raises, is swallowed, returns None).""" + calls = _install_custom_code_tokenizer(monkeypatch) + model_info = _model_info(trust_remote_code = True) + + out = render_native_template( + model_info = model_info, + active_model_name = "acme/custom-tokenizer-model", + messages = _MESSAGES, + tools = _TOOLS, + ) + + assert out is not None, "native fallback should render the tools prompt with consent" + assert "[AVAILABLE_TOOLS]" in out + assert "get_weather" in out + assert calls["trust_remote_code"] is True # the stored consent was threaded through + # A successful fetch is cached so the next tool turn skips the reload. + assert model_info["native_chat_template"] == _NATIVE_TEMPLATE + + +def test_native_reload_without_consent_returns_none(monkeypatch): + """Without stored consent the custom-code reload raises, is swallowed, and + ``render_native_template`` returns None (no unconsented code execution). Proves + the stored flag -- not a hard-coded True -- drives the reload.""" + calls = _install_custom_code_tokenizer(monkeypatch) + model_info = _model_info(trust_remote_code = False) + + out = render_native_template( + model_info = model_info, + active_model_name = "acme/custom-tokenizer-model", + messages = _MESSAGES, + tools = _TOOLS, + ) + + assert out is None + assert calls["trust_remote_code"] is False + # A failed fetch must not be cached as "no template" (would pin the tool drop). + assert model_info["native_chat_template"] is None + + +def test_backend_model_info_persists_trust_remote_code(): + """Both backends must store ``trust_remote_code`` on their per-model info dict so + ``render_native_template`` can source the consent value. Guards against the read + landing on a key ``load_model`` never sets (which would silently no-op the fix).""" + inf = (Path(_BACKEND_DIR) / "core" / "inference" / "inference.py").read_text() + mlx = (Path(_BACKEND_DIR) / "core" / "inference" / "mlx_inference.py").read_text() + assert '"trust_remote_code": trust_remote_code,' in inf + assert '"trust_remote_code": trust_remote_code,' in mlx diff --git a/studio/backend/tests/test_pr5624_regressions.py b/studio/backend/tests/test_pr5624_regressions.py new file mode 100644 index 0000000000..4f5471675c --- /dev/null +++ b/studio/backend/tests/test_pr5624_regressions.py @@ -0,0 +1,1011 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Regression tests for PR #5624 (DeepSeek R1/V3.x, GLM 4.x, Kimi K2 tool +parsing). Each test pins a specific edge case surfaced during the +review: + +* GLM string-vs-JSON-encoded value coercion (template emits strings + raw and non-strings JSON-encoded; the parser must not coerce a + bare string ``"42"`` into ``42``). +* GLM ```` containing a literal ``<`` (e.g. ``if x < 10``). +* Kimi K2 dotted name ``functions.my.tool:0`` keeps its full name + (``my.tool``) after stripping only the ``functions.`` prefix and + ``:idx`` suffix, while the full id is preserved on the call. +* Kimi K2 bare-counter id (no ``functions.`` prefix, no ``:IDX``) is + dropped rather than surfaced under a numeric name. +* DeepSeek V3.1 truncated mid-stream produces an empty result without + raising. +* ``routes.inference._strip_tool_xml`` strips the DeepSeek envelope and + the Kimi section markers added by this PR. +""" + +import json + +import pytest + +from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_tool_markup, +) + + +# GLM string-vs-JSON-encoded value coercion (finding B in plan) + + +@pytest.mark.parametrize( + "raw_val, expected_python", + [ + # Bare numeric / bool / null shapes are still treated as JSON + # literals (ambiguous with strings; the template doesn't tell us). + ("42", 42), + ("true", True), + ("false", False), + ("null", None), + ("3.14", 3.14), + ("-7", -7), + ("1e3", 1000.0), + ], +) +def test_glm_numeric_and_bool_literals_are_json_decoded(raw_val, expected_python): + text = ( + "n\n" + f"v\n" + f"{raw_val}\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["v"] == expected_python + + +@pytest.mark.parametrize( + "raw_val", + [ + "hello world", # plain prose + "True", # Python literal, NOT JSON -- no longer eaten by ast.literal_eval + "None", # Python literal, NOT JSON -- no longer eaten by ast.literal_eval + "if x < 10: pass", # code with literal < (well, < not in arg_value here) + "{not valid json", # looks like an object but is malformed -- must stay raw + "[oops", # looks like an array but is malformed + ], +) +def test_glm_non_json_shapes_stay_raw(raw_val): + text = ( + "n\n" + f"v\n" + f"{raw_val}\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["v"] == raw_val + assert isinstance(args["v"], str) + + +def test_glm_json_object_arg_decoded(): + text = ( + "nest\n" + "opts\n" + '{"limit": 10}\n' + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["opts"] == {"limit": 10} + + +def test_glm_json_array_arg_decoded(): + text = ( + "nest\n" + "ids\n" + "[1, 2, 3]\n" + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["ids"] == [1, 2, 3] + + +def test_glm_arg_value_with_literal_less_than(): + text = ( + "run\n" + "code\n" + "if x < 10: pass\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "if x < 10: pass" + + +# GLM 4.7 no-newline emission shape + + +def test_glm_4_7_no_newlines_between_name_and_arg_key(): + """GLM 4.7 strips the ``\\n`` after the name (``{{- ... -}}`` in the + template) so ```` follows directly. Parser must accept both.""" + text = ( + "get_weather" + "cityLondon" + "unitscelsius" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"city": "London", "units": "celsius"} + + +def test_glm_4_7_no_newlines_multi_call(): + """Back-to-back GLM 4.7 calls without intervening newlines.""" + text = ( + "ax1" + "by2" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 2 + assert calls[0]["function"]["name"] == "a" + assert calls[1]["function"]["name"] == "b" + + +def test_glm_4_7_does_not_break_qwen_path(): + """Qwen ``{json}`` still dispatches to Qwen; GLM's + first-char ``[^\\n<{]`` excludes ``{``.""" + text = '{"name":"web_search","arguments":{"q":"x"}}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + +# Kimi K2 dotted name + bare counter (finding C in plan) + + +def test_kimi_dotted_namespace_keeps_full_dotted_name(): + # A dotted Kimi id keeps its FULL name; only the ``functions.`` prefix and ``:idx`` suffix drop (vLLM parity). + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.my.tool:0" + "<|tool_call_argument_begin|>{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "my.tool" + assert calls[0]["id"] == "functions.my.tool:0" + + +def test_kimi_two_sections_in_one_stream_both_parse(): + """Outer loop walks every ``<|tool_calls_section_begin|>...end|>`` + so vLLM / SGLang parity holds even on multi-section streams.""" + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.a:0" + '<|tool_call_argument_begin|>{"x":1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " some prose between sections " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.b:0" + '<|tool_call_argument_begin|>{"y":2}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 2 + assert calls[0]["function"]["name"] == "a" + assert calls[1]["function"]["name"] == "b" + assert calls[0]["id"] == "functions.a:0" + assert calls[1]["id"] == "functions.b:0" + + +def test_kimi_bare_counter_id_is_dropped(): + """Bare-digit id (``3``) is dropped (matches vLLM); SGLang infers + name from schema, which we don't have at parse time.""" + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>3" + "<|tool_call_argument_begin|>{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert calls == [] + + +# DeepSeek truncated mid-stream + + +def test_deepseek_v3_1_huge_truncated_body_is_linear(): + """Adversarial input: DeepSeek envelope with no JSON brace and a + 50k-char body. A regex-based ``[^\\n<]+?`` name capture is O(N^2) + here; the parser uses ``str.find`` on the sep marker so it stays + linear. Budget 1s to flag any future regression.""" + import time as _time + + text = "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>fn<｜tool▁sep｜>" + "x" * 50_000 + start = _time.time() + calls = parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"V3 path is non-linear: {elapsed:.2f}s" + assert calls == [] + + +def test_deepseek_r1_huge_fenceless_body_is_linear(): + """R1 detection used a greedy ``([^\\n]+)\\n```json`` regex that is O(N^2) on a + fence-less body of repeated ``function`` tokens. The parser now scans with + ``str.find``; budget 1s to flag any regression.""" + import time as _time + + text = "<｜tool▁calls▁begin｜>" + "function<｜tool▁sep｜>a" * 40_000 + start = _time.time() + calls = parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"R1 path is non-linear: {elapsed:.2f}s" + assert calls == [] + + +def test_glm_unclosed_body_many_arg_keys_is_linear(): + """An unclosed GLM ```` body runs to EOF; a lazy-group ``finditer`` + over many bare ```` tokens was O(N^2). The parser now walks pairs with + ``str.find``; budget 1s.""" + import time as _time + + text = "foo\n" + "k" * 40_000 + start = _time.time() + parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"GLM path is non-linear: {elapsed:.2f}s" + + +def test_deepseek_r1_fenced_json_parses(): + """R1 wraps args in a ```json fence after ``functionNAME``.""" + import json as _json + + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC","unit":"c"}\n' + "```<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + assert _json.loads(calls[0]["function"]["arguments"]) == {"city": "NYC", "unit": "c"} + + +def test_deepseek_v3_1_truncated_arguments_drops_call_without_crash(): + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"' # no closing brace, no end markers + ) + calls = parse_tool_calls_from_text(text) + assert calls == [] + + +def test_deepseek_v3_1_truncated_after_end_marker_still_yields_call(): + text = ( + "<｜tool▁calls▁begin｜>" "<｜tool▁call▁begin｜>get_time" "<｜tool▁sep｜>" '{"city":"Tokyo"}' + # neither <｜tool▁call▁end｜> nor <｜tool▁calls▁end｜> + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_time" + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "Tokyo"} + + +# Routes-layer strip across the three new families + + +def test_routes_layer_strip_removes_deepseek_envelope(): + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +def test_routes_layer_strip_removes_kimi_section(): + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +def test_routes_layer_strip_removes_glm_block(): + """``.*?`` covers GLM via the Qwen pattern.""" + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "web_search\n" + "q\nx\n" + "" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +# strip_tool_markup (parser-level finalise path) over the new families + + +def test_strip_tool_markup_handles_deepseek_envelope(): + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + stripped = strip_tool_markup(text, final = True) + assert "before" in stripped and "after" in stripped + assert "｜tool▁" not in stripped + assert "get_time" not in stripped and "Tokyo" not in stripped + + +def test_strip_tool_markup_handles_kimi_section(): + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + stripped = strip_tool_markup(text, final = True) + assert "before" in stripped and "after" in stripped + assert "tool_calls_section_begin" not in stripped + + +# Round-2 review findings: GLM quoted-string / unclosed-arg, DeepSeek +# strict terminator, nested wrapper-less Gemma strip + + +def test_glm_quoted_string_arg_keeps_its_quotes(): + # A GLM string value emitted verbatim that itself begins with a quote. + text = ( + "web_search\n" + "query\n" + '"exact phrase"\n' + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == '"exact phrase"' + + +def test_glm_unclosed_arg_value_is_rejected_in_strict_mode(): + # Closing present but a value never closes: strict mode must reject + # the whole call rather than execute it with the argument silently dropped. + text = ( + "web_search\n" + "query\n" + "Tokyo weather" # no + "" + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # With Auto-Heal the partial value is kept, not dropped to a no-arg call. + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 + args = json.loads(healed[0]["function"]["arguments"]) + assert "Tokyo weather" in args.get("query", "") + + +def test_deepseek_v3_missing_call_terminator_rejected_in_strict_mode(): + # Envelope closes but the per-call <｜tool▁call▁end｜> is absent. Strict mode + # must reject (it is truncated/merged); Auto-Heal still parses it. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁calls▁end｜>" # envelope end only, no per-call end + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 + assert healed[0]["function"]["name"] == "get_time" + + +def test_deepseek_v3_with_call_terminator_parses_in_strict_mode(): + # Sanity: a well-formed V3 call (with the per-call end marker) still parses + # under strict mode after the terminator check. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_time" + + +def test_strip_tool_markup_removes_nested_wrapperless_gemma_call(): + # Wrapper-less Gemma call with a NESTED object arg: the balanced helper must strip the whole call, not leave a trailing ``}``. + text = "answer: call:f{loc:{city:NYC},n:3} done" + stripped = strip_tool_markup(text, final = True) + assert "call:f" not in stripped + assert "}" not in stripped + assert "answer:" in stripped and "done" in stripped + + +# Pass-3 review findings: bare-Kimi streaming (non-final) strip symmetry +# and the wrapper-less Gemma route-display strip + + +def test_strip_tool_markup_non_final_removes_bare_kimi_call(): + # A bare ``<|tool_call_begin|>...<|tool_call_end|>`` (no section wrapper): the CLOSED (final=False) strip must remove it too. + text = ( + "before " + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + " after" + ) + stripped = strip_tool_markup(text, final = False) + assert "tool_call_begin" not in stripped + assert "tool_call_end" not in stripped + assert "before" in stripped and "after" in stripped + + +def test_routes_layer_strip_removes_wrapperless_gemma_call(): + # Gemma 4 (skip_special_tokens) emits a wrapper-less ``call:NAME{..}`` with no XML markers. + from routes.inference import _strip_tool_xml as _routes_strip + + text = 'before call:web_search{query:"weather in Sydney"} after' + stripped = _routes_strip(text) + assert "call:web_search" not in stripped + assert "before" in stripped and "after" in stripped + + +def test_deepseek_envelope_end_inside_arg_string_is_not_a_truncation(): + # A DeepSeek V3.1 call whose argument string contains the literal envelope-end token must not be dropped. + content = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>web_search<｜tool▁sep｜>" + '{"query":"what does <｜tool▁calls▁end｜> mean"}' + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(content) + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "what does <｜tool▁calls▁end｜> mean" + } + + +def test_glm_value_containing_literal_arg_value_close_is_preserved(): + # A GLM string argument may legitimately contain . + content = ( + "runcode" + 'print("")' + ) + calls = parse_tool_calls_from_text(content) + assert len(calls) == 1, calls + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + +def test_attribute_form_function_with_embedded_marker_runs_outer_call(): + # is a supported envelope; a DeepSeek/Kimi marker inside one of its + # parameter values is data, not a second call. + content = ( + '' + "The Kimi format is <|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|>" + "" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["respond"], calls + + +def test_wrapperless_gemma_call_gated_by_enabled_tools(): + # Once skip_special_tokens removes the <|tool_call> wrapper, call:NAME{...} is + # indistinguishable from prose documenting the Gemma syntax. + prose = "Here is an example of the syntax: call:foo{x:1}. That shows how tools work." + assert parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) == [] + # The display strip is gated the same way, so the example survives in the answer. + assert "call:foo{x:1}" in strip_tool_markup( + prose, final = True, enabled_tool_names = {"web_search"} + ) + # An enabled name is still a real call (parsed, and stripped from display). + real = "Answer. call:web_search{query:hi}" + calls = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert "call:web_search" not in strip_tool_markup( + real, final = True, enabled_tool_names = {"web_search"} + ) + + +def test_kimi_section_end_inside_arg_string_is_not_a_truncation(): + # In a multi-call Kimi section, a later call whose argument holds the literal section-end token must not truncate the section. + content = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>" + '{"q":"cats"}<|tool_call_end|>' + "<|tool_call_begin|>functions.explain:1<|tool_call_argument_begin|>" + '{"text":"the token <|tool_calls_section_end|> means end"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["search", "explain"], calls + assert json.loads(calls[1]["function"]["arguments"]) == { + "text": "the token <|tool_calls_section_end|> means end" + } + + +def test_closed_envelope_before_deepseek_block_owns_turn(): + # Document order is the contract: a CLOSED / call that precedes a + # DeepSeek/Kimi block owns the turn, even when prose frames it as an example. + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>search_web\n" + "```json\n" + '{"query":"weather in Paris"}\n' + "```" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + prose = ( + 'A Qwen call looks like {"name":"example_tool","arguments":{}}.\n' + ) + calls = parse_tool_calls_from_text(prose + deepseek) + assert [c["function"]["name"] for c in calls] == ["example_tool"], calls + + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.lookup:0" + '<|tool_call_argument_begin|>{"id":7}<|tool_call_end|><|tool_calls_section_end|>' + ) + calls_k = parse_tool_calls_from_text("Example: {} and now:\n" + kimi) + assert [c["function"]["name"] for c in calls_k] == ["demo"], calls_k + + +def test_marker_inside_closed_outer_envelope_still_runs_outer_call(): + # The guard must fire when the marker sits INSIDE a closed outer / envelope's arguments: the OUTER call wins. + outer = ( + "what does <｜tool▁calls▁begin｜> mean" + ) + calls = parse_tool_calls_from_text(outer) + # The outer envelope is the real call; the embedded DeepSeek marker must not + # hijack the parse into a spurious tool. + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "q": "what does <｜tool▁calls▁begin｜> mean" + } + + +def test_truncated_outer_envelope_with_embedded_marker_heals_outer_call(): + # A TRUNCATED outer call embedding a DeepSeek/Kimi marker in its argument still Auto-Heals as the outer call. + trunc = 'x = "<｜tool▁calls▁begin｜>sample"' + calls = parse_tool_calls_from_text(trunc) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_python_tag_call_with_embedded_marker_runs_outer_call(): + # ``<|python_tag|>`` is Llama-3's tool-call envelope, so a DeepSeek/Kimi example quoted + # in its argument is data: the OUTER python_tag call (``web_search``) must run, not the + # embedded marker (``delete_all``). + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + ) + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>delete_all<｜tool▁sep｜>{}" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + for embedded in (kimi, deepseek): + builtin = '<|python_tag|>web_search.call(query="explain ' + embedded + '")' + calls = parse_tool_calls_from_text(builtin, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + custom = ( + '<|python_tag|>{"name":"web_search","parameters":' + '{"query":"explain ' + embedded + '"}}' + ) + calls = parse_tool_calls_from_text(custom, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # A bare ``<|python_tag|>`` prose mention (no call shape) must NOT be treated as an + # envelope: a real Kimi call after it still parses (the call-shaped lookahead guard). + prose = "The token <|python_tag|> is used. " + kimi + calls = parse_tool_calls_from_text(prose) + assert [c["function"]["name"] for c in calls] == ["delete_all"], calls + + +def test_gemma_wrapperless_quoted_value_with_comma_not_split(): + # A wrapper-less Gemma call whose quoted value contains ``, key:``. + text = 'call:web_search{query:"weather, location: Boston", limit:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "weather, location: Boston", + "limit": 3, + } + + +def test_literal_close_tag_in_xml_arg_before_marker_runs_outer_call(): + # A literal ```` inside an outer XML argument (before a marker) is not the envelope close: the span reaches the REAL final close. + text = ( + 'x = " ' + "<|tool_call_begin|>functions.delete_all:0<|tool_call_argument_begin|>{}" + '<|tool_call_end|>"' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_literal_tool_call_close_in_qwen_json_before_marker_runs_outer_call(): + # A Qwen/Hermes whose JSON argument holds a literal then a marker must run the OUTER call. + text = ( + '{"name":"search","arguments":{"query":"explain then ' + "<|tool_call_begin|>functions.delete_all:0<|tool_call_argument_begin|>{}" + '<|tool_call_end|>"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["search"], calls + # Back-to-back Qwen calls still parse independently (real-close span must keep the + # negative-lookahead that separates adjacent calls). + bb = ( + '{"name":"a","arguments":{}}' + '{"name":"b","arguments":{}}' + ) + assert [c["function"]["name"] for c in parse_tool_calls_from_text(bb)] == ["a", "b"] + + +def test_r1_heal_keeps_later_call_when_first_omits_close_fence(): + # DeepSeek R1 multi-call where the FIRST call has balanced JSON but omits its close + # fence/terminator, followed by a well-formed second call. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n```json\n" + '{"city":"SF"}\n```' # no <｜tool▁call▁end｜> + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time\n```json\n" + '{"tz":"UTC"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + heal = [c["function"]["name"] for c in parse_tool_calls_from_text(text)] + assert "get_time" in heal, heal + # Strict keeps the later well-formed call; heal must be a superset. + strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] + assert set(strict) <= set(heal), (strict, heal) + + +def test_wrapperless_gemma_nested_call_in_arg_is_not_a_second_call(): + # A wrapper-less Gemma call whose quoted argument mentions another enabled tool must not execute that nested name. + text = 'call:web_search{query:"explain call:delete_all{target:files}"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "explain call:delete_all{target:files}" + } + # Two genuinely separate calls still both parse. + two = "call:web_search{query:hi}call:get_time{tz:UTC}" + assert [ + c["function"]["name"] + for c in parse_tool_calls_from_text(two, enabled_tool_names = {"web_search", "get_time"}) + ] == ["web_search", "get_time"] + + +def test_leading_bare_json_call_owns_quoted_gemma_snippet(): + # Document order: a leading Llama-3.2 bare-JSON call with trailing prose owns the turn. + text = ( + '{"name":"lookup","parameters":{"note":"use call:web_search{query:cats} for this"}}\n' + "That is the call I would make." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "note": "use call:web_search{query:cats} for this" + } + + # Same with the ``;`` inter-call separator: both real calls parse, the + # quoted snippet still does not. + two = ( + '{"name":"lookup","parameters":{"note":"see call:web_search{query:cats}"}};' + '{"name":"lookup","parameters":{"q":"second"}}' + ) + calls_two = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "web_search"}) + assert [c["function"]["name"] for c in calls_two] == ["lookup", "lookup"], calls_two + + +def test_leading_gemma_call_still_wins_over_trailing_json_example(): + # Reverse control: a real leading Gemma call followed by a bare-JSON example keeps the Gemma call (bare JSON matches only a LEADING object). + text = 'call:web_search{query:cats} Example JSON: {"name":"demo_tool","parameters":{}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "demo_tool"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # And prose-only enabled Gemma syntax (no leading JSON) still promotes: the + # markerless by-design behaviour is unchanged. + prose = "You can run call:web_search{query:cats} to search." + calls_p = parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls_p] == ["web_search"], calls_p + + +def test_leading_gemma_call_owns_quoted_mistral_trigger(): + # A leading wrapper-less Gemma call whose argument quotes a Mistral trigger must win: the [TOOL_CALLS] literal is data. + text = 'call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "docs say [TOOL_CALLS]delete_all{}" + } + + # Reverse control: a real leading Mistral call still parses normally. + real = '[TOOL_CALLS]delete_all{"x":1}' + calls_m = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls_m] == ["delete_all"], calls_m + + # A DISABLED Gemma example quoting the trigger is dropped as prose and a + # real call after it still parses (drop-the-span recursion). + mixed = ( + 'Example: call:demo{note:"see [TOOL_CALLS]delete_all{}"}\n' + '[TOOL_CALLS]web_search{"q":"real"}' + ) + calls_d = parse_tool_calls_from_text(mixed, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls_d] == ["web_search"], calls_d + + +def test_chained_bare_json_owns_kimi_marker_in_later_call(): + # Document order: two ;-chained bare-JSON calls own the turn even when the second's argument quotes a complete Kimi snippet. + kimi = ( + "<|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|>" + ) + two = ( + '{"name":"lookup","parameters":{"q":"first"}};' + '{"name":"lookup","parameters":{"note":"' + kimi + '"}}' + ) + calls = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls + + # Reverse control: prose followed by a real Kimi block still parses. + real = "Let me check.\n<|tool_calls_section_begin|>" + kimi + "<|tool_calls_section_end|>" + calls_k = parse_tool_calls_from_text(real, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls_k] == ["delete_all"], calls_k + + # A closed leading Mistral call preceding a trailing Kimi example owns the + # turn too (same closed-call-precedes-marker rule). + mistral = '[TOOL_CALLS]lookup{"q":"first"} then example ' + kimi + calls_m = parse_tool_calls_from_text(mistral, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls_m] == ["lookup"], calls_m + + +def test_nested_gemma_values_keep_commas_and_parens(): + # Nested wrapper-less Gemma mappings/arrays use the top-level delimiter rules, so nested arguments are not split. + calls = parse_tool_calls_from_text( + "call:python{opts:{code:print(1,2),lang:py}}", enabled_tool_names = {"python"} + ) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "opts": {"code": "print(1,2)", "lang": "py"} + } + + arr = parse_tool_calls_from_text( + "call:python{opts:[1,2,{a:f(1,2)}]}", enabled_tool_names = {"python"} + ) + assert json.loads(arr[0]["function"]["arguments"]) == {"opts": [1, 2, {"a": "f(1,2)"}]} + + prose_comma = parse_tool_calls_from_text( + "call:python{opts:{note:hello, world}}", enabled_tool_names = {"python"} + ) + assert json.loads(prose_comma[0]["function"]["arguments"]) == {"opts": {"note": "hello, world"}} + + quoted = parse_tool_calls_from_text( + 'call:python{opts:{q:say "a, b" now,n:3}}', enabled_tool_names = {"python"} + ) + assert json.loads(quoted[0]["function"]["arguments"]) == { + "opts": {"q": 'say "a, b" now', "n": 3} + } + + # Controls: nested quoted values and multi-key mappings are unchanged, and + # a truncated nested value still falls back to the raw string. + nested_q = parse_tool_calls_from_text( + 'call:python{loc:{city:"New York"}}', enabled_tool_names = {"python"} + ) + assert json.loads(nested_q[0]["function"]["arguments"]) == {"loc": {"city": "New York"}} + multi = parse_tool_calls_from_text( + "call:python{opts:{a:1,b:2},n:3}", enabled_tool_names = {"python"} + ) + assert json.loads(multi[0]["function"]["arguments"]) == {"opts": {"a": 1, "b": 2}, "n": 3} + trunc = parse_tool_calls_from_text( + "call:python{opts:{code:print(1,2}}", enabled_tool_names = {"python"} + ) + assert json.loads(trunc[0]["function"]["arguments"]) == {"opts": "{code:print(1,2}"} + + +def test_multi_gemma_calls_own_turn_over_signal_in_later_call(): + # Document order: when the first enabled Gemma call closes before the first foreign signal, the leading call still owns the turn. + en = {"get_time", "web_search", "delete_all"} + both = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in both] == ["get_time", "web_search"], both + assert json.loads(both[1]["function"]["arguments"]) == { + "query": "docs say [TOOL_CALLS]delete_all{}" + } + + # XML and Kimi markers in the later call's strings stay data too. + xml = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"see delete_all"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in xml] == ["get_time", "web_search"], xml + kimi = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"see <|tool_call_begin|>' + 'functions.delete_all:0<|tool_call_argument_begin|>{}<|tool_call_end|>"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in kimi] == ["get_time", "web_search"], kimi + + # A trailing prose example after the closed leading call defers the same way. + prose = parse_tool_calls_from_text( + "call:get_time{} Example: [TOOL_CALLS]delete_all{}", enabled_tool_names = en + ) + assert [c["function"]["name"] for c in prose] == ["get_time"], prose + + +def test_multi_gemma_ownership_reverse_controls(): + # A real leading Mistral/XML call with a trailing Gemma example keeps the leading call; a signal before every Gemma call keeps normal order. + en = {"get_time", "web_search", "delete_all"} + mistral = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"delete_all","arguments":{}}] Example: call:web_search{query:cats}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in mistral] == ["delete_all"], mistral + xml_first = parse_tool_calls_from_text( + '{"name":"delete_all","arguments":{}} call:web_search{query:cats}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in xml_first] == ["delete_all"], xml_first + agnostic = parse_tool_calls_from_text( + 'call:foo{} {"name":"delete_all","arguments":{}}' + ) + assert [c["function"]["name"] for c in agnostic] == ["delete_all"], agnostic + + +def test_disabled_leading_bare_json_does_not_hide_later_marker_call(): + # A leading bare-JSON object with a NOT-enabled name is prose: the real DeepSeek/Kimi call after it still parses. + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"cats"}<|tool_call_end|><|tool_calls_section_end|>' + ) + calls = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + kimi, enabled_tool_names = {"web_search"} + ) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"} + + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"q":"cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + calls_ds = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + deepseek, enabled_tool_names = {"web_search"} + ) + assert [c["function"]["name"] for c in calls_ds] == ["web_search"], calls_ds + + +def test_disabled_leading_bare_json_ownership_controls(): + kimi_delete = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + ) + # ENABLED leading name still owns the turn (document order, the shipped + # inside-or-after rule). + owns = parse_tool_calls_from_text( + '{"name":"web_search","parameters":{"q":"first"}} ' + kimi_delete, + enabled_tool_names = {"web_search", "delete_all"}, + ) + assert [c["function"]["name"] for c in owns] == ["web_search"], owns + # A marker INSIDE the disabled object's own strings stays data: the span + # is prose, the tail holds no call, so nothing parses. + inside = parse_tool_calls_from_text( + '{"name":"draft","parameters":{"note":"see <|tool_call_begin|>functions.delete_all:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}}\nsome trailing prose', + enabled_tool_names = {"web_search", "delete_all"}, + ) + assert inside == [], inside + # Nameless leading JSON answers keep recursing to the real call. + nameless = parse_tool_calls_from_text( + '{"answer":42} ' + kimi_delete, enabled_tool_names = {"delete_all"} + ) + assert [c["function"]["name"] for c in nameless] == ["delete_all"], nameless + # Name-agnostic path unchanged: the leading object is the call. + agnostic = parse_tool_calls_from_text('{"name":"draft","parameters":{}} ' + kimi_delete) + assert [c["function"]["name"] for c in agnostic] == ["draft"], agnostic + + +def test_leading_json_answer_with_prose_keeps_quoted_gemma_snippet_as_data(): + # A LEADING JSON answer followed by prose is data (same contract as the whole-content JSON exemption). + obj = '{"summary":"use call:web_search{query:cats} to search"}\nHope that helps!' + assert parse_tool_calls_from_text(obj, enabled_tool_names = {"web_search"}) == [] + arr = '["use call:web_search{query:cats} to search"]\nHope that helps!' + assert parse_tool_calls_from_text(arr, enabled_tool_names = {"web_search"}) == [] + assert strip_tool_markup(obj, enabled_tool_names = {"web_search"}) == obj + + # A REAL call in the tail after the answer still parses (and strips). + tail = '{"summary":"done"}\ncall:web_search{query:cats}' + calls = parse_tool_calls_from_text(tail, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # A leading brace run that is NOT valid JSON gets no exemption. + not_json = "{not json} call:web_search{query:cats}" + calls_nj = parse_tool_calls_from_text(not_json, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls_nj] == ["web_search"], calls_nj + + +def test_glm_heal_bounds_unclosed_value_at_tool_call_close(): + # Auto-Heal: a value missing only its before the block's heals to the + # value text, not the close tag and everything after it swallowed into the argument. + one = "get_weathercityNYC" + calls = parse_tool_calls_from_text(one, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["get_weather"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "NYC"} + + # Trailing prose after the close stays out of the healed value. + two = one + "\nLet me check that for you." + calls_two = parse_tool_calls_from_text(two, allow_incomplete = True) + assert json.loads(calls_two[0]["function"]["arguments"]) == {"city": "NYC"} + + # Strict mode still rejects the unclosed value outright. + assert parse_tool_calls_from_text(one, allow_incomplete = False) == [] + + # A value truncated at EOF (no structural tag follows) keeps the partial heal, and a proper + # close whose value holds a literal is untouched by the bounding. + eof = "get_weathercityNew York Ci" + calls_eof = parse_tool_calls_from_text(eof, allow_incomplete = True) + assert json.loads(calls_eof[0]["function"]["arguments"]) == {"city": "New York Ci"} + lit = ( + "get_weathercity" + 'print("")' + ) + calls_lit = parse_tool_calls_from_text(lit, allow_incomplete = True) + assert json.loads(calls_lit[0]["function"]["arguments"]) == {"city": 'print("")'} + + +def test_prose_mentioning_ds_kimi_markers_survives_final_strip(): + # False-alarm literals: the trailing strip arms require a call-shaped + # lookahead, so an answer documenting a marker keeps its tail. + from core.inference.tool_call_parser import strip_tool_markup + + for text in [ + "The Kimi marker <|tool_calls_section_begin|> starts a section.", + "DeepSeek uses <｜tool▁calls▁begin｜> to open calls.", + "See <|tool_call_begin|> in the docs.", + ]: + assert strip_tool_markup(text, final = True) == text + + # Truncated REAL calls still drop, and a bare marker at EOF is a fragment. + truncated_kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q' + ) + assert strip_tool_markup(truncated_kimi, final = True) == "" + assert strip_tool_markup("prefix <|tool_calls_section_begin|>", final = True) == "prefix" diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index ce5688be3e..89a0b3879b 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -1990,8 +1990,8 @@ class TestTranslatedMessagesValidate: ChatMessage(**m.model_dump(exclude_none = True)) -# reasoning_prefilled: Qwen3/GLM enable_thinking templates prefill an unclosed , so generation -# begins inside the think block and emits only the closing ; extractor starts in reasoning. +# reasoning_prefilled mode: Qwen3/GLM enable_thinking templates prefill an unclosed , so +# generation begins inside the think block and emits only the closing ; the extractor starts in reasoning. class TestReasoningPrefilledExtractor: def test_prefilled_single_feed_splits_lone_close(self): # T1: reasoning...answer with a prefilled (unseen) open tag. @@ -2055,7 +2055,8 @@ class TestReasoningPrefilledExtractor: assert visible == "\n\nanswer" def test_prefilled_stray_open_tag_is_suppressed(self): - # T7: a re-emitted literal inside prefilled reasoning is dropped, not leaked. + # T7: a re-emitted literal inside prefilled reasoning is dropped, + # not leaked into the drawer (covers enable_thinking_effort full-tag output). reasoning, visible = _extract_responses_reasoning( "abc", parse_think_markers = True, @@ -2076,7 +2077,9 @@ class TestReasoningPrefilledExtractor: assert visible == "hi" def test_not_prefilled_lone_close_preserves_current_behavior(self): - # T9: without prefilled, a lone keeps pre-fix behavior (reasoning stays visible, tag dropped). + # T9: GGUF-parity guard -- WITHOUT prefilled, a lone keeps the + # pre-fix behavior (reasoning stays visible, tag dropped). Ensures GGUF and + # every existing caller are byte-identical. reasoning, visible = _extract_responses_reasoning( "reasoningans", parse_think_markers = True, @@ -2096,7 +2099,8 @@ class TestReasoningPrefilledExtractor: assert visible == "v" def test_prefilled_ignored_when_markers_not_parsed(self): - # T11: a non-reasoning model (parse_think_markers False) passes text straight through. + # T11: a non-reasoning model (parse_think_markers False) still passes text + # straight through even if reasoning_prefilled were mistakenly set False. reasoning, visible = _extract_responses_reasoning( "just an answer", parse_think_markers = False, diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 643d64af7a..3701a00dd2 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -11,6 +11,8 @@ from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) @@ -127,8 +129,8 @@ def test_detect_safetensors_features_gptoss_disables_tools(): assert flags["supports_tools"] is False -# Llama-3 / Mistral / Gemma 4 tool-call formats are parser-supported, so supports_tools stays True; -# only templates matching none of the known markers are suppressed. +# Llama-3 / Mistral / Gemma 4 tool-call formats are now parser-supported, so supports_tools=True +# must hold for all of them; only templates matching none of the five known markers are suppressed. LLAMA3_TEMPLATE = """ {%- if tools %} @@ -198,6 +200,86 @@ def test_detect_safetensors_features_gemma4_template_keeps_tools_on(): assert flags["supports_tools"] is True +# DeepSeek V3 / V3.1 / R1 emit ``<｜tool▁calls▁begin｜>...`` blocks. +# Note the full-width pipe (U+FF5C) and lower-1/8-block (U+2581). +DEEPSEEK_TEMPLATE = """ +{%- if tools %} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'assistant' and message.tool_calls %} + {%- for tc in message.tool_calls %} + {{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tc.function.name + + '<｜tool▁sep｜>' + tc.function.arguments + '<｜tool▁call▁end｜>' }} + {%- endfor %} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_deepseek_template_keeps_tools_on(): + """DeepSeek emits ``<｜tool▁calls▁begin｜>...``; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1") + flags = _detect_safetensors_features(backend, DEEPSEEK_TEMPLATE) + assert flags["supports_tools"] is True + + +# GLM 4.5 / 4.6 / 4.7 emit ``NAME\n...... +GLM_TEMPLATE = """ +{%- if tools %} + For each function call, output the function name and arguments within + the following XML format: + {function-name} + {arg-key} + {arg-value} + + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" + + +def test_detect_safetensors_features_glm_template_keeps_tools_on(): + """GLM 4.x emits ``NAME\\n...``; parser handles it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/GLM-4.6") + flags = _detect_safetensors_features(backend, GLM_TEMPLATE) + assert flags["supports_tools"] is True + + +# Kimi K2 / Moonshot uses ``<|tool_calls_section_begin|>...`` blocks +# with ``functions.NAME:IDX`` as the per-call id. +KIMI_TEMPLATE = """ +{%- if tools %} + <|im_system|>tool_declare<|im_middle|>{{ tools | tojson }}<|im_end|> +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'assistant' and message.tool_calls %} + <|tool_calls_section_begin|> + {%- for tc in message.tool_calls %} + <|tool_call_begin|>{{ tc.id }}<|tool_call_argument_begin|>{{ tc.function.arguments | tojson }}<|tool_call_end|> + {%- endfor %} + <|tool_calls_section_end|> + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_kimi_template_keeps_tools_on(): + """Kimi K2 emits ``<|tool_calls_section_begin|>...``; parser handles it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Kimi-K2-Instruct") + flags = _detect_safetensors_features(backend, KIMI_TEMPLATE) + assert flags["supports_tools"] is True + + LLAMA3_2_BARE_JSON_TEMPLATE = """ {%- if tools %} {{- 'Given the following functions, respond with JSON for a function call.' }} @@ -534,7 +616,34 @@ def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors(): assert flags["supports_preserve_thinking"] is True -# Templates advertising tools whose ``{"name":`` example is pretty-printed or JSON-escaped. +@pytest.mark.parametrize( + "opener", + [ + "<｜tool▁calls▁begin｜>", # canonical + "<｜tool_calls_begin｜>", # ASCII underscores + "<｜tool▁calls｜>", # short form + "<｜tool calls begin｜>", # spaces + "<｜tool\\_calls\\_begin｜>", # escaped underscores + ], +) +def test_detect_safetensors_features_deepseek_opener_variants_keep_tools_on(opener): + # Every DeepSeek opener the parser accepts must keep supports_tools on; the route gate derives + # its markers from the parser's TOOL_XML_SIGNALS so it can no longer drift behind the parser ... + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}tools{%- endif %}" + + opener + + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time{}" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1") + flags = _detect_safetensors_features(backend, tpl) + assert flags["supports_tools"] is True + + +# Templates that advertise tools ({%- if tools %}) and prompt the bare-JSON +# call form, but whose ``{"name":`` example is pretty-printed or JSON-escaped. _WHITESPACE_BARE_JSON_TEMPLATE = ( "{%- if tools %}\n" "To call a tool, output JSON of the form:\n" @@ -554,7 +663,8 @@ _TOOLS_ADVERTISED_NO_PARSEABLE_FORM = ( def test_detect_safetensors_features_keeps_tools_for_pretty_printed_bare_json(): - # Pretty-printed bare-JSON (``{ "name" :``) keeps supports_tools: parser accepts the whitespace. + # A pretty-printed bare-JSON example (``{ "name" :``) must keep supports_tools since the parser + # accepts that whitespace via raw_decode. from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") @@ -571,7 +681,8 @@ def test_detect_safetensors_features_keeps_tools_for_escaped_bare_json(): def test_detect_safetensors_features_drops_tools_when_no_parseable_form(): - # Negative control: tools advertised but no parser-recognised emission form -> pill dropped. + # Negative control: tools advertised but no parser-recognised emission form at + # all -> the pill is still dropped (the gate is not now matching everything). from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") @@ -580,7 +691,8 @@ def test_detect_safetensors_features_drops_tools_when_no_parseable_form(): def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json(): - # The {"function":...} bare-JSON alias keeps supports_tools, mirroring {"name":...}. + # A template documenting the parser-supported {"function":...} bare-JSON alias + # must keep supports_tools, mirroring the {"name":...} form. from routes.inference import _detect_safetensors_features tpl = ( @@ -594,9 +706,10 @@ def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json(): assert flags["supports_tools"] is True -# _sf_reasoning_prefill_mode gates the prefilled- extractor for enable_thinking models. +# _sf_reasoning_prefill_mode gates the prefilled- extractor so safetensors/MLX reach +# GGUF reasoning-block parity for enable_thinking models. class TestSafetensorsReasoningPrefillGate: - # Qwen3-style template with the standard / markers. + # A minimal Qwen3-style template with the standard / markers. _QWEN_TPL = "{% if enable_thinking %}{% endif %}......" # gemma-style bespoke reasoning channel -- no standard markers. _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" @@ -650,8 +763,8 @@ class TestSafetensorsReasoningPrefillGate: assert _sf_reasoning_prefill_mode(feats, False, self._QWEN_TPL) is True def test_g8_gemma_bespoke_channel_excluded(self): - # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled (else the - # whole answer is swallowed as reasoning). Regression guard. + # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled + # (would otherwise swallow the whole answer as reasoning). Regression guard. from routes.inference import _sf_reasoning_prefill_mode assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py index 9158d1ad5e..4a5423fa87 100644 --- a/studio/backend/tests/test_safetensors_reasoning_stream.py +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -34,7 +34,7 @@ def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict: visible_deltas: list[str] = [] monitor: list[str] = [] tool_starts: list[dict] = [] - order: list[str] = [] # "reasoning" | "visible" | "tool_start" sequence + order: list[str] = [] # sequence of ("reasoning"|"visible"|"tool_start") events def _flush(): fr, fv = extractor.finish() @@ -155,8 +155,11 @@ _THINK_TPL = "...{% if enable_thinking %}{% endif %}......" def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): - # GLM-5.2 enable_thinking_effort + reasoning_effort="none" disables thinking like - # enable_thinking=False, so prefilled must be OFF (else the answer is swallowed into reasoning). + # GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and + # enable_thinking omitted) disables thinking exactly like enable_thinking=False, so + # prefilled mode must be OFF. Otherwise the model emits no and a plain + # answer is swallowed whole into reasoning_content, leaving the visible response + # empty (the exact bug: prefilled=True below eats the whole answer). feats = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "none") is False # Thinking on (effort level or default) still prefills. @@ -171,7 +174,8 @@ def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): plain = {"reasoning_style": "enable_thinking", "supports_reasoning": True} assert _sf_reasoning_prefill_mode(plain, None, _THINK_TPL, "none") is True - # End-to-end: with prefilled=False, a plain no- answer stays visible. + # End-to-end: with the corrected prefilled=False, a plain no- answer is + # emitted as visible content rather than swallowed into the thinking drawer. events = [{"type": "content", "text": "The capital of France is Paris."}] out = _replay_sf_reasoning_stream(events, prefilled = False) assert out["visible"] == "The capital of France is Paris." diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 984d5f8ae9..38b30fe8f6 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -139,7 +139,7 @@ class TestParser: assert "print('hi')" in result[0]["function"]["arguments"] def test_xml_param_preserves_leading_indentation(self): - # Only the wrapping newline is trimmed, so code indentation survives. + # Only the wrapping newline is trimmed, so code-argument indentation survives (str.strip() destroyed it). text = ( "\n" " indented = 1\n" @@ -204,7 +204,9 @@ class TestParser: assert strip_tool_markup(text) == "before after" def test_strip_named_mistral_call_consumes_trailing_eos(self): - # The named [TOOL_CALLS]name{json} shape must eat the optional trailing . + # The named ``[TOOL_CALLS]name{json}`` shape must eat the optional + # trailing ```` like the array shape, so the EOS marker is not left + # behind as visible content. text = '[TOOL_CALLS]web_search{"query":"cats"}' assert strip_tool_markup(text) == "" text = '[TOOL_CALLS]web_search{"query":"cats"} and then' @@ -235,8 +237,27 @@ class TestParser: == "before " ) + def test_streaming_strip_handles_nested_mistral_json(self): + # The non-greedy [TOOL_CALLS]name{...} pattern truncates nested JSON at the first }; the + # balanced helper must remove the whole call so no trailing brace leaks to the streaming ... + raw = 'ok [TOOL_CALLS]foo{"a":{"b":1}} tail' + out = strip_tool_markup_streaming(raw) + assert "[TOOL_CALLS]" not in out + assert "}" not in out + assert "ok " in out and "tail" in out + + def test_streaming_strip_handles_nested_wrapperless_gemma(self): + # Same class of bug for the wrapper-less Gemma call:NAME{...} form with a + # nested object argument. + raw = "ok call:f{loc:{city:NYC},n:3} tail" + out = strip_tool_markup_streaming(raw) + assert "call:f" not in out + assert "}" not in out + assert "ok " in out and "tail" in out + def test_streaming_strip_keeps_prose_after_function_xml_with_literal_marker(self): - # A literal in a value is data: the strip closes at the REAL , keeping prose. + # A literal ```` in a value is data: the strip must close at the REAL + # ```` and keep trailing prose (the open-ended regex ate to EOF). raw = ( "pref " 'print("") tail' @@ -246,15 +267,19 @@ class TestParser: assert strip_tool_markup_streaming(raw) == strip_tool_markup(raw, final = True) def test_streaming_strip_drops_leading_magistral_reasoning(self): - # Magistral reasoning is a leading [THINK]...[/THINK] block; the streaming strip must drop it. + # Magistral emits reasoning as a leading ``[THINK]...[/THINK]`` bracket block + # (not the ```` the reasoning channel renders). The streaming display + # strip must drop it so the raw chain-of-thought does not leak into the + # safetensors content; GGUF routes it to reasoning_content natively. closed = "[THINK]Let me think. 2+2 is 4.[/THINK]The answer is 4." assert strip_tool_markup_streaming(closed) == "The answer is 4." assert strip_tool_markup_streaming(closed) == strip_tool_markup(closed, final = True) - # Unclosed mid-stream reasoning is held; cleaned text grows only after [/THINK]. + # Unclosed mid-stream reasoning is held from the marker on (nothing leaks, and + # the cleaned text only grows as the answer streams in after ``[/THINK]``). assert strip_tool_markup_streaming("[THINK]still thinking") == "" assert strip_tool_markup_streaming("[THINK]r[/THINK]The") == "The" assert strip_tool_markup_streaming("[THINK]r[/THINK]The answer") == "The answer" - # A non-leading [THINK] is ordinary prose, left untouched. + # A non-leading ``[THINK]`` is ordinary prose and is left untouched. assert strip_tool_markup_streaming("hi [THINK] later") == "hi [THINK] later" @@ -294,7 +319,7 @@ class TestParserMultiFormat: assert args == {"query": "hi", "n": 5} def test_llama3_python_tag_json_form_with_eom(self): - # Llama-3 emits <|eom_id|> after the JSON; must not break parsing. + # Llama-3 emits ``<|eom_id|>`` after the JSON; must not break parsing. import json text = '<|python_tag|>{"name":"python","parameters":{"code":"print(2+2)"}}<|eom_id|>' @@ -307,10 +332,22 @@ class TestParserMultiFormat: text = '<|python_tag|>brave_search.call(query="x")' assert strip_tool_markup(text, final = True) == "" - # Llama-3.2 bare JSON ``custom_tools`` + def test_llama3_python_tag_json_form_non_scalar_args_skipped(self): + # Should NOT fabricate ``{"value": args}`` when the JSON form + # has a non-dict / non-string ``arguments`` value. + for bad in ( + '<|python_tag|>{"name":"foo","arguments":42}', + '<|python_tag|>{"name":"foo","arguments":[1,2,3]}', + '<|python_tag|>{"name":"foo","arguments":null}', + '<|python_tag|>{"name":"foo","arguments":true}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + + # ── Llama-3.2 bare JSON ``custom_tools`` ───────────────────── def test_llama3_2_bare_json_parameters(self): - # Llama-3.2-Instruct emits bare JSON directly as content, no <|python_tag|> prefix. + # Llama-3.2-Instruct emits bare JSON directly as content; no + # <|python_tag|> prefix per its training template. import json text = '{"name":"web_search","parameters":{"query":"Tokyo weather"}}' @@ -330,7 +367,7 @@ class TestParserMultiFormat: assert args == {"a": 1, "b": 2} def test_llama3_2_bare_json_multi_call(self): - # Llama-3 may chain calls with "; " per training template. + # Llama-3 may chain calls with ``; `` per training template. text = '{"name":"a","parameters":{}}; {"name":"b","parameters":{}}' result = parse_tool_calls_from_text(text) assert len(result) == 2 @@ -356,7 +393,8 @@ class TestParserMultiFormat: assert parse_tool_calls_from_text(text) == [] def test_llama3_2_bare_json_embedded_in_prose_does_not_fire(self): - # Defensive: JSON embedded in prose must NOT fire (content must START with `{`). + # Defensive: JSON embedded in prose must NOT fire (parser is + # strict about content STARTING with `{`). text = 'The tool result was: {"name":"foo"}' assert parse_tool_calls_from_text(text) == [] @@ -373,12 +411,14 @@ class TestParserMultiFormat: assert parse_tool_calls_from_text(text) == [] def test_llama3_2_bare_json_string_parameters_does_not_fire(self): - # Llama-3 spec: parameters must be a dict; a string value must NOT trigger. + # Llama-3 spec: parameters must be a dict. Prose like + # ``{"name":"foo","parameters":"a sentence"}`` must NOT trigger. text = '{"name":"foo","parameters":"this is a sentence"}' assert parse_tool_calls_from_text(text) == [] def test_llama3_2_bare_json_string_arguments_not_json_does_not_fire(self): - # OpenAI arguments may be a JSON-string of a dict, but a plain non-JSON string must not pass. + # OpenAI ``arguments`` may be a JSON-string of a dict, but a + # plain non-JSON string must not pass the guard. text = '{"name":"foo","arguments":"not json"}' assert parse_tool_calls_from_text(text) == [] @@ -417,7 +457,8 @@ class TestParserMultiFormat: def test_mistral_array_parameters_key_alias(self): import json - # Array object keyed on parameters (not arguments) must keep its payload. + # Array object keyed on ``parameters`` (not ``arguments``) must keep its + # payload, matching the JSON/XML paths and SGLang's base detector. text = '[TOOL_CALLS] [{"name":"get_weather","parameters":{"city":"Paris"}}]' result = parse_tool_calls_from_text(text) assert len(result) == 1 @@ -435,7 +476,7 @@ class TestParserMultiFormat: assert result[1]["function"]["name"] == "b" def test_mistral_pre_v11_unclosed_array(self): - # Closing ] truncated: parser must heal off individual objects. + # Closing ``]`` truncated -- parser must heal off individual objects. text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"},"id":"id"}' result = parse_tool_calls_from_text(text) assert len(result) == 1 @@ -444,7 +485,7 @@ class TestParserMultiFormat: # Mistral v11+ def test_mistral_v11_single(self): - # Magistral / Mistral Small 3.1: bare name{json} after trigger. + # Magistral / Mistral Small 3.1: bare ``name{json}`` after trigger. import json text = '[TOOL_CALLS]add{"a":3.5,"b":4}' @@ -454,7 +495,7 @@ class TestParserMultiFormat: assert json.loads(result[0]["function"]["arguments"]) == {"a": 3.5, "b": 4} def test_mistral_v11_parallel(self): - # v11+ parallel: [TOOL_CALLS]a{...}[TOOL_CALLS]b{...}. + # v11+ parallel: ``[TOOL_CALLS]a{...}[TOOL_CALLS]b{...}``. text = '[TOOL_CALLS]add{"a":1}[TOOL_CALLS]sub{"b":2}' result = parse_tool_calls_from_text(text) assert len(result) == 2 @@ -462,7 +503,7 @@ class TestParserMultiFormat: assert result[1]["function"]["name"] == "sub" def test_mistral_v11_with_args_marker(self): - # Ministral / Mistral Large 3: [TOOL_CALLS]name[ARGS]{json}. + # Ministral / Mistral Large 3: ``[TOOL_CALLS]name[ARGS]{json}``. import json text = '[TOOL_CALLS]add[ARGS]{"a":1,"b":2}' @@ -476,7 +517,9 @@ class TestParserMultiFormat: assert strip_tool_markup(text, final = True) == "" def test_mistral_call_id_form(self): - # Mistral Small 3.2: the [CALL_ID] segment must be skipped, not treated as a stop (llama.cpp test-chat.cpp:4785). + # Mistral Small 3.2: ``[TOOL_CALLS]name[CALL_ID][ARGS]{json}``. + # The ``[CALL_ID]`` segment must be skipped, not treated as a stop + # (llama.cpp test-chat.cpp:4785 parses this to one call). import json text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' @@ -501,7 +544,9 @@ class TestParserMultiFormat: assert strip_tool_markup(text, final = True) == "" def test_mistral_think_reasoning_ignored(self): - # A [TOOL_CALLS] inside [THINK]...[/THINK] is reasoning; only the call after [/THINK] counts (llama.cpp test-chat.cpp:2285). + # Magistral wraps reasoning in ``[THINK]...[/THINK]``. A ``[TOOL_CALLS]`` + # inside the reasoning is chain-of-thought, not a real call; only the + # call after ``[/THINK]`` counts (llama.cpp test-chat.cpp:2285). import json text = ( @@ -514,12 +559,14 @@ class TestParserMultiFormat: assert json.loads(result[0]["function"]["arguments"]) == {"y": 2} def test_mistral_think_reasoning_no_real_call(self): - # Reasoning that mentions a call but emits none after [/THINK] yields no calls. + # Reasoning that merely mentions a tool call but does not emit one + # after ``[/THINK]`` yields no calls. text = '[THINK]I might call [TOOL_CALLS]fake[ARGS]{"x":1}[/THINK]Done.' assert parse_tool_calls_from_text(text) == [] def test_mistral_think_literal_in_argument_preserved(self): - # A literal [THINK] inside a real tool argument must not be stripped or corrupt the parse. + # A literal ``[THINK]`` inside a real tool argument (after the call) + # must not be stripped or corrupt the parse. import json text = '[TOOL_CALLS]search[ARGS]{"q":"explain the [THINK] token"}' @@ -554,7 +601,7 @@ class TestParserMultiFormat: assert args == {"enabled": True, "attempts": 5, "threshold": 1.5, "nickname": None} def test_gemma4_nested_args(self): - # Gemma 4 nests dicts / lists with bare keys and <|"|> strings. + # Gemma 4 nests dicts / lists with bare keys and ``<|"|>`` strings. import json text = ( @@ -585,7 +632,65 @@ class TestParserMultiFormat: text = "<|tool_call>call:foo{x:1}" assert strip_tool_markup(text, final = True) == "" - # Cross-format sentinels + # ── Gemma 4 wrapper-less (skip_special_tokens stripped) ─────────── + + def test_gemma4_bare_stripped_call(self): + # skip_special_tokens removes <|tool_call>/ and <|"|>, + # leaving a bare call:NAME{...} with an unquoted value. + import json + + text = "call:web_search{query:weather in San Francisco right now}" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "weather in San Francisco right now"} + + def test_gemma4_bare_code_with_commas(self): + # A code value with commas must not truncate at the first comma. + import json + + text = ( + "call:python{code:def f(n):\n a, b = 0, 1\n" + " for _ in range(2, n+1):\n a, b = b, a + b\n" + " return b\n\nprint(f(30))}" + ) + result = parse_tool_calls_from_text(text) + assert result[0]["function"]["name"] == "python" + code = json.loads(result[0]["function"]["arguments"])["code"] + assert "a, b = 0, 1" in code and "print(f(30))" in code + + def test_gemma4_bare_quotes_normalized(self): + # The same value quoted vs unquoted must parse identically so the + # agentic loop can collapse a looping model's repeated calls. + import json + + a = parse_tool_calls_from_text('call:web_search{query:"foo bar"}') + b = parse_tool_calls_from_text("call:web_search{query:foo bar}") + assert json.loads(a[0]["function"]["arguments"]) == {"query": "foo bar"} + assert json.loads(a[0]["function"]["arguments"]) == json.loads( + b[0]["function"]["arguments"] + ) + + def test_gemma4_bare_multi_arg(self): + import json + + text = "call:web_search{query:pytorch latest, url:https://pytorch.org}" + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "pytorch latest", "url": "https://pytorch.org"} + + def test_gemma4_bare_not_matched_in_prose(self): + # A word ending in "call:" must not trigger a bare tool call. + text = "I will recall:that the function{ } is helpful." + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_gemma4_bare_strip_markup_final(self): + text = "Here you go: call:web_search{query:weather today}" + assert "call:web_search" not in strip_tool_markup(text, final = True) + + # ── Cross-format sentinels ──────────────────────────────────── def test_all_markers_in_tool_xml_signals(self): # Streaming buffer wakes up on every emission marker. @@ -703,6 +808,553 @@ def _make_loop( ), exec_fn +class TestParserDeepSeek: + """DeepSeek R1 / V3 / V3.1 coverage. Markers use full-width pipes + (U+FF5C) and lower-one-eighth-block (U+2581). R1 wraps args in a + Markdown ``` ```json ``` ``` fence; V3 / V3.1 emit bare JSON.""" + + def test_r1_simple_call_with_code_fence(self): + import json as _json + + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>function" + "<｜tool▁sep｜>special_function\n" + "```json\n" + '{"arg1": 1}\n' + "```" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "special_function" + assert _json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_r1_short_form_outer_marker(self): + # llama.cpp accepts ``<｜tool▁calls｜>`` as the short-form opener. + import json as _json + + text = ( + "<｜tool▁calls｜>function" + "<｜tool▁sep｜>get_time\n" + "```json\n" + '{"city": "Paris"}\n' + "```" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + + def test_v3_1_bare_json(self): + # V3 / V3.1 omit the ``function`` prefix and the code fence. + import json as _json + + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + assert _json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_v3_1_multi_call_shares_envelope(self): + # Parallel calls share one outer envelope; each inner call has + # its own ``<｜tool▁call▁begin｜>...<｜tool▁call▁end｜>``. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "get_time" + assert result[1]["function"]["name"] == "get_weather" + + def test_v3_1_with_reasoning(self): + # Reasoning ... precedes the tool block. + text = ( + "I'm thinking\n" + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + + def test_v3_1_strict_rejects_unclosed_envelope(self): + # Envelope truncated mid-stream (no <｜tool▁calls▁end｜>): healed by + # default, rejected with Auto-Heal off. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + ) + assert len(parse_tool_calls_from_text(text)) == 1 + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_v3_1_multi_call_recovers_when_first_end_marker_missing(self): + # First inner call omits its <｜tool▁call▁end｜>; the second must still be parsed. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["get_time", "get_weather"] + + def test_v3_1_strict_recovers_after_missing_call_end(self): + # Strict mode (Auto-Heal off): the FIRST inner call is missing its <｜tool▁call▁end｜> + # terminator, so it is skipped -- but the parser must keep scanning and still return the ... + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "SF"}' + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"tz": "PST"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + # Auto-Heal keeps both; strict skips the truncated first, keeps the second. + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == [ + "get_weather", + "get_time", + ] + strict = parse_tool_calls_from_text(text, allow_incomplete = False) + assert [c["function"]["name"] for c in strict] == ["get_time"] + + def test_r1_strict_recovers_after_missing_close_fence(self): + # R1 form. + text = ( + "<｜tool▁calls▁begin｜>" + "function<｜tool▁sep｜>get_weather\n```json\n" + '{"city": "SF"}' + "function<｜tool▁sep｜>get_time\n```json\n" + '{"tz": "PST"}' + "\n```<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + strict = parse_tool_calls_from_text(text, allow_incomplete = False) + assert [c["function"]["name"] for c in strict] == ["get_time"] + + def test_deepseek_strip_markup(self): + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>foo" + "<｜tool▁sep｜>" + "{}" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_deepseek_signal_wakes_streaming(self): + # The streaming buffer state machine must wake on the DeepSeek opener so the rest of the + # section is drained instead of leaked. + text = "<｜tool▁calls▁begin｜>..." + assert has_tool_signal(text) + + def test_deepseek_short_opener_is_stripped(self): + # The short ``<｜tool▁calls｜>`` opener is parsed, so its markup must also be stripped (the + # strip patterns used to require ...calls_begin and left the short-opener markup leaking to ... + text = ( + "before " + "<｜tool▁calls｜>" + "<｜tool▁call▁begin｜>foo" + "<｜tool▁sep｜>" + "{}" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + +class TestParserGLM: + """GLM 4.5 / 4.6 / 4.7 coverage. Marker collides with Qwen's + ```` but the body shape is XML kv pairs instead of JSON, + so the dispatch order keeps both formats working.""" + + def test_glm_simple_call(self): + import json as _json + + text = ( + "web_search\n" + "query\n" + "weather Tokyo\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = _json.loads(result[0]["function"]["arguments"]) + # Strings come through raw; the parser does not double-quote. + assert args == {"query": "weather Tokyo"} + + def test_glm_mixed_types_decode_correctly(self): + # Per the chat_template.jinja, strings are emitted raw and non-strings are JSON-encoded. + import json as _json + + text = ( + "complex_function\n" + "name\nJohn Doe\n" + "age\n30\n" + "active\ntrue\n" + "score\n95.5\n" + "" + ) + result = parse_tool_calls_from_text(text) + args = _json.loads(result[0]["function"]["arguments"]) + assert args == {"name": "John Doe", "age": 30, "active": True, "score": 95.5} + + def test_glm_multi_call_back_to_back(self): + # GLM emits parallel calls as consecutive ``... + # `` blocks with no outer envelope. + text = ( + "a\nx\n1\n" + "b\ny\n2\n" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_glm_unclosed_tool_call_does_not_lose_value(self): + # Truncated mid-stream (no ) -- the parser must + # still surface what it found rather than dropping the call. + text = "web_search\nquery\npartial" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_glm_does_not_break_qwen_path(self): + # Real Qwen emission must still be parsed by the Qwen branch, + # not silently misrouted to GLM (the marker is shared). + text = '{"name":"web_search","arguments":{"q":"x"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_glm_strip_markup(self): + text = ( + "before " + "a\nx\n1\n" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_glm_zero_arg_inline_call(self): + # GLM 4.7 emits a no-argument call inline as ``name`` (name followed + # straight by the close tag, no \n / ). + import json as _json + + text = "get_current_date" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_current_date" + assert _json.loads(result[0]["function"]["arguments"]) == {} + + def test_glm_zero_arg_call_in_parallel_batch(self): + # A no-arg call alongside a normal one must not make either vanish. + text = ( + "get_current_date" + "get_weather\ncity\n" + "Tokyo" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "get_current_date" + assert result[1]["function"]["name"] == "get_weather" + + def test_glm_string_value_whitespace_preserved(self): + # The template emits string args verbatim, so significant leading / trailing whitespace + # (code, diffs) must survive. + import json as _json + + text = ( + "run\ncode\n" + " indented code " + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = _json.loads(result[0]["function"]["arguments"]) + assert args == {"code": " indented code "} + + +class TestParserKimi: + """Kimi K2 / Moonshot coverage. ASCII pipes only (NOT full-width). + Name arrives as ``functions.NAME:IDX``; the parser strips the + prefix and the index to recover the bare callable name while + preserving the full id for round-trip rendering.""" + + def test_kimi_simple_call(self): + import json as _json + + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.special_function:0" + "<|tool_call_argument_begin|>" + '{"arg1": 1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + # Bare name recovered; full id preserved verbatim. + assert result[0]["function"]["name"] == "special_function" + assert result[0]["id"] == "functions.special_function:0" + assert _json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_outer_tool_call_with_embedded_kimi_marker_parses_outer(self): + # A Qwen/Hermes whose argument contains literal Kimi markup (a user asking + # about that syntax) must execute the OUTER call, not the embedded marker via the ... + text = ( + '{"name":"web_search","arguments":{"query":' + '"explain <|tool_call_begin|>functions.evil:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}}' + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_genuine_kimi_call_without_envelope_still_parses(self): + # Control: a real Kimi call with no leading envelope must + # still go through the pre-pass. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"query":"x"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_kimi_multi_call_with_index(self): + # Multiple consecutive calls inside a single section, each + # with its own monotonically incrementing ``:IDX``. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.read_file:0" + "<|tool_call_argument_begin|>" + '{"path":"a"}' + "<|tool_call_end|>" + "<|tool_call_begin|>functions.web_search:1" + "<|tool_call_argument_begin|>" + '{"query":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "read_file" + assert result[0]["id"].endswith(":0") + assert result[1]["function"]["name"] == "web_search" + assert result[1]["id"].endswith(":1") + + def test_kimi_dotted_name_keeps_full_dotted_name(self): + # A dotted Kimi id keeps its FULL name after stripping only the ``functions.`` prefix and + # ``:idx`` suffix -- matching current vLLM ... + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>a.b.c:2" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "a.b.c" + + def test_kimi_dotted_mcp_name_with_functions_prefix(self): + # ``functions.mcp.server-list:0`` must resolve to ``mcp.server-list`` + # (only the ``functions.`` prefix and ``:idx`` are removed). + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.mcp.server-list:0" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp.server-list" + + def test_kimi_multi_call_recovers_when_first_end_marker_missing(self): + # First call omits its <|tool_call_end|>; the second must still parse. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.read_file:0" + "<|tool_call_argument_begin|>" + '{"path":"a"}' + "<|tool_call_begin|>functions.web_search:1" + "<|tool_call_argument_begin|>" + '{"query":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["read_file", "web_search"] + + def test_kimi_handles_unclosed_section(self): + # End marker missing -- the parser must still extract the call. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.foo:0" + "<|tool_call_argument_begin|>" + '{"a":1}' + "<|tool_call_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + + def test_kimi_strip_markup(self): + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.x:0" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_kimi_signal_wakes_streaming(self): + text = "<|tool_calls_section_begin|>..." + assert has_tool_signal(text) + + def test_kimi_call_without_section_wrapper(self): + # llama.cpp makes the ``<|tool_calls_section_begin|>`` wrapper optional -- Kimi K2 can emit + # a bare ``<|tool_call_begin|>`` call. + import json as _json + + text = ( + "<|tool_call_begin|>functions.execute_command:0" + "<|tool_call_argument_begin|>" + '{"cmd":"ls"}' + "<|tool_call_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "execute_command" + assert _json.loads(result[0]["function"]["arguments"]) == {"cmd": "ls"} + + def test_kimi_malformed_json_recovers_later_calls(self): + # A call with malformed / truncated JSON must not drop the valid calls that follow it in + # the same section (the bad call is skipped, the good one is recovered). + import json as _json + + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.a:0" + '<|tool_call_argument_begin|>{"city":"Beijing"' # missing closing brace + "<|tool_call_end|>" + "<|tool_call_begin|>functions.b:1" + '<|tool_call_argument_begin|>{"city":"Shanghai"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "b" + assert _json.loads(result[0]["function"]["arguments"]) == {"city": "Shanghai"} + + +class TestParserCrossFormatRouting: + """Ensure the per-format dispatch order doesn't misroute any + family. Real emissions for each new family + every old family + must still parse correctly when intermixed.""" + + def test_dispatch_routes_each_family_correctly(self): + cases = [ + ( + "Qwen", + '{"name":"a","arguments":{"x":1}}', + "a", + ), + ( + "DeepSeek V3.1", + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>", + "get_time", + ), + ( + "GLM", + "web_search\n" + "q\nx\n" + "", + "web_search", + ), + ( + "Kimi", + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.add:0" + "<|tool_call_argument_begin|>" + '{"a":1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>", + "add", + ), + ] + for label, text, expected_name in cases: + result = parse_tool_calls_from_text(text) + assert len(result) == 1, f"{label}: parser missed the call" + assert result[0]["function"]["name"] == expected_name, ( + f"{label}: got {result[0]['function']['name']!r}, " f"expected {expected_name!r}" + ) + + def test_all_new_markers_in_tool_xml_signals(self): + # The safetensors / MLX streaming buffer must wake on every supported emission marker -- + # otherwise the BUFFERING state leaks tool content to the user before parse. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + for marker in ( + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>", + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>", + ): + assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + + def test_active_tools_are_passed_to_single_turn_after_render_html_success(): captured_tool_names: list[list[str]] = [] exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) @@ -739,7 +1391,8 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success(): def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): - # A late unclosed heals only with Auto-Heal on; off, it must not execute. + # A late call caught by the safety net: an unclosed ```` heals only with Auto-Heal on; + # off, the safety net must not pass ``allow_incomplete=True`` and execute a truncated call. prose = "Sure, let me look that up for you right now. " incomplete = '{"name":"web_search","arguments":{"query":"weather in Sydney"}}' @@ -764,7 +1417,9 @@ def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): def test_bare_json_tool_call_is_not_streamed_as_content(): - # Llama-3.2 bare form carries no XML signal: BUFFER until the object closes, never leak the JSON. + # Llama-3.2 ``custom_tools`` bare form ``{"name":..,"parameters":..}`` carries no + # XML signal. The loop must BUFFER it until the object closes and execute it via + # the safety net, never leaking the raw JSON to streaming clients as content. bare = '{"name":"web_search","parameters":{"query":"cats"}}' loop, exec_fn = _make_loop( turns = [[bare], ["Here are the results."]], @@ -779,7 +1434,9 @@ def test_bare_json_tool_call_is_not_streamed_as_content(): def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(): - # Markerless JSON whose "name" is not an enabled tool must be shown, not dropped. + # Markerless JSON whose "name" is not an enabled tool (e.g. a person record + # ``{"name":"Alice",...}``) must be shown as the answer, not misread as a call + # to a disabled tool and dropped. _make_loop enables web_search/python/terminal. answer = '{"name":"Alice","parameters":{"age":30}}' loop, exec_fn = _make_loop(turns = [[answer]], max_tool_iterations = 1) events = _collect_events(loop) @@ -789,7 +1446,8 @@ def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(): def test_bare_json_tool_call_split_across_chunks_is_not_streamed(): - # Same as above but the bare object arrives split mid-key, held across chunks until it balances. + # Same as above but the bare object arrives split mid-key, so the buffer is + # held open across chunks before it balances. loop, exec_fn = _make_loop( turns = [ ['{"name":"web_', 'search","parameters":{"query":"cats"}}'], @@ -804,8 +1462,68 @@ def test_bare_json_tool_call_split_across_chunks_is_not_streamed(): assert not any('"name"' in t or "web_search" in t for t in contents), contents +def test_gemma_wrapperless_call_is_not_streamed_as_content(): + # Gemma 4 wrapper-less ``call:NAME{...}`` has no XML signal; the loop must hold + # it (BUFFERING) and execute it, never streaming the raw call text. + loop, exec_fn = _make_loop( + turns = [["call:web_search{query:cats}"], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call:web_search" in t for t in contents), contents + + +def test_gemma_wrapperless_call_with_whitespace_is_suppressed_when_streamed(): + # Gemma may emit ``call : NAME{...}`` with whitespace around the colon, split across stream + # chunks. + loop, exec_fn = _make_loop( + turns = [["call", " : ", "web_search", "{query:cats}"], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call" in t for t in contents), contents + + +def test_long_gemma_tool_name_is_not_streamed_as_content(): + # A tool name longer than the small buffer cap (OpenAI 64 chars, MCP longer) + # must still be held: the ``call:NAME`` prefix keeps buffering until ``{`` + # instead of leaking ``call:longname`` as visible text. + long_name = "mcp__github__list_repository_issues" # 35 chars + turns = iter([list('call:%s{repo:"octo/hello"}' % long_name), ["Done."]]) + + def _gen(_messages): + try: + chunks = next(turns) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(["RESULT"]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": long_name}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [(long_name, {"repo": "octo/hello"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call:" in t for t in contents), contents + + def test_leading_json_answer_is_not_dropped(): - # A leading {...} that is NOT a call must still surface; the hold only delays it. + # A leading ``{...}`` that is NOT a tool call must still surface as content: + # the bare-JSON hold can only ever delay it to end-of-object, never drop it. obj = '{"answer": 42, "note": "done"}' loop, exec_fn = _make_loop( turns = [[obj]], @@ -844,7 +1562,8 @@ def _reprompt_loop(*, auto_heal_tool_calls): def test_reprompt_names_only_active_tools_not_hardcoded(): - # The nudge must name the tools actually enabled, not hardcoded web_search/python. + # The plan-without-action nudge must name the tools actually enabled, never the + # old hardcoded ``web_search``/``python`` (which a restricted set would reject). captured, _events = _reprompt_loop(auto_heal_tool_calls = True) assert len(captured) >= 2, "intent prose should have triggered a re-prompt turn" reprompt = captured[1][-1] @@ -855,7 +1574,8 @@ def test_reprompt_names_only_active_tools_not_hardcoded(): def test_reprompt_suppressed_when_auto_heal_disabled(): - # With Auto-Heal off the nudge stays silent for GGUF parity, so only the initial generation runs. + # With Auto-Heal off the safetensors nudge must stay silent for backend parity + # with the GGUF loop, so only the single initial generation runs. captured, events = _reprompt_loop(auto_heal_tool_calls = False) assert len(captured) == 1, captured contents = [e["text"] for e in events if e["type"] == "content"] @@ -922,7 +1642,8 @@ class TestLoopBasic: assert "Result: 1" in contents[-1]["text"] def test_llama3_python_tag_form(self): - # The loop must recognise Llama-3's <|python_tag|> marker, drain the turn, and execute the call. + # The agentic loop must recognise Llama-3's <|python_tag|> + # marker, drain the rest of the turn, and execute the call. loop, exec_fn = _make_loop( turns = [ [ @@ -940,8 +1661,12 @@ class TestLoopBasic: assert "sunny" in contents[-1]["text"].lower() def test_llama3_bare_json_form_fires_tool(self): - # Llama-3.1/3.2 bare-JSON calls carry no XML signal; the safety-net parse must still fire - # the tool. Regression for the has_tool_signal gate that dropped these. + # Llama-3.1 / 3.2 emit a bare-JSON tool call + # ``{"name":..,"parameters":..}`` with NO XML signal. The loop's + # safety-net parse must still fire the tool instead of treating the + # turn as "planned without calling tools" and re-prompting the model + # into giving up. Regression for the has_tool_signal gate that + # dropped these; GGUF's llama-server parses them natively. loop, exec_fn = _make_loop( turns = [ ['{"name": "web_search", "parameters": {"query": "weather in SF"}}'], @@ -955,7 +1680,7 @@ class TestLoopBasic: assert "sunny" in contents[-1]["text"].lower() def test_mistral_pre_v11_form(self): - # Pre-v11 Mistral emission: [TOOL_CALLS] [{...}]. + # Pre-v11 Mistral emission: ``[TOOL_CALLS] [{...}]``. loop, exec_fn = _make_loop( turns = [ [ @@ -973,7 +1698,7 @@ class TestLoopBasic: assert tool_start["tool_call_id"] == "abc" def test_mistral_v11_form(self): - # v11+ Mistral emission: bare name{json} after the trigger. + # v11+ Mistral emission: bare ``name{json}`` after the trigger. loop, exec_fn = _make_loop( turns = [ ['[TOOL_CALLS]web_search{"query":"hi"}'], @@ -985,7 +1710,7 @@ class TestLoopBasic: assert exec_fn.calls == [("web_search", {"query": "hi"})] def test_gemma4_form(self): - # Gemma 4 emission: <|tool_call>call:NAME{...}. + # Gemma 4 emission: ``<|tool_call>call:NAME{...}``. loop, exec_fn = _make_loop( turns = [ [ @@ -1000,6 +1725,70 @@ class TestLoopBasic: events = _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "weather"})] + def test_deepseek_v3_1_form(self): + # DeepSeek V3.1 emission inside the agentic loop -- the buffer state machine must wake on + # ``<｜tool▁calls▁begin｜>`` and the parser must extract the V3.1 bare-JSON body. + loop, exec_fn = _make_loop( + turns = [ + [ + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>web_search", + "<｜tool▁sep｜>", + '{"query":"Tokyo weather"}', + "<｜tool▁call▁end｜>", + "<｜tool▁calls▁end｜>", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "Tokyo weather"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + def test_glm_form(self): + # GLM 4.x emission: ``NAME\n...``. + loop, exec_fn = _make_loop( + turns = [ + [ + "web_search\n", + "query\n", + "Tokyo\n", + "", + ], + ["found"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "Tokyo"})] + + def test_kimi_form(self): + # Kimi K2 emission ``<|tool_calls_section_begin|>...``. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>functions.web_search:0", + "<|tool_call_argument_begin|>", + '{"query":"Tokyo"}', + "<|tool_call_end|>", + "<|tool_calls_section_end|>", + ], + ["done"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + # The bare name must reach execute_tool, even though the model + # emitted ``functions.web_search:0`` as the formatted id. + assert exec_fn.calls == [("web_search", {"query": "Tokyo"})] + # tool_start carries the original full id so the conversation + # roundtrip can replay it verbatim. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "functions.web_search:0" + def test_render_html_emits_provisional_tool_start(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -1360,8 +2149,12 @@ class TestLoopBehaviour: assert captured_tool_names[2] == ["web_search", "python"] def test_duplicate_noop_does_not_consume_budget_at_small_cap(self): - # A duplicate no-op turn must NOT spend the tool budget: only turns that execute a tool - # count (GGUF parity), so a distinct call can still follow at max_tool_iterations=2. + # A duplicate/disabled no-op turn is a correction turn and must NOT spend the + # caller's tool budget, so with max_tool_iterations=2 the model can still make a + # DISTINCT valid call after repeating one. Only turns that actually execute a + # tool count -- matching the GGUF loop. (The budget used to be charged per + # non-re-prompt iteration, so the duplicate burned the second slot and the third + # turn was sent with no tools, dropping the ``python`` call.) captured_tool_names: list[list[str]] = [] turns = iter( [ @@ -1657,7 +2450,8 @@ class TestLoopRePrompt: assert contents and contents[-1]["text"].strip() == "4" def test_max_reprompts_capped_at_three(self): - # Model keeps stalling with intent -- after 3 re-prompts the loop must give up. + # Model keeps stalling with intent -- after 3 re-prompts the + # loop must give up rather than burn forever. turns = [["Let me search for that."]] * 6 # well over the cap loop, exec_fn = _make_loop( turns = turns, @@ -1670,7 +2464,9 @@ class TestLoopRePrompt: assert statuses and statuses[-1]["text"] == "" def test_short_intent_below_buffer_threshold_triggers_reprompt(self): - # Short emission that never exits BUFFERING must still trigger the intent re-prompt. + # Short emission that never exits BUFFERING (< 32 chars + no + # marker prefix). The unified buffer-end path must still + # trigger the intent re-prompt, not silently terminate. loop, exec_fn = _make_loop( turns = [ ["Let me check."], @@ -1683,7 +2479,9 @@ class TestLoopRePrompt: assert exec_fn.calls == [("web_search", {"query": "x"})] def test_reprompt_does_not_consume_tool_budget(self): - # max_tool_iterations=1: the re-prompt must not eat the slot, so the real call still runs. + # max_tool_iterations=1: one re-prompt, then one real tool call, + # then the budget-exhausted final answer must still fire. If the + # re-prompt ate the slot the tool call would never run. loop, exec_fn = _make_loop( turns = [ # 1. Intent stall (re-prompt 1/3). @@ -1714,7 +2512,8 @@ class TestLoopCanonicalHealKey: exec_results = ["1\n"], ) events = _collect_events(loop) - # The bare string must heal to {"code": ...}, not {"query": ...}, so the python sandbox runs it. + # The bare string must heal to {"code": "print(1)"}, not + # {"query": ...}, so the python sandbox actually executes it. assert exec_fn.calls == [("python", {"code": "print(1)"})] def test_terminal_bare_string_heals_to_command(self): @@ -1744,7 +2543,10 @@ class TestGGUFSafetensorsHealingParity: """Pin GGUF vs safetensors/MLX loop parity so a regression on either side breaks CI.""" def test_gguf_imports_shared_signal_markers(self): - # The GGUF BUFFERING machine must wake on every shared emission marker, else calls slip past as prose. + # The GGUF BUFFERING state machine must wake on every emission + # marker the shared parser knows -- otherwise Llama-3 / Mistral + # / Gemma 4 emissions slip past as plain prose when the + # llama-server structured channel fails. import inspect from core.inference.llama_cpp import LlamaCppBackend @@ -1756,7 +2558,10 @@ class TestGGUFSafetensorsHealingParity: ) def test_gguf_uses_shared_strip_helper(self): - # The GGUF stream-cleanup must delegate to the shared strip_tool_markup for every family. + # The GGUF stream-cleanup function must delegate to the shared + # strip_tool_markup so closed-pair markup is removed for every + # emission family (Llama-3 <|python_tag|>, Mistral [TOOL_CALLS], + # Gemma 4 <|tool_call>...). import inspect from core.inference.llama_cpp import LlamaCppBackend @@ -1767,7 +2572,11 @@ class TestGGUFSafetensorsHealingParity: ), "GGUF stream cleanup must delegate to the shared strip_tool_markup helper" def test_gguf_uses_canonical_heal_keys(self): - # GGUF and safetensors heal a bare-string argument to the same canonical key via the shared coerce_tool_arguments. + # GGUF and safetensors heal a bare-string ``arguments`` to the same + # per-tool canonical key -- ``code`` for python, ``command`` for + # terminal, ``query`` for everything else. The mapping is centralised in + # the shared ToolLoopController (both backends route bare-string args + # through ``coerce_tool_arguments``), so the two paths cannot drift. from core.inference.tool_loop_controller import ( _CANONICAL_HEAL_ARG, coerce_tool_arguments, @@ -1786,7 +2595,9 @@ class TestGGUFSafetensorsHealingParity: } def test_intent_regex_matches_same_phrases_as_gguf(self): - # The intent re-prompt regex must match the SAME phrases on both backends. + # The intent re-prompt regex must match the SAME forward-looking + # phrases on both backends so behaviour is the same on Mac (MLX + # / safetensors) and on Linux (GGUF). from core.inference.llama_cpp import _INTENT_SIGNAL as gguf_re from core.inference.safetensors_agentic import ( _INTENT_SIGNAL as sf_re, @@ -1811,7 +2622,8 @@ class TestGGUFSafetensorsHealingParity: "I can help with that.", "I should mention", "Let's go.", - # Negated intent is a refusal, not a plan: neither backend may re-prompt on it. + # Negated intent is a refusal, not a plan: neither backend may + # force a tool-call re-prompt on it. "I will not search the web for that.", "I'll never call that tool.", ): @@ -2240,6 +3052,28 @@ class TestGuardrails: and event.get("type") in {"tool_start", "tool_end"} ] + def test_same_turn_distinct_calls_are_capped(self): + # >_MAX_TOOL_CALLS_PER_TURN DISTINCT calls in one turn must be capped so a runaway turn + # cannot fan out into many executions (the GGUF path is held back by llama-server's lazy ... + from core.inference.safetensors_agentic import _MAX_TOOL_CALLS_PER_TURN + + n = _MAX_TOOL_CALLS_PER_TURN + 4 + turn = "".join( + '{"name":"web_search","arguments":{"query":"q%d"}}' % i + for i in range(n) + ) + loop, exec_fn = _make_loop( + turns = [[turn], ["final"]], + exec_results = ["r"] * n, + max_tool_iterations = 2, + ) + _collect_events(loop) + assert len(exec_fn.calls) == _MAX_TOOL_CALLS_PER_TURN + # The first N distinct queries executed, in document order. + assert [a["query"] for _name, a in exec_fn.calls] == [ + "q%d" % i for i in range(_MAX_TOOL_CALLS_PER_TURN) + ] + def test_coerce_string_args_python_uses_code_key(self): assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"} @@ -2283,7 +3117,8 @@ class TestRoutesPythonTagStrip: """``_TOOL_XML_RE`` must consume multi-line code, embedded JSON, and bare ``<`` (earlier ``[^\n<]*`` / ``[^\n]*`` revisions leaked tails); the streaming route-level strip is the regression-prone path.""" def _strip(self, text: str) -> str: - # Import inside the test so a routes-module import error doesn't fail collection. + # Import inside the test so a routes-module import error does + # not blow up the entire test file at collection time. from routes.inference import _strip_tool_xml return _strip_tool_xml(text) @@ -2293,7 +3128,8 @@ class TestRoutesPythonTagStrip: assert self._strip(text) == "" def test_python_tag_with_less_than_in_code(self): - # 5615 regression: a literal < inside code must NOT terminate the strip early. + # 5615 regression: literal ``<`` inside code must NOT terminate + # the strip early. text = '<|python_tag|>python.call(code="if x < 10: pass")' assert self._strip(text) == "" @@ -2303,7 +3139,7 @@ class TestRoutesPythonTagStrip: assert self._strip(text) == "" def test_python_tag_multiline_with_less_than(self): - # Combined: multi-line code AND literal < in code. + # Combined: multi-line code AND literal ``<`` in code. text = ( '<|python_tag|>python.call(code="for i in range(10):\n' " if i < 5:\n" @@ -2312,7 +3148,8 @@ class TestRoutesPythonTagStrip: assert self._strip(text) == "" def test_python_tag_stops_at_eom_sentinel(self): - # Strip stops at the next Llama-3 <| sentinel so trailing assistant content survives. + # Strip stops at the next Llama-3 ``<|`` sentinel so any + # trailing assistant content survives. text = '<|python_tag|>python.call(code="multi\nline")' "<|eom_id|>final answer text" assert self._strip(text) == "<|eom_id|>final answer text" @@ -2326,20 +3163,25 @@ class TestRoutesPythonTagStrip: assert self._strip(text) == "" def test_python_tag_with_eom_then_trailing_python_tag(self): - # Two python_tag emissions back-to-back across a sentinel: both strip independently. + # Two python_tag emissions back-to-back across a sentinel: both + # should strip independently. text = ( '<|python_tag|>brave_search.call(query="a")' "<|eom_id|>" '<|python_tag|>python.call(code="x=1")' ) - # <|eom_id|> between the two strips remains; both python_tag blocks are consumed. + # ``<|eom_id|>`` between the two strips remains; both + # python_tag blocks are fully consumed. assert self._strip(text) == "<|eom_id|>" # Robustness fixes uncovered while validating against vLLM / sglang. class TestParserRobustness: def test_tool_call_json_accepts_parameters_key(self): - # Hermes wrapper using parameters instead of arguments; this path now accepts both keys. + # Hermes wrapper around a Llama-3.2 bare-JSON object that uses + # ``parameters`` instead of ``arguments``. The bare-JSON and + # python_tag paths already accept both keys; this path now does + # too. Was extracting name only and silently dropping the args. import json text = "\n" '{"name": "search", "parameters": {"q": "ramen"}}\n' "" @@ -2349,7 +3191,8 @@ class TestParserRobustness: assert json.loads(result[0]["function"]["arguments"]) == {"q": "ramen"} def test_function_xml_attribute_form(self): - # MiniCPM-5 / MiniMax-M2 attribute syntax: v. + # MiniCPM-5 / MiniMax-M2 attribute syntax: + # ``v``. import json text = '' 'Tokyo' "" @@ -2373,7 +3216,8 @@ class TestParserRobustness: assert args == {"city": "Tokyo", "unit": "celsius"} def test_function_xml_legacy_equals_form_still_works(self): - # Regression guard: the old v syntax must keep parsing after the regex broadening. + # Regression guard: the old ``v`` + # syntax must keep parsing after the regex broadening. import json text = "Tokyo" @@ -2383,17 +3227,24 @@ class TestParserRobustness: assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} def test_function_attribute_form_has_tool_signal(self): - # The standalone form must flip the streaming buffer, else the call is dropped. + # The standalone ```` attribute form must flip + # the streaming buffer; otherwise the end-of-turn safety-net parse in + # the agentic loop is gated off and the real call is dropped. assert has_tool_signal('') is True def test_function_attribute_form_strip_markup(self): - # The attribute form must also be stripped from displayed text, like . + # The attribute form must also be stripped from displayed text, like + # the legacy ```` form. text = 'result X' assert strip_tool_markup(text, final = True) == "result" def test_llama3_chat_template_round_trip(self): - # Llama-3.x prefixes assistant turns with <|start_header_id|>...<|end_header_id|>; the - # sentinel-strip must reach past the role label to the JSON body, else history calls drop. + # Meta's official Llama-3.x chat template prefixes every + # assistant turn with + # ``<|start_header_id|>assistant<|end_header_id|>\n\n``. The + # sentinel-strip in ``_parse_llama3_bare_json`` must reach past + # the role label to the JSON body, else every round-tripped + # tool call in history silently drops. import json text = ( @@ -2418,7 +3269,8 @@ class TestParserRobustness: assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} def test_llama3_round_trip_with_eot_prefix(self): - # Prior turn closes with <|eot_id|>, then the new header opens; both sentinels + role must be consumed. + # Prior assistant turn closes with ``<|eot_id|>``, then the + # new header opens. Both sentinels + the role must be consumed. import json text = ( @@ -2430,7 +3282,10 @@ class TestParserRobustness: assert result[0]["function"]["name"] == "f" def test_function_xml_followed_by_prose(self): - # Body must terminate at even without a wrapper, else prose leaks into the value. + # Models routinely follow a tool call with explanatory prose. + # Body must terminate at ```` even without a + # ```` wrapper, else trailing prose leaks into the + # last parameter value. import json text = ( @@ -2456,8 +3311,236 @@ class TestParserRobustness: assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} +def test_render_with_native_template_returns_render_only_when_tools_emitted(): + # The native-template fallback re-renders with the model's repo template when an override drops + # the tools schema. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_native_template + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + model_info = { + "native_chat_template": "TPL", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|TOOLS=" + ",".join(t["function"]["name"] for t in tools) if tools else "") + + def ignoring(tokenizer, msgs, *, tools, **_kw): + return "".join(m["content"] for m in msgs) # never reflects tools + + out = render_native_template( + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + assert out == "hi|TOOLS=web_search" + # The native template must be restored on the live tokenizer after probing. + assert model_info["tokenizer"].chat_template == "OVERRIDE" + + assert ( + render_native_template( + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = ignoring, + ) + is None + ) + + # No tokenizer and no processor -> return None instead of an AttributeError. + no_tok = {"native_chat_template": "TPL"} + assert ( + render_native_template( + model_info = no_tok, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + is None + ) + + +def test_render_with_native_template_does_not_mutate_shared_tokenizer(): + # The shared tokenizer must never carry the temporary native template, even mid-render: this + # runs outside the generation lock, so a concurrent request could otherwise render with the ... + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_native_template + + shared = SimpleNamespace(chat_template = "OVERRIDE") + seen = [] + + def capture(tokenizer, msgs, *, tools, **_kw): + seen.append((tokenizer is shared, shared.chat_template)) + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + model_info = {"native_chat_template": "TPL", "tokenizer": shared} + render_native_template( + model_info = model_info, + active_model_name = "x", + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + apply_fn = capture, + ) + # Rendering happened on a copy, and the shared tokenizer stayed "OVERRIDE" + # throughout (never the temporary "TPL"). + assert seen and all(not is_shared for is_shared, _ in seen) + assert all(tpl == "OVERRIDE" for _, tpl in seen) + assert shared.chat_template == "OVERRIDE" + + +def test_native_template_loads_from_base_model_for_lora(monkeypatch): + # For a LoRA adapter the chat template lives on the base model; active_model_name + # is the adapter id and may ship no template. The loader must read base_model. + from types import SimpleNamespace + + import transformers + + from core.inference.chat_template_helpers import render_native_template + + captured = {} + + def fake_from_pretrained(name, *args, **kwargs): + captured["source"] = name + return SimpleNamespace(chat_template = "BASE_TPL") + + monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", fake_from_pretrained) + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + model_info = { + "base_model": "base/model-id", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + out = render_native_template( + model_info = model_info, + active_model_name = "adapter/path", + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + apply_fn = emitting, + ) + assert captured["source"] == "base/model-id" + assert out == "hi|T" + + +def test_render_with_native_template_fallback_swaps_when_override_drops_tools(): + # The shared gate (used by the transformers and MLX backends): when the live render is + # identical with and without tools, re-render with the native template and return it. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + # apply_fn that IGNORES tools -> live render drops the schema. + def ignoring(tokenizer, msgs, *, tools, **_kw): + return "".join(m["content"] for m in msgs) + + model_info = { + "native_chat_template": "TPL", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + + # Native render emits the tools, so the fallback swaps to it. + def native_emits(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|TOOLS" if tools else "") + + out = render_with_native_template_fallback( + formatted_prompt = ignoring(None, messages, tools = tools), + tokenizer = SimpleNamespace(), + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = lambda tok, msgs, *, tools, **kw: ( + native_emits(tok, msgs, tools = tools) + if getattr(tok, "chat_template", None) == "TPL" + else ignoring(tok, msgs, tools = tools) + ), + ) + assert out == "hi|TOOLS", out + + +def test_render_with_native_template_fallback_keeps_prompt_when_tools_emitted(): + # Live render already differs with vs without tools -> no fallback, returned + # unchanged. Also a no-tools call is a passthrough. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + kept = render_with_native_template_fallback( + formatted_prompt = emitting(None, messages, tools = tools), + tokenizer = SimpleNamespace(), + model_info = {"native_chat_template": "TPL", "tokenizer": SimpleNamespace()}, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + assert kept == "hi|T", kept + + # No tools -> passthrough (native template never consulted). + passthrough = render_with_native_template_fallback( + formatted_prompt = "hi", + tokenizer = SimpleNamespace(), + model_info = {}, + active_model_name = "x", + messages = messages, + tools = None, + apply_fn = emitting, + ) + assert passthrough == "hi" + + +def test_render_with_native_template_fallback_keeps_prompt_when_no_tools_probe_raises(): + # A template that REQUIRES tools can raise on the no-tools probe. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def raises_without_tools(tokenizer, msgs, *, tools, **_kw): + if not tools: + raise RuntimeError("template requires tools") + return "".join(m["content"] for m in msgs) + "|T" + + out = render_with_native_template_fallback( + formatted_prompt = "hi|T", + tokenizer = SimpleNamespace(), + model_info = {"native_chat_template": "TPL", "tokenizer": SimpleNamespace()}, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = raises_without_tools, + ) + assert out == "hi|T", out + + def test_truncated_bare_json_at_eof_is_not_leaked(): - # Stream ends mid bare-JSON: the held fragment must be dropped at EOF, not flushed as content. + # Stream ends mid bare-JSON object: the held fragment must be dropped at the + # EOF resolver, not flushed as plain assistant content (GGUF parity). loop, _exec = _make_loop( turns = [['{"name":"web_search","parameters":{"query":"weather in S']], max_tool_iterations = 1, @@ -2468,7 +3551,9 @@ def test_truncated_bare_json_at_eof_is_not_leaked(): def test_oversized_bare_json_call_is_not_leaked_and_executes(): - # A bare-JSON call exceeding _MAX_BARE_JSON_BUFFER must DRAIN, not stream the prefix, and still execute. + # A bare-JSON call whose arguments exceed _MAX_BARE_JSON_BUFFER must DRAIN + # (suppress) rather than stream the raw JSON prefix, and still execute once + # the full object is parsed by the safety net. from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) @@ -2483,7 +3568,8 @@ def test_oversized_bare_json_call_is_not_leaked_and_executes(): def test_oversized_plain_json_answer_still_streams(): - # A giant plain JSON answer (no "name" key) is NOT a call and must still stream. + # A giant plain JSON answer (no "name" key) is NOT a tool call and must still + # stream -- the oversized DRAIN route is gated on a "name" key. from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) @@ -2496,7 +3582,9 @@ def test_oversized_plain_json_answer_still_streams(): def test_oversized_disabled_name_json_answer_still_streams(): - # A giant still-open JSON answer whose "name" is NOT an enabled tool must stream, not drain. + # A giant still-open JSON answer whose "name" is NOT an enabled tool must stream: + # the oversized DRAIN branch was gated only on the presence of a "name" key, so a + # large ordinary record ({"name":"Alice",...}) was drained instead of shown. from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) @@ -2510,7 +3598,8 @@ def test_oversized_disabled_name_json_answer_still_streams(): def test_truncated_disabled_name_json_is_shown_at_eof(): - # A truncated JSON answer whose name is not an enabled tool must be shown at EOF. + # A truncated ordinary JSON answer whose name is not an enabled tool, held to EOF, + # must be shown -- the EOF bare-JSON DRAIN branch was gated only on a "name" key. truncated = '{"name":"Alice","parameters":{"age":' loop, exec_fn = _make_loop(turns = [[truncated]], max_tool_iterations = 1) events = _collect_events(loop) @@ -2520,7 +3609,9 @@ def test_truncated_disabled_name_json_is_shown_at_eof(): def test_truncated_plain_json_with_nested_enabled_name_is_visible(): - # A truncated answer with only a NESTED "name" must be shown: the gate uses the TOP-LEVEL name. + # A truncated ordinary JSON answer with a NESTED ``"name"`` matching an enabled + # tool ({"result":{"name":"web_search",...) must be shown, not suppressed: the + # gate now extracts the TOP-LEVEL name only, so the nested field is just data. loop, exec_fn = _make_loop( turns = [['{"result":{"name":"web_search","age":']], max_tool_iterations = 1, @@ -2532,7 +3623,8 @@ def test_truncated_plain_json_with_nested_enabled_name_is_visible(): def test_bare_json_call_not_replayed_in_next_turn_content(): - # After a bare-JSON call executes, the next-turn assistant content must not contain the raw call. + # After a complete bare-JSON call executes, the assistant content fed to the + # next turn must not contain the raw call (next-turn contamination). captured: list[list[dict]] = [] exec_fn = FakeExecuteTool(["RESULT"]) @@ -2562,7 +3654,10 @@ if __name__ == "__main__": def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): - # With Auto-Heal OFF a truncated enabled-name bare-JSON fragment stays visible; with it ON, suppressed. + # F3: with Auto-Heal OFF, a truncated ENABLED-name bare-JSON fragment that did + # not parse must stay visible (disabled-Auto-Heal contract: malformed markup is + # preserved), matching the XML strip in the same drain branch. With Auto-Heal ON + # the same fragment is suppressed. trunc = '{"name":"web_search","parameters":{"query":"weather' off, exec_off = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = False) events_off = _collect_events(off) @@ -2578,7 +3673,9 @@ def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): def test_looks_like_enabled_bare_json_accepts_function_alias(): - # The buffering gate must recognise the "function" bare-JSON alias, so it is buffered, not streamed. + # The safetensors buffering gate must recognise the "function" bare-JSON alias + # the parser accepts, so a truncated/complete {"function":} call is + # buffered/healed instead of streaming as visible content. from core.inference.safetensors_agentic import _looks_like_enabled_bare_json enabled = {"web_search"} @@ -2591,7 +3688,8 @@ def test_looks_like_enabled_bare_json_accepts_function_alias(): class TestFalseAlarmMarkerProse: def test_leading_marker_prose_streams_intact(self): - # An answer starting with a literal marker is a false alarm: the full prose must reach the client. + # An answer that starts with a literal marker is a false alarm: the + # drain finds no calls and the full prose must reach the client. text = "[TOOL_CALLS] is the Mistral tool marker. More prose after." loop, exec_fn = _make_loop(turns = [[text]]) events = _collect_events(loop) @@ -2600,7 +3698,8 @@ class TestFalseAlarmMarkerProse: assert texts and texts[-1] == text def test_chained_bare_json_calls_not_replayed_in_history(self): - # Both chained calls execute; the next-turn history must not contain the second call's raw JSON. + # Both chained calls execute; the kept content (next-turn assistant + # history) must not contain the second call's raw JSON. chained = ( '{"name":"web_search","parameters":{"q":"first"}};' '{"name":"python","parameters":{"code":"x"}}' diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index fded2a8443..7f47140b8d 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -72,10 +72,8 @@ class TestFunctionStyleTrailingText: assert call == {"name": "python", "arguments": {"code": 'print("")'}} def test_closed_function_with_trailing_prose_heal_path(self): - # Regression: the heal / finalize path (allow_incomplete=True) used to fold - # and the trailing prose into the argument and drop - # the prose from visible content. It must now match the strict path -- keep a - # clean argument and leave the trailing prose outside the call span. + # Regression: the heal path (allow_incomplete=True) must match the strict path -- + # keep a clean argument and leave trailing prose outside the call span. text = "cats trailing words" calls = parse_tool_calls_from_text(text, allow_incomplete = True) assert len(calls) == 1 @@ -103,7 +101,8 @@ class TestFunctionStyleTrailingText: assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] def test_attribute_form_literal_close_tag_is_preserved(self): - # Attribute form ends at the LAST , so a literal close inside code survives. + # The attribute form (MiniCPM-5 / MiniMax-M2) also ends at the + # LAST , so a literal close tag inside a code argument survives. text = ( '' 'print("")' @@ -113,7 +112,8 @@ class TestFunctionStyleTrailingText: assert call == {"name": "python", "arguments": {"code": 'print("")'}} def test_closed_zero_param_attribute_call_is_accepted_in_strict_mode(self): - # A closed zero-param call is valid; strict mode must not treat it as truncated. + # A closed call with no parameters is a valid zero-argument call; strict + # mode must not treat the empty parameter list as a truncated call. assert _only('') == {"name": "ping", "arguments": {}} # A no-arg call that never closes is still rejected as truncated. assert parse_tool_calls_from_text('', allow_incomplete = False) == [] @@ -231,9 +231,8 @@ class TestHealingPathUnaffected: assert calls[0]["function"]["name"] == "web_search" def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self): - # allow_incomplete exists for truncated output; a call that DID close - # must parse identically to strict mode, leaving prose after - # out of the last parameter and out of the removal span. + # A call that DID close must parse identically to strict mode, leaving prose after + # out of the last parameter and the removal span. from core.tool_healing import parse_tool_calls_from_text as parse_with_spans text = "cats trailing" @@ -246,7 +245,8 @@ class TestHealingPathUnaffected: ) def test_wrapperless_fallback_calls_carry_spans(self): - # The wrapperless fallback must report spans so consumers strip exactly the markup. + # The wrapperless function-XML fallback must report spans too, so with_spans + # consumers strip exactly the promoted markup (through when closed). from core.tool_healing import parse_tool_calls_from_text as parse_with_spans closed = "before cats after" @@ -266,6 +266,50 @@ class TestHealingPathUnaffected: assert healed[span[0] : span[1]] == "dogs" +class TestGlmStrict: + def test_closed_glm_call_is_accepted(self): + text = ( + "get_weather\n" + "city\nParis\n" + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + + def test_unclosed_glm_call_is_rejected(self): + # No close: truncated, reject with Auto-Heal off. + text = "get_weather\ncity\nParis" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestKimiStrict: + _SB = "<|tool_calls_section_begin|>" + _KB = "<|tool_call_begin|>" + _AB = "<|tool_call_argument_begin|>" + _KE = "<|tool_call_end|>" + _SE = "<|tool_calls_section_end|>" + + def test_full_kimi_call_is_accepted(self): + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + self._SE + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "x" + + def test_kimi_call_without_call_end_is_rejected(self): + # Section closed but the call lacks <|tool_call_end|>: reject in strict. + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._SE + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + def test_kimi_without_section_end_is_rejected(self): + # No <|tool_calls_section_end|>: truncated section, reject in strict. + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + class TestParserLinearity: """Llama-3 ``.call`` kwargs and Mistral-array healing must stay linear (a regex-per-offset blew up on long truncated bodies).""" @@ -293,6 +337,27 @@ class TestParserLinearity: parse_tool_calls_from_text(text, allow_incomplete = True) assert time.perf_counter() - t0 < 2.0 + def test_gemma_wrapperless_deep_nesting_is_linear(self): + # Wrapper-less Gemma ``call:f{a:{a:{...}}}`` deep nesting must parse in linear time (no quadratic re-scan). + import time + + def nested(d): + return "call:f{a:" + "{a:" * d + "x:1" + "}" * d + "}" + + def best_ms(depth): + text = nested(depth) + best = float("inf") + for _ in range(5): + t0 = time.perf_counter() + calls = parse_tool_calls_from_text(text) + best = min(best, time.perf_counter() - t0) + assert calls and json.loads(calls[0]["function"]["arguments"]), "nested args dropped" + return best + + t200 = best_ms(200) + t400 = best_ms(400) + assert t400 < t200 * 3.0, (t200, t400) + def test_llama3_call_kwargs_still_parse(self): text = '<|python_tag|>do.call(s="hi 😀", n=42, f=1.5, b=true, z=null)' calls = parse_tool_calls_from_text(text, allow_incomplete = True) @@ -334,7 +399,8 @@ class TestLlamaBuiltinChainAndNesting: assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} def test_nested_python_tag_in_json_string_arg_is_not_a_call(self): - # A <|python_tag|> literal inside a code arg is data: the outer "python" call wins. + # A code arg literally containing a <|python_tag|>...call(...) string: the real call is the + # outer "python", not the nested "os" -- the scan stays anchored to the first tag. text = ( '<|python_tag|>{"name":"python","parameters":' '{"code":"<|python_tag|>os.call(\'rm -rf /\')"}}' @@ -353,6 +419,41 @@ class TestLlamaBuiltinChainAndNesting: assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} +def test_glm_open_does_not_parse_spaced_prose_as_tool_name(): + # The GLM NAME opener must reject spaced literal prose (V10); only a + # valid [\w.\-]+ name (followed by newline//) is a call. + assert parse_tool_calls_from_text("not a call") == [] + ok = parse_tool_calls_from_text( + "get_weather\ncity\nNYC\n" + ) + assert [c["function"]["name"] for c in ok] == ["get_weather"] + + +def test_deepseek_r1_missing_call_terminator_rejected_in_strict_mode(): + # R1 must reject a fenced call whose closing ``` + <｜tool▁call▁end｜> never + # arrived when Auto-Heal is off, matching V3/V3.1 strictness (V6). + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC"}' + "<｜tool▁calls▁end｜>" + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +def test_deepseek_r1_complete_call_accepted_in_strict_mode(): + # A fully-terminated R1 call (close fence + per-call end) is still accepted. + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC"}\n' + "```<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 and calls[0]["function"]["name"] == "get_weather" + + def test_strip_leading_bare_json_call_drops_complete_call(): from core.inference.tool_call_parser import strip_leading_bare_json_call @@ -386,6 +487,57 @@ def test_strip_leading_bare_json_call_preserves_plain_json_and_prose(): assert strip_leading_bare_json_call("just a sentence.") == "just a sentence." +def test_glm_literal_close_tag_in_string_arg_not_truncated(): + import json + + from core.inference.tool_call_parser import parse_tool_calls_from_text + + # A GLM string argument may legitimately contain the literal close tag ````. + text = ( + "run_code\n" + "code\n" + 'print("")\n' + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == 'print("")', args + + +def test_glm_truncated_block_rejected_in_strict_mode_but_healed_otherwise(): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + # No close: strict mode (Auto-Heal off) rejects the truncated + # block; with Auto-Heal it keeps the partial call. + text = "get_weather\ncity\nNYC" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 and healed[0]["function"]["name"] == "get_weather" + + +def test_truncated_wrapperless_gemma_call_is_stripped(): + from core.inference.tool_call_parser import strip_tool_markup + + # A wrapper-less Gemma ``call:NAME{...`` cut off mid-arguments (no closing + # brace) must not leak the raw call into the visible stream. + text = 'Sure!\ncall:web_search{"query": "weather in San Fr' + stripped = strip_tool_markup(text, final = True) + assert "call:web_search" not in stripped, repr(stripped) + assert stripped.strip() == "Sure!" + + +def test_complete_wrapperless_gemma_call_keeps_trailing_prose(): + from core.inference.tool_call_parser import strip_tool_markup + + # The truncation pattern must run AFTER the closed form, so a complete call + # followed by prose keeps the prose instead of eating to EOS. + text = 'call:web_search{"query": "cats"} Here you go.' + stripped = strip_tool_markup(text, final = True) + assert "call:web_search" not in stripped + assert stripped.strip() == "Here you go." + + def test_bare_json_gated_on_enabled_tool_names(): from core.inference.tool_call_parser import parse_tool_calls_from_text @@ -421,7 +573,8 @@ def test_strip_leading_bare_json_call_gated_on_enabled_tool_names(): def test_function_xml_strip_keeps_literal_close_tag_in_param_value(): from core.inference.tool_call_parser import strip_tool_markup - # Strip uses the LAST so a literal in a value survives; calls strip independently. + # The strip uses the LAST (like the parser) so a literal in a value doesn't + # truncate it; separate calls still strip independently. text = 'print("") done' assert strip_tool_markup(text, final = True) == "done" two = ( @@ -434,7 +587,8 @@ def test_function_xml_strip_keeps_literal_close_tag_in_param_value(): def test_function_xml_strip_keeps_trailing_text_after_literal_open_tag(): from core.inference.tool_call_parser import parse_tool_calls_from_text, strip_tool_markup - # A literal opener inside a value is data: the strip keeps " done". + # A literal ```` opener inside a parameter value is data, not a call: the scan-based + # strip keeps " done" (the old negative-lookahead regex ate the trailing prose). text = 'print("") done' assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python" assert strip_tool_markup(text, final = True) == "done" @@ -446,10 +600,11 @@ def test_function_xml_strip_keeps_trailing_text_after_literal_open_tag(): def test_final_strip_removes_magistral_think_reasoning(): from core.inference.tool_call_parser import strip_tool_markup - # Magistral reasoning is [THINK]...[/THINK]; end-of-turn must drop it. + # Magistral emits reasoning as ``[THINK]...[/THINK]`` (bracket form, not ````); + # at end-of-turn it must be dropped so it doesn't leak into display / history. text = "[THINK]The user greeted me, I should say hi.[/THINK]Hello! How can I help?" assert strip_tool_markup(text, final = True) == "Hello! How can I help?" - # A [TOOL_CALLS] living inside the reasoning goes with it. + # A ``[TOOL_CALLS]`` living inside the reasoning goes with it. with_call = '[THINK]Maybe I should search.[/THINK][TOOL_CALLS]search{"q":"x"}' assert strip_tool_markup(with_call, final = True) == "" @@ -457,7 +612,8 @@ def test_final_strip_removes_magistral_think_reasoning(): def test_streaming_strip_keeps_magistral_think_buffered(): from core.inference.tool_call_parser import strip_tool_markup - # Mid-stream (final=False) leaves the reasoning block intact; only end-of-turn removes it. + # Mid-stream (final=False) the reasoning block is left intact; only the + # end-of-turn pass removes it. text = "[THINK]still thinking" assert strip_tool_markup(text, final = False) == text @@ -465,7 +621,7 @@ def test_streaming_strip_keeps_magistral_think_buffered(): def test_final_strip_leaves_non_magistral_bracket_text_untouched(): from core.inference.tool_call_parser import strip_tool_markup - # Only a LEADING [THINK] block is reasoning; unrelated bracketed prose stays. + # Only a LEADING ``[THINK]`` block is reasoning; unrelated bracketed prose stays. text = "See [THINK about it] later" assert strip_tool_markup(text, final = True) == "See [THINK about it] later" @@ -473,7 +629,8 @@ def test_final_strip_leaves_non_magistral_bracket_text_untouched(): def test_strip_leading_bare_json_call_ignores_nested_name(): from core.inference.tool_call_parser import strip_leading_bare_json_call - # A nested "name" must NOT gate the strip; the JSON answer is kept verbatim. + # A nested ``"name"`` must NOT gate the strip (only a TOP-LEVEL enabled name is a call); the + # ordinary JSON answer is kept verbatim, truncated or complete. nested_trunc = '{"result":{"name":"web_search","age":' nested_full = '{"result":{"name":"web_search","age":1}}' assert strip_leading_bare_json_call(nested_trunc, {"web_search"}) == nested_trunc @@ -493,7 +650,8 @@ def test_mistral_single_object_call_is_stripped_for_display(): parse_tool_calls_from_text, ) - # The parser accepts single-object [TOOL_CALLS]{...}, so the strip must remove it too. + # The parser accepts the single-object [TOOL_CALLS]{...} shape, so the display + # strip must remove it too (asymmetry would leak the raw object). text = '[TOOL_CALLS]{"name":"web_search","arguments":{"filters":{"date":"2024"}}} tail' assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == ["web_search"] assert _strip_mistral_closed_calls(text) == " tail" @@ -502,7 +660,8 @@ def test_mistral_single_object_call_is_stripped_for_display(): def test_tool_call_parser_declares_future_annotations_for_py39_import(): - # PEP 604 X | None annotations need `from __future__ import annotations` on py3.9; guard it stays. + # F1: the parser is imported standalone on python >=3.9, where its PEP 604 ``X | None`` + # annotations need ``from __future__ import annotations``; guard that the import stays. from pathlib import Path src = ( Path(__file__).resolve().parent.parent / "core" / "inference" / "tool_call_parser.py" @@ -510,8 +669,23 @@ def test_tool_call_parser_declares_future_annotations_for_py39_import(): assert "from __future__ import annotations" in src +def test_glm_strip_treats_literal_close_tag_in_arg_value_as_data(): + # Core strip parity: a literal inside a GLM is argument data, so the whole call is stripped (no leaked tail). + from core.inference.tool_call_parser import strip_tool_markup + + text = ( + "web_search\nquery\n" + "see tag\n tail" + ) + assert strip_tool_markup(text, final = True) == "tail" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "see tag"} + + def test_bare_json_function_alias_parses_and_strips_symmetrically(): - # The "function" alias for the call name must parse and strip symmetrically. + # The bare-JSON parser accepts the "function" alias for the call name; + # strip_leading_bare_json_call must recognise it too (parser/strip symmetry). from core.inference.tool_call_parser import ( parse_tool_calls_from_text, strip_leading_bare_json_call, @@ -585,6 +759,92 @@ class TestHealerSignalAlignment: assert not list(healer.finalize()) or all(k == "text" for k, _v in healer.finalize()) +class TestGemmaWrapperlessLiteralMarkers: + """Wrapper-less Gemma calls whose ARGUMENTS mention Gemma's own markup. + + The tool_healing deferral must key on an actual wrapped opener + (``<|tool_call>call:...``), not the wrapper literal anywhere in content: + a query about the marker has nothing tool_healing can parse, and deferring + it loses the call entirely (not executed AND stripped from display).""" + + def test_marker_literal_in_argument_still_parses(self): + text = 'call:web_search{query:"what does <|tool_call> mean"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what does <|tool_call> mean" + + def test_real_wrapped_call_still_deferred_to_tool_healing(self): + from core.inference.tool_call_parser import _parse_gemma_tool_calls + + # An actual wrapped opener present: the Gemma fallback must keep + # deferring to the shared tool_healing parser that owns that form. + text = '<|tool_call>call:web_search{query:<|"|>cats<|"|>}' + assert _parse_gemma_tool_calls(text, id_offset = 0) == [] + + def test_single_quoted_brace_does_not_truncate_code(self): + text = "call:python{code:print('}')}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "print('}')" + + def test_single_quoted_brace_strip_span_covers_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = "call:python{code:print('}')} Done." + stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"python"}) + assert "call:python" not in stripped + assert "')}" not in stripped + assert stripped.strip() == "Done." + + +class TestGlmEmbeddedClosePair: + """A GLM value whose string literal embeds the full close-tag pair + ```` (code documenting the GLM format) must not be + truncated at the embedded pair: a structural close sits at balanced quote + state, an embedded one is inside an open string literal.""" + + def test_embedded_pair_inside_quoted_value_not_structural(self): + text = ( + "python\n" + "code\n" + 'print("")\nx = 1\n' + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == 'print("")\nx = 1' + + def test_strip_covers_the_full_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = ( + "python\n" + "code\n" + 'print("")\nx = 1\n' + " Done." + ) + stripped = strip_tool_markup(text, final = True) + assert "arg_value" not in stripped + assert stripped.strip() == "Done." + + def test_unbalanced_apostrophe_falls_back_to_first_candidate(self): + # Prose-like value with an apostrophe: no candidate reaches balanced + # quote state, so the first token-valid close wins (prior behavior). + text = ( + "web_search\n" + "query\n" + "it's fine\n" + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "it's fine" + + class TestPythonTagLiteralInsideMistralArgs: """A python_tag LITERAL inside a leading Mistral call's arguments is data; the outer call executes.""" @@ -719,6 +979,60 @@ class TestMagistralThinkRehearsal: assert parse_tool_calls_from_text(text) == [] +class TestGemmaUnquotedApostrophes: + """Quotes open strings only at value-start context: an apostrophe inside + an unquoted wrapper-less value (contractions, possessives) is prose, and + treating it as an opener swallowed the closing brace and lost the call.""" + + def test_contraction_in_unquoted_query_parses(self): + text = "call:web_search{query:what's the weather}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what's the weather" + + def test_contraction_does_not_swallow_next_key(self): + text = "call:web_search{query:what's up, n:3}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what's up" + assert args["n"] == 3 + + def test_contraction_strip_span_covers_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = "call:web_search{query:what's the weather} Done." + stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) + assert "call:web_search" not in stripped + assert stripped.strip() == "Done." + + def test_quoted_values_still_hide_delimiters(self): + text = 'call:web_search{query:"weather, location: Boston", n:2}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "weather, location: Boston" + assert args["n"] == 2 + + +class TestGlmKeyWithoutValue: + """A GLM with no tag: strict mode rejects the call + (same contract as an unclosed value) instead of executing it with the + argument silently dropped; Auto-Heal keeps the lenient skip.""" + + def test_strict_rejects_key_without_value(self): + text = "web_search\nquery\n" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_heal_keeps_the_lenient_skip(self): + text = "web_search\nquery\n" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == {} + + class TestDisabledBareJsonLiteralNotPromoted: """A leading non-enabled-name object is content: nothing inside promotes, and a call after it still parses.""" @@ -742,6 +1056,38 @@ class TestDisabledBareJsonLiteralNotPromoted: assert [c["function"]["name"] for c in calls] == ["web_search"] +class TestDeepSeekMarkerInsideLeadingEnvelopes: + """A DeepSeek/Kimi marker quoted inside a leading bare-JSON or Mistral + call's argument strings is data: the pre-pass must not promote the + embedded no-arg literal and drop the real outer call.""" + + def test_marker_inside_leading_json_call_stays_data(self): + text = ( + '{"name": "web_search", "arguments": ' + '{"query": "what is <｜tool▁calls▁begin｜>...{}..."}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert "tool▁calls▁begin" in args["query"] + + def test_marker_inside_leading_mistral_call_stays_data(self): + text = ( + '[TOOL_CALLS] [{"name": "web_search", "arguments": ' + '{"query": "docs on <｜tool▁calls▁begin｜> markers"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_standalone_deepseek_call_still_parses(self): + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"query": "cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + class TestMistralLiteralInsideLeadingJson: """A [TOOL_CALLS] literal quoted inside a leading JSON object must not be promoted over it.""" @@ -776,6 +1122,28 @@ class TestGemmaWrappedWhitespace: assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] +class TestDisabledJsonBeforeDeepSeekCall: + """A disabled leading bare-JSON object whose strings mention a + DeepSeek/Kimi marker is dropped and the tail parsed, so a REAL + DeepSeek/Kimi call after the object still executes instead of the whole + message skipping the pre-pass.""" + + _DS = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"query": "cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + + def test_real_deepseek_call_after_disabled_json_parses(self): + text = '{"name": "Alice", "note": "<｜tool▁calls▁begin｜>"} ' + self._DS + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_disabled_json_with_marker_alone_stays_data(self): + text = '{"name": "Alice", "note": "<｜tool▁calls▁begin｜>"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + class TestGemmaDottedArgumentKeys: """Dotted Gemma keys (namespaced schemas) must survive key-quoting or the call is lost.""" @@ -787,6 +1155,29 @@ class TestGemmaDottedArgumentKeys: assert args == {"user.name": "bob", "query": "x"} +class TestLeadingWrapperlessGemmaOverEmbeddedMarkers: + """A leading wrapper-less Gemma call to an enabled tool owns the turn: a + quoted foreign literal inside its argument (a query citing another tool + syntax) is data, and tool_healing must not promote it before the Gemma + fallback runs. Foreign markup leading keeps the normal order.""" + + def test_leading_gemma_wins_over_quoted_xml_literal(self): + text = ( + 'call:web_search{query:"explain ' + '{"name":"evil","arguments":{}}"}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_leading_keeps_normal_order(self): + text = ( + '{"name":"web_search","arguments":' + '{"query":"call:evil{x:1} example"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + class TestLeadingMistralCallOwnsTheTurn: """A leading Mistral call wins in document order over literal XML in trailing prose.""" @@ -798,7 +1189,7 @@ class TestLeadingMistralCallOwnsTheTurn: calls = parse_tool_calls_from_text(text) assert [c["function"]["name"] for c in calls] == ["web_search"] - def test_xml_leading_keeps_normal_order(self): + def test_function_xml_leading_keeps_normal_order(self): text = ( "x " "[TOOL_CALLS]evil[ARGS]{}" @@ -816,6 +1207,63 @@ class TestGemmaDottedKeyAfterBareValue: assert args == {"query": "foo", "user.name": "bob"} +class TestJsonAnswersAreDataForMarkerlessScans: + """A whole-content JSON value is a structured answer: a quoted example of + an enabled tool's syntax inside it must not execute the tool, and the + display strip must not mutilate the answer.""" + + def test_gemma_example_inside_json_answer_not_promoted(self): + text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_gemma_example_inside_json_answer_not_stripped(self): + from core.inference.tool_call_parser import strip_tool_markup + text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' + assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + + def test_kimi_marker_inside_json_answer_not_promoted(self): + text = ( + '{"answer":"<|tool_call_begin|>functions.web_search:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaNestedQuotedLeaves: + def test_nested_object_and_array_values_are_unquoted(self): + text = 'call:f{loc:{city:"New York"},items:["a","b"],n:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"f"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"loc": {"city": "New York"}, "items": ["a", "b"], "n": 3} + + +class TestEarliestEnvelopeWinsAcrossDeepSeekKimi: + """The DeepSeek/Kimi pre-pass dispatches by earliest envelope opener: a + leading real call wins over a trailing example of the sibling format in + either direction (document order, like the other leading guards).""" + + _DS = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>evil\n" + '```json\n{"x": 1}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + _KIMI = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"query": "cats"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + def test_leading_kimi_wins_over_trailing_deepseek_example(self): + text = self._KIMI + " For reference: " + self._DS + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_deepseek_wins_over_trailing_kimi_example(self): + text = self._DS + " Kimi format: " + self._KIMI + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["evil"] + + class TestNamelessLeadingJsonAnswerIsData: """A nameless leading JSON answer is an envelope: quoted markup stays data, and a call after it parses.""" @@ -832,6 +1280,140 @@ class TestNamelessLeadingJsonAnswerIsData: assert [c["function"]["name"] for c in calls] == ["web_search"] +class TestClosedCallPrecedesMarkerPrePass: + """A closed non-DeepSeek/Kimi call that precedes the first DS/Kimi marker + owns the turn: a trailing example (or an example quoted inside a wrapped + Gemma argument) must not be promoted by the pre-pass.""" + + _KIMI_EVIL = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.evil:0" + '<|tool_call_argument_begin|>{"x": 1}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + def test_kimi_example_inside_wrapped_gemma_arg_stays_data(self): + text = ( + '<|tool_call>call:web_search{query:<|"|>explain ' + + self._KIMI_EVIL + + '<|"|>}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_xml_call_wins_over_trailing_kimi_example(self): + text = ( + '{"name":"web_search","arguments":{"query":"cats"}}' + " For reference: " + self._KIMI_EVIL + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_standalone_kimi_call_still_parses(self): + calls = parse_tool_calls_from_text(self._KIMI_EVIL) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestTruncatedWrapperlessGemmaStopsScan: + def test_call_quoted_inside_truncated_arg_not_promoted(self): + text = 'call:python{code:example("call:web_search{query:hi}") and then it cut' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + + +class TestGemmaQuotedNestedDelimiters: + def test_comma_inside_quoted_nested_string_not_a_split(self): + text = 'call:f{loc:{city:"New, York"},n:1}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"f"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"loc": {"city": "New, York"}, "n": 1} + + +class TestGemmaStringMarkerLiteralInArgs: + def test_string_marker_literal_does_not_lose_the_call(self): + text = "call:web_search{query:'what does <|\"|> mean in Gemma'}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == 'what does <|"|> mean in Gemma' + + +class TestGemmaMidValueQuotedPhrase: + def test_quoted_phrase_mid_value_hides_delimiters(self): + text = 'call:web_search{query:find "weather, location: Boston", limit:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": 'find "weather, location: Boston"', "limit": 3} + + def test_apostrophes_still_prose_mid_value(self): + text = "call:web_search{query:what's on at the museum, n:2}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "what's on at the museum", "n": 2} + + +class TestGlmStrictRefusesInQuoteFallback: + """A truncated GLM value whose only close candidates sit inside a string + literal must reject in strict mode instead of executing truncated + arguments; Auto-Heal keeps the lenient partial value.""" + + _TRUNC = ( + 'python\ncode\nprint("")' + ) + + def test_strict_rejects_truncated_in_string_close(self): + assert parse_tool_calls_from_text(self._TRUNC, allow_incomplete = False) == [] + + def test_heal_keeps_partial_value(self): + calls = parse_tool_calls_from_text(self._TRUNC, allow_incomplete = True) + assert len(calls) == 1 and calls[0]["function"]["name"] == "python" + + +class TestGemmaGuardCoversPreambles: + def test_preamble_then_gemma_call_quoting_xml_wins(self): + text = ( + "Sure, searching now. call:web_search{query:" + '"explain {"name":"evil","arguments":{}}"}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestGlmStrictAcceptsApostrophes: + def test_apostrophe_value_parses_in_strict_mode(self): + text = ( + "web_search\nquery\n" + "what's the weather\n" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "what's the weather"} + + +class TestDisabledGemmaCallLiteralsAreData: + def test_literal_inside_disabled_call_not_promoted(self): + text = 'call:foo{query:"x"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + + def test_real_call_after_disabled_example_still_parses(self): + text = ( + 'call:foo{query:"x"}' + " call:web_search{query:hi}" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestLeadingJsonArrayAnswerIsData: + def test_kimi_marker_inside_json_array_answer_not_promoted(self): + text = ( + '[{"answer": "<|tool_call_begin|>functions.web_search:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}]' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + class TestLeadingBareJsonOwnsTurnOverTrailingXml: """Document order: a leading closed bare-JSON call owns the turn even when tool XML appears AFTER it (inside-or-after, mirroring the Mistral rule).""" @@ -855,7 +1437,8 @@ class TestLeadingBareJsonOwnsTurnOverTrailingXml: assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls def test_non_call_leading_object_defers_to_trailing_real_call(self): - # Nameless/disabled-name objects decline: dropped, and the real trailing call still parses. + # Nameless answers and disabled-name objects take the decline path: + # the object is dropped and the real trailing call still parses. for lead in ('{"answer": 42}', '{"name":"draft","parameters":{}}'): text = lead + ' {"name":"delete_all","arguments":{}}' calls = parse_tool_calls_from_text(text, enabled_tool_names = {"delete_all"}) @@ -891,7 +1474,8 @@ class TestProseCloseTagAfterClosedFunctionCall: assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} def test_attribute_form_arguments_do_not_swallow_prose(self): - # The attribute form shares the first-balanced-close rule: prose closes never fold in. + # The attribute form shares the first-balanced-close + # rule: prose mentioning a literal close tag never folds into arguments. text = ( 'cats' " Done. The tag closes a call." diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index 7fe52a664d..d50c27130f 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -24,19 +24,39 @@ import re as _re _src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() _m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) assert _m, "could not extract _TOOL_XML_RE source" -# Provide both helpers so the extracted _strip_tool_xml_for_display resolves. -from core.inference.tool_call_parser import _strip_function_xml_calls, _strip_mistral_closed_calls +# The lazy ``(.*?)\n\)`` could grab a shorter expression if an arm is ever wrapped; +# pin the DeepSeek + bare-Kimi arms so a silent truncation fails loudly here. +assert "_DS_OPEN_SRC" in _m.group(1) and "tool_call_begin" in _m.group( + 1 +), "extracted _TOOL_XML_RE is missing expected arms (extraction truncated?)" +# The regex reuses the parser's shared DeepSeek opener alternation; provide it so the extracted +# ``_re.compile`` expression resolves the same source. +from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC +from core.inference.tool_call_parser import ( + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) + +from typing import Optional as _Optional _ns = { "_re": _re, + "_DS_OPEN_SRC": _DS_OPEN_SRC, + "Optional": _Optional, "_strip_mistral_closed_calls": _strip_mistral_closed_calls, + "_strip_gemma_wrapperless_calls": _strip_gemma_wrapperless_calls, + "_strip_glm_calls": _strip_glm_calls, "_strip_function_xml_calls": _strip_function_xml_calls, } exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) _TOOL_XML_RE = _ns["_TOOL_XML_RE"] +# Signatures may span multiple lines and now carry the enabled_tool_names gate; match +# the whole (possibly multi-line) signature up to ``-> str:`` then the indented body. _xml_helper = _re.search( - r"def _strip_tool_xml\(text: str\) -> str:\n(?: .+\n)+", + r"def _strip_tool_xml\((?:.|\n)*?\) -> str:\n(?: .+\n)+", _src, ) assert _xml_helper, "could not extract _strip_tool_xml source" @@ -44,17 +64,27 @@ assert "_strip_mistral_closed_calls" in _xml_helper.group( 0 ), "extracted _strip_tool_xml no longer runs the Mistral balanced strip" exec(_xml_helper.group(0), _ns) +_strip_tool_xml = _ns["_strip_tool_xml"] _helper = _re.search( - r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n" - r"(?: .+\n)+", + r"def _strip_tool_xml_for_display\((?:.|\n)*?\) -> str:\n(?: .+\n)+", _src, ) assert _helper, "could not extract _strip_tool_xml_for_display source" +# After the V1 fix the display helper delegates to _strip_tool_xml; confirm the +# extracted body actually reached that call rather than truncating early. assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates" exec(_helper.group(0), _ns) _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] +_gate_src = _re.search( + r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?: .+\n)+", + _src, +) +assert _gate_src, "could not extract _gemma_strip_gate source" +exec(_gate_src.group(0), _ns) +_gemma_strip_gate = _ns["_gemma_strip_gate"] + # ── Well-formed pairs ───────────────────────────────────────────── @@ -66,7 +96,8 @@ def test_route_display_strip_respects_disabled_auto_heal_contract(): def test_route_display_strip_removes_mistral_tool_calls_with_nested_json(): - # [TOOL_CALLS] with nested JSON needs the Mistral balanced-brace strip, not the regex. + # _TOOL_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral + # balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON). text = 'ok [TOOL_CALLS]web_search{"filters":{"date":"2024"},"query":"cats"} tail' assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) @@ -102,7 +133,8 @@ def test_strips_function_only_well_formed(): def test_strips_function_attribute_form(): - # Attribute form must strip from the route too; dotted/hyphenated names included. + # Attribute form ```` (MiniCPM-5 / MiniMax-M2) must strip from the route too + # (it previously leaked into the UI); a dotted/hyphenated name also strips. text = ( 'Sure.\n\n' "\nSydney\n\n\nDone." @@ -330,9 +362,84 @@ def test_no_catastrophic_backtracking_on_orphan_opening_spam(): assert "" not in cleaned +# ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ── + + +def test_strips_deepseek_space_opener_variant(): + # The space-separated opener is parsed by the parser, so the display strip + # must remove it too (the shared opener alternation is reused here). + text = ( + "pre <｜tool calls begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "tool" not in cleaned.replace("post", "").replace("pre", "") + assert cleaned == "pre post" + + +def test_strips_deepseek_escaped_underscore_opener_variant(): + text = ( + "pre <｜tool\\_calls\\_begin｜><｜tool▁call▁begin｜>get_y<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "pre post" + + +def test_strips_bare_kimi_call_without_section_wrapper(): + # Kimi can emit a bare <|tool_call_begin|>...<|tool_call_end|> with no + # section wrapper; the parser accepts it, so the strip must cover it. + text = ( + "pre <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>" + '{"a":1}<|tool_call_end|> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "tool_call_begin" not in cleaned + assert cleaned == "pre post" + + +@pytest.mark.parametrize( + "text", + [ + # Prose that merely names a Kimi/DeepSeek marker (no real call follows) must + # survive: the call-shaped lookahead fires only on a real call or a bare EOF + # fragment, so an answer discussing the protocol is never truncated. + "See <|tool_call_begin|> in the docs. More prose after it.", + "The <|tool_calls_section_begin|> marker opens a batch. Read on.", + "DeepSeek uses <｜tool▁calls▁begin｜> to start a call block, then continues.", + ], +) +def test_deepseek_kimi_false_alarm_prose_is_kept(text): + # Regression for the route arm truncating a prose answer that references a marker + # without a following call (parser _TOOL_ALL_PATS already had this lookahead). + assert _TOOL_XML_RE.sub("", text) == text + + +def test_deepseek_kimi_real_calls_still_strip_after_false_alarm_fix(): + # The lookahead must not weaken real-call stripping: closed, truncated, and bare + # EOF-fragment forms all still get removed. + closed = ( + "answer <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>" + '{"a":1}<|tool_call_end|> tail' + ) + assert _TOOL_XML_RE.sub("", closed) == "answer tail" + eof_fragment = "prefix <|tool_call_begin|>" + assert _TOOL_XML_RE.sub("", eof_fragment) == "prefix " + deepseek = ( + "reply <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + assert _TOOL_XML_RE.sub("", deepseek) == "reply " + + +# ── Llama-3 <|python_tag|> arm bounds on REAL sentinels only ────── + + # Llama-3 <|python_tag|> arm bounds on REAL sentinels only def test_python_tag_strip_consumes_literal_sentinel_in_arg(): - # A literal <|...|> token inside the arg must not end the strip early. + # A <|python_tag|> tool call whose JSON argument carries a literal <|...|> + # token (here <|cite|>) must be stripped whole. The old `<(?!\|)` arm stopped + # at any `<|`, leaking the call tail (e.g. `<|cite|> here"}}`) into display. text = '<|python_tag|>{"name": "send", "parameters": {"text": "use <|cite|> here"}}' cleaned = _TOOL_XML_RE.sub("", text) assert cleaned == "", f"python_tag call leaked at literal sentinel: {cleaned!r}" @@ -348,7 +455,8 @@ def test_python_tag_strip_consumes_literal_sentinel_in_arg(): ], ) def test_python_tag_strip_stops_at_real_sentinel(sentinel): - # A real control sentinel bounds the strip so following text survives. + # A genuine Llama control sentinel still bounds the strip so following + # assistant text is preserved (the arm must not swallow past it). text = f'<|python_tag|>{{"name": "x", "parameters": {{}}}}{sentinel}visible answer' cleaned = _TOOL_XML_RE.sub("", text) assert ( @@ -357,14 +465,37 @@ def test_python_tag_strip_stops_at_real_sentinel(sentinel): def test_python_tag_strip_restarts_on_second_python_tag(): - # A second <|python_tag|> opens a new region; both are stripped. + # A second <|python_tag|> opens a new tool-call region, so the whole pair is + # stripped (the arm bounds the first, then the next match consumes the rest). text = '<|python_tag|>{"name": "a"}<|python_tag|>{"name": "b"}' cleaned = _TOOL_XML_RE.sub("", text) assert cleaned == "", f"second python_tag region leaked: {cleaned!r}" +def test_glm_call_with_literal_close_tag_in_arg_value_is_stripped_whole(): + # GLM 4.x emits NAMEkv .... + text = ( + "web_search\nquery\n" + "find here\n done" + ) + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out + assert "" not in out + assert out.strip() == "done" + + +def test_glm_normal_and_qwen_calls_still_stripped_by_route(): + # Regression: a normal GLM call (no literal close tag) and a Qwen + # {json} are still stripped; trailing prose is kept. + glm = "get_time\ntz\nUTC\n ok" + assert _strip_tool_xml_for_display(glm, auto_heal_tool_calls = True).strip() == "ok" + qwen = '{"name":"web_search","arguments":{"q":"x"}} after' + assert _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after" + + def test_route_strip_removes_param_alias_close_tag(): - # Orphan (attribute-form alias of ) must strip too. + # The parser accepts the ... attribute-form alias of + # ; the route tail cleanup must strip an orphan close too. assert _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " assert ( _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " @@ -372,13 +503,44 @@ def test_route_strip_removes_param_alias_close_tag(): def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup(): - # A literal in a value must not truncate the strip. + # A literal in a value must not truncate the strip: the route runs the + # parser's guarded function-XML scan before the regex, matching the core strip. text = " tail" assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail" +def test_route_strip_gates_wrapperless_gemma_by_enabled_tools(): + # The route strip must gate the markerless Gemma call:NAME{...} form on the enabled tool names, + # like the parser/loop, so a disabled/example name in prose is preserved in ... + prose = "To document syntax you write call:foo{query:example}. That shows the format." + assert "call:foo{query:example}" in _strip_tool_xml(prose, {"web_search"}) + # An enabled name is still a real call and stripped. + assert "call:web_search" not in _strip_tool_xml( + "Answer. call:web_search{query:x}", {"web_search"} + ) + # No gate (legacy) strips every closed call. + assert "call:foo" not in _strip_tool_xml(prose) + + +def test_gemma_strip_gate_empty_tools_preserves_prose(): + # With NO tools enabled the gate must return an EMPTY set (strip nothing), not None: None falls + # back to strip-all and deletes an answer that documents the call:NAME{...} syntax. + assert _gemma_strip_gate([]) == set() + assert _gemma_strip_gate(None) == set() + assert _gemma_strip_gate([{"function": {"name": "web_search"}}]) == {"web_search"} + prose = "To document syntax you write call:foo{query:example}. That shows the format." + assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate([])) + assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate(None)) + # An enabled tool's real call is still stripped. + assert "call:web_search" not in _strip_tool_xml( + "Answer. call:web_search{query:x}", + _gemma_strip_gate([{"function": {"name": "web_search"}}]), + ) + + def test_strip_keeps_prose_after_closed_function_call_with_literal_close(): - # The call ends at its first non-data close; prose after (even a literal ) survives. + # The call ends at its first non-data close: prose after it survives the + # strip even when it mentions a literal . from core.inference.tool_call_parser import strip_tool_markup text = ( "cats" @@ -388,7 +550,8 @@ def test_strip_keeps_prose_after_closed_function_call_with_literal_close(): def test_final_strip_keeps_prose_mentioning_bare_markers(): - # A false-alarm marker in prose must not drop trailing text; only call-start-shaped text drops. + # A false-alarm marker in a normal answer must not lose everything after + # it; only text that looks like that family's call start drops. from core.inference.tool_call_parser import strip_tool_markup for text in ( "See [TOOL_CALLS] docs for details. More prose after.", @@ -413,7 +576,8 @@ def test_final_strip_still_drops_truncated_marker_calls(): def test_chained_bare_json_strip_consumes_all_calls(): - # Next-turn history must not keep an executed call, else it replays. + # The loops keep this text as next-turn history: a leftover executed call + # would be replayed alongside the structured tool_calls. from core.inference.tool_call_parser import strip_leading_bare_json_call enabled = {"web_search", "python"} From 233949cc9c04f80d7878f9973e0c9b4dd24c01bb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 18:34:18 -0700 Subject: [PATCH 025/113] scan_packages: baseline transitive-dep drift in the supply-chain scan (#6917) The pip scan-packages gate (SCAN_ENFORCE=1) blocks on non-baselined CRITICAL/HIGH findings. Recent upstream releases of transitive dependencies added new files/loops that trip the pattern scanner, so all three shards (extras, hf-stack, studio) red-failed on legitimate library code. Add the 7 reviewed findings to scripts/scan_packages_baseline.json. Each entry is genuine upstream code from the official PyPI archive: - huggingface-hub huggingface_hub/_sandbox.py (staged dropper + C2 loop): the HF Jobs sandbox bootstrap string and its host-pool reservation loop. New in huggingface_hub 1.x (pulled via huggingface_hub>=0.34.0). - huggingface-hub huggingface_hub/hf_api.py, utils/_http.py (C2 loop): standard polling / retry while True loops. - fastapi fastapi/routing.py (C2 loop): websocket receive loop. - fastmcp-slim fastmcp/cli/apps_dev.py (fs enum + network): the FastMCP dev CLI (PrefectHQ) making httpx/socket calls. - cffi cffi/_cffi_gen_src.py (compile + exec): cffi generating and running C extension source, its core purpose. Additive only: no existing baseline entry is changed or removed. Verified by re-running the scanner over the full closure on Python 3.12.13 (the CI interpreter); it now exits 0 with only MEDIUM findings remaining. --- scripts/scan_packages_baseline.json | 56 +++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 046566d148..d42225e205 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -1489,6 +1489,62 @@ "severity": "HIGH", "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" + }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", + "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" + }, + { + "package": "fastmcp-slim", + "file": "fastmcp/cli/apps_dev.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/_sandbox.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", + "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/_sandbox.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", + "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", + "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/utils/_http.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", + "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" + }, + { + "package": "cffi", + "file": "cffi/_cffi_gen_src.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", + "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" } ] } From f109e7f0e6a1fdfde8bb38e73f7359cb777bcc50 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 18:52:13 -0700 Subject: [PATCH 026/113] Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes (#5704) * Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes Extends the rescue parsers in core/tool_healing.py and core/inference/tool_call_parser.py to recognise two extra serialisations local models commonly emit when bypassing native function calling: * [TOOL_CALLS]name{json_args} (Devstral-Small-2, Mistral-Small-3.x). * name[ARGS]{json_args} (reasoning-model rehearsal). Both extractors use a brace-balance scan that honours escapes and quoted strings so nested JSON args stay intact. Also pre-strips ... and [THINK]...[/THINK] blocks before matching so calls emitted after a reasoning preamble are recognised regardless of position. Streaming gates (TOOL_XML_SIGNALS, llama_cpp.py _TOOL_XML_SIGNALS) and the SSE strip regex (routes/inference.py _TOOL_XML_RE) gain the new sentinels so the parser is actually invoked and the raw markup never leaks to the UI. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Strip unclosed think blocks and catch rehearsal [ARGS] mid-buffer The pre-existing ``_THINK_TAG_RE`` only matched closed thinking blocks (``...`` or ``[THINK]...[/THINK]``). During streaming the model is still inside the open block when the parser runs, so any tool-shaped markup the model is REHEARSING inside that block survived the strip and could be executed as a real call. Switch both copies of the regex (parser + healing) to accept the trailing block being terminated by end-of-string in addition to the explicit closer. The ``_TOOL_XML_SIGNALS`` list on the llama_cpp streaming buffer included ``[ARGS]`` to catch rehearsal syntax, but the gate used a ``startswith`` check against the buffer head -- rehearsal is shaped ``name[ARGS]{json}``, so the buffer never STARTS with ``[ARGS]`` and the signal had no effect. Add a substring fallback for the bracket-style signals so the BUFFERING window can still divert the stream into DRAINING when rehearsal markup arrives mid-buffer. Adds three regression tests covering rehearsal inside unclosed ```` / ``[THINK]`` blocks (must yield no calls) and the positive case after a closed think block (still parsed). * [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: harden bracket-tag tool-call parsing and streaming strip Address review findings on the Mistral [TOOL_CALLS] / rehearsal [ARGS] paths: - Accept hyphenated tool names in the bracket parsers and strip patterns. _MISTRAL_BRACKET_RE and _REHEARSAL_RE used \w+, which dropped or truncated MCP function names containing dashes (mcp__srv__list-issues). Use [\w-]+ to match the XML and Gemma parsers. - Strip a partial bracket marker streamed before its opening brace. The trailing-unclosed patterns required the {, so a [TOOL_CALLS]web_search or python[ARGS] split across deltas leaked the raw marker to the UI. Match the bare marker to end-of-text, mirroring how the bare open tags are stripped. Closed pairs are unchanged so in-progress markup stays buffered until parsed. - Strip a truncated bracket tail in the route-level display regex. _TOOL_XML_RE required a balanced JSON object; a tool call truncated by EOS now strips up to \Z, like the orphan-opening XML shapes. Complete calls still strip only their balanced JSON so following prose survives. Add regression tests for hyphenated names, the streaming partial-marker strip, and the unclosed-tail route strip. * Studio: preserve XML parameter indentation in tool_healing The chat template emits \nVALUE\n; the parameter-start regex consumed the wrapping newline AND the value's first-line indentation via a trailing \s*, then str.strip() removed the rest, corrupting code/diff arguments. Narrow the trailing class to horizontal whitespace and trim exactly one wrapping newline (_trim_param_value), preserving indentation. Matches SGLang's qwen3_coder detector and the same fix on the multi-format parser. Add a regression test. * Studio: tighten Mistral/rehearsal tool-call comments Compress the comments in the Mistral [TOOL_CALLS] / rehearsal [ARGS] healing shim and its callers to one or two lines, keeping the bracket-tag stripping rationale, the thinking-block handling note, and the forge attribution intact. Comment-only: no code or behavior change (verified with comment_tools.py check --strip-docstrings; tests green). * Studio: fix think-strip arg corruption and nested bracket-JSON strip Review follow-up for the Mistral/rehearsal healing shim: - The /[THINK] strip ran unconditionally over the whole content before parsing, so a real tool argument that legitimately contained a / [THINK] literal was silently corrupted. Don't delete the blocks: compute the reasoning-block spans and skip any tool-call candidate that STARTS inside one, across all parse paths (JSON, Gemma, XML, bracket, rehearsal). A rehearsed call inside reasoning is still ignored; a real call after still parses. - The bracket-tag display strip used a fixed one-level-nesting regex, so a call with two-level-nested JSON args either leaked raw markup or, in final mode, let the catch-all eat the trailing prose. Add a balanced-brace _strip_bracket_tag_calls pass (any nesting depth) used by strip_tool_call_markup and the route display strip. Add regressions: /[THINK] literal inside a real argument, rehearsal-inside- think with a real call after, and two-level-nested bracket/rehearsal strip keeping trailing prose. * Studio: correct think-block comments to match span-skip behavior The think-strip fix replaced the unconditional think-block strip with a span-skip (the block is kept and any tool-call candidate starting inside it is ignored), but two comments still described the old strip-first behavior. Update the _THINK_TAG_RE comment and the parse_tool_calls_from_text docstring. * Studio: parse Mistral arrays and call-ids, unify bracket parse/strip, keep it linear - Parse the canonical Mistral array form (TOOL_CALLS followed by a JSON list of calls) and emit every call; parse the v11 shape that carries an opaque CALL_ID token between the name and ARGS (the function name is the token after TOOL_CALLS, never the call-id); and parse a Mistral call plus a rehearsal call in one message (the second was dropped yet still stripped from display). - One shared balanced forward scan (_iter_bracket_spans) backs both the parser and the strip path, so they no longer diverge. It is linear: each regex is re-searched only once its cached match falls behind the cursor, replacing the per-match full-tail re-scan that was O(n^2) (O(n^3) over a stream). A length cap before the scan is a backstop. - strip_tool_call_markup preserves think/reasoning blocks verbatim (the parser skips tool markup inside them), stripping only the visible text around them. - _in_think uses bisect over the sorted think spans (was a linear scan per candidate). - GGUF streaming strip runs the balanced bracket pre-pass before the regex patterns so nested-arg calls do not leak or eat trailing prose, and the BUFFERING ARGS detector requires the rehearsal name-ARGS shape. - Tests: canonical array, array string-args, array strip keeps prose, Mistral plus rehearsal multi-call, v11 call-id name, think-rehearsal strip preservation, and bracket-strip linearity. * Studio: preserve reasoning blocks in the route and streaming strip paths too Addresses Gemini/Codex review: making strip_tool_call_markup preserve think blocks left the route display strip and the GGUF streaming strip inconsistent, so a rehearsed call inside a reasoning block was still deleted from the visible text on those paths. - Extract the think-block segmentation into one shared helper (strip_outside_think) and route all three strip paths through it: strip_tool_call_markup, _strip_tool_xml_for_display, and the GGUF _strip_tool_markup_streaming closure. - Add a route-strip regression test that a rehearsal inside a reasoning block is preserved while a real call outside it is still stripped. * Studio: fix bracket-tag strip/buffer review findings Address the live code-review findings on the Mistral bracket-tag / rehearsal tool-call rescue path: - tool_healing: a literal think block inside a tool-call argument is no longer treated as a reasoning block. strip_outside_think now excludes think spans that sit inside a complete tool-call span, so the call is stripped whole instead of the split hiding its open/close pair and leaking the raw call. - tool_healing: the rehearsal trailing-strip pattern requires a following brace or end-of-text, so prose that merely mentions name[ARGS] is not truncated as a phantom call. The bracket strip patterns are aligned with the parser regexes (whitespace, v11 [CALL_ID]/[ARGS] metadata, and the [CALL_ID] lookbehind). - routes: strip a truncated canonical Mistral array ([TOOL_CALLS] [{... with no closing bracket) that the balanced scan cannot remove, align the display regex with the parser regexes, and apply the same rehearsal-prose guard. - safetensors loop: mirror the GGUF [ARGS] rehearsal-substring check during BUFFERING so a rehearsal name does not stream before its [ARGS] arrives. Adds regression tests for each; existing parser suite stays green. * Studio: hold split rehearsal tool-name prefix in both streaming loops A reasoning-model rehearsal call can stream the tool name and its [ARGS] arm in separate chunks (web_search then [ARGS]{...}). The buffering detector only recognised the rehearsal once [ARGS] was present, so the bare tool name was emitted as visible content before the call drained and executed. Add _is_rehearsal_prefix (mirrored in the safetensors loop and the GGUF loop): when a no-signal buffer is a bare active-tool name -- or a partial prefix of NAME[ARGS] -- hold it as a prefix instead of streaming it, so the next chunk's [ARGS] flips it to a drain. A whitespace in the buffer means prose, not a split call, so ordinary text still streams. Adds regression tests for the split rehearsal in both loops and a guard that a plain non-tool word still streams. * Studio: route Anthropic tool-call cleanup through the protected display strip The Anthropic stream, non-stream, and passthrough paths cleaned content with raw _TOOL_XML_RE.sub instead of _strip_tool_xml_for_display, so a rehearsal call inside was deleted from the reasoning and a nested [TOOL_CALLS] call dropped its trailing prose (the OpenAI-compatible paths already use the helper). Route all four sites (prior-assistant cleanup, streaming content events, non-stream aggregation, passthrough conversion) through the protected helper, and add a source-level guard test so raw _TOOL_XML_RE.sub stays confined to the helper itself. * Studio: stop split rehearsal tool names leaking once streaming, uncapped, or unrestricted The split-rehearsal guard (NAME in one chunk, [ARGS]{...} in the next) only held the name in the initial BUFFERING state. Three gaps remained where the bare tool name still streamed as visible content before the call drained: - STREAMING: after prose had already streamed, both loops emitted a trailing active-tool-name token (and the GGUF/safetensors [ARGS] boundary was not pulled back over the name). Hold the trailing rehearsal token and release it on the next chunk, with an end-of-stream flush so a plain answer that merely ends on a tool-name word is never dropped. - Buffer cap: a realistic MCP name longer than the 32-char _MAX_BUFFER_CHARS cap defeated the BUFFERING hold. A rehearsal prefix is self-bounding (it stops matching once it grows past NAME[ARGS]), so the generic cap no longer applies to it. - Unrestricted mode (tools=[]): with no declared tool list, any bare identifier may be a NAME[ARGS] rehearsal, so the prefix check now recognises one instead of leaking the name and mis-parsing the call. Regression tests cover the streaming, long-name, and unrestricted cases plus the plain-prose paths that must not be held or corrupted. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio tools: protect think blocks in safetensors streaming, hold split rehearsal on initial flush, advertise Mistral tools Pass-3 review follow-ups on the Mistral [TOOL_CALLS] / rehearsal [ARGS] work: - Safetensors streaming display strip now preserves think / [THINK] reasoning verbatim (routes through strip_outside_think like the GGUF path). A call rehearsed inside a reasoning block was stripped mid-stream and then restored by the final strip, a non-monotonic shrink/grow that corrupted append-by-length stream consumers and the visible reasoning. - The first flush out of BUFFERING (safetensors and GGUF) now applies the same trailing-name hold the STREAMING branch uses, so a split rehearsal (prose plus a trailing active tool name in one chunk, [ARGS]{...} in the next) no longer leaks the bare name before the call drains. - Safetensors capability gate no longer suppresses tools for Mistral [TOOL_CALLS] templates, which the shared bracket-tag parser now handles end to end. Llama python_tag stays suppressed (still unparseable). - Route display strip applies the open-ended / bare-marker tail arms only on the segment after the last reasoning block (closed-only regex before it), matching strip_tool_call_markup, so a bare foo[ARGS] before a reasoning block is preserved while complete calls are still removed in every segment. Adds regression tests for each and updates the now-stale Mistral capability test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix tool-call think-marker and bracket-wrapper edge cases Round-1 review follow-ups on the Mistral/rehearsal tool-call healing: - tool_healing: a reasoning marker that opens INSIDE a tool call's arguments is argument data, not a reasoning block. Add _think_spans_outside_tool_markup (start-inside test) and use it in both parse_tool_calls_from_text and strip_outside_think so a literal marker in one call's args no longer hides a later call (parse) or leaks the raw markup (strip) when the greedy match runs past the call's closer. - tool_healing: strip the orphan Mistral v11 [/TOOL_CALLS] closer left behind after the balanced scan removes the call body. Add a route arm for the same closer in _TOOL_XML_RE / _TOOL_XML_CLOSED_RE. - safetensors + llama_cpp streaming strip: run the open-ended (EOS anchored) tail patterns only on the last segment; segments before a reasoning block use the closed-only patterns, matching the final strip and the route strip. A bare foo[ARGS] before a reasoning block is prose, not a truncated call. - safetensors streaming detector: validate each [ARGS] hit before draining. A bare foo[ARGS] in prose (no active tool name in front) no longer drains the rest of the turn; a later real NAME[ARGS] call is still found and the prose in between is preserved. Regression tests added for each case across the parser, strip helpers, and both streaming loops. * Strip incomplete-XML tool markup with literal think tags; widen render-html detector Round-2 review follow-ups. - tool_healing: an UNCLOSED / /[THINK] reasoning block, but the provisional render_html detector scanned raw content. A render_html rehearsed inside followed by a real non-render_html call emitted a provisional render_html tool_start (reusing the later call's id) that the loop never executed. Drop candidates that start inside a think span and use the first marker of each shape outside the blocks. Also resolve the [TOOL_CALLS] [{...}] array shape through the parser so a nested "name" argument key no longer fires a false provisional card ahead of the real top-level tool name. Adds regression tests for both loops: inactive-name foo[ARGS]{...} is not drained into a disabled no-op or a retry turn, a think-block render_html rehearsal emits no provisional card, and the array top-level name is read correctly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate ambiguous bare-rehearsal parse and strip on the active tool list A bare NAME[ARGS]{json} is a genuine rehearsal call only when NAME is an active tool; otherwise it is prose. The earlier round gated only detection (so an inactive foo[ARGS] no longer drained the buffer or forced a retry turn), but the parse and strip stayed unrestricted, which produced two regressions: 1. An inactive foo[ARGS]{...} placed immediately before a real web_search[ARGS]{...} in the same content span made the real call fail to execute (parse consumed the phantom foo call). 2. An inactive foo[ARGS]{...} in a prose answer had its markup stripped from the visible text, corrupting the sentence to " is just syntax." Thread enabled_tool_names through the shared parser/strip so parse and strip apply the SAME active-tool gate as detection: - core/tool_healing.py: _iter_bracket_spans skips an inactive rehearsal span; parse_tool_calls_from_text, _strip_bracket_tag_calls, _strip_markup_segment and strip_tool_call_markup accept and thread the gate; apply_tool_strip_patterns keeps an inactive rehearsal match. - core/inference/tool_call_parser.py: wrappers forward the gate. - core/inference/safetensors_agentic.py and core/inference/llama_cpp.py: compute the gate from the active tool list (None when unrestricted, to keep the legacy strip-all behavior) and thread it into every parse and streaming/final strip site. - routes/inference.py: _strip_tool_xml_for_display accepts the gate and keeps an inactive rehearsal via a capture group on its rehearsal arm, so the display cleanup does not re-strip the already-correct loop output. The [TOOL_CALLS] control-token arms still strip unconditionally. Wire the current turn's active tool names into the GGUF and safetensors content-display sites. Tests: parse and strip gate coverage in test_tool_call_parser_strict.py, test_tool_xml_strip.py and test_safetensors_tool_loop.py; end-to-end GGUF coverage for the real-call-after-inactive-rehearsal case and a strengthened assertion that the inactive rehearsal prose survives intact. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: render the reasoning block for safetensors and MLX like GGUF enable_thinking chat templates (Qwen3/Qwen3.5/GLM) prefill an unclosed into the generation prompt, so the model emits only the closing then the answer. The safetensors/MLX chat stream emitted that as plain content, so the reasoning showed inline with no collapsible thinking block, while GGUF (which surfaces reasoning via reasoning_content) rendered one. This brings safetensors and MLX to parity. - _ResponsesReasoningExtractor gains a reasoning_prefilled mode that starts inside the reasoning block and splits on the first ; default False keeps GGUF and every existing caller byte-identical. It suppresses a stray re-emitted and holds partial markers back across chunk boundaries. - _sf_reasoning_prefill_mode gates the mode on reasoning being enabled for the request, an enable_thinking or enable_thinking_effort style, and the template actually using the standard / markers. Models with a bespoke reasoning channel (e.g. gemma's <|think|>/<|channel>) are excluded so their answer is never swallowed; gpt-oss (Harmony) and thinking-off requests are excluded too. - sf_tool_stream and stream_chunks (the latter also serves MLX) feed text through the extractor, emitting reasoning_content then content deltas, with a per-turn reset in the tool loop and a flush before each tool_start; only the visible delta reaches the monitor reply. The two non-streaming drains split reasoning_content the same way. - Tests: extractor prefilled mode (streaming and edge cases), the gate matrix including the gemma-style exclusion, and a route-replay of the tool-loop reasoning stream. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: skip tool calls rehearsed in prefilled reasoning Reasoning models (Qwen3.5 enable_thinking) open in the prompt, so the generated text starts inside the thought and emits only a closing with no opener. _think_spans_outside_tool_markup only found spans with an explicit opener, so a NAME[ARGS]{...} or [TOOL_CALLS] call rehearsed in that leading thought was parsed and executed as a real call. Add a leading think span (offset 0 through the first close marker) when the content opens with a bare close, so the rehearsed call is skipped and the reasoning is preserved by strip_outside_think. Guarded by the existing call-span check: a literal inside a real call's arguments does not trigger the span, so a genuine leading call still fires. Tests for both cases. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: do not start prefilled reasoning mode when reasoning_effort is none enable_thinking_effort models (e.g. GLM-5.2) express thinking-off via reasoning_effort="none" rather than enable_thinking=False, but _sf_reasoning_prefill_mode only looked at enable_thinking, so such a request started the extractor in prefilled mode. With thinking off the model never emits , so the whole answer was captured as reasoning_content and the visible content/stream came back empty. Thread reasoning_effort through and return False when it is "none". Tests for none vs a real effort level. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: only treat a leading bare as prefilled reasoning when a real call follows The prefilled-reasoning virtual span fired on any unmatched leading close marker, so a non-prefilled turn that emits a real call before a stray (for example "Now web_search[ARGS]{...} answer") had the call swallowed by the span and dropped. Require that a real tool call also appear after the close (the actual turn that follows the thought) before adding the span, so a stray close in a normal answer no longer suppresses a genuine leading call. The rehearse-then- call case still skips the rehearsal. Test for the stray-close case. * Studio: trim redundant comments (comment-only, AST-verified) * studio: keep tool_healing importable on Python 3.9 _balanced_json_span was annotated -> int | None. With no from __future__ import annotations, that PEP 604 union is evaluated at import time, so on Python 3.9 (which the package still supports, requires-python >=3.9, and where external inference servers import this module standalone) the def raises TypeError and the whole module fails to import before any parsing runs. Add from __future__ import annotations so annotations stay lazy strings, matching the prevailing convention across studio/backend. No behavior change: the module has no runtime annotation introspection. * Studio: gate the Anthropic tool-stream display strip on declared tools The Anthropic streaming and non-streaming tool paths called _strip_tool_xml_for_display without enabled_tool_names, so with the default strip-all behavior a final answer that literally contains an inactive-name NAME[ARGS]{json} (prose, not a call) lost those bytes in the delivered text. The GGUF and safetensors paths already pass _display_tool_name_gate(tools); these two sites were missed when that gate was threaded through. Compute the gate from the declared tools and pass it at both sites (threading openai_tools into _anthropic_tool_non_streaming and its caller), so an inactive-name rehearsal survives while an active-name one is still stripped. Add a regression test. * Studio: hold a split unrestricted rehearsal prefix at the bracket In unrestricted tool mode (tools=[]) the rehearsal-prefix regex required [A after the bracket, so a chunk boundary landing right after NAME[ (e.g. web_search[ then ARGS]{...}) failed the prefix check and streamed the partial tool markup web_search[ to the client before the call drained. Restricted mode already holds this via a startswith check. Make the bracket and each ARGS letter individually optional so NAME[ is held too, matching the documented intent. Add a regression test. * Studio: gate rehearsal detection and history strip on the original tool set Two display/loop gate fixes so a spent one-shot tool is handled consistently: - Rehearsal DETECTION (safetensors and GGUF loops) now uses the ORIGINAL tool list, matching the strip gate, instead of the post-removal active_tools. After a one-shot tool (render_html) runs it is dropped from active_tools; a repeat render_html[ARGS]{...} while another tool is still active was stripped from display yet never detected, so it was not routed to the render_html_repeat no-op and the turn ended as a blank continuation. Detection now fires for it. - The GGUF assistant-history sanitiser forwards the enabled-tool-name gate (like the live-response strip), so a prior turn documenting an inactive foo[ARGS]{...} shape is preserved in the replayed prompt context instead of being deleted. Add regression tests for both loops and the history strip. * Studio: thread the tool-name gate through the remaining rehearsal/history sites Follow-up to the rehearsal-detection and history-strip gate fixes, covering the sibling sites that were missed: - GGUF loop: the rehearsal-prefix and trailing-name hold checks now use the original tool list (_detect_tools) like the detection path, so a spent one-shot's split repeat (bare render_html then [ARGS]{...}) is held instead of flushed as visible text. - The safetensors and Anthropic assistant-history sanitisers and the Anthropic non-streaming passthrough now forward the enabled-tool-name gate to _strip_tool_xml_for_display, matching the GGUF history sanitiser and the live strips, so a prior turn documenting an inactive foo[ARGS]{...} example is preserved in the replayed prompt / final text instead of deleted. Add regression tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tile bracket-call spans per array item and include the v11 closer Two with_spans fixes for the Mistral bracket parser, both hit through the client-tool passthrough healers: - A multi-call [TOOL_CALLS] array carried its whole markup span on the first call and zero-width spans after, so a consumer that filters promotions by the declared tool set either re-emitted the full raw array as text next to the promoted call or silently dropped a filtered call's bytes. The region is now tiled across the call-producing items (each call's span covers its own JSON object plus the separator bytes before it; the last span runs to the region end), so promoted markup strips exactly once and a skipped call's bytes stay visible. - The v11 wrapper closer [/TOOL_CALLS] sat outside the reported span and leaked as stray text after promotion; the region now extends over an immediately-following closer. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: decouple healer signals from the loop signal set The passthrough healer buffered on every TOOL_XML_SIGNALS entry, so the bare [ARGS] rehearsal marker this branch adds for the loops (where it is gated on active tool names) put legitimate prose like 'Use foo[ARGS] in templates' into the holding state and stalled the stream until finalization. The healer can never promote a bare rehearsal call, so it now buffers only on formats its parser promotes: , <|tool_call>, in the template, including markup that only renders PAST assistant history (Kimi-K2-Thinking) while the generation prompt opens no . Starting the reasoning extractor in prefilled mode there captured a normal answer entirely as reasoning_content and returned blank visible content. Prefill only when rendering the generation prompt actually leaves open (DeepSeek-R1 / QwQ / Qwen3-Thinking); history-only templates start the extractor in normal mode and parse the model's own .... Adds a Kimi-shape regression test. * Keep bare scalar Mistral array arguments raw instead of double-encoding A scalar string argument in the canonical Mistral [TOOL_CALLS] array (for example [TOOL_CALLS][{"name":"web_search","arguments":"weather"}]) was run through json.dumps, turning weather into the JSON string "weather". The downstream argument healer then wrapped that quoted form, so a single-string tool like web_search searched for the literal "weather" with quotes. The path already keeps a scalar argument raw; mirror it here so only a dict is serialized. Add a regression test asserting both paths yield the same healed arguments. * Tighten tool-call rescue and reasoning-prefill 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 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 208 ++++- .../core/inference/passthrough_healing.py | 64 +- .../core/inference/safetensors_agentic.py | 394 +++++++-- .../core/inference/tool_call_parser.py | 116 ++- studio/backend/core/tool_healing.py | 521 +++++++++-- studio/backend/routes/inference.py | 277 ++++-- .../backend/tests/test_anthropic_messages.py | 18 + .../tests/test_gemma_tool_parse_edge_cases.py | 14 + .../backend/tests/test_llama_cpp_tool_loop.py | 306 +++++++ .../backend/tests/test_passthrough_healing.py | 89 ++ .../tests/test_responses_tool_passthrough.py | 11 +- .../test_safetensors_capability_advertise.py | 37 +- .../test_safetensors_reasoning_stream.py | 45 +- .../tests/test_safetensors_tool_loop.py | 818 ++++++++++++++++++ .../tests/test_tool_call_parser_strict.py | 141 ++- studio/backend/tests/test_tool_xml_strip.py | 324 ++++++- 16 files changed, 3114 insertions(+), 269 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 455d1d084c..aab3470193 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -44,7 +44,8 @@ from core.inference.llama_server_args import ( from core.inference.tool_call_parser import ( _GEMMA_BARE_TC_PREFIX_RE, _GEMMA_BARE_TC_RE, - _TOOL_ALL_PATS, + _TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS, + _TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS, _balanced_brace_end, _strip_function_xml_calls, _strip_gemma_wrapperless_calls, @@ -58,6 +59,16 @@ from core.inference.tool_call_parser import ( strip_llama3_leading_sentinels, strip_tool_markup as _shared_strip_tool_markup, ) + +# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated +# pattern lists, so the GGUF streaming strip stays aligned with the parser. +from core.tool_healing import ( + _REHEARSAL_TAIL_STRIP_RE, + _strip_bracket_tag_calls, + apply_tool_strip_patterns, + strip_outside_think, + strip_tool_call_markup, +) from utils.native_path_leases import child_env_without_native_path_secret from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( @@ -256,6 +267,72 @@ _FINAL_ANSWER_SIGNAL = re.compile( ) +def _gguf_active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in (active_tools or []) + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +# Rehearsal NAME chars (word + hyphen, matching the parser); the lookbehind excludes the +# Mistral [CALL_ID]...[ARGS] shape. +_GGUF_REHEARSAL_ARGS_RE = re.compile(r"(? int: + """Index of the first ``NAME[ARGS]`` whose NAME is an active tool, else -1. A + bare/inactive-name ``foo[ARGS]`` in prose is not a call; mirrors the safetensors + ``_earliest_tool_signal`` name-gating (no unrestricted GGUF mode).""" + active = set(_gguf_active_tool_names(active_tools)) + if not active: + return -1 + for m in _GGUF_REHEARSAL_ARGS_RE.finditer(text): + if m.group(1) in active: + return m.start() + return -1 + + +def _gguf_has_genuine_tool_signal(text: str, signals, active_tools: list[dict]) -> bool: + """True when ``text`` holds a genuine tool-call boundary for one of ``signals``. + + Unambiguous markers (````, ``[TOOL_CALLS]``, ``= 0: + return True + continue + if sig in text: + return True + return False + + +def _is_rehearsal_prefix(stripped: str, active_tools: list[dict]) -> bool: + """True if ``stripped`` is a (possibly partial) prefix of ``NAME[ARGS]`` for an + active tool -- the bare tool name arriving in its own chunk before ``[ARGS]{...}``. + Mirrors the safetensors loop so the split rehearsal call is not streamed.""" + if not stripped or any(ch.isspace() for ch in stripped): + return False + for name in _gguf_active_tool_names(active_tools): + if stripped == name or f"{name}[ARGS]".startswith(stripped): + return True + return False + + +def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int: + """Length of a trailing bare tool-name token that may be a split rehearsal call + (``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it + instead of leaking the name. Returns 0 for ordinary prose. Mirrors safetensors.""" + i = len(text) + while i > 0 and not text[i - 1].isspace(): + i -= 1 + tail = text[i:] + return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0 + + def _is_short_intent_without_action(text: str) -> bool: stripped = text.strip() return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None @@ -8418,6 +8495,13 @@ class LlamaCppBackend: _reasoning_started_at: Optional[float] = None _reasoning_summary_emitted = False + # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the + # ORIGINAL tools list so a spent one-shot still reads as a tool name. None = no gate. + _enabled_names_gate = set(_gguf_active_tool_names(tools)) if tools else None + # Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent + # one-shot), else its repeat is stripped but never drained and the turn ends blank. + _detect_tools = list(tools or []) + def _reasoning_summary_event(started_at: float) -> dict: return { "type": "reasoning_summary", @@ -8436,25 +8520,42 @@ class LlamaCppBackend: ) -> str: if not (auto_heal_tool_calls or force): return text + # Delegate to the shared parser-side strip so the GGUF cleanup covers every family the + # parser promotes (Llama <|python_tag|>, Mistral [TOOL_CALLS], bare rehearsal, function + # XML, Gemma) and stays aligned with detection; tool_healing's strip omits the loop-only + # forms (python_tag / Mistral name) and would leak them into display. return _shared_strip_tool_markup( - text, final = final, enabled_tool_names = _enabled_tool_names + text, final = final, enabled_tool_names = _enabled_names_gate ) def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text - # Shared parser patterns (not the legacy tool_healing set) so textual - # Mistral/python_tag calls entering DRAINING never leak. Balanced strips - # first (nested JSON removed whole); no final trim so length compares hold. - text = _strip_mistral_closed_calls(text) - text = _strip_gemma_wrapperless_calls(text, _enabled_tool_names) - # Parser-accurate scans close at each call's REAL terminator before - # the regex arms: literal markup inside a value is data. - text = _strip_function_xml_calls(text, final = True) - text = _strip_glm_calls(text, final = True) - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + + def _seg(segment: str, is_last: bool) -> str: + # Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced + # strips first (nested JSON removed whole; literal markup inside a value is that + # call's data), then the guarded function-XML / GLM scans, then the regex arms + # (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last + # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal + markerless + # strips are name-gated on the ORIGINAL list (strip/detect aligned). + seg = _strip_mistral_closed_calls(segment) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = _enabled_names_gate) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, _enabled_names_gate) + seg = _strip_function_xml_calls(seg, final = is_last) + seg = _strip_glm_calls(seg, final = is_last) + pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if is_last: + seg = apply_tool_strip_patterns( + seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = _enabled_names_gate + ) + return seg + + # Preserve think blocks verbatim (a rehearsed call inside one must not be deleted). + return strip_outside_think(text, _seg) def _build_metadata_event(usage, timings, finish_reason): """Final usage+timings metadata event for the given pass, merging its @@ -8814,12 +8915,18 @@ class LlamaCppBackend: in_thinking = False cumulative_display += token cleaned = _strip_tool_markup_streaming(cumulative_display) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned + # Hold a trailing bare active-tool-name (split rehearsal) + # until [ARGS] arrives; released by later prose or stream end. + _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _emit = ( + cleaned[: len(cleaned) - _hold] if _hold else cleaned + ) + if len(_emit) > len(_last_emitted): + _last_emitted = _emit if not _suppress_visible_output: yield { "type": "content", - "text": cleaned, + "text": _emit, } elif detect_state == _S_BUFFERING: @@ -8828,7 +8935,8 @@ class LlamaCppBackend: if not stripped_buf: continue - # Check tool signal prefixes. + # Bracket tags arrive mid-buffer, so substring-check too; + # ``[ARGS]`` counts only as a regex-matched NAME[ARGS]. is_prefix = False is_match = False for sig in _tool_xml_signals: @@ -8838,6 +8946,31 @@ class LlamaCppBackend: if sig.startswith(stripped_buf): is_prefix = True break + if sig == "[ARGS]": + # Active NAME[ARGS] only; inactive-name prose + # is gated out, not drained/parsed. + if ( + _gguf_rehearsal_signal_pos( + stripped_buf, _detect_tools + ) + >= 0 + ): + is_match = True + break + elif sig.startswith("[") and sig in stripped_buf: + is_match = True + break + + # Split rehearsal: hold the bare name until + # its [ARGS] arrives and matches above. + is_rehearsal_prefix = False + if ( + not is_match + and not is_prefix + and _is_rehearsal_prefix(stripped_buf, _detect_tools) + ): + is_prefix = True + is_rehearsal_prefix = True # Signal-less call shapes (mirror the safetensors # loop): Llama-3.2 bare {"name":..} and Gemma @@ -8884,9 +9017,14 @@ class LlamaCppBackend: # Tool signal -- flush any visible # prefix before DRAINING so the # route sends it before tool_start. + # Use the final strip (all families incl. Llama + # <|python_tag|> / Mistral name): the buffer holds + # the whole call, so a streaming closed-only strip + # would leak its open-ended markup as display text. _flush_reasoning_and_buffer() - cleaned = _strip_tool_markup_streaming( + cleaned = _strip_tool_markup( cumulative_display, + final = True, force = True, ) if len(cleaned) > len(_last_emitted): @@ -8898,8 +9036,14 @@ class LlamaCppBackend: } detect_state = _S_DRAINING elif _hold_buffer or ( - is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS + is_prefix + and ( + is_rehearsal_prefix + or len(stripped_buf) < _MAX_BUFFER_CHARS + ) ): + # A rehearsal prefix is self-bounded; the buffer + # cap must not cut long MCP names short. pass # keep buffering else: # Not a tool -- flush buffer @@ -8910,12 +9054,20 @@ class LlamaCppBackend: cleaned = _strip_tool_markup( cumulative_display, ) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned + # Same trailing-name hold as STREAMING for this + # first flush out of BUFFERING. + _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _emit = ( + cleaned[: len(cleaned) - _hold] + if _hold + else cleaned + ) + if len(_emit) > len(_last_emitted): + _last_emitted = _emit if not _suppress_visible_output: yield { "type": "content", - "text": cleaned, + "text": _emit, } except json.JSONDecodeError: @@ -8933,7 +9085,9 @@ class LlamaCppBackend: _is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json( _bare_eos, _enabled_tool_names ) - if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals): + if stripped_buf and _gguf_has_genuine_tool_signal( + stripped_buf, _tool_xml_signals, _detect_tools + ): detect_state = _S_DRAINING elif _is_bare_tc: detect_state = _S_DRAINING @@ -9066,6 +9220,12 @@ class LlamaCppBackend: "type": "content", "text": forced_visible_text, } + elif not _suppress_visible_output: + # Turn ended as a plain answer (no [ARGS] followed): the held + # rehearsal tail is real prose, release it. + _final_clean = _strip_tool_markup_streaming(cumulative_display) + if len(_final_clean) > len(_last_emitted): + yield {"type": "content", "text": _final_clean} # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index fe1aca0e4a..a444431f8d 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -32,16 +32,15 @@ from typing import Any, Optional from core.inference.tool_loop_controller import coerce_tool_arguments from core.tool_healing import parse_tool_calls_from_text -# Signals limited to the formats parse_tool_calls_from_text (core.tool_healing) -# actually promotes. The parser module's broader signal list also covers Llama -# <|python_tag|> and Mistral [TOOL_CALLS] for the streaming DRAIN buffers whose -# full parser handles them; buffering those here would hold a streamed -# client-tool call until finalization and then flush it as prose (this healer -# cannot promote them), so the passthrough keeps its own aligned list. +# Only the formats this healer's parser can promote -- narrower than the loops' +# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare +# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a +# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in. _HEAL_SIGNALS = ( "", "<|tool_call>", " list: diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 8e86d09754..aa732f47e4 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -14,6 +14,7 @@ parses tool calls from the cumulative text and dispatches via ``core.inference.tools``. """ +import bisect import re import threading from typing import Callable, Generator, Optional @@ -23,7 +24,8 @@ from loggers import get_logger from core.inference.tool_call_parser import ( _GEMMA_BARE_TC_PREFIX_RE, _GEMMA_BARE_TC_RE, - _TOOL_ALL_PATS, + _TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS, + _TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS, _balanced_brace_end, _strip_function_xml_calls, _strip_gemma_wrapperless_calls, @@ -39,6 +41,16 @@ from core.inference.tool_call_parser import ( strip_llama3_leading_sentinels, strip_tool_markup, ) + +# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated +# pattern lists, so the safetensors streaming strip stays aligned with the parser. +from core.tool_healing import ( + _REHEARSAL_TAIL_STRIP_RE, + _strip_bracket_tag_calls, + _think_spans_outside_tool_markup, + apply_tool_strip_patterns, + strip_outside_think, +) from core.inference.tool_loop_controller import ( ToolLoopController, coerce_tool_arguments, @@ -93,6 +105,147 @@ def _active_tool_names(active_tools: list[dict]) -> list[str]: return [name for name in names if name] +def _active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in active_tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +# Unrestricted mode has no tool list, so any identifier may open a NAME[ARGS] rehearsal; +# ``[`` and each ARGS letter stay optional so a chunk split after ``NAME[`` is still held. +_UNRESTRICTED_REHEARSAL_RE = re.compile(r"[\w-]+(?:\[(?:A(?:R(?:G(?:S)?)?)?)?)?") + + +def _is_rehearsal_prefix( + stripped: str, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> bool: + """True if ``stripped`` is a (possibly partial) prefix of a ``NAME[ARGS]`` + rehearsal split across chunks (``web_search`` then ``[ARGS]{...}``). A space + means prose. Unrestricted mode accepts any identifier; else NAME must be active.""" + if not stripped or any(ch.isspace() for ch in stripped): + return False + if unrestricted: + return _UNRESTRICTED_REHEARSAL_RE.fullmatch(stripped) is not None + for name in _active_tool_names(active_tools): + if stripped == name or f"{name}[ARGS]".startswith(stripped): + return True + return False + + +def _held_rehearsal_tail_len( + text: str, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """Length of a trailing bare tool-name token that may be a split rehearsal call + (``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it + instead of leaking the name. Returns 0 for ordinary prose.""" + i = len(text) + while i > 0 and not text[i - 1].isspace(): + i -= 1 + tail = text[i:] + return ( + len(tail) + if tail and _is_rehearsal_prefix(tail, active_tools, unrestricted = unrestricted) + else 0 + ) + + +def _rehearsal_name_start( + candidate: str, + signal_pos: int, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """For an ``[ARGS]`` signal at ``signal_pos``, return the start of the preceding + bare tool-name token (``NAME[ARGS]``), else ``signal_pos`` unchanged when the + signal is not ``[ARGS]`` or NAME is not an active tool (restricted mode).""" + if not candidate.startswith("[ARGS]", signal_pos): + return signal_pos + j = signal_pos + while j > 0 and (candidate[j - 1].isalnum() or candidate[j - 1] in "_-"): + j -= 1 + if j < signal_pos and ( + unrestricted or candidate[j:signal_pos] in _active_tool_names(active_tools) + ): + return j + return signal_pos + + +def _earliest_tool_signal( + candidate: str, + signals, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """Index where the turn's first genuine tool-call boundary begins, or -1. + + Non-``[ARGS]`` markup wins on first occurrence. An ``[ARGS]`` hit is a rehearsal + only when an active tool name (any name in unrestricted mode) precedes it, so a + literal ``foo[ARGS]`` in prose is skipped rather than draining the turn; for a + real ``NAME[ARGS]`` the boundary is pulled back to NAME.""" + best = -1 + for sig in signals: + if sig != "[ARGS]": + p = candidate.find(sig) + if p >= 0 and (best < 0 or p < best): + best = p + continue + from_idx = 0 + while True: + p = candidate.find("[ARGS]", from_idx) + if p < 0: + break + name_start = _rehearsal_name_start( + candidate, p, active_tools, unrestricted = unrestricted + ) + if name_start < p: + # Genuine ``NAME[ARGS]``: the boundary is the start of NAME. + if best < 0 or name_start < best: + best = name_start + break + # Bare/prose [ARGS]: skip it so a later real call in the same chunk is still found. + from_idx = p + len("[ARGS]") + return best + + +def _has_genuine_tool_signal( + candidate: str, + signals, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> bool: + """True when ``candidate`` holds a genuine tool-call boundary for one of ``signals``. + + Non-``[ARGS]`` markers count on a substring hit; an ``[ARGS]`` hit is genuine only + when an active tool name (any in unrestricted mode) precedes it. Mirrors the + ``_earliest_tool_signal`` name-gating so BUFFERING / end-of-stream checks do not + drain inactive-name prose.""" + for sig in signals: + if sig == "[ARGS]": + if ( + _earliest_tool_signal( + candidate, ("[ARGS]",), active_tools, unrestricted = unrestricted + ) + >= 0 + ): + return True + continue + if sig in candidate: + return True + return False + + def strip_tool_markup_streaming( text: str, *, @@ -101,25 +254,46 @@ def strip_tool_markup_streaming( enabled_tool_names: Optional[set] = None, ) -> str: """Strip open-ended tool XML from display text without trimming whitespace. - ``enabled_tool_names`` gates the markerless Gemma ``call:NAME{...}`` strip so a - disabled/example name in prose is kept (mirrors the parser gate).""" + + Mirrors the parser-side ``strip_tool_markup`` segment scan (minus the final trim) so + streaming and final display agree: balanced strips first (nested JSON removed whole), + then the guarded function-XML / GLM scans that close at each call's REAL terminator so + literal markup inside argument values is data and trailing prose survives. Reasoning + ```` / ``[THINK]`` blocks are preserved verbatim (a rehearsed call inside one must + not be deleted, else the cumulative text shrinks then regrows). ``enabled_tool_names`` + keeps an inactive-name ``foo[ARGS]{..}`` / ``call:NAME{..}`` example visible (it is prose, + not a call), matching the parse / detection active-tool gate.""" if not (auto_heal_tool_calls or tool_protocol_active): return text - # Mirror the final strip's scan order so streaming and final display agree: - # balanced strips first (nested JSON removed whole), then the guarded - # function-XML/GLM scans that close at each call's REAL terminator, so literal - # markup inside argument values is data and trailing prose survives. No final - # trim so streaming length comparisons hold. Leading Magistral [THINK]...[/THINK] - # is dropped (bracket form, not the reasoning channel's ); an unclosed - # [THINK] holds until [/THINK] so the cleaned text stays monotonic. + + # Drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket reasoning form, not the + # ```` channel) so raw reasoning does not leak into streamed display; an unclosed + # leading block is held (dropped to EOF) until its closer streams in. text = _strip_mistral_reasoning(text) - text = _strip_mistral_closed_calls(text) - text = _strip_gemma_wrapperless_calls(text, enabled_tool_names) - text = _strip_function_xml_calls(text, final = True) - text = _strip_glm_calls(text, final = True) - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + + def _seg(segment: str, is_last: bool) -> str: + # Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced + # strips first, then the guarded function-XML / GLM scans, then the regex arms + # (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last + # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal strips are name-gated. + seg = _strip_mistral_closed_calls(segment) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + seg = _strip_function_xml_calls(seg, final = is_last) + seg = _strip_glm_calls(seg, final = is_last) + pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if is_last: + seg = apply_tool_strip_patterns( + seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = enabled_tool_names + ) + return seg + + # Preserve think blocks verbatim: stripping a rehearsed call inside one shrinks then + # regrows the cumulative text, corrupting append-by-length consumers. + return strip_outside_think(text, _seg) def _strip_tool_markup_final( @@ -149,23 +323,66 @@ def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) _FUNCTION_SIGNAL_RE = re.compile(r"") _TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"') +# Mistral name/v11 and rehearsal forms, aligned with the parser so the provisional +# render-html card fires for bracket-tag serializations too. +_MISTRAL_RENDER_NAME_RE = re.compile( + r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)" +) +_REHEARSAL_RENDER_NAME_RE = re.compile(r"(? bool: - """Return True when the first drained tool call is clearly render_html.""" - function_match = _FUNCTION_SIGNAL_RE.search(content) - tool_call_index = content.find("") - if not function_match and tool_call_index < 0: + """Return True when the FIRST tool call in ``content`` is clearly render_html. + + Covers every serialization the loop executes (XML ```` / ````, + Mistral ``[TOOL_CALLS]``, rehearsal ``NAME[ARGS]``); the earliest marker wins so a + render_html marker inside another call's argument is treated as data. Markers inside + a ```` / ``[THINK]`` block are dropped since the parser skips them.""" + think_spans = _think_spans_outside_tool_markup(content) + _think_starts = [s for s, _e in think_spans] + + def _in_think(pos: int) -> bool: + if not think_spans: + return False + i = bisect.bisect_right(_think_starts, pos) - 1 + return i >= 0 and think_spans[i][0] <= pos < think_spans[i][1] + + def _first_outside(start: int, finder) -> int: + # First occurrence at/after ``start`` that is not inside a think span. + pos = finder(start) + while pos >= 0 and _in_think(pos): + pos = finder(pos + 1) + return pos + + candidates: list[tuple[int, str]] = [] + for fm in _FUNCTION_SIGNAL_RE.finditer(content): + if not _in_think(fm.start()): + candidates.append((fm.start(), fm.group(1))) + break + tc = _first_outside(0, lambda i: content.find("", i)) + if tc >= 0: + nm = _TOOL_CALL_NAME_RE.search(content[tc:]) + candidates.append((tc, nm.group(1) if nm else "")) + mt = _first_outside(0, lambda i: content.find("[TOOL_CALLS]", i)) + if mt >= 0: + mm = _MISTRAL_RENDER_NAME_RE.match(content, mt) + if mm: + candidates.append((mt, mm.group(1))) + else: + # Array shape: a bare ``"name"`` search can latch onto an argument key, so resolve the + # first call through the parser (it reads top-level names). + arr_calls = parse_tool_calls_from_text(content[mt:]) + if arr_calls: + candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or "")) + for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content): + if not _in_think(rm.start(1)): + candidates.append((rm.start(1), rm.group(1))) + break + + if not candidates: return False - - if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index): - return function_match.group(1) == "render_html" - - if tool_call_index >= 0: - name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:]) - return bool(name_match and name_match.group(1) == "render_html") - - return False + _pos, name = min(candidates, key = lambda c: c[0]) + return name == "render_html" def _coerce_arguments_with_provenance( @@ -256,6 +473,12 @@ def run_safetensors_tool_loop( conversation.extend(_auto["messages"]) unrestricted_tools = not tools + # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the + # ORIGINAL tools list so a spent one-shot still reads as a tool name. None = unrestricted. + _enabled_names_gate = None if unrestricted_tools else set(_active_tool_names(tools)) + # Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent + # one-shot), else its repeat is stripped but never drained and the turn ends blank. + _detect_tools = [] if unrestricted_tools else list(tools or []) tool_controller = ToolLoopController( tools = None if unrestricted_tools else tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -381,18 +604,18 @@ def run_safetensors_tool_loop( if detect_state == _state_streaming: candidate = cumulative_display + delta - signal_pos = -1 - for sig in tool_xml_signals: - p = candidate.find(sig) - if p >= 0 and (signal_pos < 0 or p < signal_pos): - signal_pos = p + # Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is + # pulled back to NAME so the name is not flushed. + signal_pos = _earliest_tool_signal( + candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools + ) if signal_pos >= 0: before_tool = candidate[:signal_pos] cleaned_before = strip_tool_markup_streaming( before_tool, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) if len(cleaned_before) > len(last_emitted): last_emitted = cleaned_before @@ -423,11 +646,20 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) - if len(cleaned) > len(last_emitted): - last_emitted = cleaned - yield {"type": "content", "text": cleaned} + # Hold a trailing bare active-tool-name (split rehearsal) until its [ARGS] arrives; + # released by later prose or the end-of-stream flush. + if tool_protocol_active: + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools, unrestricted = unrestricted_tools + ) + emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned + else: + emit = cleaned + if len(emit) > len(last_emitted): + last_emitted = emit + yield {"type": "content", "text": emit} continue # BUFFERING: hold until we know it is not a tool call. @@ -445,6 +677,34 @@ def run_safetensors_tool_loop( if sig.startswith(stripped): is_prefix = True break + # Bracket-tag forms arrive mid-buffer, so substring-check too (mirrors GGUF); [ARGS] + # counts only with an active NAME so prose is not drained into a no-op. + if sig == "[ARGS]": + if ( + _earliest_tool_signal( + stripped, + ("[ARGS]",), + _detect_tools, + unrestricted = unrestricted_tools, + ) + >= 0 + ): + is_match = True + break + elif sig.startswith("[") and sig in stripped: + is_match = True + break + + # Split rehearsal: hold the bare name until its [ARGS] arrives and matches above. + is_rehearsal_prefix = False + if ( + not is_match + and not is_prefix + and tool_protocol_active + and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools) + ): + is_prefix = True + is_rehearsal_prefix = True # Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML # signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses @@ -512,7 +772,7 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) if len(cleaned) > len(last_emitted): last_emitted = cleaned @@ -536,7 +796,8 @@ def run_safetensors_tool_loop( "arguments": {}, "provenance": _tool_event_provenance(provisional = True), } - elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS: + elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS): + # A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short. continue else: detect_state = _state_streaming @@ -545,24 +806,38 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) - if len(cleaned) > len(last_emitted): - last_emitted = cleaned - yield {"type": "content", "text": cleaned} + # Same trailing-name hold as STREAMING for this first flush out of BUFFERING. + if tool_protocol_active: + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools, unrestricted = unrestricted_tools + ) + emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned + else: + emit = cleaned + if len(emit) > len(last_emitted): + last_emitted = emit + yield {"type": "content", "text": emit} # Stream finished -- resolve what we collected. if cancel_event is not None and cancel_event.is_set(): return if detect_state == _state_buffering: - # Buffer never resolved -- tool XML or plain content? + # Buffer never resolved: [ARGS] is name-gated so a prose answer with a literal + # ``foo[ARGS]{...}`` is not parsed. stripped = content_buffer.lstrip() _bare_eos = strip_llama3_leading_sentinels(stripped) if ( stripped and tool_protocol_active - and any(sig in stripped for sig in tool_xml_signals) + and _has_genuine_tool_signal( + stripped, + tool_xml_signals, + _detect_tools, + unrestricted = unrestricted_tools, + ) ): detect_state = _state_draining elif tool_protocol_active and _looks_like_enabled_bare_json( @@ -629,6 +904,17 @@ def run_safetensors_tool_loop( # in full; route-level cleanup still applies the Auto-Heal policy. if content_accum and any(sig in content_accum for sig in tool_xml_signals): yield {"type": "content", "text": content_accum} + else: + # Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real + # prose, release it. + final_clean = strip_tool_markup_streaming( + cumulative_display, + auto_heal_tool_calls = auto_heal_tool_calls, + tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, + ) + if len(final_clean) > len(last_emitted): + yield {"type": "content", "text": final_clean} yield {"type": "status", "text": ""} return tool_calls = safety_tc @@ -636,19 +922,23 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) logger.info( "Safetensors safety net: parsed %d tool call(s) from streamed content", len(tool_calls), ) else: - # DRAINING: parse tool calls out of full content. + # DRAINING: parse tool calls out of full content. Gate the bare rehearsal on the + # ORIGINAL tool list (``_enabled_names_gate``), the same names detection/strip used to + # drain here: a spent one-shot (render_html) is off the active list but its re-emitted + # ``render_html[ARGS]{..}`` must still parse so it routes to the repeat no-op instead of + # being dropped into a blank continuation. tool_calls = parse_tool_calls_from_text( content_accum, id_offset = next_call_id, allow_incomplete = auto_heal_tool_calls, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) if not tool_calls: # Parser found nothing. Auto-Heal-enabled display cleanup @@ -682,7 +972,7 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) if tool_calls: diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 08a6bf418a..70115e5744 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -41,6 +41,9 @@ TOOL_XML_SIGNALS = ( "<|python_tag|>", "[TOOL_CALLS]", "<|tool_call>", + # Bare reasoning-rehearsal marker (``name[ARGS]{...}``, no leading [TOOL_CALLS]); + # keeps a rehearsed call held in the stream so it is promoted, not leaked as prose. + "[ARGS]", # DeepSeek R1 / V3 / V3.1 -- 5 opener variants llama.cpp keeps. "<｜tool▁calls▁begin｜>", "<｜tool▁call▁begin｜>", @@ -90,7 +93,7 @@ _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ # follows; a prose mention (``See [TOOL_CALLS] docs...``) keeps its tail. Bare marker at EOF drops. re.compile(r"<\|tool_call>(?=\s*call\s*:|\s*$).*$", re.DOTALL), re.compile( - r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*[\[{])|\s*$).*$", + r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*(?:[\[{]|\s*$))|\s*$).*$", re.DOTALL, ), re.compile( @@ -572,26 +575,51 @@ def strip_tool_markup( enabled_tool_names: Optional[set] = None, ) -> str: """Strip tool-call markup. ``final=False`` keeps in-progress markup buffered; - ``final=True`` also drops trailing unclosed runs and trims. ``enabled_tool_names`` - gates the markerless Gemma ``call:NAME{...}`` strip so a disabled/example name in - prose is kept (mirrors the parser gate); ``None`` strips every closed call.""" + ``final=True`` also drops trailing unclosed runs and trims. + + ``enabled_tool_names`` gates the name-conditioned forms so a disabled/example name in + prose is kept (mirrors the parser gate): the bare reasoning-rehearsal ``name[ARGS]{...}`` + and the markerless Gemma ``call:NAME{...}`` strip. ``None`` strips every closed call. + """ if final: # Drop a leading Magistral ``[THINK]...[/THINK]`` at end-of-turn; its bracket # form is not the ```` the reasoning channel renders. text = _strip_mistral_reasoning(text) - text = _strip_mistral_closed_calls(text) - if final: - text = _strip_gemma_wrapperless_calls(text, enabled_tool_names) - # Scan-strip the function-XML form (a literal ```` inside a value is - # data). The regex arms below cover the other formats but no-op on function calls here. - text = _strip_function_xml_calls(text, final = final) - # GLM 4.x: scan to the call's real so a literal one inside a value is data, - # not a leak. Qwen {json} is left to the regex arms. - text = _strip_glm_calls(text, final = final) - pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in pats: - text = pat.sub("", text) - return text.strip() if final else text + + def _strip_segment(segment: str, is_last: bool) -> str: + seg_final = final and is_last + seg = _strip_mistral_closed_calls(segment) + # Bare reasoning-rehearsal ``name[ARGS]{json}`` and the Mistral name form promote through + # the shared balanced scan, so strip them the same way (any nesting depth removed whole). + # The rehearsal arm is name-gated: an inactive ``foo[ARGS]{..}`` is prose and is kept. + seg = _tool_healing._strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if seg_final: + # Markerless Gemma ``call:NAME{...}`` (name-gated, mirrors the parse gate); end-of-turn only. + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + # Scan-strip the function-XML form (parser-accurate: a literal ```` in a + # value is data, not a call); the regex arms below cover the other formats. + seg = _strip_function_xml_calls(seg, final = seg_final) + # GLM 4.x: scan to the call's real so a literal one inside a value is data, + # not a leak. Qwen {json} is left to the regex arms. + seg = _strip_glm_calls(seg, final = seg_final) + pats = _TOOL_ALL_PATS if seg_final else _TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if seg_final: + # Drop a trailing partial bare rehearsal (``name[ARGS]`` with a truncated or absent + # body) the balanced scan cannot close; gated so prose ``foo[ARGS] ...`` survives. + seg = _tool_healing.apply_tool_strip_patterns( + seg, + [_tool_healing._REHEARSAL_TAIL_STRIP_RE], + enabled_tool_names = enabled_tool_names, + ) + return seg + + # ```` / ``[THINK]`` reasoning is preserved verbatim (a rehearsed call inside it is + # not executed, so it must not be stripped from display either); a literal think marker + # inside a real call's arguments is that call's data and is stripped with the call. + result = _tool_healing.strip_outside_think(text, _strip_segment) + return result.strip() if final else result def has_tool_signal(text: str) -> bool: @@ -711,6 +739,41 @@ def _xml_signal_inside_leading_mistral(content: str) -> bool: return _mistral_region_end(content, trig) is not None +def _parse_bare_rehearsals( + content: str, + *, + id_offset: int = 0, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Promote bare reasoning-rehearsal ``name[ARGS]{json}`` calls that a leading [TOOL_CALLS] + owns-the-turn parse would miss. Only the ``rehearsal`` kind is taken (a Mistral + ``[TOOL_CALLS]name[ARGS]{..}`` yields ``name`` and is not double-counted), and a rehearsal + inside a ```` / ``[THINK]`` block is reasoning, so it is skipped.""" + out: list[dict] = [] + think_spans = _tool_healing._think_spans_outside_tool_markup(content) + for start, end, kind, m in _tool_healing._iter_bracket_spans( + content, enabled_tool_names = enabled_tool_names + ): + if kind != "rehearsal": + continue + if any(s <= start < e for s, e in think_spans): + continue + try: + payload = json.loads(content[m.end() : end]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(payload, dict): + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": m.group(1), "arguments": json.dumps(payload)}, + } + ) + return out + + _ATTR_FUNC_OPEN_RE = re.compile(r'`` quoted-string handling the GGUF path relies on). + # Qwen/Hermes, Qwen3.5 XML, Gemma 4, plus Mistral [TOOL_CALLS] / bare rehearsal + # ``name[ARGS]{json}`` use the shared tool_healing parser (strict/Auto-Heal contract + + # nested-marker, trailing-prose, and ``<|"|>`` quoted-string handling the GGUF path + # relies on). ``enabled_tool_names`` gates the ambiguous bare-rehearsal form so an + # inactive ``foo[ARGS]{..}`` stays prose. calls = _tool_healing.parse_tool_calls_from_text( content, id_offset = id_offset, allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, ) if calls: return calls diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index b91403ed57..1b6b05768a 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -1,38 +1,91 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# +# Bracket-tag, rehearsal, and thinking-block-strip logic adapted from forge +# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026 +# Antoine Zambelli, used under the MIT License. -"""Lightweight tool-call XML parsing and stripping helpers. +"""Lightweight tool-call parsing and stripping helpers. External inference servers import this module without pulling in the inference -orchestrator, structlog, httpx, or the rest of the studio backend. +orchestrator, structlog, httpx, or the rest of the studio backend. Kept in +lockstep with ``core/inference/tool_call_parser.py`` so those servers +(llama-server wrappers, llama-swap, custom shims) reuse the same logic. Any +change here must also land there. + +Handles these serializations (see ``parse_tool_calls_from_text``): + +* ``{json}`` +* ``<|tool_call>call:name{...}`` (Gemma) +* ``v`` +* ``[TOOL_CALLS]name{json}`` (Mistral / Devstral fallback) +* ``name[ARGS]{json}`` (reasoning-model rehearsal) """ +# PEP 604 annotations must stay import-safe on Python 3.9 (requires-python >=3.9). +from __future__ import annotations + +import bisect import json import re -# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs -# strip first so a closed call goes as a unit before any to-EOF sweep reaches -# nested markup; only the final list adds the .*$ EOF sweeps. +# One nesting level in the strip regexes; deeper may leak markup (still parsed). +_BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}" + +# Rehearsal ``name[ARGS]{..}`` strips; group 1 = name for tool-list gating. Closed = +# complete body, tail = truncated; ``(?.*?`` rescans to EOF from every opener +# (quadratic on a stream of unclosed openers). Also reused by the quote-aware Gemma pre-pass. _TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) _TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) -_TC_GEMMA_END_PAT = re.compile(r"") _TOOL_CLOSED_PATS = [ _TC_JSON_CLOSED_PAT, _TC_GEMMA_CLOSED_PAT, + re.compile(r""), _TC_FUNC_CLOSED_PAT, - _TC_GEMMA_END_PAT, + # Mirror the parser regexes: tolerate whitespace and v11 [CALL_ID]/[ARGS] metadata. + re.compile( + r"\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*" + + _BRACKETED_JSON_ONE_LEVEL, + re.DOTALL, + ), + _REHEARSAL_CLOSED_STRIP_RE, + # Drop the bare v11 [/TOOL_CALLS] closer the balanced scan leaves behind. + re.compile(r"\[/TOOL_CALLS\]"), ] -_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ - re.compile(r"<\|tool_call>.*$", re.DOTALL), +# Bare open markers strip a partial call mid-stream; the rehearsal tail needs `{` or EOF +# so prose ``foo[ARGS]`` survives. The XML open-tail forms reach EOF and are reused by +# _tool_call_markup_spans (a think tag in an unclosed call's args stays argument data). +_TOOL_OPEN_XML_TAIL_PATS = [ re.compile(r".*$", re.DOTALL), + re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] -# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in -# their argument data cannot make the helper truncate the block and its tail. +_TOOL_ALL_PATS = ( + _TOOL_CLOSED_PATS + + _TOOL_OPEN_XML_TAIL_PATS + + [ + re.compile(r"\[TOOL_CALLS\].*$", re.DOTALL), + _REHEARSAL_TAIL_STRIP_RE, + ] +) + +# Rehearsal strips (name in group 1); name-gated via ``enabled_tool_names``, strip-all when None. +_REHEARSAL_STRIP_PATS = frozenset({_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE}) + +# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in argument +# data cannot make the helper truncate the block and its tail. _TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT] -# A lazy closed-pair pattern whose close token is absent rescans to EOF from -# every opener (quadratic, re-run per streamed token); skip that doomed pass. +# A lazy closed-pair pattern whose close token is absent would rescan to EOF from every +# opener; skip that doomed (quadratic) pass. Shared by both strip helpers. _PAT_REQUIRED_TOKEN = { _TC_JSON_CLOSED_PAT: "", _TC_GEMMA_CLOSED_PAT: "", @@ -50,26 +103,102 @@ def strip_tool_patterns(text: str, patterns) -> str: return text +def apply_tool_strip_patterns( + text: str, + patterns, + enabled_tool_names = None, +) -> str: + """Apply strip ``patterns`` to ``text``. A bare rehearsal ``name[ARGS]{..}`` pattern + strips only when ``name`` is an enabled tool (or when ``enabled_tool_names`` is + ``None``); every other pattern is removed unconditionally. A closed-pair pattern whose + close token is absent is skipped so an unclosed-marker stream stays linear.""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + if enabled_tool_names is not None and pat in _REHEARSAL_STRIP_PATS: + text = pat.sub(lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text) + else: + text = pat.sub("", text) + return text + + # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") -# Name class allows dots/hyphens for dotted Gemma names; whitespace-tolerant around -# ``call`` / ``:`` since drift emits ``call: name{`` and ``call : name{``. _TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") _TC_GEMMA_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -# Horizontal whitespace only so the newline + value indentation survive (_trim_param_value trims one newline). +# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it. _TC_PARAM_START_RE = re.compile(r"[^\S\n]*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _GEMMA_QUOTE = '<|"|>' _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" -# A bare (unquoted) Gemma value ends at `}` or at a comma beginning the next -# identifier-shaped `key:` pair; a comma before a non-key (`New York, NY`, -# `10:00, 11:00`) stays in the value. Dots let a dotted key end the value. +# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next +# `key:` pair. A comma NOT followed by a key token is part of the value (e.g. +# `location:New York, NY`), so it must not terminate the value. The key token +# must be identifier-shaped (start with a letter or underscore); a comma +# followed by digits-then-colon is value text such as a timestamp or ratio +# (`meet at 10:00, 11:00 tomorrow`), not a new key. _GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") +# A candidate starting inside a think block is a rehearsal (block kept so literal tags in +# real args survive); ``$`` accepts an unclosed block mid-stream. +_THINK_TAG_RE = re.compile(r".*?(?:|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL) +# Bare open/close markers for prefilled-reasoning turns (template opens in the prompt). +_THINK_OPEN_RE = re.compile(r"|\[THINK\]") +_THINK_CLOSE_RE = re.compile(r"|\[/THINK\]") + +# Mistral canonical array: [TOOL_CALLS] + JSON list of {"name","arguments"} objects. +_MISTRAL_ARRAY_RE = re.compile(r"\[TOOL_CALLS\]\s*(?=\[)") + +# Mistral name form + v11 [ARGS]/[CALL_ID] shapes; [CALL_ID] is metadata, not the name, +# and hyphens keep dashed MCP names whole. +_MISTRAL_BRACKET_RE = re.compile( + r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)" +) + +# Rehearsal ``name[ARGS]{json}`` (no [TOOL_CALLS]); the lookbehind keeps the v11 call-id +# from being taken as the function name. +_REHEARSAL_RE = re.compile(r"(? int | None: + """Return the end index of a balanced JSON object opening at ``start``, + or ``None`` if the braces don't balance. Honors escapes and strings. + """ + if start >= len(text) or text[start] != "{": + return None + depth = 0 + in_string = False + escape = False + for j in range(start, len(text)): + ch = text[j] + if escape: + escape = False + continue + if ch == "\\": + escape = True + continue + if in_string: + if ch == '"': + in_string = False + continue + if ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return j + return None + def _balanced_brace_end( content: str, @@ -134,6 +263,94 @@ def _balanced_bracket_end(src: str, start: int) -> int: return -1 +def _decode_array_items(text: str, body_start: int, body_end: int): + """Return ``(objs, ends)`` for each top-level element of the JSON array between + ``body_start`` (at or before its ``[``) and ``body_end`` (exclusive): the decoded + object and its absolute exclusive end offset. + + Decoding element-by-element with ``raw_decode`` tolerates the comma-less object + separators the repo's own Mistral/Ollama multi-call templates emit + (``[{...}{...}]``; see ollama_template_mappers.py). A single ``json.loads`` of the + whole body rejects that form and would drop every call. The ends also tile the + region across the calls' spans so a with_spans consumer strips each exactly once.""" + decoder = json.JSONDecoder() + objs: list = [] + ends: list[int] = [] + i = text.find("[", body_start) + if i < 0: + return objs, ends + i += 1 + while i < body_end: + while i < body_end and text[i] in " \t\r\n,": + i += 1 + if i >= body_end or text[i] == "]": + break + try: + obj, rel = decoder.raw_decode(text[i:body_end]) + except (json.JSONDecodeError, ValueError): + break + i += rel + objs.append(obj) + ends.append(i) + return objs, ends + + +def _iter_bracket_spans( + text: str, + start: int = 0, + enabled_tool_names = None, +): + """Yield ``(span_start, span_end, kind, match)`` for each balanced bracket-tag + call from ``start`` on, in document order; ``span_end`` exclusive. ``kind`` is + ``"array"`` ([TOOL_CALLS] [..]), ``"name"`` ([TOOL_CALLS]name{..}, incl. v11 + [CALL_ID]/[ARGS]) or ``"rehearsal"`` (name[ARGS]{..}). + + ``enabled_tool_names`` (set, or None = unrestricted) gates only the ambiguous + bare rehearsal form: name[ARGS]{..} is a call ONLY when ``name`` is enabled, so a + prose ``foo[ARGS]{..}`` (foo disabled) is neither parsed nor stripped. Explicit + [TOOL_CALLS] markers stay unconditional, keeping parse/strip/detection symmetric. + + Balance-only (no JSON validation) so strip and parse share one scan. The cursor + jumps past each consumed span, so a marker inside consumed JSON is never + re-matched and each regex re-searches only once its match falls behind: linear.""" + n = len(text) + specs = ( + ("array", _MISTRAL_ARRAY_RE), + ("name", _MISTRAL_BRACKET_RE), + ("rehearsal", _REHEARSAL_RE), + ) + nexts = {kind: rx.search(text, start) for kind, rx in specs} + cursor = start + while cursor < n: + for kind, rx in specs: + m = nexts[kind] + if m is not None and m.start() < cursor: + nexts[kind] = rx.search(text, cursor) + live = [(kind, m) for kind, m in nexts.items() if m is not None] + if not live: + return + kind, m = min(live, key = lambda km: km[1].start()) + if kind == "array": + end = _balanced_bracket_end(text, m.end()) + end = None if end < 0 else end + else: + end = _balanced_json_span(text, m.end()) + if end is None: + # Truncated body: skip and keep scanning; the caller's catch-all strips the tail. + cursor = m.end() + continue + if ( + kind == "rehearsal" + and enabled_tool_names is not None + and m.group(1) not in enabled_tool_names + ): + # Inactive-name rehearsal is prose: advance past its body without yielding. + cursor = end + 1 + continue + yield (m.start(), end + 1, kind, m) + cursor = end + 1 + + def _split_top_level_commas(src: str) -> list: """Split on commas that are not inside a nested ``[]``/``{}`` or a string.""" parts: list[str] = [] @@ -164,8 +381,14 @@ def _split_top_level_commas(src: str) -> list: def _quote_gemma_array_elements(body: str) -> str: - """Normalise a Gemma array value (``labels:[bug,ui]``) so json.loads succeeds: - quote bare strings, recurse into objects/arrays, keep quoted/JSON literals.""" + """Normalise the elements of a Gemma array value so json.loads succeeds. + + Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of + objects (``items:[{path:a}]``) whose keys/values also lack quotes; left + as-is json.loads fails and the whole call is dropped. Bare string elements + are quoted, object and nested-array elements are normalised recursively, and + quoted strings (already normalised from ``<|"|>``), numbers, and JSON + literals are preserved.""" out: list[str] = [] for element in _split_top_level_commas(body): stripped = element.strip() @@ -173,9 +396,11 @@ def _quote_gemma_array_elements(body: str) -> str: out.append(element) continue if stripped[0] == "{": + # Object element: quote its keys/bare values like a top-level object. out.append(_quote_gemma_object_keys(stripped)) continue if stripped[0] == "[": + # Nested array: normalise its elements too. inner_end = _balanced_bracket_end(stripped, 0) if inner_end == len(stripped) - 1: out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]") @@ -240,8 +465,6 @@ def _quote_gemma_object_keys(src: str) -> str: while i < len(src) and src[i].isspace(): i += 1 key_name_start = i - # Dots match the parser's key/name charset: Gemma emits dotted argument keys - # (user.name:...) for namespaced schemas. while i < len(src) and (src[i].isalnum() or src[i] in "_-."): i += 1 key_name = src[key_name_start:i] @@ -254,12 +477,15 @@ def _quote_gemma_object_keys(src: str) -> str: parts.append(src[i:colon_pos]) parts.append(":") i = colon_pos + 1 - # Quote bare string values ({unit:celsius}); JSON stays as-is. + # Gemma may emit bare string values ({unit:celsius}); quote them so + # json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is. ws = i while i < len(src) and src[i].isspace(): i += 1 parts.append(src[ws:i]) if i < len(src) and src[i] == "[": + # Array value: quote bare string elements (e.g. labels:[bug,ui]) + # so json.loads succeeds instead of dropping the call. arr_end = _balanced_bracket_end(src, i) if arr_end < 0: parts.append(src[i:]) @@ -269,7 +495,9 @@ def _quote_gemma_object_keys(src: str) -> str: i = arr_end + 1 elif i < len(src) and src[i] not in '"{': v_start = i - # Bare value: up to `}` or a comma that starts the next key:pair. + # Consume the bare value up to `}` or a comma that starts the + # next key:value pair; a comma inside the value (e.g. + # `New York, NY`) does not terminate it. while i < len(src): if src[i] == "}": break @@ -329,7 +557,9 @@ def _func_close_index(content: str, body_start: int, body: str) -> int: def _trim_param_value(val: str) -> str: - """Trim only the wrapping newline (not str.strip) so code/diff argument indentation survives.""" + """Trim the single wrapping newline the chat template adds around an XML + parameter value, preserving indentation inside VALUE (``str.strip()`` destroyed + code/diff argument indentation).""" if val.startswith("\n"): val = val[1:] if val.endswith("\n"): @@ -404,6 +634,7 @@ def parse_tool_calls_from_text( *, id_offset: int = 0, allow_incomplete: bool = True, + enabled_tool_names = None, with_spans: bool = False, ): """Parse OpenAI-format tool calls from model text. @@ -412,22 +643,36 @@ def parse_tool_calls_from_text( {"name":"web_search","arguments":{"query":"..."}} <|tool_call>call:web_search{query:"..."} ... + [TOOL_CALLS]web_search{"query":"..."} (Mistral / Devstral fallback) + web_search[ARGS]{"query":"..."} (reasoning-model rehearsal) + + A call rehearsed inside a ```` / ``[THINK]`` block is skipped, not + executed; the block is kept so a literal tag in a real argument is preserved. With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]`` is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup in ``content`` (including its close tag when present), so a caller can remove exactly the parsed markup and keep every other byte intact. """ + # Candidates starting inside a think block are rehearsals, skipped; blocks are kept, and a + # think marker opening inside a call is argument data (excluded from spans). + _think_spans = _think_spans_outside_tool_markup(content) + _think_starts = [s for s, _e in _think_spans] + + def _in_think(pos: int) -> bool: + # Spans are ordered and non-overlapping; bisect gives O(log M) per candidate. + i = bisect.bisect_right(_think_starts, pos) - 1 + return i >= 0 and _think_spans[i][0] <= pos < _think_spans[i][1] + tool_calls: list[dict] = [] call_spans: list[tuple] = [] - # Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside - # another call's coverage, or an open value, is data not executed. - markers = _build_markers(content) - coverage = _marker_coverage(content, markers) + # Collect JSON/Gemma markers; _marker_coverage decides nesting so a marker inside + # another call's coverage (even one that failed to parse) is data, not executed. A + # marker opening inside a think block is a rehearsal and is skipped. parsed_items = [] # (start, span_end, name, arguments) in document order + markers = [mk for mk in _build_markers(content) if not _in_think(mk[0])] + coverage = _marker_coverage(content, markers) for idx, (start, brace_end, kind, m) in enumerate(markers): - # A marker starting inside another's coverage is that call's data. The - # end is exclusive so a marker at a close's end is an adjacent sibling. if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx): continue if brace_end < 0: @@ -441,7 +686,7 @@ def parse_tool_calls_from_text( if kind == "json": obj = json.loads(content[m.end() - 1 : brace_end + 1]) name = obj.get("name", "") - # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes ). + # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes). arguments = obj.get("arguments") if arguments is None: arguments = obj.get("parameters", {}) @@ -452,7 +697,6 @@ def parse_tool_calls_from_text( arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) except (json.JSONDecodeError, ValueError): continue - # Span reaches through the close tag when present, else just the braces. span_end = brace_end + 1 close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE ws = len(content[span_end:]) - len(content[span_end:].lstrip()) @@ -461,14 +705,11 @@ def parse_tool_calls_from_text( span_end = close_m.end() parsed_items.append((start, span_end, name, arguments)) - # Function-XML calls promote in document order alongside marker calls (the - # #6801 contract). A inside any marker's coverage is excluded -- - # even if that marker failed to parse -- so nested XML cannot escape; one - # after a balanced close-less marker is a sibling, not swallowed to EOF. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) + and not _in_think(fm.start()) and not any(s <= fm.start() < e for s, e in coverage) ] for idx, fm in enumerate(func_starts): @@ -545,11 +786,170 @@ def parse_tool_calls_from_text( ) call_spans.append((start, span_end)) + # Patterns 3+4: Mistral [TOOL_CALLS] and bare rehearsal via one balanced scan in document + # order, so a Mistral call and a rehearsal in one message both parse. + if not tool_calls: + for start, end, kind, m in _iter_bracket_spans( + content, enabled_tool_names = enabled_tool_names + ): + if _in_think(start): + continue + # Extend the region over an immediately-following v11 closer so with_spans consumers strip it too. + closer = re.match(r"\s*\[/TOOL_CALLS\]", content[end:]) + region_end = end + closer.end() if closer else end + if kind == "array": + # Decode elements individually (comma-tolerant): one json.loads of the whole + # body rejects the comma-less multi-call arrays Mistral/Ollama templates emit. + payload, item_ends = _decode_array_items(content, m.end(), end) + if not payload: + continue + # Tile the region so every byte belongs to exactly one span; a with_spans consumer + # keeps skipped bytes visible and strips promoted markup exactly once. + tile_start = start + last_span_idx = -1 + for item_idx, item in enumerate(payload): + if not isinstance(item, dict) or "name" not in item: + continue + args = item.get("arguments", {}) + if isinstance(args, str): + # ``arguments`` may itself be a JSON string (OpenAI spec). + try: + args = json.loads(args) + except (json.JSONDecodeError, ValueError): + pass + if not isinstance(args, (dict, str)): + # ``"arguments": null`` (or any non-object scalar) becomes {} like the + # path, not the string "null" auto-heal would mangle to + # a bogus {"query":"null"}. + args = {} + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": item.get("name", ""), + # A bare scalar string stays raw (like the path); + # json.dumps would double-encode it so the arg healer wraps + # "weather" with its literal quotes. + "arguments": args if isinstance(args, str) else json.dumps(args), + }, + } + ) + item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end + last_span_idx = len(call_spans) + call_spans.append((tile_start, item_end)) + tile_start = item_end + if last_span_idx >= 0: + tile_start, _tile_end = call_spans[last_span_idx] + call_spans[last_span_idx] = (tile_start, region_end) + else: + try: + payload = json.loads(content[m.end() : end]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(payload, dict): + continue + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": m.group(1), + "arguments": json.dumps(payload), + }, + } + ) + call_spans.append((start, region_end)) + if with_spans: return tool_calls, call_spans return tool_calls +def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str: + """Strip complete [TOOL_CALLS] arrays / name / bare name[ARGS]{..} calls with one + balanced forward scan, so nested JSON args are removed whole (a fixed-depth regex + left two-level args behind). Truncated tails go to the caller's catch-all. Linear. + ``enabled_tool_names`` gates the rehearsal form (inactive-name prose kept; None + strips every span).""" + if len(text) > _MAX_BRACKET_SCAN_CHARS: + return text + out: list[str] = [] + cursor = 0 + for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names): + out.append(text[cursor:start]) + cursor = end + out.append(text[cursor:]) + return "".join(out) + + +def _tool_call_markup_spans(text: str) -> list[tuple[int, int]]: + """Spans of tool-call markup, so a literal /[THINK] inside a call's args is + stripped WITH the call, not kept as a reasoning block. Covers closed XML/bracket + calls and an unclosed XML call (run via allow_incomplete); without the open-ended + span the unclosed call's markup would leak after execution.""" + # Skip a lazy closed-pair pattern whose close token is absent: its finditer would rescan + # to EOF from every opener (quadratic on a stream of unclosed openers). + spans = [ + m.span() + for pat in _TOOL_CLOSED_PATS + if (_PAT_REQUIRED_TOKEN.get(pat) is None or _PAT_REQUIRED_TOKEN[pat] in text) + for m in pat.finditer(text) + ] + spans.extend((start, end) for start, end, _kind, _m in _iter_bracket_spans(text)) + # An unclosed opener is a real incomplete call only outside closed/bracket spans. + for pat in _TOOL_OPEN_XML_TAIL_PATS: + for m in pat.finditer(text): + if not any(s <= m.start() < e for s, e in spans): + spans.append(m.span()) + return spans + + +def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]: + """/[THINK] block spans, minus any whose opening marker sits INSIDE a + tool-call span (that tag is argument data, not reasoning). Keeping it would drop a + real call after it as rehearsed and leak the call's markup. START tested only, so + a greedy unclosed past the call is still that call's argument data.""" + think_spans = [m.span() for m in _THINK_TAG_RE.finditer(text)] + call_spans = _tool_call_markup_spans(text) + # Prefilled reasoning: the template opens in the prompt, so add a leading span + # (0..close) to skip calls rehearsed there; guarded so a stray close in a normal answer is safe. + close = _THINK_CLOSE_RE.search(text) + if close is not None: + opener = _THINK_OPEN_RE.search(text) + if ( + (opener is None or close.start() < opener.start()) + and not any(cs <= close.start() < ce for cs, ce in call_spans) + and any(cs >= close.end() for cs, ce in call_spans) + ): + think_spans = [(0, close.end())] + think_spans + if not think_spans: + return think_spans + if not call_spans: + return think_spans + return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)] + + +def strip_outside_think(text: str, strip_segment) -> str: + """Apply ``strip_segment(segment, is_last)`` to visible text around /[THINK] + blocks, preserving the blocks verbatim (tool-looking text inside is rehearsal). + ``is_last`` is True only after the final block, so trailing-tail patterns apply + only there. Shared by every strip path so they stay consistent.""" + # A think marker opening inside a complete call is argument text; excluding it lets the + # stripper see the whole call. START-tested, so an unclosed match stays argument data. + think_spans = _think_spans_outside_tool_markup(text) + if not think_spans: + return strip_segment(text, True) + pieces: list[str] = [] + prev = 0 + for s, e in think_spans: + pieces.append(strip_segment(text[prev:s], False)) + pieces.append(text[s:e]) + prev = e + pieces.append(strip_segment(text[prev:], True)) + return "".join(pieces) + + def _strip_gemma_native_spans(text: str, *, final: bool) -> str: """Remove complete Gemma-native spans, brace/quote-balanced so a literal ```` in a quoted argument cannot truncate the span. An incomplete @@ -635,26 +1035,43 @@ def _strip_closed_blocks_outside_gemma(text: str) -> str: return text -def strip_tool_markup_final(text: str) -> str: - """Final display strip, shared with the streaming wrappers so all paths order - the passes identically: Gemma-aware closed JSON/function blocks first, then - well-formed Gemma spans (quote-aware), then the regex sweeps mop up malformed - spans and drop any unclosed remainder to EOF. Whitespace is kept.""" +def _strip_markup_segment( + text: str, + *, + final: bool, + enabled_tool_names = None, +) -> str: + # Bracket-tag calls (Mistral/rehearsal) first via balanced scan (any nesting depth, + # rehearsal name-gated); then the quote-aware Gemma-native passes so a literal + # in an argument cannot truncate a block; finally the regex XML/tail sweeps. + text = _strip_bracket_tag_calls(text, enabled_tool_names = enabled_tool_names) text = _strip_closed_blocks_outside_gemma(text) - text = _strip_gemma_native_spans(text, final = True) - return strip_tool_patterns(text, _TOOL_ALL_PATS) + text = _strip_gemma_native_spans(text, final = final) + patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS + return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names) -def strip_tool_call_markup(text: str, *, final: bool = False) -> str: +def strip_tool_call_markup( + text: str, + *, + final: bool = False, + enabled_tool_names = None, +) -> str: """Strip tool-call XML markup from text. When ``final`` is False, only fully closed tool-call blocks are removed. When ``final`` is True, trailing incomplete tool-call blocks are removed too, and the result is stripped of surrounding whitespace. + + ```` / ``[THINK]`` reasoning is preserved verbatim (see + ``strip_outside_think``); the trailing-tail patterns apply only after the + last block. ``enabled_tool_names`` keeps an inactive-name ``foo[ARGS]{..}`` + example visible (it is prose, not a call) so display cleanup matches detection. """ - if final: - return strip_tool_markup_final(text).strip() - # Non-final: same ordering as the final path, but incomplete blocks are kept. - text = _strip_closed_blocks_outside_gemma(text) - text = _strip_gemma_native_spans(text, final = False) - return strip_tool_patterns(text, _TOOL_CLOSED_PATS) + result = strip_outside_think( + text, + lambda seg, is_last: _strip_markup_segment( + seg, final = final and is_last, enabled_tool_names = enabled_tool_names + ), + ) + return result.strip() if final else result diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3341a9c628..1042dda004 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1366,17 +1366,60 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: return flags +def _generation_prompt_opens_think(template: Optional[str]) -> bool: + """True when rendering the template's generation prompt ends INSIDE an unclosed ````. + + Distinguishes templates that PREFILL an open ```` in the assistant generation + prompt (DeepSeek-R1, QwQ, Qwen3-Thinking) -- where the model emits only the closing + ```` and the extractor must start in reasoning mode -- from templates that merely + render PAST assistant ``...`` history while leaving the generation prompt + open with no ```` (e.g. Kimi-K2-Thinking), where the model self-emits its own block + and the extractor must start in normal mode. Renders a single-user-message probe with the + same sandbox transformers uses; on any failure returns True, preserving the historical + always-on prefill for templates that cannot be rendered here. + """ + if not template: + return False + try: + from jinja2.sandbox import ImmutableSandboxedEnvironment + + def _raise_exception(message: str): + raise RuntimeError(message) + + env = ImmutableSandboxedEnvironment( + trim_blocks = True, + lstrip_blocks = True, + extensions = ["jinja2.ext.loopcontrols"], + ) + env.filters["tojson"] = lambda value, **kwargs: json.dumps(value, ensure_ascii = False) + env.globals["raise_exception"] = _raise_exception + rendered = env.from_string(template).render( + messages = [{"role": "user", "content": "hi"}], + add_generation_prompt = True, + bos_token = "", + eos_token = "", + ) + except Exception: + return True + # ```` is not a substring of ```` (the ``/`` breaks it), so the last open + # tag sitting after the last close tag means the prompt ends inside an open block. + return rendered.rfind("") > rendered.rfind("") + + def _sf_reasoning_prefill_mode( features: dict, enable_thinking: Optional[bool], template: Optional[str] = None, reasoning_effort: Optional[str] = None, ) -> bool: - """Whether this request begins INSIDE an unclosed ```` (Qwen3/Qwen3.5/GLM prefill it). + """Whether a safetensors/MLX generation begins INSIDE an unclosed ````. - Gated on the STANDARD ````/```` markers: a bespoke reasoning channel (e.g. gemma) - never emits ````, so prefilled mode would swallow the whole answer -- excluded, as are - gpt-oss and thinking-disabled requests. ``enable_thinking=None`` defaults ON, so plain requests prefill. + ``enable_thinking`` templates (Qwen3/GLM) prefill an open ```` so the model + emits only the closing ````, and the extractor must start in reasoning mode. + Gated on the STANDARD ````/```` markers: bespoke channels (gemma's + ``<|think|>``) never emit ```` and would swallow the answer, so they and + gpt-oss and thinking-disabled requests return False. ``enable_thinking`` None + defaults thinking ON, so a plain request still prefills. """ if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"): return False @@ -1384,16 +1427,21 @@ def _sf_reasoning_prefill_mode( if "" not in tpl and "" not in tpl: return False if features.get("reasoning_always_on"): - return True + # enable_thinking_effort + always-on: the effort mechanism (not the prompt shape) keeps + # thinking on, so always-on wins over reasoning_effort and we prefill. + if features.get("reasoning_style") == "enable_thinking_effort": + return True + # ``reasoning_always_on`` fires on paired ``...`` anywhere in the + # template, including markup that only renders PAST assistant history (Kimi-K2-Thinking) + # while the generation prompt opens none. Prefill only when the generation prompt opens + # one, else the extractor captures a normal answer as reasoning_content and returns blank. + return _generation_prompt_opens_think(tpl) if not features.get("supports_reasoning"): return False if enable_thinking is False: return False - # A reasoning_effort="none" request disables thinking for enable_thinking_effort - # (GLM-5.2) models the same way enable_thinking=False does (see - # ``_request_reasoning_kwargs``). Without this, the model emits no ```` and - # a plain answer is swallowed whole into reasoning_content, leaving the visible - # response empty. + # Thinking-off arrives as reasoning_effort "none" on enable_thinking_effort models; honor it + # so we don't prefill and capture the answer. Plain enable_thinking models ignore effort. if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none": return False return True @@ -1669,11 +1717,17 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str: return nudge + " " + _RAG_GROUNDING_NUDGE -# Strip leaked tool-call markup: every shared-parser format plus the four leak -# shapes llama_cpp.py's speculative buffer splits across the visible/DRAIN -# boundary. Mistral [TOOL_CALLS] uses the parser's balanced-brace helper (a -# non-greedy regex would truncate nested JSON); the DeepSeek opener alternation -# is the parser's own, so a signal we parse is never left un-stripped. +# Strip leaked tool-call markup: every shared-parser format plus the leak shapes +# llama_cpp.py's speculative buffer splits across the visible/DRAIN boundary: +# 1. well-formed `...` / `...` +# 2. orphan opening to EOF (close was DRAINED) +# 3. bare orphan close (open was DRAINED) +# 4. tail-only `` (outer close truncated by EOS); anchored to +# `\Z` so mid-text `` in user code samples survives. +# 5. Mistral `[TOOL_CALLS]name{json}` / rehearsal `name[ARGS]{json}`: the balanced +# scan removes the whole call (a non-greedy regex would truncate nested JSON). +# DeepSeek/GLM/Kimi envelopes are covered by the parser's own arms/scans, so a signal +# we parse is never left un-stripped; the DeepSeek opener alternation is the parser's own. from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC _TOOL_XML_RE = _re.compile( @@ -1694,6 +1748,17 @@ _TOOL_XML_RE = _re.compile( r"|" r"|" r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*" + r"|\[/TOOL_CALLS\]" + # Truncated canonical array (closing ``]`` lost to EOS): the balanced scan cannot remove + # it, so strip its tail here. + r"|\[TOOL_CALLS\]\s*\[.*\Z" + # Named / v11 forms and bare rehearsal; arms aligned with the parser regexes. + r"|\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|.*?\Z)" + # Rehearsal: balanced/truncated body or bare marker at EOS only (prose ``foo[ARGS]`` + # survives); NAME captured as ``reh`` for the inactive-name display gate. + r"|(?[\w-]+)\[ARGS\]\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|\{.*\Z|\Z)" + # DeepSeek envelopes (all opener variants), Kimi section blocks, and bare Kimi calls; + # each arm carries a call-shaped lookahead so prose merely mentioning a marker survives. r"|" + _DS_OPEN_SRC + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*?(?:<｜tool▁calls▁end｜>|\Z)" @@ -1705,6 +1770,17 @@ _TOOL_XML_RE = _re.compile( _re.DOTALL, ) +# Closed-only variant for segments before the last think block: the ``\Z``-anchored arms +# would treat a segment boundary as EOS and strip prose ``foo[ARGS]``. +_TOOL_XML_CLOSED_RE = _re.compile( + r"<(?:tool_call|function=[\w-]+)>.*?" + r"|<\|tool_call>.*?" + r"|" + r"|" + r"|\[/TOOL_CALLS\]", + _re.DOTALL, +) + def _gemma_strip_gate(tools) -> set: """Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the @@ -1720,18 +1796,18 @@ def _gemma_strip_gate(tools) -> set: return names -def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str: - """Combine the parser's scan-based strips (Mistral balanced-brace, gated - Gemma wrapper-less, GLM real-close, guarded function-XML) with - ``_TOOL_XML_RE`` -- the scan strips close at each call's REAL terminator so - literal markup inside argument values is data, not a leaked tail. - ``enabled_tool_names`` gates the Gemma strip; ``None`` strips every closed call.""" - cleaned = _strip_glm_calls( - _strip_gemma_wrapperless_calls(_strip_mistral_closed_calls(text), enabled_tool_names), - final = True, - ) - cleaned = _strip_function_xml_calls(cleaned, final = True) - return _TOOL_XML_RE.sub("", cleaned) +def _display_tool_name_gate(active_tools): + """Active tool NAMES for gating the rehearsal display strip, or None when no tools + are enabled. ``None`` keeps the legacy strip-all behavior, mirroring the loop gate: + a bare ``NAME[ARGS]`` is a call only when NAME is active; without a tool list every + identifier stays ambiguous, so strip.""" + names = { + (t.get("function") or {}).get("name") + for t in (active_tools or []) + if isinstance(t, dict) and isinstance(t.get("function"), dict) + } + names.discard(None) + return names or None def _strip_tool_xml_for_display( @@ -1740,12 +1816,56 @@ def _strip_tool_xml_for_display( auto_heal_tool_calls: bool, enabled_tool_names: Optional[set] = None, ) -> str: - """Route-level leak cleanup (Auto-Heal only). Delegates to ``_strip_tool_xml`` - so the Mistral balanced-brace pass runs too (``_TOOL_XML_RE`` alone has no - ``[TOOL_CALLS]`` arm). ``enabled_tool_names`` gates the Gemma strip.""" + """Apply route-level XML leak cleanup only when Auto-Heal is enabled. + + Mirrors the parser-side segment scan: balanced strips first (Mistral, gated Gemma + wrapper-less, GLM real-close, guarded function-XML close at each call's REAL terminator + so literal markup inside a value is data), then the ``_TOOL_XML_RE`` arms cover the + DeepSeek / Kimi / orphan forms. ```` blocks are preserved verbatim and the + ``\\Z``-anchored tail arms run only on the last segment (prose ``foo[ARGS]`` before a + block survives). ``enabled_tool_names`` (when not None) gates the ambiguous bare-rehearsal + ``NAME[ARGS]{...}`` and wrapper-less Gemma ``call:NAME{...}`` strips on the active tool + list; an inactive NAME is prose and is kept. The ``[TOOL_CALLS]`` control-token arms strip + unconditionally regardless of NAME.""" if not auto_heal_tool_calls: return text - return _strip_tool_xml(text, enabled_tool_names) + from core.tool_healing import _strip_bracket_tag_calls, strip_outside_think + + def _keep_inactive_rehearsal(m) -> str: + # Only the bare-rehearsal arm captures ``reh``; with a tool list an inactive + # NAME[ARGS]{...} is prose -- keep it. + if enabled_tool_names is not None: + name = m.groupdict().get("reh") + if name is not None and name not in enabled_tool_names: + return m.group(0) + return "" + + def _strip_segment(seg: str, is_last: bool) -> str: + # Scan strips close at each call's REAL terminator (a literal ```` or a + # nested marker quoted inside a value cannot truncate the strip); the regex arms below + # cover the attribute form and the DeepSeek / Kimi / orphan families. + seg = _strip_mistral_closed_calls(seg) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + seg = _strip_glm_calls(seg, final = is_last) + seg = _strip_function_xml_calls(seg, final = is_last) + if is_last: + return _TOOL_XML_RE.sub(_keep_inactive_rehearsal, seg) + return _TOOL_XML_CLOSED_RE.sub("", seg) + + return strip_outside_think(text, _strip_segment) + + +def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str: + # Mistral balanced-brace pre-strip (kept explicit so the regression guards see it), then + # the shared think-aware display strip -- the one raw _TOOL_XML_RE.sub lives inside + # _strip_tool_xml_for_display, so every route cleanup site shares it. ``enabled_tool_names`` + # gates the Gemma wrapper-less strip; ``None`` strips every closed call. + text = _strip_mistral_closed_calls(text) + return _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = enabled_tool_names + ) logger = get_logger(__name__) @@ -6010,14 +6130,18 @@ async def openai_chat_completions( _gguf_auto_heal_tool_calls = ( payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True ) + # Active tool names gating the bare-rehearsal strip, matching the loop gate. + _gguf_display_tool_names = _display_tool_name_gate(tools_to_use) # ── Strip stale tool-call XML from conversation history ─ for _msg in gguf_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + # Gate on enabled tool names, like the live strip, so a documented inactive + # ``foo[ARGS]{...}`` survives in the replayed prompt context. _msg["content"] = _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _gguf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(tools_to_use), + enabled_tool_names = _gguf_display_tool_names, ).strip() def gguf_generate_with_tools(): @@ -6151,7 +6275,7 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(tools_to_use), + enabled_tool_names = _gguf_display_tool_names, ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative @@ -6258,7 +6382,7 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _gguf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(tools_to_use), + enabled_tool_names = _gguf_display_tool_names, ) return full_text, usage, finish finally: @@ -6631,14 +6755,17 @@ async def openai_chat_completions( _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") _sf_features = _detect_safetensors_features(backend, _sf_tpl) - # Split prefilled-```` output into reasoning_content deltas (GGUF parity) so the UI - # renders the thinking block for safetensors AND MLX. + # GGUF parity: enable_thinking templates prefill an unclosed ; split into + # reasoning_content deltas so the UI renders the block for safetensors and MLX. _sf_parse_think = bool( _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on") ) - # Prefilled-open only for prefill styles with thinking on this request; gpt-oss excluded. + # Prefilled-open only for prefill styles with thinking on; gpt-oss uses the normal mode. _sf_reasoning_prefilled = _sf_reasoning_prefill_mode( - _sf_features, payload.enable_thinking, _sf_tpl, payload.reasoning_effort + _sf_features, + payload.enable_thinking, + _sf_tpl, + reasoning_effort = payload.reasoning_effort, ) def _new_sf_reasoning_extractor(): @@ -6722,6 +6849,8 @@ async def openai_chat_completions( _sf_auto_heal_tool_calls = ( payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True ) + # Active tool names gating the bare-rehearsal strip, matching the loop gate. + _sf_display_tool_names = _display_tool_name_gate(_sf_tools_to_use) # Strip stale tool-call XML from prior assistant turns. _sf_chat_messages = [] @@ -6733,7 +6862,7 @@ async def openai_chat_completions( "content": _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _sf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), + enabled_tool_names = _sf_display_tool_names, ).strip(), } ) @@ -6792,7 +6921,7 @@ async def openai_chat_completions( reasoning_extractor = _new_sf_reasoning_extractor() def _sf_flush_reasoning(): - # Drain the extractor at a turn boundary / stream end (GGUF parity); only visible text reaches the monitor. + # Drain the extractor at turn/stream end (mirrors GGUF); only visible text hits the monitor. fr, fv = reasoning_extractor.finish() out = [] if fr: @@ -6818,7 +6947,7 @@ async def openai_chat_completions( if event["type"] == "status": if not event["text"]: - # Iteration boundary: flush reasoning, then start a fresh extractor for the next turn. + # Iteration boundary: flush reasoning, then a fresh prefilled extractor for the next turn. for _c in _sf_flush_reasoning(): yield _c prev_text = "" @@ -6834,7 +6963,7 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": - # Flush reasoning before the tool_start line so the thinking block closes ahead of the tool card. + # Flush reasoning before tool_start so the thinking block closes ahead of the card. for _c in _sf_flush_reasoning(): yield _c prev_text = "" @@ -6847,7 +6976,7 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _sf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), + enabled_tool_names = _sf_display_tool_names, ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative @@ -6939,12 +7068,12 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _sf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), + enabled_tool_names = _sf_display_tool_names, ) return full_text content_text = await asyncio.to_thread(_drain_to_text) - # Split prefilled reasoning out of the visible answer (GGUF parity); monitor gets visible text only. + # Split prefilled out of the visible answer (GGUF parity); the monitor gets visible text only. _reasoning_text, _visible_text = _extract_responses_reasoning( content_text, parse_think_markers = _sf_parse_think, @@ -7043,7 +7172,7 @@ async def openai_chat_completions( yield _chat_role_chunk(completion_id, created, model_name) prev_text = "" - # Split prefilled into reasoning_content deltas (GGUF parity). Single turn (no per-turn reset); also serves MLX. + # Split prefilled into reasoning_content deltas (GGUF parity); single turn, serves MLX. reasoning_extractor = _new_sf_reasoning_extractor() # Run the sync generator in a thread pool to avoid blocking the # event loop. Critical for compare mode: two SSE requests arrive @@ -7150,7 +7279,7 @@ async def openai_chat_completions( for token in generate(): full_text = token - # Split prefilled reasoning from the visible answer (GGUF parity); also covers MLX. + # Split prefilled reasoning (GGUF parity); also covers MLX via the shared generate(). _reasoning_text, _visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _sf_parse_think, @@ -8000,8 +8129,8 @@ class _ResponsesReasoningExtractor: reasoning_prefilled: bool = False, ) -> None: self._buffer = "" - # ``reasoning_prefilled``: output begins INSIDE an unclosed ```` (Qwen3/GLM prefill), - # so start in reasoning to capture leading text until the first ````. Callers default False. + # reasoning_prefilled: the template inserts an unclosed , so output begins inside + # the block; start in reasoning until the first close tag. Existing callers pass False. self._in_reasoning = reasoning_prefilled # Splitting requires marker parsing; a prefilled open implies it. self._parse_think_markers = parse_think_markers or reasoning_prefilled @@ -8033,8 +8162,8 @@ class _ResponsesReasoningExtractor: self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] self._in_reasoning = False continue - # Hold back a trailing partial of EITHER marker: the close (clean chunk-boundary split) - # and a stray open (so a re-emitted ```` isn't leaked into the reasoning drawer). + # Hold back a trailing partial of either marker: the close (clean split across chunks) + # and a stray open (a re-emitted is suppressed, not leaked). keep = _responses_marker_holdback( self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN) ) @@ -9919,11 +10048,15 @@ async def anthropic_messages( else: openai_messages.insert(0, {"role": "system", "content": _nudge}) - # Strip stale tool-call XML from conversation + # Strip stale tool-call XML via the protected display helper (think rehearsal and [TOOL_CALLS] + # prose survive), gated on enabled tool names so documented inactive examples are kept. + _anthropic_history_gate = _display_tool_name_gate(openai_tools) for _msg in openai_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): - _msg["content"] = _strip_tool_xml( - _msg["content"], _gemma_strip_gate(openai_tools) + _msg["content"] = _strip_tool_xml_for_display( + _msg["content"], + auto_heal_tool_calls = True, + enabled_tool_names = _anthropic_history_gate, ).strip() def _run_tool_gen(): @@ -10023,6 +10156,10 @@ async def _anthropic_tool_stream( """Streaming response for the tool-calling path.""" _sentinel = object() + # Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final + # answer is prose and must survive in the delivered text. + _display_names = _display_tool_name_gate(openai_tools) + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens # makes blocking HTTP calls to llama-server, so run it off the event loop. # Pass the tools so tool-schema tokens are counted (the generator renders @@ -10074,9 +10211,15 @@ async def _anthropic_tool_stream( captured_finish_reason = _fr # Strip leaked tool-call XML from content events first, so a # content event that was purely tool XML doesn't count as text. + # Protected helper preserves rehearsal and balanced + # [TOOL_CALLS] trailing prose (raw _TOOL_XML_RE.sub corrupts both). if etype == "content": event = dict(event) - event["text"] = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools)) + event["text"] = _strip_tool_xml_for_display( + event["text"], + auto_heal_tool_calls = True, + enabled_tool_names = _display_names, + ) # disable_parallel_tool_use: keep only the first tool_use block, # dropping every later tool_start and its paired tool_end (robust # to empty tool-call ids — tracked by state, not id matching). @@ -10250,6 +10393,9 @@ async def _anthropic_tool_non_streaming( usage = {} prev_text = "" captured_finish_reason = None + # Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final + # answer is prose and must survive in the delivered text. + _display_names = _display_tool_name_gate(openai_tools) # Pending client tool_use; cleared by tool_end (server execution) or # trailing text. See the stop_reason mapping below. ends_on_tool_use = False @@ -10259,8 +10405,10 @@ async def _anthropic_tool_non_streaming( for event in events: etype = event.get("type", "") if etype == "content": - # Strip leaked tool-call XML - clean = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools)) + # Strip leaked tool XML (protected helper keeps think rehearsal and trailing prose). + clean = _strip_tool_xml_for_display( + event["text"], auto_heal_tool_calls = True, enabled_tool_names = _display_names + ) new = clean[len(prev_text) :] prev_text = clean if new: @@ -10730,13 +10878,16 @@ async def _anthropic_passthrough_non_streaming( text = message.get("content") or "" if text: # Keep unpromoted bytes when healing is active; legacy stripping is - # only for opted-out or no-client-tool requests. Use the full - # _strip_tool_xml pass so Mistral [TOOL_CALLS] and guarded - # function-XML leaks are cleaned too, not just _TOOL_XML_RE forms, - # with the Gemma display gate so a disabled/example call:NAME{...} - # in prose survives. + # only for opted-out or no-client-tool requests. Protected helper (not + # raw _TOOL_XML_RE.sub): preserves rehearsal and balanced + # [TOOL_CALLS] trailing prose, gated on the declared tools so an + # inactive NAME[ARGS]{...} example in the final text is kept. if not healing_active: - text = _strip_tool_xml(text, _gemma_strip_gate(openai_tools)) + text = _strip_tool_xml_for_display( + text, + auto_heal_tool_calls = True, + enabled_tool_names = _display_tool_name_gate(openai_tools), + ) text = text.strip() if text: content_blocks.append(AnthropicResponseTextBlock(text = text)) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index a6c1fcda9c..170b456eac 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -889,6 +889,24 @@ class TestAnthropicToolNonStreaming: assert tool_blocks[0]["name"] == "render_html" assert tool_blocks[0]["input"] == {"code": ""} + def test_display_strip_gates_on_declared_tools(self): + # A final answer containing NAME[ARGS]{json} is gated on the declared tools: undeclared + # ``foo`` markup is prose and survives, the declared web_search rehearsal strips. + def _run_gen(): + yield { + "type": "content", + "text": 'Try foo[ARGS]{"x": 1} but not web_search[ARGS]{"q": "hi"} here.', + } + + tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}] + response = asyncio.run( + _anthropic_tool_non_streaming(_run_gen, "msg_1", "m", openai_tools = tools) + ) + body = json.loads(response.body) + text = "".join(b["text"] for b in body["content"] if b["type"] == "text") + assert 'foo[ARGS]{"x": 1}' in text # inactive name preserved as prose + assert "web_search[ARGS]" not in text # active name stripped from display + # ===================================================================== # Pass-through emitter tests (client-side tool execution path) diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index 7b653f47aa..e3055d2127 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -160,6 +160,20 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): assert [c["function"]["name"] for c in calls] == ["python"], calls +def test_unclosed_think_literal_inside_tool_argument_does_not_hide_later_call(): + # A literal inside a completed call's arguments is argument data; both calls must parse. + text = '[TOOL_CALLS]a{"x":"literal marker"} b[ARGS]{"y":2}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["a", "b"], calls + + +def test_real_think_block_with_rehearsal_inside_still_skips_only_the_rehearsal(): + # A genuine reasoning block still hides its rehearsal while a real call after it parses. + text = 'web_search[ARGS]{"q":"draft"}real[ARGS]{"q":"go"}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["real"], calls + + def test_wrapperless_nested_object_argument_is_parsed(): # skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare. calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}") diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index dcc759a210..9e16be2160 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -2130,6 +2130,312 @@ def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch): assert "prompt_tokens_details" not in metadata[-1]["usage"] +def test_gguf_rehearsal_name_split_before_args_is_not_leaked(monkeypatch): + """Finding 6: a rehearsal call whose name (``web_search``) and ``[ARGS]{...}`` + arrive in separate content deltas must hold the bare name in the buffer until + ``[ARGS]`` flips it to a drain. Without _is_rehearsal_prefix the GGUF path + streams the tool name as visible content before the call executes.""" + + first_stream = [ + _sse({"content": "web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all("[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_initial_buffer_flush_holds_split_rehearsal_name(monkeypatch): + """The first flush out of BUFFERING (prose plus a trailing active-tool-name in + the first delta, ``[ARGS]{...}`` in the next) must apply the same trailing-name + hold the STREAMING branch uses. The first delta has spaces so it is not a + rehearsal prefix and falls to the initial flush, which previously emitted the + bare name before the call drained.""" + + first_stream = [ + _sse({"content": "I will use web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all("[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_rehearsal_name_after_prose_in_streaming_is_not_leaked(monkeypatch): + """Finding 9: the BUFFERING guard only covers a rehearsal at the turn start. + When prose has already streamed (STREAMING state) and the model then emits the + tool name and ``[ARGS]{...}`` in later deltas, the bare name must still be held, + not flushed as visible content before the call drains.""" + + first_stream = [ + _sse({"content": "Let me think. "}), + _sse({"content": "I will search "}), + _sse({"content": "web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + + +def test_gguf_plain_answer_ending_with_tool_name_word_is_preserved(monkeypatch): + """End-of-stream flush: a plain answer that ENDS on a tool-name word with no + ``[ARGS]`` following is real prose and must not be dropped by the streaming + rehearsal hold.""" + + first_stream = [ + _sse({"content": "I think "}), + _sse({"content": "you should "}), + _sse({"content": "web_search"}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "advise"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any(t.rstrip().endswith("web_search") for t in content_texts), content_texts + + +def test_gguf_long_tool_name_split_rehearsal_is_not_capped_and_executes(monkeypatch): + """Finding 11: a realistic MCP name longer than the 32-char buffer cap split as + NAME then [ARGS]{...} must still be held (a rehearsal prefix is self-bounding), + so the name does not leak and the call executes.""" + name = "mcp__github__create_pull_request" + assert len(name) >= 32, len(name) + + first_stream = [ + _sse({"content": name}), + _sse({"content": '[ARGS]{"x":1}'}), + _done(), + ] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": name}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [(name, {"x": 1})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert not any(name in t for t in content_texts), content_texts + + +def test_gguf_streaming_keeps_bare_args_before_think_block(monkeypatch): + """F4: the GGUF streaming strip must run its open-ended ``[ARGS]`` tail cleanup + only on the LAST segment. A bare ``foo[ARGS]`` (no JSON body, ``foo`` not a tool) + before a block is prose, not a truncated call, so the final visible text + must keep it verbatim instead of dropping ``foo[ARGS]`` and corrupting the + sentence.""" + + first_stream = [ + _sse({"content": "Please pass foo[ARGS] "}), + _sse({"content": "pause "}), + _sse({"content": "to the template."}), + _done(), + ] + backend = _make_backend(monkeypatch, [first_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert content_texts, events + assert content_texts[-1] == "Please pass foo[ARGS] pause to the template." + + +def test_gguf_inactive_name_args_in_prose_is_not_drained(monkeypatch): + """BUG A: an inactive-name ``foo[ARGS]{...}`` in a prose answer must not be treated + as a tool call. The BUFFERING and end-of-stream safety-net ``[ARGS]`` checks gate on + active tool names (like the safetensors loop and the mid-stream path), so ``foo`` + (``web_search`` is the only enabled tool) is neither drained/parsed into a disabled + no-op nor forced into another generation turn.""" + first_stream = [ + _sse({"content": 'foo[ARGS]{"x":1} is just syntax.'}), + _done(), + ] + backend = _make_backend(monkeypatch, [first_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + # No tool executed for the inactive name; a spurious no-op re-prompt would exhaust the + # single supplied stream and error. + assert calls == [], calls + assert not any(e.get("type") in ("tool_start", "tool_end") for e in events), events + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + # The inactive ``foo[ARGS]{...}`` is prose: the name-gated strip keeps the whole sentence. + assert any('foo[ARGS]{"x":1} is just syntax.' in t for t in content_texts), content_texts + + +def test_gguf_inactive_rehearsal_before_active_call_executes_and_keeps_prose(monkeypatch): + """BUG X (#5704): an inactive ``foo[ARGS]{...}`` before a real ``web_search[ARGS]{...}`` + in one delta must NOT swallow the real call; web_search executes while the inactive + rehearsal stays visible as prose.""" + first_stream = [ + _sse({"content": 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + backend = _make_backend(monkeypatch, [first_stream, final_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + # The real call runs; ``foo`` is not executed as a phantom disabled call. + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + # The inactive rehearsal is preserved as prose; the active one is stripped. + assert any('foo[ARGS]{"a":1}' in t for t in content_texts), content_texts + assert all("web_search[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_rehearsal_detection_recognises_spent_one_shot_with_original_tools(): + # Rehearsal detection is fed the ORIGINAL tool list, so a spent one-shot's re-emitted + # repeat is still detected (matching the strip gate) instead of blanking the turn. + from core.inference.llama_cpp import _gguf_has_genuine_tool_signal + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + + repeat = 'render_html[ARGS]{"code":"x"}' + active_only = [{"type": "function", "function": {"name": "web_search"}}] + original = active_only + [{"type": "function", "function": {"name": "render_html"}}] + assert not _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, active_only) + assert _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, original) + + +def test_gguf_rehearsal_prefix_and_tail_hold_recognise_spent_one_shot(): + # The BUFFERING prefix check and STREAMING/flush tail-holds use the ORIGINAL tool list, + # so a spent one-shot's split repeat is held rather than leaked as visible text. + from core.inference.llama_cpp import _held_rehearsal_tail_len, _is_rehearsal_prefix + + active_only = [{"type": "function", "function": {"name": "web_search"}}] + original = active_only + [{"type": "function", "function": {"name": "render_html"}}] + assert not _is_rehearsal_prefix("render_html", active_only) + assert _is_rehearsal_prefix("render_html", original) + assert _held_rehearsal_tail_len("answer render_html", active_only) == 0 + assert _held_rehearsal_tail_len("answer render_html", original) == len("render_html") + + def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): """An oversized bare-JSON call drains rather than streams, and still executes via the safety net.""" diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py index 06316e2243..83bcc5864a 100644 --- a/studio/backend/tests/test_passthrough_healing.py +++ b/studio/backend/tests/test_passthrough_healing.py @@ -280,6 +280,56 @@ class TestStreamHealer: assert [c["id"] for c in calls] == ["call_0", "call_1"] assert _events_text(events).strip() == "then" + def test_mistral_array_multiple_calls_all_promoted_in_stream(self): + # A canonical Mistral [TOOL_CALLS] array carries several calls under a + # SINGLE signal. Draining only the first call would leave the residue + # starting at ",{...}]" (no signal), so later calls in the same array + # must be promoted in the same pass, not flushed as raw text. + healer = StreamToolCallHealer({"get_weather", "get_time"}) + array = ( + '[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},' + '{"name":"get_time","arguments":{"tz":"UTC"}}]' + ) + events = healer.feed(array) + healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"] + assert [c["id"] for c in calls] == ["call_0", "call_1"] + assert _events_text(events) == "" + + def test_mistral_array_multiple_calls_promoted_char_by_char(self): + healer = StreamToolCallHealer({"get_weather", "get_time"}) + array = ( + '[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},' + '{"name":"get_time","arguments":{"tz":"UTC"}}]' + ) + events = [] + for ch in array: + events += healer.feed(ch) + events += healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"] + assert _events_text(events) == "" + + def test_mistral_array_undeclared_middle_kept_as_text_others_promoted(self): + # A mid-array element for a tool that is not declared must survive as + # text while the declared neighbours on either side still promote in + # document order. + healer = StreamToolCallHealer({"a", "c"}) + array = ( + '[TOOL_CALLS][{"name":"a","arguments":{}},' + '{"name":"b","arguments":{}},{"name":"c","arguments":{}}]' + ) + events = healer.feed(array) + healer.finalize() + assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "c"] + assert '"b"' in _events_text(events) + + def test_mistral_array_then_trailing_prose(self): + healer = StreamToolCallHealer({"a", "b"}) + array = '[TOOL_CALLS][{"name":"a","arguments":{}},{"name":"b","arguments":{}}]' + events = healer.feed(f"{array} all done") + healer.finalize() + assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "b"] + assert "all done" in _events_text(events) + def test_incomplete_call_healed_at_finalize(self): healer = StreamToolCallHealer({"Bash"}) events = healer.feed('{"name":"Bash","arguments":{"cmd":"ls"}}') @@ -1356,3 +1406,42 @@ class TestOpenaiStreamingRoute: assert chunks[0] == line + "\n\n" # byte-for-byte relay asyncio.run(_run()) + + +class TestHealerSignalAlignment: + """The passthrough healer buffers only formats its parser can promote. + The loops' bare [ARGS] rehearsal signal is gated on active tool names + there; ungated in the healer it would stall legitimate prose until + finalization without ever producing a promotable call.""" + + def test_heal_signals_are_promotable_formats_only(self): + from core.inference.passthrough_healing import _HEAL_SIGNALS + assert set(_HEAL_SIGNALS) == { + "", + "<|tool_call>", + ", so -# generation begins inside the think block and emits only the closing ; the extractor starts in reasoning. +# reasoning_prefilled: enable_thinking templates prefill an unclosed , so +# generation begins inside the block; the extractor must start in reasoning. class TestReasoningPrefilledExtractor: def test_prefilled_single_feed_splits_lone_close(self): # T1: reasoning...answer with a prefilled (unseen) open tag. @@ -2077,9 +2077,7 @@ class TestReasoningPrefilledExtractor: assert visible == "hi" def test_not_prefilled_lone_close_preserves_current_behavior(self): - # T9: GGUF-parity guard -- WITHOUT prefilled, a lone keeps the - # pre-fix behavior (reasoning stays visible, tag dropped). Ensures GGUF and - # every existing caller are byte-identical. + # T9: without prefilled, a lone close tag keeps the pre-fix behavior (parity guard). reasoning, visible = _extract_responses_reasoning( "reasoningans", parse_think_markers = True, @@ -2099,8 +2097,7 @@ class TestReasoningPrefilledExtractor: assert visible == "v" def test_prefilled_ignored_when_markers_not_parsed(self): - # T11: a non-reasoning model (parse_think_markers False) still passes text - # straight through even if reasoning_prefilled were mistakenly set False. + # T11: a non-reasoning model passes text through even with reasoning_prefilled False. reasoning, visible = _extract_responses_reasoning( "just an answer", parse_think_markers = False, diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 3701a00dd2..9fd1535f22 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -183,7 +183,9 @@ def test_detect_safetensors_features_llama3_template_keeps_tools_on(): def test_detect_safetensors_features_mistral_template_keeps_tools_on(): - """Mistral emits [TOOL_CALLS]; parser now supports it.""" + """Mistral emits [TOOL_CALLS]name{json}, which the safetensors loop now parses + (the shared bracket-tag parser). The gate must no longer suppress it, or the + PR's Mistral tool support is unreachable through normal capability detection.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3") @@ -706,13 +708,28 @@ def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json(): assert flags["supports_tools"] is True -# _sf_reasoning_prefill_mode gates the prefilled- extractor so safetensors/MLX reach -# GGUF reasoning-block parity for enable_thinking models. +# _sf_reasoning_prefill_mode gates the prefilled- extractor (GGUF reasoning parity). class TestSafetensorsReasoningPrefillGate: # A minimal Qwen3-style template with the standard / markers. _QWEN_TPL = "{% if enable_thinking %}{% endif %}......" # gemma-style bespoke reasoning channel -- no standard markers. _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" + # always-on template whose GENERATION PROMPT opens an unclosed (DeepSeek-R1 / QwQ / + # Qwen3-Thinking shape): the model emits only the closing , so prefill. + _ALWAYS_ON_OPEN_TPL = ( + "{% for m in messages %}{{ m['content'] }}{% endfor %}" + "{% if add_generation_prompt %}<|assistant|>\n{% endif %}" + ) + # always-on template that renders PAST assistant ... history but leaves the + # generation prompt open with no (Kimi-K2-Thinking shape): the model self-emits its + # own block, so prefill mode would blank a normal answer. + _ALWAYS_ON_HISTORY_TPL = ( + "{% for m in messages %}" + "{% if m['role'] == 'assistant' %}{{ m.get('reasoning_content', '') }}" + "{{ m['content'] }}{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}<|im_assistant|>assistant<|im_middle|>{% endif %}" + ) def _features(self, **over): base = { @@ -756,11 +773,19 @@ class TestSafetensorsReasoningPrefillGate: feats = self._features(supports_reasoning = False, reasoning_style = None) assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False - def test_g7_reasoning_always_on(self): - # G7: hardcoded- template -> prefilled regardless of the flag. + def test_g7_reasoning_always_on_prompt_opens_think(self): + # G7: always-on template whose generation prompt opens -> prefilled regardless of the flag. from routes.inference import _sf_reasoning_prefill_mode feats = self._features(reasoning_always_on = True) - assert _sf_reasoning_prefill_mode(feats, False, self._QWEN_TPL) is True + assert _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True + + def test_g7b_reasoning_always_on_history_only_not_prefilled(self): + # G7b (#5704): always-on classification from rendered assistant HISTORY + # (Kimi-K2-Thinking) whose generation prompt opens no . Prefill mode would capture a + # normal answer entirely as reasoning_content and blank the visible answer, so it must be off. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_always_on = True) + assert _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) is False def test_g8_gemma_bespoke_channel_excluded(self): # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py index 4a5423fa87..4e708139b7 100644 --- a/studio/backend/tests/test_safetensors_reasoning_stream.py +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -3,9 +3,11 @@ """Safetensors/MLX reasoning-block parity with GGUF. -enable_thinking templates prefill an unclosed ````, so the stream must split the leading -text into ``reasoning_content`` deltas (per turn, monitor gets visible text only). Replays a copy -of ``sf_tool_stream``'s reasoning loop from routes/inference.py against synthetic events. +enable_thinking templates (Qwen3/GLM) prefill an unclosed ```` so the model +emits only the closing ```` then the answer; the safetensors stream must +split the leading text into ``reasoning_content`` deltas (plain stream and tool +loop), resetting per turn and appending only visible text to the monitor. Replays a +copy of ``sf_tool_stream``'s reasoning loop against synthetic events. """ from __future__ import annotations @@ -24,8 +26,40 @@ from routes.inference import ( ) +_THINK_TPL = "........." +_ETHINK = {"reasoning_style": "enable_thinking", "supports_reasoning": True} +_ETHINK_EFFORT = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} + + +def test_prefill_mode_on_for_enable_thinking_default(): + assert _sf_reasoning_prefill_mode(_ETHINK, None, _THINK_TPL) is True + + +def test_prefill_mode_off_when_thinking_disabled(): + assert _sf_reasoning_prefill_mode(_ETHINK, False, _THINK_TPL) is False + + +def test_prefill_mode_off_for_reasoning_effort_none(): + # enable_thinking_effort turns thinking off via reasoning_effort="none"; prefilled mode + # would capture the whole answer as reasoning_content. + assert ( + _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none") + is False + ) + assert ( + _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high") + is True + ) + + +def test_prefill_mode_off_without_think_markers(): + assert _sf_reasoning_prefill_mode(_ETHINK, None, "no markers here") is False + + def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict: - """Mirror sf_tool_stream's reasoning loop: diff cumulative snapshots, reset (flushing) on turn end.""" + """Mirror sf_tool_stream's reasoning loop: diff each cumulative ``content`` + snapshot, feed the delta through the extractor, and reset (flushing first) on + ``tool_start`` / empty ``status`` so each turn splits independently.""" prev_text = "" extractor = _ResponsesReasoningExtractor( parse_think_markers = True, reasoning_prefilled = prefilled @@ -151,9 +185,6 @@ def test_s5_thinking_off_no_reasoning_deltas(): assert out["monitor"] == "Just the plain answer, no thinking." -_THINK_TPL = "...{% if enable_thinking %}{% endif %}......" - - def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): # GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and # enable_thinking omitted) disables thinking exactly like enable_thinking=False, so diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 38b30fe8f6..f826f3cddf 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -115,6 +115,22 @@ class TestParser: assert result[0]["function"]["name"] == "python" assert "print('hi')" in result[0]["function"]["arguments"] + def test_xml_param_preserves_leading_indentation(self): + import json + + # Only the wrapping newline is trimmed; code-argument indentation survives. + text = ( + "\n" + " indented = 1\n" + " more\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "code": " indented = 1\n more" + } + def test_xml_unclosed(self): # Closing tags omitted; parser must still extract the value. text = "ls -la" @@ -183,6 +199,8 @@ class TestParser: assert has_tool_signal("blah x") assert has_tool_signal("blah <|tool_call>call:terminal") assert has_tool_signal("hi ...") + assert has_tool_signal("ok [TOOL_CALLS]web_search{...") + assert has_tool_signal("fine python[ARGS]{...") assert not has_tool_signal("hello world") def test_render_html_start_detector_uses_first_tool(self): @@ -197,6 +215,44 @@ class TestParser: '{"name":"python","arguments":{"code":""}}' ) + def test_render_html_start_detector_covers_mistral_and_rehearsal_forms(self): + # The provisional render-html card must fire for bracket-tag forms too, not only XML. + assert _detect_render_html_tool_start('[TOOL_CALLS]render_html{"code":""}') + assert _detect_render_html_tool_start('[TOOL_CALLS]render_html[ARGS]{"code":"x"}') + assert _detect_render_html_tool_start( + '[TOOL_CALLS] [{"name":"render_html","arguments":{}}]' + ) + assert _detect_render_html_tool_start('render_html[ARGS]{"code":""}') + # A different first tool (or a prose mention with no JSON body) must not fire. + assert not _detect_render_html_tool_start('[TOOL_CALLS]web_search{"q":"x"}') + assert not _detect_render_html_tool_start('web_search[ARGS]{"q":"x"}') + assert not _detect_render_html_tool_start('python[ARGS]{"code":"render_html[ARGS]{}"}') + assert not _detect_render_html_tool_start("use render_html[ARGS] to render") + + def test_render_html_start_detector_skips_think_block_rehearsal(self): + # A render_html rehearsed inside think must not fire the card; the outside-think call decides. + assert not _detect_render_html_tool_start( + 'draft render_html[ARGS]{"code":"x"}python[ARGS]{"code":"print(1)"}' + ) + assert not _detect_render_html_tool_start( + '[THINK]render_html[ARGS]{"code":"x"}[/THINK]web_search[ARGS]{"q":"y"}' + ) + # A real render_html AFTER a rehearsed non-render_html inside think still fires. + assert _detect_render_html_tool_start( + 'web_search[ARGS]{"q":"x"}render_html[ARGS]{"code":""}' + ) + # A render_html rehearsed inside think with no real call after does not fire. + assert not _detect_render_html_tool_start('render_html[ARGS]{"code":"x"}') + + def test_render_html_start_detector_reads_top_level_array_name(self): + # Array form: the name is the object's top-level ``"name"``, not an argument key. + assert not _detect_render_html_tool_start( + '[TOOL_CALLS] [{"arguments":{"name":"render_html"},"name":"python"}]' + ) + assert _detect_render_html_tool_start( + '[TOOL_CALLS] [{"arguments":{"name":"python"},"name":"render_html"}]' + ) + def test_strip_markup_closed(self): text = "before {} after" assert strip_tool_markup(text) == "before after" @@ -237,6 +293,376 @@ class TestParser: == "before " ) + # Mistral [TOOL_CALLS] bracket-tag. + + def test_mistral_bracket_basic(self): + # Devstral / Mistral-Small fallback when bypassing native FC. + text = '[TOOL_CALLS]web_search{"query":"weather"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + assert isinstance(result[0]["function"]["arguments"], str) + assert "weather" in result[0]["function"]["arguments"] + + def test_rehearsal_inside_unclosed_think_is_ignored(self): + """Rehearsal-shaped markup inside an unclosed block must + not be executed as a real tool call. Mid-stream the + tag has not arrived yet, so the strip regex has to accept + end-of-string as a terminator. Regression for the Gemini + high-severity flag on this PR.""" + text = ( + "I should call web_search[ARGS]" '{"query":"weather"} next to find the answer.' + ) + result = parse_tool_calls_from_text(text) + # Inside an unclosed think block no calls are yielded. + assert result == [] + + def test_rehearsal_inside_unclosed_bracket_think_is_ignored(self): + text = "[THINK]planning to use python[ARGS]" '{"code":"print(1)"} but not yet.' + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_rehearsal_after_closed_think_still_parsed(self): + text = "planning" 'python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_rehearsal_inside_prefilled_think_is_ignored(self): + """Reasoning models (Qwen3.5 enable_thinking) open in the PROMPT, + so generated content starts inside the thought and carries only a closing + . A call rehearsed in that leading thought must be skipped, while a + real call after the close still fires.""" + text = 'planning web_search[ARGS]{"query":"draft"}python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_literal_close_think_in_leading_argument_not_prefill(self): + """A literal inside a real leading call's arguments must not be + read as a prefilled-reasoning close (which would skip the call).""" + text = 'web_search[ARGS]{"query":"what is "}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_stray_close_after_real_call_not_treated_as_prefill(self): + """A real leading call followed by a stray and no further call is + a normal answer, not prefilled reasoning; the call must still fire (the + virtual span only applies when a real call follows the close).""" + text = 'Now web_search[ARGS]{"query":"x"} answer' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_mistral_bracket_with_whitespace(self): + # Optional whitespace (incl. newlines) between the name and the opening brace. + text = '[TOOL_CALLS]python \n {"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print(1)" in result[0]["function"]["arguments"] + + def test_mistral_bracket_nested_json(self): + # Brace-balance scan handles nested objects and braces inside string literals. + text = "[TOOL_CALLS]web_search" '{"query":"a {nested} brace","opts":{"limit":5}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + import json as _json + + args = _json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "a {nested} brace" + assert args["opts"] == {"limit": 5} + + def test_mistral_bracket_with_prose(self): + # Bracket-tag surrounded by prose is still recognised. + text = ( + "Sure, I will look that up.\n" + '[TOOL_CALLS]web_search{"query":"weather"}\n' + "Calling now." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_mistral_bracket_bad_json_dropped(self): + text = "[TOOL_CALLS]web_search{not valid}" + result = parse_tool_calls_from_text(text) + # No usable tool call; callers fall back to text. + assert result == [] + + def test_mistral_bracket_object_with_array_value(self): + # Args must be a JSON object; a dict wrapping an array value is accepted. + text = '[TOOL_CALLS]web_search{"opts":[1,2,3]}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # Rehearsal syntax name[ARGS]{json}. + + def test_rehearsal_basic(self): + text = 'python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print(1)" in result[0]["function"]["arguments"] + + def test_rehearsal_with_prose(self): + text = "I should call the python tool. Like this: " 'python[ARGS]{"code":"x = 1"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_rehearsal_bad_json_dropped(self): + text = "python[ARGS]{not valid json}" + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_mistral_bracket_hyphenated_mcp_name(self): + # Dashed MCP names must be captured whole, not truncated at the first dash. + text = '[TOOL_CALLS]mcp__srv__list-issues{"q":"x"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__list-issues" + + def test_rehearsal_hyphenated_mcp_name(self): + text = 'mcp__srv__list-issues[ARGS]{"q":"x"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__list-issues" + + def test_streaming_strip_removes_partial_bracket_marker(self): + # A bracket tag streamed before its opening brace must strip on the final pass, not leak. + assert strip_tool_markup("answer [TOOL_CALLS]web_search", final = True) == "answer" + assert strip_tool_markup("text python[ARGS]", final = True) == "text" + # Non-final must keep the in-progress tag buffered (not yet stripped). + partial = "answer [TOOL_CALLS]web_search" + assert strip_tool_markup(partial, final = False) == partial + + def test_strip_removes_two_level_nested_bracket_call_keeps_prose(self): + # Two-level-nested args must be removed whole; the balanced scan handles any depth. + text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after' + assert strip_tool_markup(text, final = False) == "before after" + assert strip_tool_markup(text, final = True) == "before after" + + def test_strip_removes_call_with_literal_think_in_argument(self): + # A literal think block inside arguments strips with the call, not as a reasoning block. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + assert strip_tool_markup(text, final = True) == "" + + def test_strip_preserves_real_think_but_strips_call_with_literal_think(self): + text = ( + "planning ok " + '{"name":"w","arguments":{"t":"x"}} done' + ) + out = strip_tool_markup(text, final = True) + assert "planning" in out + assert "" not in out and '"name"' not in out + assert "ok" in out and "done" in out + + def test_prose_mentioning_args_marker_is_not_truncated(self): + # ``foo[ARGS] to the template`` is prose; the catch-all must not delete the sentence. + text = "Please pass foo[ARGS] to the template and continue reading." + assert strip_tool_markup(text, final = True) == text + + def test_streaming_strip_handles_mistral_v11_call_id_args(self): + # The streaming strip uses the regex patterns directly, so they must cover the v11 + # [CALL_ID]/[ARGS] metadata (aligned with the parser). + raw = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' + out = strip_tool_markup_streaming(raw) + assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert "before" in out and "after" in out + + # pre-strip. + + def test_think_block_stripped_before_xml(self): + # The think block is stripped before matching so the post-thinking call is recognised. + text = ( + "I will use web_search to find the weather." + '{"name":"web_search","arguments":{"query":"sf"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_think_block_stripped_before_bracket_tag(self): + text = ( + "Let me search for that.\n" '[TOOL_CALLS]web_search{"query":"weather"}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_uppercase_think_tag_stripped(self): + # Some templates use [THINK]...[/THINK] instead of . + text = "[THINK]planning my next call[/THINK]" '[TOOL_CALLS]python{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_think_block_hides_inner_tool_call(self): + # A call mentioned inside think is a rehearsal; the wrapper strip removes the inner markup. + text = ( + "I might call " + '{"name":"web_search","arguments":{}} ' + "but I am not sure\n" + "Let me just answer directly." + ) + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_think_literal_inside_real_tool_argument_is_preserved(self): + # A real call whose argument contains a literal think tag must not be corrupted. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["text"] == ( + "compare and tags" + ) + + def test_bracket_tag_argument_with_think_literal_is_preserved(self): + text = '[TOOL_CALLS]search{"q":"explain [THINK] blocks"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["q"] == "explain [THINK] blocks" + + def test_real_call_after_think_with_rehearsal_inside(self): + # A rehearsal inside is skipped, but the real call after the close tag parses. + text = 'plan: search[ARGS]{"q":"x"}search[ARGS]{"q":"real"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["q"] == "real" + + # XML takes precedence over bracket-tag. + + def test_xml_wins_over_bracket(self): + # When a model emits both forms in one message, the XML form is canonical and wins. + text = ( + '{"name":"primary","arguments":{}}' + '[TOOL_CALLS]secondary{"k":"v"}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "primary" + + # Strip patterns include bracket-tag and rehearsal. + + def test_strip_bracket_tag_closed(self): + text = 'before [TOOL_CALLS]web_search{"q":"hi"} after' + assert "[TOOL_CALLS]" not in strip_tool_markup(text) + assert "before" in strip_tool_markup(text) + assert "after" in strip_tool_markup(text) + + def test_strip_rehearsal_closed(self): + text = 'prose python[ARGS]{"code":"x"} more prose' + cleaned = strip_tool_markup(text) + assert "[ARGS]" not in cleaned + assert "prose" in cleaned + assert "more prose" in cleaned + + def test_strip_bracket_tag_unclosed_final(self): + text = 'before [TOOL_CALLS]web_search{"q":"part' + # Final-mode strip drops the trailing unclosed run. + cleaned = strip_tool_markup(text, final = True) + assert "TOOL_CALLS" not in cleaned + assert cleaned == "before" + + # Canonical Mistral array, v11 [CALL_ID], unified multi-call (PR review fixes). + + def test_mistral_canonical_array_is_parsed(self): + # Canonical multi-call array: every call must parse (was dropped then deleted to EOS). + text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}},{"name":"b","arguments":{"y":2}}]' + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["a", "b"] + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + assert json.loads(result[1]["function"]["arguments"]) == {"y": 2} + + def test_mistral_array_string_arguments_are_decoded(self): + # OpenAI-spec arguments arrive as a JSON string; decode to an object. + text = '[TOOL_CALLS] [{"name":"a","arguments":"{\\"x\\":1}"}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + + def test_mistral_array_scalar_string_argument_not_double_encoded(self): + # A bare scalar string argument in the Mistral array form must be kept + # raw, exactly like the path, so the downstream argument + # healer wraps ``weather`` into the single-string tool's key -- not + # ``"weather"`` with literal quotes from a redundant json.dumps. + array = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"web_search","arguments":"weather"}]' + ) + xml = parse_tool_calls_from_text( + '{"name":"web_search","arguments":"weather"}' + ) + assert array[0]["function"]["arguments"] == xml[0]["function"]["arguments"] == "weather" + healed = _coerce_arguments( + array[0]["function"]["arguments"], heal = True, tool_name = "web_search" + ) + assert healed == {"query": "weather"} + + def test_mistral_array_strip_keeps_trailing_prose(self): + # The array form must be removed whole, not deleted to end-of-string. + text = 'answer [TOOL_CALLS] [{"name":"a","arguments":{}}] tail' + assert strip_tool_markup(text, final = True) == "answer tail" + + def test_mistral_and_rehearsal_in_one_message_both_parse(self): + # A Mistral call and a rehearsal call together: both must parse. + text = '[TOOL_CALLS]a{"x":1} then b[ARGS]{"y":2}' + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["a", "b"] + + def test_mistral_v11_call_id_is_not_the_function_name(self): + # v11 shape: the function name is ``name``, never the opaque call-id token. + result = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[CALL_ID]abc123[ARGS]{"q":"x"}') + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"q": "x"} + # v11 without a call-id parses the same name. + r2 = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[ARGS]{"q":"y"}') + assert r2[0]["function"]["name"] == "get_weather" + + def test_strip_preserves_rehearsal_inside_think(self): + # A rehearsal inside is reasoning; strip keeps it verbatim. + text = 'plan: search[ARGS]{"q":"x"} A' + out = strip_tool_markup(text, final = True) + assert out == text + assert "search[ARGS]" in out + + def test_streaming_strip_preserves_rehearsal_inside_think(self): + # The streaming strip must also preserve a think rehearsal: a mid-stream strip shrinks + # then regrows the cumulative text (corrupts append-by-length consumers). Matches GGUF. + text = 'plan: search[ARGS]{"q":"x"} A' + assert strip_tool_markup_streaming(text) == text + assert strip_tool_markup_streaming(text, tool_protocol_active = True) == text + # An unclosed block during streaming is preserved too (the parser keeps it). + partial = 'plan: search[ARGS]{"q":"x"}' + assert strip_tool_markup_streaming(partial, tool_protocol_active = True) == partial + + def test_streaming_strip_still_removes_real_call_outside_think(self): + # The think guard must not stop the streaming strip removing a call outside the block. + text = 'reason web_search[ARGS]{"q":"x"}' + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert "web_search[ARGS]" not in out + assert "reason" in out + + def test_strip_bracket_calls_is_linear(self): + # Many complete bracket calls must strip in ~linear time (was O(n^2) per match). + import time + + text = '[TOOL_CALLS]f{"a":1}' * 4000 # ~80KB, 4000 complete calls + t0 = time.perf_counter() + out = strip_tool_markup(text, final = True) + elapsed = time.perf_counter() - t0 + assert "[TOOL_CALLS]" not in out + assert elapsed < 1.0, f"strip took {elapsed * 1000:.0f}ms on 4000 bracket calls" + def test_streaming_strip_handles_nested_mistral_json(self): # The non-greedy [TOOL_CALLS]name{...} pattern truncates nested JSON at the first }; the # balanced helper must remove the whole call so no trailing brace leaks to the streaming ... @@ -1390,6 +1816,254 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success(): assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) +def test_spent_one_shot_rehearsal_repeat_is_detected_not_blank_continuation(): + # A spent one-shot (render_html) stays in the ORIGINAL tool list; detection is gated on + # that list (matching the strip gate) so a re-emitted repeat is drained and routed to the + # repeat no-op instead of stripped into a blank continuation. + exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) + turns = iter( + [ + [ + '{"name":"render_html","arguments":{"code":"one"}}' + ], + ['render_html[ARGS]{"code":"two"}'], # spent one-shot rehearsal + ["The chart is above."], + ] + ) + + def gen(_messages, *, active_tools = None): + try: + chunks = next(turns) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = gen, + messages = [{"role": "user", "content": "make a chart"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "web_search"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 5, + ) + ) + contents = [e["text"] for e in events if e["type"] == "content"] + # render_html ran exactly once; the repeat was a no-op, not a second execution. + assert exec_fn.calls == [("render_html", {"code": "one"})], exec_fn.calls + # The loop continued past the repeat to the real answer (not a blank continuation). + assert any("The chart is above." in t for t in contents), contents + # The raw rehearsal markup never leaked as visible content. + assert not any("render_html[ARGS]" in t for t in contents), contents + + +def test_rehearsal_call_name_is_not_streamed_before_args(): + # A rehearsal whose name and [ARGS] arrive together must drain, not stream the bare name. + loop, exec_fn = _make_loop( + turns = [['web_search[ARGS]{"query":"cats"}'], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_rehearsal_call_name_split_before_args_is_not_streamed(): + # Finding 5: name and [ARGS] in separate chunks -- the bare name is held until [ARGS] arrives. + loop, exec_fn = _make_loop( + turns = [["web_search", '[ARGS]{"query":"cats"}'], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_plain_word_matching_no_tool_still_streams(): + # The prefix guard must not swallow prose: a non-tool bare word streams. + loop, _exec = _make_loop( + turns = [["weather", " is nice today."]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "weather is nice today." in contents, contents + + +def test_rehearsal_name_after_prose_in_streaming_is_not_streamed(): + # After prose has streamed (STREAMING state), a split rehearsal name must still be held. + loop, exec_fn = _make_loop( + turns = [ + # _make_loop accumulates these deltas into cumulative snapshots. + ["Let me think. ", "I will search ", "web_search", '[ARGS]{"query":"cats"}'], + ["Found."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_rehearsal_name_after_prose_same_chunk_in_streaming_is_not_streamed(): + # Prose then ``web_search[ARGS]{...}`` in one chunk: the boundary is pulled back over the name. + loop, exec_fn = _make_loop( + turns = [ + ["Sure. ", 'now web_search[ARGS]{"query":"cats"}'], + ["Found."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_initial_buffer_flush_holds_split_rehearsal_name(): + # First flush out of BUFFERING applies the same trailing-name hold as STREAMING. + loop, exec_fn = _make_loop( + turns = [["I will use python", '[ARGS]{"code":"print(1)"}'], ["done"]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("python", {"code": "print(1)"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("python" in t for t in contents), contents + + +def test_think_rehearsal_streams_monotonically_and_keeps_reasoning(): + # A think rehearsal streams the same text the final strip keeps: cumulative content is + # monotonically non-decreasing and ends with the markup intact. + loop, exec_fn = _make_loop( + turns = [["plan ", 'search[ARGS]{"q":"x"}', " visible"]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert exec_fn.calls == [], exec_fn.calls + assert all(len(b) >= len(a) for a, b in zip(contents, contents[1:])), contents + final = contents[-1] if contents else "" + assert 'search[ARGS]{"q":"x"}' in final, contents + assert "visible" in final, contents + + +def test_plain_answer_ending_with_tool_name_word_is_preserved(): + # End-of-stream flush: a plain answer ending on a tool-name word is prose, not dropped. + loop, exec_fn = _make_loop( + turns = [["I think ", "you should ", "web_search"]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert any(t.rstrip().endswith("web_search") for t in contents), contents + + +def test_long_tool_name_split_rehearsal_is_not_capped_and_executes(): + # Finding 10/11: an MCP name longer than the buffer cap, split before [ARGS], is still + # held (self-bounding prefix); no leak and the call executes. + from core.inference.safetensors_agentic import _MAX_BUFFER_CHARS + + name = "mcp__github__create_pull_request" + assert len(name) >= _MAX_BUFFER_CHARS, len(name) + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([[name, name + '[ARGS]{"x":1}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": name}}], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [(name, {"x": 1})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any(name in t for t in contents), contents + + +def test_unrestricted_mode_split_rehearsal_name_is_not_streamed(): + # Finding 6: unrestricted mode treats any bare identifier as a possible rehearsal NAME. + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([["web_search", 'web_search[ARGS]{"q":"x"}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [], # unrestricted + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_unrestricted_mode_split_after_bracket_is_not_streamed(): + # Unrestricted mode: a chunk split right after ``NAME[`` is still held (parity with the + # restricted-mode startswith hold). + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([["web_search[", 'web_search[ARGS]{"q":"x"}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [], # unrestricted + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search[" in t for t in contents), contents + + +def test_unrestricted_mode_plain_prose_still_streams(): + # The unrestricted hold releases a held identifier once the rest of the sentence follows. + def st(_messages, active_tools = None): + for snap in ("Hello", "Hello there friend."): + yield snap + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = FakeExecuteTool([]), + max_tool_iterations = 1, + ) + ) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Hello there friend." in contents, contents + + def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): # A late call caught by the safety net: an unclosed ```` heals only with Auto-Heal on; # off, the safety net must not pass ``allow_incomplete=True`` and execute a truncated call. @@ -1977,6 +2651,42 @@ class TestLoopBasic: assert tool_starts[0]["tool_name"] == "python" assert exec_fn.calls == [("python", {"code": "print('')"})] + def test_render_html_rehearsed_in_think_block_emits_no_provisional_start(self): + # BUG B: a render_html rehearsed inside think before a real python call must not emit a + # provisional render_html card; only the outside-think call fires. + exec_fn = FakeExecuteTool(["ok"]) + turn_iter = iter( + [ + [ + 'draft render_html[ARGS]{"code":"x"}', + 'python[ARGS]{"code":"print(1)"}', + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "run code"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + assert [e["tool_name"] for e in tool_starts] == ["python"], tool_starts + assert exec_fn.calls == [("python", {"code": "print(1)"})] + def test_render_html_success_blocks_second_canvas_call(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -3653,6 +4363,114 @@ if __name__ == "__main__": pytest.main([__file__, "-v"]) +def test_streaming_strip_keeps_bare_args_before_think_block(): + # F3: a bare ``foo[ARGS]`` before a think block is prose; EOS-anchored tail arms run only + # on the last segment. + text = "Please pass foo[ARGS] pause to the template." + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert out == text + + +def test_streaming_strip_still_removes_complete_call_before_think_block(): + # A complete bracket call before a think block still strips in the non-last segment. + text = 'go web_search[ARGS]{"q":"x"} z done' + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert "web_search[ARGS]" not in out + assert "z" in out + assert "go" in out and "done" in out + + +def test_prose_args_marker_before_real_call_does_not_drain_the_prose(): + # F5: an inactive ``foo[ARGS]`` in prose is not a call boundary; the prose streams in + # full and the later real call still executes. + loop, exec_fn = _make_loop( + turns = [ + ["Intro ", "foo[ARGS] syntax. ", 'web_search[ARGS]{"query":"cats"}'], + ["Cats are great."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + # The prose between the bogus marker and the real call must survive. + assert any("foo[ARGS] syntax." in t for t in contents), contents + # The real call markup is never shown as content. + assert not any("web_search[ARGS]" in t for t in contents), contents + + +def test_inactive_name_args_with_body_is_not_parsed_into_disabled_noop(): + # BUG A: a prose answer with an inactive ``foo[ARGS]{...}`` is not drained into a + # disabled no-op extra turn; the [ARGS] checks are name-gated. + turns = [['foo[ARGS]{"x":1} is just syntax.']] + turn_calls: list[int] = [] + + def _gen(_messages): + turn_calls.append(1) + chunks = turns[len(turn_calls) - 1] if len(turn_calls) <= len(turns) else [] + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "explain"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + assert not any(e["type"] in ("tool_start", "tool_end") for e in events), events + # Exactly one generation turn -- no disabled ``foo`` no-op re-prompt. + assert len(turn_calls) == 1, turn_calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert any("is just syntax." in t for t in contents), contents + + +class TestEnabledToolNameGate: + """The safetensors loop passes the active tool names into parse/strip so the + ambiguous bare-rehearsal ``NAME[ARGS]{json}`` is treated as a call only when NAME + is an active tool (#5704). Without the gate an inactive ``foo[ARGS]{...}`` in prose + was parsed into a disabled no-op call and stripped from the visible text.""" + + def _names(self, calls): + return [c["function"]["name"] for c in calls] + + def test_parse_inactive_rehearsal_does_not_swallow_active_call(self): + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_parse_inactive_rehearsal_alone_is_prose(self): + assert ( + parse_tool_calls_from_text('foo[ARGS]{"a":1}', enabled_tool_names = {"web_search"}) == [] + ) + + def test_streaming_strip_keeps_inactive_rehearsal(self): + raw = 'answer foo[ARGS]{"x":1} tail' + assert strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) == raw + + def test_streaming_strip_removes_active_rehearsal(self): + raw = 'answer web_search[ARGS]{"q":1} tail' + out = strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) + assert "web_search[ARGS]" not in out + assert out == "answer tail" + + def test_final_strip_keeps_inactive_rehearsal(self): + text = 'foo[ARGS]{"x":1} is just syntax.' + assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + + def test_gate_none_preserves_legacy_strip_and_parse(self): + text = 'foo[ARGS]{"x":1} tail' + assert self._names(parse_tool_calls_from_text(text)) == ["foo"] + assert strip_tool_markup_streaming(text) == " tail" + + def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): # F3: with Auto-Heal OFF, a truncated ENABLED-name bare-JSON fragment that did # not parse must stay visible (disabled-Auto-Heal contract: malformed markup is diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 7f47140b8d..c6da1e90e7 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -266,6 +266,133 @@ class TestHealingPathUnaffected: assert healed[span[0] : span[1]] == "dogs" +class TestEnabledToolNameGate: + """``enabled_tool_names`` disambiguates the ambiguous bare-rehearsal + ``NAME[ARGS]{json}`` form (#5704): NAME is a call only when it is an active tool, + otherwise it is prose. ``None`` (the default) keeps the legacy unrestricted parse + so existing callers are unaffected.""" + + def _names(self, calls): + return [c["function"]["name"] for c in calls] + + def test_inactive_rehearsal_before_active_call_does_not_swallow_it(self): + # P1: an inactive ``foo[ARGS]{...}`` before a real call must not consume the real call. + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_inactive_rehearsal_alone_is_not_a_call(self): + text = 'foo[ARGS]{"a":1}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_active_rehearsal_is_still_parsed(self): + text = 'web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + + def test_unrestricted_gate_none_preserves_legacy_behavior(self): + # Without a gate every ``NAME[ARGS]{...}`` is parsed, as before the gate landed. + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + assert self._names(parse_tool_calls_from_text(text)) == ["foo", "web_search"] + assert self._names(parse_tool_calls_from_text(text, enabled_tool_names = None)) == [ + "foo", + "web_search", + ] + + +class TestBracketCallSpans: + """with_spans tiling for Mistral bracket calls: promoted markup strips + exactly once, filtered calls' bytes stay visible, closers strip too.""" + + def test_mixed_array_filtered_first_keeps_its_bytes_only(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"bad","arguments":{"x":1}},' + '{"name":"lookup","arguments":{"q":"cats"}}]' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + kinds = [k for k, _v in events] + assert kinds == ["text", "tool_call"] + text = events[0][1] + assert '"bad"' in text + # The promoted call's markup must not survive in the text event. + assert '"lookup"' not in text + + def test_mixed_array_filtered_second_stays_visible(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"lookup","arguments":{"q":"cats"}},' + '{"name":"bad","arguments":{"x":1}}]' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + assert events[0][0] == "tool_call" + trailing = "".join(v for k, v in events if k == "text") + assert '"bad"' in trailing + + def test_v11_closer_inside_span(self): + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + text = '[TOOL_CALLS]web_search[ARGS]{"query":"cats"}[/TOOL_CALLS] after' + calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) + (call,) = calls + assert call["function"]["name"] == "web_search" + (span,) = spans + assert text[span[0] : span[1]].endswith("[/TOOL_CALLS]") + assert text[span[1] :] == " after" + + def test_fully_promoted_array_strips_whole_region(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"lookup","arguments":{"q":"a"}},' + '{"name":"lookup","arguments":{"q":"b"}}] after' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + assert [k for k, _v in events] == ["tool_call", "tool_call", "text"] + assert events[2][1] == " after" + + +class TestMistralArrayHealing: + """Draining the whole [TOOL_CALLS] array for the shapes the repo's own + Mistral/Ollama templates emit.""" + + def test_comma_less_multi_call_array_parses_all_calls(self): + # ollama_template_mappers.py renders multi-call turns as [{...}{...}] with no + # comma separator; a single json.loads of the body rejects it and dropped every + # call. The element-by-element decode must recover all of them. + text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}}{"name":"b","arguments":{"y":2}}]' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["a", "b"] + assert json.loads(calls[0]["function"]["arguments"]) == {"x": 1} + assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} + + def test_comma_separated_and_single_arrays_still_parse(self): + both = parse_tool_calls_from_text( + '[TOOL_CALLS] [{"name":"a","arguments":{}},{"name":"b","arguments":{}}]' + ) + assert [c["function"]["name"] for c in both] == ["a", "b"] + one = parse_tool_calls_from_text('[TOOL_CALLS] [{"name":"a","arguments":{}}]') + assert [c["function"]["name"] for c in one] == ["a"] + + def test_mistral_array_null_arguments_normalized_to_empty_object(self): + # ``"arguments": null`` is a no-arg call; it must become {} (as the + # path does), not the string "null" that auto-heal turns into {"query":"null"}. + calls = parse_tool_calls_from_text('[TOOL_CALLS][{"name":"get_time","arguments":null}]') + assert calls[0]["function"]["arguments"] == "{}" + + class TestGlmStrict: def test_closed_glm_call_is_accepted(self): text = ( @@ -740,22 +867,26 @@ class TestMistralOuterOverXmlLiteral: class TestHealerSignalAlignment: - """The healer buffers only promotable formats; Mistral/Llama text calls stream through.""" + """The healer buffers only formats its shared parser can promote. Mistral's + ``[TOOL_CALLS]`` is promotable (rescued), so it is a heal signal; the loop-only + text-call markers (Llama ``<|python_tag|>``, bare ``[ARGS]``) are not, so they + stream through instead of stalling as prose that never yields a call.""" def test_heal_signals_subset_of_promotable_formats(self): from core.inference.passthrough_healing import _HEAL_SIGNALS - assert set(_HEAL_SIGNALS) == {"", "<|tool_call>", "", "<|tool_call>", " is not a healer-promotable format, so it streams through as text. + events = list(healer.feed('<|python_tag|>web_search.call(query="cats")')) text_out = "".join(v for k, v in events if k == "text") - assert "[TOOL_CALLS]" in text_out # streamed through, not buffered + assert "<|python_tag|>" in text_out # streamed through, not buffered assert not list(healer.finalize()) or all(k == "text" for k, _v in healer.finalize()) diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index d50c27130f..f7792a2a71 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -52,6 +52,11 @@ _ns = { } exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) _TOOL_XML_RE = _ns["_TOOL_XML_RE"] +# The display helper uses the closed-only variant before the last think block; keep it in scope. +_mc = _re.search(r"_TOOL_XML_CLOSED_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) +assert _mc, "could not extract _TOOL_XML_CLOSED_RE source" +exec(f"_TOOL_XML_CLOSED_RE = _re.compile({_mc.group(1)})", _ns) +_TOOL_XML_CLOSED_RE = _ns["_TOOL_XML_CLOSED_RE"] # Signatures may span multiple lines and now carry the enabled_tool_names gate; match # the whole (possibly multi-line) signature up to ``-> str:`` then the indented body. @@ -66,16 +71,19 @@ assert "_strip_mistral_closed_calls" in _xml_helper.group( exec(_xml_helper.group(0), _ns) _strip_tool_xml = _ns["_strip_tool_xml"] +# Extract the gate helper and display strip up to the next top-level ``logger =``. _helper = _re.search( - r"def _strip_tool_xml_for_display\((?:.|\n)*?\) -> str:\n(?: .+\n)+", + r"def _display_tool_name_gate\(.*?(?=\nlogger = get_logger)", _src, + _re.DOTALL, ) -assert _helper, "could not extract _strip_tool_xml_for_display source" -# After the V1 fix the display helper delegates to _strip_tool_xml; confirm the -# extracted body actually reached that call rather than truncating early. +assert _helper, "could not extract display strip helper source" +# The extracted block spans _display_tool_name_gate through _strip_tool_xml (defined before +# ``logger =``); confirm the shared _strip_tool_xml delegate is present. assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates" exec(_helper.group(0), _ns) _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] +_display_tool_name_gate = _ns["_display_tool_name_gate"] _gate_src = _re.search( r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?: .+\n)+", @@ -95,6 +103,56 @@ def test_route_display_strip_respects_disabled_auto_heal_contract(): assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) +def test_route_display_strip_preserves_rehearsal_inside_think(): + # A rehearsed bracket call inside think is reasoning: the block is preserved while a real + # call outside it still strips. + text = 'plan: search[ARGS]{"q":"x"} answer [TOOL_CALLS]web_search{"q":"y"} tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert 'plan: search[ARGS]{"q":"x"}' in out + assert "[TOOL_CALLS]web_search" not in out + assert "answer" in out and "tail" in out + + +def test_route_display_strip_keeps_bare_args_before_think_block(): + # A bare ``foo[ARGS]`` before a think block is prose: EOS-anchored tail arms run only on + # the last segment (earlier segments use the closed-only regex). + text = "Please pass foo[ARGS] pause to the template." + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) == text + + +def test_route_display_strip_removes_complete_call_before_think_block(): + # A complete bracket call before a think block still strips (balanced scan runs on every segment). + text = 'before search[ARGS]{"q":"x"} pause after' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "search[ARGS]" not in out + assert "pause" in out + assert "before" in out and "after" in out + + +def test_route_display_strip_removes_closed_xml_before_think_block(): + # A closed before a think block is removed in the non-last segment. + text = 'pre {"name":"x"} p tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out + assert "p" in out + assert "pre" in out and "tail" in out + + +def test_all_route_cleanup_sites_use_protected_display_helper(): + # Every route cleanup site must use _strip_tool_xml_for_display (think-preserving, + # balanced); raw _TOOL_XML_RE.sub corrupted think rehearsal and trailing prose. The only + # legitimate raw sub lives inside the helper itself. + raw_sub_lines = [ + (i, line) + for i, line in enumerate(_src.splitlines(), 1) + if "_TOOL_XML_RE.sub(" in line and not line.lstrip().startswith("#") + ] + assert len(raw_sub_lines) == 1, ( + "raw _TOOL_XML_RE.sub must appear only inside _strip_tool_xml_for_display; " + f"found extra call sites: {raw_sub_lines!r}" + ) + + def test_route_display_strip_removes_mistral_tool_calls_with_nested_json(): # _TOOL_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral # balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON). @@ -234,6 +292,32 @@ def test_strips_tail_only_parameter_orphan_no_trailing_ws(): assert "Final answer." in cleaned +def test_strips_complete_bracket_tag_keeps_trailing_prose(): + # A complete Mistral call strips only its balanced JSON, leaving following prose intact. + cleaned = _TOOL_XML_RE.sub("", '[TOOL_CALLS]web_search{"q":"x"} and then prose') + assert "[TOOL_CALLS]" not in cleaned + assert "and then prose" in cleaned + + +def test_strips_unclosed_bracket_tail(): + # Close brace lost to EOS: the truncated tail strips to the end instead of leaking. + cleaned = _TOOL_XML_RE.sub("", 'here [TOOL_CALLS]web_search{"query":"weather"') + assert "[TOOL_CALLS]" not in cleaned + assert cleaned.strip() == "here" + + +def test_strips_unclosed_rehearsal_tail(): + cleaned = _TOOL_XML_RE.sub("", 'text python[ARGS]{"code":"print(1)"') + assert "[ARGS]" not in cleaned + assert cleaned.strip() == "text" + + +def test_strips_hyphenated_mcp_bracket_name(): + cleaned = _TOOL_XML_RE.sub("", 'x [TOOL_CALLS]mcp__srv__list-issues{"q":"x"}') + assert "list-issues" not in cleaned + assert cleaned.strip() == "x" + + def test_preserves_mid_string_parameter_in_code_sample(): # Tail-anchor on `` so doc/example prose survives. text = ( @@ -362,6 +446,238 @@ def test_no_catastrophic_backtracking_on_orphan_opening_spam(): assert "" not in cleaned +# ── Two-level-nested bracket JSON (balanced-scan strip) ────────── + + +def test_route_strip_two_level_nested_bracket_keeps_trailing_prose(): + # Two-level-nested args must be removed whole so the trailing prose survives. + text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after' + cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert cleaned == "before after" + assert "[TOOL_CALLS]" not in cleaned + + +def test_route_strip_two_level_nested_rehearsal_keeps_trailing_prose(): + text = 'note python[ARGS]{"a":{"b":{"c":1}}} done' + cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert cleaned == "note done" + assert "[ARGS]" not in cleaned + + +def test_route_strip_removes_call_with_literal_think_in_argument(): + # A literal inside a call argument strips with the call, not as reasoning. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out and '"name"' not in out + + +def test_route_strip_removes_truncated_mistral_array(): + # A canonical array truncated by EOS is stripped by the route fallback like other orphans. + text = 'before [TOOL_CALLS] [{"name":"a","arguments":{"x":1}}' # missing ] + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "{" not in out + assert "before" in out + + +def test_route_strip_keeps_prose_mentioning_args_marker(): + # ``foo[ARGS] in a sentence`` is prose; the rehearsal arm must not truncate the line. + text = "Please pass foo[ARGS] to the template and continue reading." + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert out == text + + +def test_route_strip_handles_mistral_v11_call_id_args_shape(): + # v11 [CALL_ID]/[ARGS] shape (Mistral Small 3.2) must strip whole. + text = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert "before" in out and "after" in out + + +# ── Mistral [/TOOL_CALLS] closer + literal inside a call ─────────────── + +from core.tool_healing import strip_tool_call_markup as _strip_tool_call_markup + + +def test_core_strip_removes_orphan_tool_calls_closer_array_form(): + # The bare v11 [/TOOL_CALLS] closer left by the balanced scan must not leak as content. + text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]' + assert _strip_tool_call_markup(text, final = True) == "" + + +def test_core_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail(): + text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail' + assert _strip_tool_call_markup(text, final = True) == "tail" + + +def test_core_strip_removes_call_with_literal_think_in_argument(): + # An unclosed literal inside call arguments strips with the call (argument data). + text = 'before {"name":"write","arguments":{"text":"literal marker"}} after' + assert _strip_tool_call_markup(text, final = True) == "before after" + + +def test_route_display_strip_removes_orphan_tool_calls_closer_array_form(): + text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert out.strip() == "" + + +def test_route_display_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail(): + text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[/TOOL_CALLS]" not in out + assert out.strip() == "tail" + + +def test_incomplete_xml_call_with_literal_think_in_arg_is_stripped(): + # An incomplete holding a literal strips to EOS, not as a reasoning + # block (the unclosed tail _tool_call_markup_spans previously missed). + from core.tool_healing import parse_tool_calls_from_text as _parse + from core.tool_healing import strip_tool_call_markup as _strip + + text = 'before {"name":"write","arguments":{"text":"literal marker"}} after' + assert [c["function"]["name"] for c in _parse(text)] == ["write"] + assert _strip(text, final = True) == "before" + + # A real reasoning block with no tool call is still preserved verbatim. + assert ( + _strip("answer real done", final = True) == "answer real done" + ) + + # A complete call followed by a real reasoning block: call stripped, block kept. + mixed = '{"name":"a","arguments":{}} mid r end' + assert _strip(mixed, final = True) == "mid r end" + + +# ── enabled-tool gate for the ambiguous bare-rehearsal strip (#5704) ── + + +def test_display_tool_name_gate_returns_active_names_or_none(): + # Empty / no tools -> None (unrestricted; keep the legacy strip-all behavior). + assert _display_tool_name_gate([]) is None + assert _display_tool_name_gate(None) is None + # OpenAI-shaped tool dicts -> set of function names, malformed entries dropped. + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "run_python"}}, + {"type": "function"}, # no name + {"nope": 1}, # no function + ] + assert _display_tool_name_gate(tools) == {"web_search", "run_python"} + + +def test_route_display_strip_keeps_inactive_rehearsal_when_gated(): + # P1 #5704: an inactive ``foo[ARGS]{...}`` is prose; the gated strip leaves the sentence intact. + gate = {"web_search"} + text = 'foo[ARGS]{"x":1} is just syntax.' + assert ( + _strip_tool_xml_for_display(text, auto_heal_tool_calls = True, enabled_tool_names = gate) + == text + ) + # A bare marker with no JSON body is likewise prose when inactive. + assert ( + _strip_tool_xml_for_display( + "use foo[ARGS] here", auto_heal_tool_calls = True, enabled_tool_names = gate + ) + == "use foo[ARGS] here" + ) + + +def test_route_display_strip_removes_active_rehearsal_when_gated(): + # Mirror case: an active tool name is a real rehearsal and still strips. + gate = {"web_search"} + out = _strip_tool_xml_for_display( + 'web_search[ARGS]{"query":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + assert "web_search[ARGS]" not in out + assert out.strip() == "done" + + +def test_route_display_strip_ungated_strips_all_rehearsal_unchanged(): + # Backwards-compat: with no gate (None) the bare rehearsal strips as before. + text = 'foo[ARGS]{"x":1} is just syntax.' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "is just syntax." + assert ( + _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = None + ).strip() + == "is just syntax." + ) + + +def test_route_display_strip_control_token_stripped_regardless_of_gate(): + # [TOOL_CALLS] is a control token: stripped even when its NAME is not in the gate. + gate = {"web_search"} + out = _strip_tool_xml_for_display( + '[TOOL_CALLS]foo[ARGS]{"x":1} keep', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + assert "[TOOL_CALLS]" not in out and "foo[ARGS]" not in out + assert out.strip() == "keep" + + +def test_core_strip_gates_bare_rehearsal_on_enabled_tools(): + # P1 (#5704): the shared strip gate mirrors the parse gate -- inactive names are prose + # and preserved, active names strip, ``None`` keeps legacy strip-all. + from core.tool_healing import strip_tool_call_markup as _strip + + text = 'foo[ARGS]{"x":1} is just syntax.' + assert _strip(text, final = True, enabled_tool_names = {"web_search"}) == text + assert ( + _strip('web_search[ARGS]{"q":1} done', final = True, enabled_tool_names = {"web_search"}) + == "done" + ) + assert _strip(text, final = True).strip() == "is just syntax." + assert _strip(text, final = True, enabled_tool_names = None).strip() == "is just syntax." + + +def test_route_display_strip_gate_preserves_inactive_history_rehearsal(): + # The GGUF history sanitiser passes the gate, so a documented inactive shape survives in + # the replayed prompt context. + gate = _display_tool_name_gate([{"function": {"name": "web_search"}}]) + text = 'To call it write foo[ARGS]{"x":1} in your reply.' + assert 'foo[ARGS]{"x":1}' in _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = gate + ) + # An ACTIVE name is still stripped as a real rehearsed call. + assert "web_search[ARGS]" not in _strip_tool_xml_for_display( + 'Result web_search[ARGS]{"q":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + # No gate (legacy) strips every NAME[ARGS]{...}. + assert "foo[ARGS]" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + + +def test_gguf_history_sanitizer_forwards_enabled_tool_names_gate(): + # Wiring guard: the GGUF history strip must forward the display gate like the live strip. + block = _re.search( + r"Strip stale tool-call XML from conversation history.*?\.strip\(\)", + _src, + _re.DOTALL, + ) + assert block, "could not locate GGUF history sanitizer block" + assert "enabled_tool_names" in block.group( + 0 + ), "GGUF history sanitizer must pass enabled_tool_names to _strip_tool_xml_for_display" + + +def test_route_history_and_passthrough_forward_the_display_gate(): + # The safetensors/Anthropic history sanitisers and the Anthropic non-stream passthrough + # must forward the gate so inactive examples survive in replayed prompt / final text. + blocks = { + "safetensors history": r"Strip stale tool-call XML from prior assistant turns.*?\.strip\(\)", + "anthropic history": r"Strip stale tool-call XML via the protected display helper.*?\.strip\(\)", + "anthropic passthrough": r"gated on the declared tools so an\n.*?\.strip\(\)", + } + for label, pat in blocks.items(): + m = _re.search(pat, _src, _re.DOTALL) + assert m, f"could not locate {label} strip block" + assert "enabled_tool_names" in m.group( + 0 + ), f"{label} must forward enabled_tool_names to _strip_tool_xml_for_display" + + # ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ── From c2a7b78f6b55ffb80d294e00a85f3d2fd5272706 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 19:40:06 -0700 Subject: [PATCH 027/113] Studio: exclude mlx-lm 0.31.3 (broke gemma4/qwen3_5 QK-norm load on Apple Silicon) (#6803) * Studio: exclude mlx-lm 0.31.3 (broke gemma4/qwen3_5 QK-norm load) mlx-lm 0.31.3 regressed the QK-norm archs: its strict load_weights rejects the q_norm/k_norm tensors with "Received N parameters not in model", so gemma4 and qwen3_5 checkpoints fail to load. Studio installs the MLX stack unpinned at latest, which pulls 0.31.3. Verified on a real macos-14 runner: gemma4 fails to load on 0.31.3 but loads and generates coherently on 0.31.2 and on git-main (future 0.31.4). See mlx-lm #1242. Exclude just that release (!=0.31.3) in the installer and the self-heal floor so --upgrade still resolves to the newest good build, and treat an already-installed 0.31.3 as unsatisfied so the self-heal replaces it. * Studio MLX: cover fresh-install path + robust bad-version compare Address PR review: - Fresh install.sh (Apple Silicon) runs the base 'uv pip install unsloth' with SKIP_STUDIO_BASE=1, skipping the guarded MLX-stack step, so transitive resolution could still pull mlx-lm 0.31.3. install.sh already exports UV_OVERRIDE -> overrides-darwin-arm64.txt before that install, so exclude mlx-lm 0.31.3 there too; this also strengthens the self-heal (same override). - Match the known-bad version with parsed packaging.Version so 0.31.3 == 0.31.3.0 (trailing-zero normalization) instead of raw string equality. * Studio: exclude mlx-lm 0.31.3 on the fresh Apple Silicon install too The overrides file only applies via UV_OVERRIDE when it exists relative to the script, which is not true for a curl-piped install, and the guarded MLX step in install_python_stack.py is skipped there (SKIP_STUDIO_BASE=1). So the base install could still resolve the transitive mlx-lm to the broken 0.31.3. Append mlx-lm!=0.31.3 to the base install on Apple Silicon (empty elsewhere), so the fresh path pins away from 0.31.3 without waiting for the runtime self-heal. * Studio: exclude mlx-lm 0.31.3 on the migrated install; keep the >=0.22.0 floor The with-deps migrated install did not append ${_MLX_LM_EXCLUDE_ARG:-}, so a curl-piped Apple Silicon migration (no repo overrides file, UV_OVERRIDE unset) could resolve mlx-lm 0.31.3 transitively. Append the exclusion there, matching the fresh install path. The no-torch migration is left alone since --no-deps never resolves mlx-lm (same as the fresh no-torch path). Also restore the >=0.22.0 floor in overrides-darwin-arm64.txt: a uv override replaces the transitive constraint, so a bare !=0.31.3 could let the resolver drop below the supported minimum that mlx_repair.py enforces at runtime. * Triage huggingface_hub 1.22.0 / fastapi / multiprocess scanner false positives The scan-packages gate red-failed on all three shards after transitive deps bumped. Every new CRITICAL is a benign false positive, verified against upstream: - huggingface_hub 1.22.0 added _sandbox.py for the remote HF sandbox feature. Its job-startup bootstrap string (fetch sbx-server into the container /tmp and exec it) and the SandboxPool host-reservation loop trip the staged-dropper and C2-loop heuristics; that script runs inside a remote HF container, not on the user machine. The bump also re-hashed the already-reviewed benign polling loops in hf_api.py and utils/_http.py. The PyPI artifact is byte-identical to the official v1.22.0 tag. - fastapi 0.139.0 routing.py re-hashed the websocket keepalive while-True loop; byte-identical to upstream 0.139.0. - multiprocess 0.70.19 forkserver.py and tests/__init__.py re-hashed the AF_UNIX fork-server IPC and fd-inheritance tests; genuine uqfoundation release, local IPC not network. Added 7 reviewed allowlist entries (no blind regenerate). All three shards (hf-stack, studio, extras) exit 0 locally. * Tighten mlx-lm 0.31.3 exclusion comments * Trim mlx-lm 0.31.3 exclusion comments --- install.sh | 12 ++++++++-- scripts/scan_packages_baseline.json | 16 +++++++++++++ .../single-env/overrides-darwin-arm64.txt | 6 +++++ studio/backend/tests/test_mlx_repair.py | 23 +++++++++++++++++++ studio/backend/utils/mlx_repair.py | 21 +++++++++++++++-- studio/install_python_stack.py | 8 ++++++- 6 files changed, 81 insertions(+), 5 deletions(-) diff --git a/install.sh b/install.sh index 0370559540..81c50bc899 100755 --- a/install.sh +++ b/install.sh @@ -1442,8 +1442,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then SKIP_TORCH=true fi +# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression for +# gemma4 / qwen3_5; mlx-lm #1242). A curl-piped install has no overrides file +# and skips the guarded MLX step (SKIP_STUDIO_BASE=1), so this is the only cover. +_MLX_LM_EXCLUDE_ARG="" + # Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file). if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _MLX_LM_EXCLUDE_ARG="mlx-lm!=0.31.3" _OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt" if [ -f "$_OVERRIDES_FILE" ]; then # uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace @@ -2679,9 +2685,11 @@ if [ "$_MIGRATED" = true ]; then run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi else + # Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no + # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2912,7 +2920,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth -- "$PACKAGE_NAME" + --upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-} fi # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index d42225e205..1d34cfb66d 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -1545,6 +1545,22 @@ "severity": "HIGH", "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" + }, + { + "package": "multiprocess", + "file": "multiprocess/forkserver.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L6: import socket sha256:6c707119169286c9a798e2c8d13a48614e481d8a503950916fd4ffb4c94d3182", + "evidence_hash": "50fec0f0522a8e4e636bf348b752002d7935d8455af31fb78c6f11e2eba19f6d" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3569: os.dup2(conn.fileno(), i) | L3601: \"test needs os.dup2()\") | L3619: os.dup2(fd, newfd) | L20: import socket sha256:c824dc0f409f242420c3fbb324790c53cb3078d2c8b07ee8f2a05694b01c2946", + "evidence_hash": "3878a2b430c175dbc5877a95195bfe52f9588ff73fb74e2261ed5e33087915ad" } ] } diff --git a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt index 56b948644f..43f37b3183 100644 --- a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt +++ b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt @@ -10,3 +10,9 @@ transformers>=4.57.6 # anyio that also ImportErrors on TaskHandle and 500s the server. An override # wins the fight, so force one consistent <4.14 here too. anyio<4.14.0 + +# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights +# rejects q_norm/k_norm, so those checkpoints fail to load. mlx-lm #1242. +# The override also drops it from transitive resolution; keep the >=0.22.0 floor +# (mirrors mlx_repair.py _MLX_MIN_VERSIONS) or the resolver could go below it. +mlx-lm>=0.22.0,!=0.31.3 diff --git a/studio/backend/tests/test_mlx_repair.py b/studio/backend/tests/test_mlx_repair.py index 1b0cbf9df1..365cc46410 100644 --- a/studio/backend/tests/test_mlx_repair.py +++ b/studio/backend/tests/test_mlx_repair.py @@ -271,6 +271,29 @@ def test_stack_available_requires_runtime_imports_and_versions(monkeypatch): assert imported == list(mr._MLX_RUNTIME_IMPORTS) +def test_mlx_packages_exclude_known_bad_mlx_lm(): + # mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5); the install spec + # must exclude it so the resolver picks 0.31.2 or >=0.31.4. See mlx-lm #1242. + (mlx_lm_spec,) = [p for p in mr.MLX_PACKAGES if p.startswith("mlx-lm")] + assert mlx_lm_spec == "mlx-lm>=0.22.0,!=0.31.3" + + +@pytest.mark.parametrize("bad_form", ["0.31.3", "0.31.3.0"]) +def test_known_bad_installed_mlx_lm_triggers_repair(monkeypatch, bad_form): + # An installed 0.31.3 counts as unsatisfied so the self-heal replaces it; + # parsed-Version compare also catches the trailing-zero form 0.31.3.0. + import importlib.metadata as metadata + + def _version(name): + return bad_form if name == "mlx-lm" else mr._MLX_MIN_VERSIONS[name] + + monkeypatch.setattr(metadata, "version", _version) + monkeypatch.setattr( + mr.importlib, "import_module", lambda _n: pytest.fail("versions must gate imports") + ) + assert mr.mlx_stack_available() is False + + def test_no_op_off_apple_silicon(monkeypatch): monkeypatch.setattr(mr, "is_apple_silicon", lambda: False) called = {"n": 0} diff --git a/studio/backend/utils/mlx_repair.py b/studio/backend/utils/mlx_repair.py index 520c11c3b1..7e1c9864c9 100644 --- a/studio/backend/utils/mlx_repair.py +++ b/studio/backend/utils/mlx_repair.py @@ -45,9 +45,21 @@ DISABLE_ENV_VAR = "UNSLOTH_DISABLE_MLX_AUTOREPAIR" # deps). mlx-vlm especially must be >=0.4.4: an older one still imports but # breaks VLM Train/Export, so installing it would wrongly clear chat-only. _MLX_MIN_VERSIONS = {"mlx": "0.22.0", "mlx-lm": "0.22.0", "mlx-vlm": "0.4.4"} +# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights +# rejects q_norm/k_norm, so a self-heal must not pull it. mlx-lm #1242. +_MLX_BAD_VERSIONS = {"mlx-lm": ("0.31.3",)} _MLX_PACKAGE_NAMES = tuple(_MLX_MIN_VERSIONS) _MLX_RUNTIME_IMPORTS = ("mlx.core", "mlx_lm", "mlx_lm.sample_utils", "mlx_vlm") -MLX_PACKAGES = tuple(f"{name}>={version}" for name, version in _MLX_MIN_VERSIONS.items()) + + +def _mlx_spec(name: str, version: str) -> str: + spec = f"{name}>={version}" + for bad in _MLX_BAD_VERSIONS.get(name, ()): + spec += f",!={bad}" + return spec + + +MLX_PACKAGES = tuple(_mlx_spec(name, version) for name, version in _MLX_MIN_VERSIONS.items()) _MLX_REINSTALL_ARGS = tuple( arg for name in _MLX_PACKAGE_NAMES for arg in ("--reinstall-package", name) ) @@ -140,7 +152,12 @@ def _mlx_versions_satisfy_minimums() -> bool: return False for name, minimum in _MLX_MIN_VERSIONS.items(): try: - if Version(_dist_version(name)) < Version(minimum): + installed = Version(_dist_version(name)) + if installed < Version(minimum): + return False + # A known-broken build counts as unsatisfied so the self-heal + # reinstalls a good one; Version compare matches 0.31.3(.0/+local). + if any(installed == Version(bad) for bad in _MLX_BAD_VERSIONS.get(name, ())): return False except PackageNotFoundError: return False diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 439d3ffe7b..e033a56a0a 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1532,6 +1532,10 @@ LOCAL_DD_UNSTRUCTURED_PLUGIN = ( ) LOCAL_DD_GITHUB_PLUGIN = SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed" +# mlx-lm 0.31.3 broke gemma4 / qwen3_5 loading (strict load_weights rejects the +# QK-norm q_norm/k_norm tensors); exclude just that release. See mlx-lm #1242. +MLX_LM_BAD_VERSION_EXCLUSION = "!=0.31.3" + # Apple Silicon: override mlx-vlm/mlx-lm's transformers pin (see overrides). # _uv_safe_path: uv truncates UV_OVERRIDE at the first space too (issue #6503). _MLX_OVERRIDES = SINGLE_ENV / "overrides-darwin-arm64.txt" @@ -2092,6 +2096,8 @@ def install_python_stack() -> int: # macOS arm64: install MLX stack at latest (UV_OVERRIDE relaxes the # mlx-vlm / mlx-lm transformers pin -- set at module load). + # Exclude mlx-lm 0.31.3 (see MLX_LM_BAD_VERSION_EXCLUSION); it broke + # gemma4 / qwen3_5 QK-norm loading. mlx-lm #1242. if IS_MAC_ARM and not skip_base: _progress("MLX stack (Apple Silicon)") pip_install( @@ -2100,7 +2106,7 @@ def install_python_stack() -> int: "--upgrade", "mlx", "mlx-metal", - "mlx-lm", + f"mlx-lm{MLX_LM_BAD_VERSION_EXCLUSION}", "mlx-vlm", ) From 9dabe96786da7317148dd6e96137904dcee04105 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 19:41:19 -0700 Subject: [PATCH 028/113] Studio chat: tool-call nudging on by default (API stays opt-in) (#6883) * Studio chat: tool-call nudging on by default (API stays opt-in) Healing is already default-on everywhere and the nudge retry from the client-tool passthrough is opt-in on the API. Studio chat had neither signal: the frontend never sent nudge_tool_calls, and the safetensors and MLX server-side loop lacked the GGUF loop's plan-without-action re-prompt entirely. Backend: the re-prompt helpers move from llama_cpp.py into tool_call_parser.py (shared, cycle-free; the GGUF loop imports them under its old names with zero behavior change) and run_safetensors_tool_loop now re-prompts once at the streaming no-tool-call exit, gated on Auto-Heal, active tools, nothing executed yet, and short forward-looking text. Re-prompts do not consume tool iterations. Frontend: the chat adapter sends nudge_tool_calls from a new nudgeToolCalls runtime setting (default true) with the same persistence, hydration, and settings toggle plumbing as Auto-Heal. Request-model defaults are untouched, so raw API callers stay opt-in. * Address review: persist the nudge setting, consume the flag in the loops, skip the re-prompt after RAG autoinject ChatSettingsPayload uses extra forbid, so a settings patch containing nudgeToolCalls failed to persist any settings; the field is now typed and round-trips. nudge_tool_calls now plumbs into both server-side tool loops and gates the plan-without-action re-prompt with None meaning on, so API callers keep today's behavior, explicit false disables it, and Studio's default-on flag actually controls the path Studio chat runs. The safetensors loop no longer re-prompts after RAG autoinject: the injected retrieval bypasses the tool controller, so the nothing-executed gate saw an empty history and re-asked after a successful retrieval. * Safetensors loop: the plan-without-action retry requires an explicit nudge flag The retry is new on this loop, so an omitted nudge_tool_calls must not change existing API behavior; Studio opts in explicitly. The GGUF loop keeps None as on because its re-prompt predates the flag. * Suppress the plan-without-action re-prompt after a denied tool confirmation A denial appends TOOL_REJECTED_MESSAGE but records nothing in the tool controller history, so the nothing-executed gate re-prompted the model to call the tool the user had just rejected, producing another confirmation prompt. A denial now suppresses the re-prompt for the rest of the request, mirroring the RAG autoinject handling. * Tighten plan-without-action re-prompt comments * Tighten plan-without-action re-prompt comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: match unified plan-without-action nudge cap to GGUF default of 3 The shared MAX_ACT_REPROMPTS was set to 1, but GGUF's established default (llama_cpp.py) has re-prompted a stalling model up to 3 times since #5620. Restore the GGUF-matched cap so safetensors and MLX inherit the same behavior, and update the safetensors cap test to assert the cap dynamically. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/inference.py | 2 + studio/backend/core/inference/llama_cpp.py | 51 ++-- studio/backend/core/inference/orchestrator.py | 2 + .../core/inference/safetensors_agentic.py | 62 ++--- .../core/inference/tool_call_parser.py | 35 +++ studio/backend/routes/chat_history.py | 1 + studio/backend/routes/inference.py | 3 + .../backend/tests/test_chat_history_routes.py | 11 + .../backend/tests/test_llama_cpp_tool_loop.py | 42 +++ .../tests/test_nudge_tool_calls_wiring.py | 88 +++++++ .../tests/test_safetensors_tool_loop.py | 240 ++++++++++++++++-- .../src/features/chat/api/chat-adapter.ts | 1 + .../features/chat/api/chat-settings-api.ts | 1 + .../src/features/chat/chat-settings-sheet.tsx | 25 ++ .../chat/stores/chat-runtime-store.ts | 14 + .../frontend/src/features/chat/types/api.ts | 1 + .../chat/utils/chat-settings-storage.ts | 10 + .../features/settings/tabs/general-tab.tsx | 1 + 18 files changed, 507 insertions(+), 83 deletions(-) create mode 100644 studio/backend/tests/test_nudge_tool_calls_wiring.py diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 164f202681..064be30c06 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -828,6 +828,7 @@ class InferenceBackend: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -877,6 +878,7 @@ class InferenceBackend: execute_tool = execute_tool, cancel_event = cancel_event, auto_heal_tool_calls = auto_heal_tool_calls, + nudge_tool_calls = nudge_tool_calls, max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index aab3470193..5e6287f528 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -75,6 +75,12 @@ from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs +from core.inference.tool_call_parser import ( + MAX_ACT_REPROMPTS as _MAX_REPROMPTS, + REPROMPT_MAX_CHARS as _REPROMPT_MAX_CHARS, + is_short_intent_without_action as _is_short_intent_without_action, + reprompt_to_act_message as _reprompt_to_act_message, +) from core.inference.tool_loop_controller import ( ToolLoopController, tool_event_provenance, @@ -223,25 +229,8 @@ def _wsl_system_rocm_lib_dirs() -> "list[str]": return out -# ── Pre-compiled patterns for plan-without-action re-prompt ── -# Forward-looking intent signals: the model is describing what it *will* -# do rather than giving a final answer. -_INTENT_SIGNAL = re.compile( - r"(?i)(" - # Direct intent ("I'll ...", "Let me ...", straight + curly apostrophes). - # Excludes "I can"/"I should"/"I want to"/"let's" (common in answers). - # Negative lookahead drops negated forms ("I will not") so a refusal - # doesn't trigger a re-prompt. - r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" - r"|" - # Step/plan framing: "First ...", "Step 1:", "Here's my plan" - r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))" - r"|" - # "Now I" / "Next I" patterns - r"\b(?:now i|next i)\b" - r")" -) -_MAX_REPROMPTS = 3 +# Plan-without-action re-prompt state (intent signal, caps, message) now lives +# in tool_call_parser, imported above under its old aliases. # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. @@ -252,7 +241,6 @@ _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min # is exempt because it needs immediate artifact feedback. _PROVISIONAL_ARGS_MIN_CHARS = 256 _DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min -_REPROMPT_MAX_CHARS = 2000 # Cap tool calls from a single TEXTUAL-fallback turn (mirrors the safetensors # loop). Structured delta.tool_calls are grammar-bounded by llama-server; text # parsed from content is not, so one runaway turn could fan out unbounded. @@ -333,11 +321,6 @@ def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int: return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0 -def _is_short_intent_without_action(text: str) -> bool: - stripped = text.strip() - return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None - - def _should_suppress_forced_no_tool_output(text: str) -> bool: """Suppress only repeated forced-turn planning text, not final answers.""" stripped = text.strip() @@ -8456,6 +8439,7 @@ class LlamaCppBackend: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -8626,11 +8610,9 @@ class LlamaCppBackend: _kb_search_count = 0 # ── Re-prompt on plan-without-action ───────────────── - # When the model describes what it intends to do (forward-looking - # language) without calling a tool, re-prompt once. Only triggers on - # responses signaling intent/planning -- a direct answer like "4" or - # "Hello!" won't match. Pattern compiled at module level - # (_INTENT_SIGNAL). + # Model describes intent without calling a tool: re-prompt once. A + # direct answer ("4", "Hello!") won't match. Pattern shared with the + # safetensors loop (tool_call_parser.INTENT_SIGNAL). _reprompt_count = 0 # Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved # re-prompt slots don't extend the budget. Mirrors the safetensors guard. @@ -9153,8 +9135,10 @@ class LlamaCppBackend: r"(?i)\brender[_\s-]?html\b", _stripped, ) + # None keeps the default-on re-prompt; False disables it. if ( auto_heal_tool_calls + and (nudge_tool_calls is None or nudge_tool_calls) and active_tools and not _render_html_already_done_intent and _reprompt_count < _MAX_REPROMPTS @@ -9183,12 +9167,7 @@ class LlamaCppBackend: conversation.append( { "role": "user", - "content": ( - "You have access to enabled tools. If a tool is needed to satisfy " - "the user's request or complete the action you described, call " - f"{tool_hint} now. If no tool is needed, provide the final answer " - "and follow the user's requested format." - ), + "content": _reprompt_to_act_message(tool_hint), } ) # Accumulate tokens and timing from this iteration. diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 5dbd5fb479..675bd9f3ea 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -945,6 +945,7 @@ class InferenceOrchestrator: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -1007,6 +1008,7 @@ class InferenceOrchestrator: execute_tool = execute_tool, cancel_event = cancel_event, auto_heal_tool_calls = auto_heal_tool_calls, + nudge_tool_calls = nudge_tool_calls, max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index aa732f47e4..81c25b777e 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -33,10 +33,13 @@ from core.inference.tool_call_parser import ( _strip_mistral_closed_calls, _strip_mistral_reasoning, BUDGET_EXHAUSTED_NUDGE, + MAX_ACT_REPROMPTS, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, + is_short_intent_without_action, parse_tool_calls_from_text, + reprompt_to_act_message, strip_leading_bare_json_call, strip_llama3_leading_sentinels, strip_tool_markup, @@ -75,21 +78,6 @@ _MAX_BUFFER_CHARS = 32 # Memory bound for holding a leading bare-JSON object whose top-level "{" never balances. _MAX_BARE_JSON_BUFFER = 16384 -# Forward-looking intent ("I'll", "First,", "Step 1:") = planning, not answering; nudge a call. -# Negative lookahead drops negated forms ("I will not") so a refusal doesn't trigger it. Mirrors GGUF. -_INTENT_SIGNAL = re.compile( - r"(?i)(" - r"\b(i['’](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" - r"|\b(?:first\b|step \d+:?|here['’]?s (?:my |the |a )?(?:plan|approach))" - r"|\b(?:now i|next i)\b" - r")" -) -_MAX_REPROMPTS = 3 -_REPROMPT_MAX_CHARS = 2000 -# Templated so the nudge names the caller's enabled tools, not a hardcoded set. Mirrors GGUF tool_hint. -_REPROMPT_INSTRUCTION_TEMPLATE = ( - "STOP. Do NOT write code or explain. You MUST call a tool NOW. Call {tool_hint} immediately." -) # No grammar constraint here (unlike llama-server's lazy grammar): collapse # exact-duplicate calls and cap the count so a runaway turn cannot fan out. @@ -432,6 +420,7 @@ def run_safetensors_tool_loop( execute_tool: Callable[..., str], cancel_event: Optional[threading.Event] = None, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, @@ -471,6 +460,9 @@ def run_safetensors_tool_loop( for _ev in _auto["events"]: yield _ev conversation.extend(_auto["messages"]) + # Autoinject ran a KB search outside the controller, so it counts as an + # executed tool for the plan-without-action gate. + rag_autoinjected = bool(_auto) unrestricted_tools = not tools # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the @@ -488,6 +480,9 @@ def run_safetensors_tool_loop( final_attempt_done = False next_call_id = 0 reprompt_count = 0 + # A denied tool confirmation must not be answered with a plan-without-action + # re-prompt (which would raise the confirmation gate again). + tool_denied = False # Real tool-call turns completed. Only turns that actually executed a tool count # against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a # plan-without-action re-prompt) must not consume budget, matching the GGUF loop. @@ -510,7 +505,7 @@ def run_safetensors_tool_loop( _state_draining = 2 # Reserve re-prompt slots so they don't eat the caller's tool budget. - _extra_iters = _MAX_REPROMPTS if max_tool_iterations > 0 else 0 + _extra_iters = MAX_ACT_REPROMPTS if max_tool_iterations > 0 else 0 for iteration in range(max_tool_iterations + _extra_iters + 1): if cancel_event is not None and cancel_event.is_set(): return @@ -869,33 +864,39 @@ def run_safetensors_tool_loop( enabled_tool_names = _enabled_tool_names, ) if not safety_tc: - # Re-prompt only when the model planned without acting (intent - # signal); "4" / "Hello!" never trigger. Mirrors GGUF. - _stripped = content_accum.strip() + # Re-prompt once on plan-without-action, before any tool runs + # (GGUF loop parity). The retry is gated on nudge_tool_calls so + # Studio callers (which send True) always nudge, while API callers + # who omit the flag keep today's no-reprompt behavior (opt-in). + stripped_answer = content_accum.strip() if ( - tools - and auto_heal_tool_calls - and reprompt_count < _MAX_REPROMPTS - and 0 < len(_stripped) < _REPROMPT_MAX_CHARS - and _INTENT_SIGNAL.search(_stripped) - and not final_attempt_done + auto_heal_tool_calls + and nudge_tool_calls + and active_tools + and reprompt_count < MAX_ACT_REPROMPTS + and not rag_autoinjected + and not tool_denied + and not any(record.executed for record in tool_controller.history) + and is_short_intent_without_action(stripped_answer) ): reprompt_count += 1 logger.info( - "Safetensors re-prompt %d/%d: model planned without " + "Safetensors re-prompt %d/%d: model responded without " "calling tools (%d chars)", reprompt_count, - _MAX_REPROMPTS, - len(_stripped), + MAX_ACT_REPROMPTS, + len(stripped_answer), ) + conversation.append({"role": "assistant", "content": stripped_answer}) tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" - conversation.append({"role": "assistant", "content": _stripped}) conversation.append( { "role": "user", - "content": _REPROMPT_INSTRUCTION_TEMPLATE.format(tool_hint = tool_hint), + "content": reprompt_to_act_message(tool_hint), } ) + # Empty status clears the badge and resets the route's + # per-turn text cursor before the re-prompted turn streams. yield {"type": "status", "text": ""} continue @@ -1085,6 +1086,7 @@ def run_safetensors_tool_loop( "result": TOOL_REJECTED_MESSAGE, "provenance": decision.provenance, } + tool_denied = True denied_message = { "role": "tool", "name": decision.tool_name, diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 70115e5744..1ab1142eba 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -163,6 +163,41 @@ RAG_SEARCH_CAP_NUDGE = ( ) +# ── Plan-without-action re-prompt (shared by the GGUF and safetensors loops) ── +# Forward-looking intent: the model says what it *will* do, not a final answer. +INTENT_SIGNAL = re.compile( + r"(?i)(" + # Direct intent ("I'll", "Let me"); lookahead drops negated forms + # ("I will not") so a refusal does not re-prompt. + r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" + r"|" + # Step/plan framing: "First ...", "Step 1:", "Here's my plan" + r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))" + r"|" + r"\b(?:now i|next i)\b" + r")" +) +# Matches GGUF's established default (llama_cpp.py has re-prompted up to 3 +# times since #5620); safetensors and MLX inherit the same cap from here. +MAX_ACT_REPROMPTS = 3 +REPROMPT_MAX_CHARS = 2000 + + +def is_short_intent_without_action(text: str) -> bool: + stripped = text.strip() + return 0 < len(stripped) < REPROMPT_MAX_CHARS and INTENT_SIGNAL.search(stripped) is not None + + +def reprompt_to_act_message(tool_hint: str) -> str: + """The user message appended when re-prompting a plan-without-action turn.""" + return ( + "You have access to enabled tools. If a tool is needed to satisfy " + "the user's request or complete the action you described, call " + f"{tool_hint} now. If no tool is needed, provide the final answer " + "and follow the user's requested format." + ) + + # Qwen / Hermes ``{json}``. _TC_JSON_START_RE = re.compile(r"\s*\{") # Qwen3.5 ```` and the attribute form ```` diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 2c87ce8c6e..963d584303 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -177,6 +177,7 @@ class ChatSettingsPayload(BaseModel): collapseHtmlArtifacts: Optional[bool] = None allowArtifactNetworkAccess: Optional[bool] = None autoHealToolCalls: Optional[bool] = None + nudgeToolCalls: Optional[bool] = None maxToolCallsPerMessage: Optional[int] = Field(default = None, ge = 1) toolCallTimeout: Optional[int] = Field(default = None, ge = 1) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1042dda004..032a6e874a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6162,6 +6162,7 @@ async def openai_chat_completions( reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, max_tool_iterations = payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25, @@ -6888,6 +6889,7 @@ async def openai_chat_completions( reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, auto_heal_tool_calls = _sf_auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, max_tool_iterations = _sf_tool_budget, tool_call_timeout = payload.tool_call_timeout if payload.tool_call_timeout is not None @@ -10074,6 +10076,7 @@ async def anthropic_messages( cancel_event = cancel_event, max_tool_iterations = 25, auto_heal_tool_calls = True, + nudge_tool_calls = payload.nudge_tool_calls, tool_call_timeout = 300, session_id = payload.session_id, # Anthropic passthrough has no rag_scope field (RAG is local-only). diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index 2a6ebe244f..a60ac700bf 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -91,6 +91,17 @@ def test_chat_settings_payload_accepts_fast_mode_presets(): assert dumped["customPresets"][0]["params"]["fastMode"] is True +def test_chat_settings_payload_accepts_nudge_tool_calls(): + # extra="forbid" 400s PUT /api/chat/settings on unknown keys, so the + # frontend's persisted nudgeToolCalls needs a payload field (like + # autoHealToolCalls). + payload = chat_history.ChatSettingsPayload.model_validate( + {"autoHealToolCalls": True, "nudgeToolCalls": False} + ) + dumped = payload.model_dump(exclude_unset = True) + assert dumped == {"autoHealToolCalls": True, "nudgeToolCalls": False} + + def test_chat_inference_settings_covers_frontend_persisted_fields(): # Drift guard: every InferenceParams field the UI persists (all but # checkpoint) must exist on ChatInferenceSettings, else extra="forbid" diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 9e16be2160..fb1b0e52b7 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -1168,6 +1168,48 @@ def test_internal_reprompt_disabled_when_auto_heal_disabled(monkeypatch): assert len(payloads) == 1 +def test_internal_reprompt_disabled_when_nudge_tool_calls_false(monkeypatch): + # Explicit nudge_tool_calls=False disables the plan-without-action + # re-prompt even with Auto-Heal on (None keeps the default-on behavior). + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + auto_heal_tool_calls = True, + nudge_tool_calls = False, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now."] + assert len(payloads) == 1 + + def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatch): streams = [ [ diff --git a/studio/backend/tests/test_nudge_tool_calls_wiring.py b/studio/backend/tests/test_nudge_tool_calls_wiring.py new file mode 100644 index 0000000000..2c27b220ba --- /dev/null +++ b/studio/backend/tests/test_nudge_tool_calls_wiring.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Wiring guard for the plan-without-action ``nudge_tool_calls`` policy. + +Decided policy: the re-prompt is ALWAYS ON for the Studio inference paths +(safetensors, GGUF/llama_cpp, MLX) and OPT-IN for the API (/v1 OpenAI-compat + +Anthropic-compat, controlled by the request's ``nudge_tool_calls``, default off). + +Mechanism (verified here without loading a model): + + * every backend tool-loop entry point accepts and forwards ``nudge_tool_calls`` + (safetensors -> ``InferenceBackend``; MLX -> ``InferenceOrchestrator``; both + call the shared ``run_safetensors_tool_loop``; GGUF -> ``LlamaCppBackend``); + * the safetensors/MLX loop gates the retry on a truthy flag (new retry -> + opt-in), while the GGUF loop keeps its pre-existing default-on behaviour + (``None`` keeps nudging) so an omitted flag never disables GGUF; + * the API request models default the flag to ``None`` (opt-in / off); + * the Studio-facing routes forward the request's flag, and the Studio frontend + sends ``nudge_tool_calls: true`` -- exercised behaviourally in + ``test_safetensors_tool_loop.py`` and ``test_llama_cpp_tool_loop.py``. +""" + +import inspect + +from core.inference.inference import InferenceBackend +from core.inference.llama_cpp import LlamaCppBackend +from core.inference.orchestrator import InferenceOrchestrator +from core.inference.safetensors_agentic import run_safetensors_tool_loop + + +def _params(fn): + return inspect.signature(fn).parameters + + +def test_shared_loop_accepts_nudge_flag(): + assert "nudge_tool_calls" in _params(run_safetensors_tool_loop) + + +def test_all_three_backends_accept_the_flag(): + for method in ( + InferenceBackend.generate_chat_completion_with_tools, + InferenceOrchestrator.generate_chat_completion_with_tools, + LlamaCppBackend.generate_chat_completion_with_tools, + ): + assert "nudge_tool_calls" in _params(method), method.__qualname__ + + +def test_delegating_backends_forward_the_flag_to_the_shared_loop(): + # safetensors (in-process transformers) and MLX (parent-process orchestrator) + # both delegate to run_safetensors_tool_loop; GGUF runs its own in-file loop + # and consumes the flag directly (asserted separately by the gate test). + for method in ( + InferenceBackend.generate_chat_completion_with_tools, + InferenceOrchestrator.generate_chat_completion_with_tools, + ): + src = inspect.getsource(method) + assert "nudge_tool_calls = nudge_tool_calls" in src, method.__qualname__ + + +def test_safetensors_loop_is_opt_in_while_gguf_stays_default_on(): + # Safetensors/MLX: the retry is new here, so it requires a truthy flag. + sf_src = inspect.getsource(run_safetensors_tool_loop) + assert "and nudge_tool_calls" in sf_src + # GGUF: pre-existing nudge must not be accidentally disabled -- an omitted + # (None) flag keeps nudging; only an explicit False turns it off. + gguf_src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert "nudge_tool_calls is None or nudge_tool_calls" in gguf_src + + +def test_api_request_models_default_the_flag_off(): + from models.inference import AnthropicMessagesRequest, ChatCompletionRequest + for model in (ChatCompletionRequest, AnthropicMessagesRequest): + field = model.model_fields["nudge_tool_calls"] + assert field.default is None, model.__name__ + + +def test_studio_routes_forward_the_request_flag(): + # The Studio chat frontend posts to /v1/chat/completions and /v1/messages + # with nudge_tool_calls=true; the route handlers forward the request value + # (external API clients that omit it fall back to the opt-in default). + from routes import inference as routes_inference + for handler in ( + routes_inference.openai_chat_completions, + routes_inference.anthropic_messages, + ): + src = inspect.getsource(handler) + assert "nudge_tool_calls = payload.nudge_tool_calls" in src, handler.__name__ diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index f826f3cddf..a8546b82c4 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -2229,6 +2229,9 @@ def _reprompt_loop(*, auto_heal_tool_calls): tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], execute_tool = exec_fn, auto_heal_tool_calls = auto_heal_tool_calls, + # Studio always nudges (always-on for the Studio inference paths); the + # API opts in per request. Model the Studio caller here. + nudge_tool_calls = True, max_tool_iterations = 3, ) ) @@ -3100,7 +3103,7 @@ class TestLoopBehaviour: class TestLoopRePrompt: - """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``_MAX_REPROMPTS`` extra slots.""" + """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Studio always nudges, so these drive the loop with ``nudge_tool_calls=True``.""" def test_intent_signal_triggers_reprompt(self): # Turn 1: intent signal, no tool call. @@ -3116,6 +3119,7 @@ class TestLoopRePrompt: ["The sky is blue."], ], exec_results = ["Blue (Rayleigh scattering)"], + nudge_tool_calls = True, ) events = _collect_events(loop) # web_search must have been called once (after the re-prompt). @@ -3159,13 +3163,14 @@ class TestLoopRePrompt: contents = [e for e in events if e["type"] == "content"] assert contents and contents[-1]["text"].strip() == "4" - def test_max_reprompts_capped_at_three(self): - # Model keeps stalling with intent -- after 3 re-prompts the - # loop must give up rather than burn forever. + def test_max_reprompts_capped(self): + # Model keeps stalling with intent -- after MAX_ACT_REPROMPTS re-prompts + # the loop must give up rather than burn forever. turns = [["Let me search for that."]] * 6 # well over the cap loop, exec_fn = _make_loop( turns = turns, exec_results = [], + nudge_tool_calls = True, ) events = _collect_events(loop, max_events = 500) # No tool ever ran, but the loop terminated cleanly. @@ -3184,6 +3189,7 @@ class TestLoopRePrompt: ["found"], ], exec_results = ["..."], + nudge_tool_calls = True, ) events = _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "x"})] @@ -3194,7 +3200,7 @@ class TestLoopRePrompt: # re-prompt ate the slot the tool call would never run. loop, exec_fn = _make_loop( turns = [ - # 1. Intent stall (re-prompt 1/3). + # 1. Intent stall (re-prompt). ["Let me search for that."], # 2. Real tool call (uses the budget slot). ['{"name":"web_search","arguments":{"query":"weather"}}'], @@ -3203,6 +3209,7 @@ class TestLoopRePrompt: ], exec_results = ["sunny"], max_tool_iterations = 1, + nudge_tool_calls = True, ) events = _collect_events(loop) assert exec_fn.calls == [("web_search", {"query": "weather"})] @@ -3305,13 +3312,22 @@ class TestGGUFSafetensorsHealingParity: } def test_intent_regex_matches_same_phrases_as_gguf(self): - # The intent re-prompt regex must match the SAME forward-looking - # phrases on both backends so behaviour is the same on Mac (MLX - # / safetensors) and on Linux (GGUF). - from core.inference.llama_cpp import _INTENT_SIGNAL as gguf_re - from core.inference.safetensors_agentic import ( - _INTENT_SIGNAL as sf_re, + # The intent re-prompt regex is now a single shared source of truth + # (tool_call_parser.INTENT_SIGNAL) consumed by both the GGUF and the + # safetensors/MLX loops, so behaviour is identical on Mac and Linux. + # Both backends must resolve to that one shared helper. + from core.inference.llama_cpp import ( + _is_short_intent_without_action as gguf_fn, ) + from core.inference.safetensors_agentic import ( + is_short_intent_without_action as sf_fn, + ) + from core.inference.tool_call_parser import ( + INTENT_SIGNAL as shared_re, + is_short_intent_without_action as shared_fn, + ) + + assert gguf_fn is shared_fn and sf_fn is shared_fn for phrase in ( "I'll search for that", @@ -3322,8 +3338,8 @@ class TestGGUFSafetensorsHealingParity: "Here's my plan", "Now I need to call web_search", ): - assert gguf_re.search(phrase), f"GGUF missed {phrase!r}" - assert sf_re.search(phrase), f"safetensors missed {phrase!r}" + assert shared_re.search(phrase), f"missed {phrase!r}" + assert shared_fn(phrase), f"helper missed {phrase!r}" for plain in ( "4", @@ -3337,13 +3353,16 @@ class TestGGUFSafetensorsHealingParity: "I will not search the web for that.", "I'll never call that tool.", ): - assert not gguf_re.search(plain), f"GGUF wrongly fired on {plain!r}" - assert not sf_re.search(plain), f"safetensors wrongly fired on {plain!r}" + assert not shared_re.search(plain), f"wrongly fired on {plain!r}" + assert not shared_fn(plain), f"helper wrongly fired on {plain!r}" def test_max_reprompts_equal_on_both_backends(self): + # Both loops draw the cap from the shared constant, so they stay equal. from core.inference.llama_cpp import _MAX_REPROMPTS as gguf_cap - from core.inference.safetensors_agentic import _MAX_REPROMPTS as sf_cap - assert gguf_cap == sf_cap == 3 + from core.inference.safetensors_agentic import MAX_ACT_REPROMPTS as sf_cap + from core.inference.tool_call_parser import MAX_ACT_REPROMPTS as shared_cap + + assert gguf_cap == sf_cap == shared_cap class TestLoopControl: @@ -3822,6 +3841,193 @@ class TestGptOssNameDetection: assert is_gpt_oss_model_name(cast(str, None)) is False +# ──────────────────────────────────────────────────────────────────── +# Plan-without-action re-prompt (GGUF loop parity) +# ──────────────────────────────────────────────────────────────────── + + +class TestPlanWithoutActionReprompt: + def test_short_intent_is_reprompted_and_tool_executes(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the final answer."], + ], + exec_results = ["result-1"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search"] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("Here is the final answer." in t for t in texts) + + def test_reprompt_fires_up_to_the_cap(self): + # GGUF parity: a persistently stalling model is re-prompted up to + # MAX_ACT_REPROMPTS times, then the last stall is surrendered as the + # final answer and no further turn is generated. + from core.inference.tool_call_parser import MAX_ACT_REPROMPTS + + stall = "Let me look into it first." + turns = [["I'll search the web for that."]] + turns += [[stall]] * MAX_ACT_REPROMPTS + turns += [["SHOULD NOT APPEAR"]] + + generations = {"count": 0} + turn_iter = iter(turns) + + def _gen(_messages): + generations["count"] += 1 + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + # One initial turn plus exactly MAX_ACT_REPROMPTS re-prompted turns. + assert generations["count"] == MAX_ACT_REPROMPTS + 1 + texts = [e["text"] for e in events if e["type"] == "content"] + assert any(stall in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_long_prose_answer_is_not_reprompted(self): + long_answer = "I'll keep explaining the details of the topic. " * 60 + loop, exec_fn = _make_loop( + turns = [ + [long_answer], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_disabled_auto_heal_is_not_reprompted(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + auto_heal_tool_calls = False, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_explicit_nudge_off_is_not_reprompted(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = False, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_omitted_nudge_flag_is_not_reprompted(self): + # The retry is new on this loop: API callers who do not send the flag + # must keep today's behavior. Studio opts in explicitly. + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_rag_autoinject_counts_as_executed_tool(self, monkeypatch): + # Autoinject already ran a KB search outside the controller; a short + # post-retrieval intent must not trigger a spurious re-prompt. + import core.inference.tools as tools_mod + + def fake_autoinject(conversation, rag_scope): + return { + "events": [ + {"type": "tool_start", "tool_name": "search_knowledge_base"}, + {"type": "tool_end", "tool_name": "search_knowledge_base"}, + ], + "messages": [{"role": "tool", "content": "kb result"}], + } + + monkeypatch.setattr(tools_mod, "build_rag_autoinject", fake_autoinject) + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the docs."], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + assert any(e.get("type") == "tool_start" for e in events) + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the docs." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_no_reprompt_after_a_denied_tool_confirmation(self, monkeypatch): + # An explicit user denial must not be answered with a nudge to call + # the tool again (which would raise another confirmation prompt). + monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: "appr-1") + monkeypatch.setattr(safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object()) + monkeypatch.setattr(safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "deny") + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["I'll search again."], + ["SHOULD NOT APPEAR"], + ], + confirm_tool_calls = True, + session_id = "sess", + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search again." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_no_reprompt_after_a_tool_already_executed(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Now I'll refine the search."], + ["SHOULD NOT APPEAR"], + ], + exec_results = ["result-1"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search"] + texts = [e["text"] for e in events if e["type"] == "content"] + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + # Routes-level python_tag strip (multi-line; stop on next sentinel) class TestRoutesPythonTagStrip: """``_TOOL_XML_RE`` must consume multi-line code, embedded JSON, and bare ``<`` (earlier ``[^\n<]*`` / ``[^\n]*`` revisions leaked tails); the streaming route-level strip is the regression-prone path.""" diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 929c0385b0..df266fd749 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -2739,6 +2739,7 @@ export function createOpenAIStreamAdapter( : {}), auto_heal_tool_calls: useChatRuntimeStore.getState().autoHealToolCalls, + nudge_tool_calls: useChatRuntimeStore.getState().nudgeToolCalls, max_tool_calls_per_message: useChatRuntimeStore.getState().maxToolCallsPerMessage, tool_call_timeout: (() => { diff --git a/studio/frontend/src/features/chat/api/chat-settings-api.ts b/studio/frontend/src/features/chat/api/chat-settings-api.ts index 1ac569efc7..1e00357ea4 100644 --- a/studio/frontend/src/features/chat/api/chat-settings-api.ts +++ b/studio/frontend/src/features/chat/api/chat-settings-api.ts @@ -26,6 +26,7 @@ export interface PersistedChatSettings { collapseHtmlArtifacts?: boolean; allowArtifactNetworkAccess?: boolean; autoHealToolCalls?: boolean; + nudgeToolCalls?: boolean; maxToolCallsPerMessage?: number; toolCallTimeout?: number; } diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 74578b9438..ea6c409b40 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1732,6 +1732,7 @@ export function ChatSettingsPanel({
+ @@ -2006,6 +2007,30 @@ function AutoHealToolCallsToggle() { ); } +function NudgeToolCallsToggle() { + const nudgeToolCalls = useChatRuntimeStore((s) => s.nudgeToolCalls); + const setNudgeToolCalls = useChatRuntimeStore((s) => s.setNudgeToolCalls); + + return ( +
+
+ + Nudge Tool Calls + + + When a tool call cannot be repaired, re-ask the model once so the + intended tool still runs. API requests stay opt-in. + +
+ +
+ ); +} + function ConfirmToolCallsToggle() { const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 7c9685d6d0..ad50c4ece8 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -646,6 +646,7 @@ type ChatRuntimeStore = { toolStatus: string | null; generatingStatus: string | null; autoHealToolCalls: boolean; + nudgeToolCalls: boolean; maxToolCallsPerMessage: number; toolCallTimeout: number; kvCacheDtype: string | null; @@ -780,6 +781,7 @@ type ChatRuntimeStore = { setGeneratingStatus: (status: string | null) => void; setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void; setAutoHealToolCalls: (enabled: boolean) => void; + setNudgeToolCalls: (enabled: boolean) => void; setMaxToolCallsPerMessage: (value: number) => void; setToolCallTimeout: (value: number) => void; setKvCacheDtype: (dtype: string | null) => void; @@ -832,6 +834,7 @@ type ScalarSettingKey = | "collapseHtmlArtifacts" | "allowArtifactNetworkAccess" | "autoHealToolCalls" + | "nudgeToolCalls" | "maxToolCallsPerMessage" | "toolCallTimeout"; @@ -869,6 +872,7 @@ const SCALAR_SETTING_KEYS = [ "collapseHtmlArtifacts", "allowArtifactNetworkAccess", "autoHealToolCalls", + "nudgeToolCalls", "maxToolCallsPerMessage", "toolCallTimeout", ] as const satisfies readonly ScalarSettingKey[]; @@ -1103,6 +1107,7 @@ export const useChatRuntimeStore = create((set, get) => ({ generatingStatus: null, activeDiffusionCanvas: null, autoHealToolCalls: true, + nudgeToolCalls: true, maxToolCallsPerMessage: 25, toolCallTimeout: 5, kvCacheDtype: null, @@ -1544,6 +1549,15 @@ export const useChatRuntimeStore = create((set, get) => ({ ); return { autoHealToolCalls }; }), + setNudgeToolCalls: (nudgeToolCalls) => + set((state) => { + setScalarSettingVersion( + "nudgeToolCalls", + nudgeToolCalls, + state.nudgeToolCalls, + ); + return { nudgeToolCalls }; + }), setMaxToolCallsPerMessage: (maxToolCallsPerMessage) => set((state) => { setScalarSettingVersion( diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 94a7e1c56b..954e88e86b 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -357,6 +357,7 @@ export interface OpenAIChatCompletionsRequest { context_length?: number; }; auto_heal_tool_calls?: boolean; + nudge_tool_calls?: boolean; max_tool_calls_per_message?: number; tool_call_timeout?: number; session_id?: string; diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts index 4e457920d1..4899cb9c83 100644 --- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts @@ -21,6 +21,7 @@ import type { ReasoningEffort } from "../stores/chat-runtime-store"; const AUTO_TITLE_KEY = "unsloth_chat_auto_title"; const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls"; +const NUDGE_TOOL_CALLS_KEY = "unsloth_nudge_tool_calls"; const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message"; const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout"; const INFERENCE_PARAMS_KEY = "unsloth_chat_inference_params"; @@ -223,6 +224,7 @@ function sanitizeChatSettings(value: unknown): PersistedChatSettings { value.allowArtifactNetworkAccess, ); const autoHealToolCalls = sanitizeBool(value.autoHealToolCalls); + const nudgeToolCalls = sanitizeBool(value.nudgeToolCalls); const maxToolCallsPerMessage = sanitizeInt(value.maxToolCallsPerMessage, 1); const toolCallTimeout = sanitizeInt(value.toolCallTimeout, 1); @@ -245,6 +247,9 @@ function sanitizeChatSettings(value: unknown): PersistedChatSettings { if (autoHealToolCalls !== undefined) { settings.autoHealToolCalls = autoHealToolCalls; } + if (nudgeToolCalls !== undefined) { + settings.nudgeToolCalls = nudgeToolCalls; + } if (maxToolCallsPerMessage !== undefined) { settings.maxToolCallsPerMessage = maxToolCallsPerMessage; } @@ -305,6 +310,7 @@ export function isEmptyChatSettings(settings: PersistedChatSettings): boolean { settings.collapseHtmlArtifacts === undefined && settings.allowArtifactNetworkAccess === undefined && settings.autoHealToolCalls === undefined && + settings.nudgeToolCalls === undefined && settings.maxToolCallsPerMessage === undefined && settings.toolCallTimeout === undefined ); @@ -335,6 +341,7 @@ export function loadLegacyChatSettings(): PersistedChatSettings { const collapseHtmlArtifacts = loadBool(COLLAPSE_HTML_ARTIFACTS_KEY); const allowArtifactNetworkAccess = loadBool(ALLOW_ARTIFACT_NETWORK_ACCESS_KEY); const autoHealToolCalls = loadBool(AUTO_HEAL_TOOL_CALLS_KEY); + const nudgeToolCalls = loadBool(NUDGE_TOOL_CALLS_KEY); const maxToolCallsPerMessage = loadInt(MAX_TOOL_CALLS_KEY, 1); const toolCallTimeout = loadInt(TOOL_CALL_TIMEOUT_KEY, 1); const allCustomPresets = sanitizeCustomPresets([ @@ -361,6 +368,9 @@ export function loadLegacyChatSettings(): PersistedChatSettings { if (autoHealToolCalls !== undefined) { settings.autoHealToolCalls = autoHealToolCalls; } + if (nudgeToolCalls !== undefined) { + settings.nudgeToolCalls = nudgeToolCalls; + } if (maxToolCallsPerMessage !== undefined) { settings.maxToolCallsPerMessage = maxToolCallsPerMessage; } diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index df684b7752..7670aae5fa 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -78,6 +78,7 @@ const PREFS_KEYS: string[] = [ "unsloth_chat_auto_title", "unsloth_hf_token", "unsloth_auto_heal_tool_calls", + "unsloth_nudge_tool_calls", "unsloth_max_tool_calls_per_message", "unsloth_tool_call_timeout", "unsloth_chat_inference_params", From 8ba46b566a2740de4fba3cdf85ef503b005e8034 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 19:43:15 -0700 Subject: [PATCH 029/113] Studio: close switch/cancel races during model load (#6918) Fix six race conditions when a user switches or cancels a model while a previous load or generation is still in flight, across the inference orchestrator and the /load and /unload routes: - Cancel an in-flight generation on a safetensors/MLX model switch and serialize unload with load under the inference lifecycle gate. - Cancel an in-flight load off the lifecycle gate so a Stop-loading cancel does not wait out the multi-minute load; guard the dispatched mailbox against a racing unload. - Recheck the loading marker after spawn and again after the load response before publishing, so a load cancelled mid-flight is reaped instead of going live. - Discard the loading marker before tearing the subprocess down in cancel_load, closing a spawn-after-cancel window and an orphaned compare-mode dispatcher during unload. - Match the unload target before canceling an in-flight GGUF load and add an off-gate fast path for the still-loading GGUF case. - Run the Unsloth unload off the event loop so a paused SSE stream holding _gen_lock cannot block the loop. Adds studio/backend/tests/test_orchestrator_unload_cancel.py covering the unload/cancel/switch race paths. --- studio/backend/core/inference/orchestrator.py | 404 ++++- studio/backend/core/inference/worker.py | 57 +- studio/backend/routes/inference.py | 89 +- .../tests/test_orchestrator_unload_cancel.py | 1430 +++++++++++++++++ 4 files changed, 1885 insertions(+), 95 deletions(-) create mode 100644 studio/backend/tests/test_orchestrator_unload_cancel.py diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 675bd9f3ea..e0e6cef6c9 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -45,6 +45,10 @@ _DISPATCH_STOP_TIMEOUT = 5.0 _DISPATCH_IDLE_TIMEOUT = 30.0 _DISPATCH_DRAIN_TIMEOUT = 5.0 +# Max wait for a cancelled generation to release _gen_lock before unload_model +# tears the subprocess down. Only bounds a wedged worker. +_UNLOAD_GEN_LOCK_TIMEOUT = 15.0 + class InferenceOrchestrator: """ @@ -60,7 +64,13 @@ class InferenceOrchestrator: self._cmd_queue: Any = None self._resp_queue: Any = None self._cancel_event: Any = None # mp.Event — set to cancel generation + # Set for the whole unload; the worker never clears it (unlike _cancel_event), + # so a generate queued behind the cancelled one is skipped, not run. + self._drain_event: Any = None self._gen_lock = threading.Lock() # Serializes generation + # Set during a switch so a generation winning the _gen_lock handoff bails + # instead of starting on the outgoing model. + self._unload_pending = False # Dispatcher state for compare mode (adapter-controlled requests): # bypass _gen_lock, send commands directly, read from per-request @@ -159,6 +169,7 @@ class InferenceOrchestrator: self._cmd_queue = _CTX.Queue() self._resp_queue = _CTX.Queue() self._cancel_event = _CTX.Event() + self._drain_event = _CTX.Event() self._proc = _CTX.Process( target = run_without_native_path_secret, @@ -167,6 +178,7 @@ class InferenceOrchestrator: "cmd_queue": self._cmd_queue, "resp_queue": self._resp_queue, "cancel_event": self._cancel_event, + "drain_event": self._drain_event, "config": config, }, daemon = True, @@ -228,6 +240,7 @@ class InferenceOrchestrator: self._cmd_queue = None self._resp_queue = None self._cancel_event = None + self._drain_event = None logger.info("Inference subprocess shut down") def _cleanup(self): @@ -456,7 +469,15 @@ class InferenceOrchestrator: cancel ack from that same source so stale events don't leak into the next request. """ + # Latch this stream's subprocess/queue: if a wedged worker is torn down and a + # later load spawns a fresh one, bail rather than re-block on the new queue + # under _gen_lock (deadlock). + initial_proc = self._proc + initial_resp_queue = self._resp_queue while True: + if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: + yield f"Error: {self._subprocess_crash_message(crash_context)}" + return resp = read_one(read_timeout) if resp is None: # Check subprocess health @@ -595,8 +616,24 @@ class InferenceOrchestrator: if not self.active_model_name: yield "Error: No active model" return + # Latch the target model so the recheck below can detect a switch that completed + # between _start_dispatcher and mailbox registration (mirrors the locked path's + # expected_model check). + expected_model = self.active_model_name - # Ensure dispatcher is running + # Switch in flight (unload waiting on _gen_lock). This path bypasses the lock, + # so without this early-out a compare request would enqueue a generate on the + # outgoing model and delay the switch. + if self._unload_pending: + yield "Error: model is being unloaded" + return + + # Ensure the dispatcher runs. Track whether it was already running: if this call + # starts it and then bails on a racing unload, it must stop it again (see the + # unloading bail below). + dispatcher_preexisting = ( + self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() + ) self._start_dispatcher() request_id = str(uuid.uuid4()) @@ -624,10 +661,42 @@ class InferenceOrchestrator: preserve_thinking = preserve_thinking, ) - # Create mailbox BEFORE sending command + # Create the mailbox BEFORE sending, rechecking _unload_pending under + # _mailbox_lock: an unload sets _unload_pending before _wait_dispatcher_idle + # reads _mailboxes under the same lock, so either the idle check sees this + # mailbox (and tears the dispatcher down) or we see the unload and bail. + # Registering after would orphan the mailbox and hang the compare stream forever. mailbox: queue.Queue = queue.Queue() with self._mailbox_lock: - self._mailboxes[request_id] = mailbox + # _unload_pending alone is not enough: an unload that ran fully since + # _start_dispatcher clears it in its finally and stops the dispatcher, so it + # reads False here though the dispatcher is gone and the model swapped. Also + # bail when the active model changed or the dispatcher died: a mailbox with no + # dispatcher to route gen_done/gen_error hangs the compare stream. + dispatcher_alive = ( + self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() + ) + unloading = ( + self._unload_pending + or self.active_model_name != expected_model + or not dispatcher_alive + ) + if not unloading: + self._mailboxes[request_id] = mailbox + # When bailing without a mailbox, note whether any OTHER compare request still + # routes through the dispatcher; if none and this call started it, stop it below. + orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes + if unloading: + # A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was + # stopped, then set _unload_pending. The one we just started would otherwise + # linger with no mailboxes, race unload_model's _wait_response for the "unloaded" + # reply off resp_queue, and drop it as unroutable -- hanging the unload 300s. Stop + # it here so the unload stays the sole resp_queue reader. Outside _mailbox_lock: + # _stop_dispatcher joins the dispatcher, which itself takes that lock. + if orphaned_dispatcher: + self._stop_dispatcher() + yield "Error: model is being unloaded" + return try: self._send_cmd(cmd) @@ -676,14 +745,18 @@ class InferenceOrchestrator: return logger.warning("Timed out draining mailbox after cancel") - def _wait_dispatcher_idle(self) -> None: + def _wait_dispatcher_idle(self) -> bool: """Wait for all dispatched requests to complete, then stop dispatcher. - Called by _generate_inner before the _gen_lock path so the dispatcher - thread isn't competing for resp_queue reads. + Returns True if the dispatcher was stopped (all mailboxes drained, or no + dispatcher was running), and False if it was left running because compare + requests were still active after _DISPATCH_IDLE_TIMEOUT. + + Called before the _gen_lock path so the dispatcher thread isn't competing + for resp_queue reads. """ if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive(): - return + return True # Wait for all mailboxes to be emptied (dispatched requests complete) deadline = time.monotonic() + _DISPATCH_IDLE_TIMEOUT @@ -704,8 +777,9 @@ class InferenceOrchestrator: "leaving dispatcher running for compare requests", len(self._mailboxes), ) - else: - self._stop_dispatcher() + return False + self._stop_dispatcher() + return True # ------------------------------------------------------------------ # Public API — same interface as InferenceBackend @@ -772,6 +846,19 @@ class InferenceOrchestrator: ) for attempt in range(2): + # Stop-loading (/unload -> cancel_load) aborts a load by discarding this + # model's loading marker. cancel_load only kills a live child; if the cancel + # lands before any child exists (GPU placement, or between retries) there is + # nothing to kill, and without this check the loop would spawn a worker and + # load the model after /unload reported it unloaded. Observe removal and stop. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled before spawn; not starting a worker", + model_name, + ) + self.active_model_name = None + self.models.clear() + return False logger.info( "Spawning fresh inference subprocess for '%s' " "(transformers %s.x, attempt %d/2%s)", @@ -783,6 +870,22 @@ class InferenceOrchestrator: sub_config["disable_xet"] = disable_xet self._spawn_subprocess(sub_config) + # A cancel can land after the pre-spawn recheck but while _spawn_subprocess + # is still creating the queues/process. cancel_load runs off the lifecycle + # gate, so its _shutdown_subprocess can see _proc still None and no-op, + # orphaning this fresh worker; the load would then wait for "loaded" and + # publish a model /unload reported unloaded, over a live subprocess nothing + # reaps. Recheck now the child exists and tear it down before publishing. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled during spawn; tearing the worker down", + model_name, + ) + self._shutdown_subprocess(timeout = 5) + self.active_model_name = None + self.models.clear() + return False + try: resp = self._wait_response("loaded") except DownloadStallError: @@ -803,8 +906,31 @@ class InferenceOrchestrator: ) if resp.get("success"): + # A cancel can land while we were parked in _wait_response above. + # cancel_load (off the lifecycle gate) discards this model's loading + # marker BEFORE its teardown, so a Stop-loading that fired after the + # worker queued "loaded" (which we can still consume during cancel_load's + # shutdown window) shows up here only as the marker's removal. Without + # this recheck we would publish active_model_name/models for a model + # /unload reported cancelled, over a subprocess cancel_load just killed; + # its post-teardown re-clear cannot undo a publish that lands after it + # returns. Observe the removal and abort; cancel_load owns teardown. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled while waiting for 'loaded'; " + "not publishing the cancelled model", + model_name, + ) + self.active_model_name = None + self.models.clear() + return False model_info = resp.get("model_info", {}) self.active_model_name = model_info.get("identifier", model_name) + # A load always spawns a fresh subprocess holding only this model, so + # mirror that. A lingering stale name would pass unload_model's "not in + # self.models" guard, and the worker's absent-name fallback would unload + # its *active* model, not the already-gone one. + self.models = {} self.models[self.active_model_name] = { "is_vision": model_info.get("is_vision", False), "is_lora": model_info.get("is_lora", False), @@ -837,17 +963,65 @@ class InferenceOrchestrator: self.models.clear() raise - def unload_model(self, model_name: str) -> bool: - """Unload a model from the subprocess.""" - if model_name in self.loading_models: - logger.info( - "Cancelling in-flight load for model '%s' by terminating subprocess", + def cancel_load(self, model_name: str) -> bool: + """Abort an in-flight load by terminating its subprocess. + + Returns True if a load for ``model_name`` (matched case-insensitively) was + cancelled, False if nothing was loading under that name. This only tears the + loading subprocess down -- it sends no command to a worker -- so, unlike the + rest of ``unload_model``, it is safe to run WITHOUT the inference lifecycle + gate. ``/unload`` calls it off-gate so the "stop loading" button can interrupt + a safetensors load that holds the gate for its whole (multi-minute) duration; + a gated cancel could never preempt that load. + """ + target = model_name + if target not in self.loading_models: + target = next( + (m for m in self.loading_models if m.lower() == model_name.lower()), model_name, ) - self._shutdown_subprocess(timeout = 0.5) - self.loading_models.discard(model_name) - self.active_model_name = None - self.models.clear() + if target not in self.loading_models: + return False + logger.info( + "Cancelling in-flight load for model '%s' by terminating subprocess", + target, + ) + # Discard the loading marker (and clear local state) BEFORE the teardown, not + # after. cancel_load runs off the lifecycle gate, alongside a load_model that + # rechecks this marker before each spawn. But _shutdown_subprocess can block (~1s + # tearing a live child down and joining the dispatcher), so clearing only after + # leaves a window where load_model reads the marker still set, passes its pre-spawn + # recheck, and loads the model after /unload reported it cancelled. Clear first. + self.loading_models.discard(target) + self.active_model_name = None + self.models.clear() + self._shutdown_subprocess(timeout = 0.5) + # Clear the local mirrors again AFTER the teardown. A racing off-gate load_model + # may still be parked in _wait_response("loaded"): its worker already queued a + # "loaded" reply, so during the shutdown window above (the 0.5s settle before the + # response queue is drained and nulled) that thread can consume it and repopulate + # active_model_name/models, undoing the pre-teardown clear. _shutdown_subprocess + # nulls the queue but not the mirrors, so without this second clear /unload reports + # success while the backend still advertises a killed model. The nulled queue lets + # no further "loaded" through, so re-clearing here wipes any repopulation. + self.active_model_name = None + self.models.clear() + return True + + def unload_model(self, model_name: str) -> bool: + """Unload a model from the subprocess.""" + # active_model_name can differ in case from the client's raw /unload name (the + # load path canonicalizes casing). Match case-insensitively and use the canonical + # spelling so the guard, unload command, and cleanup below hit the loaded model. + if ( + self.active_model_name is not None + and model_name != self.active_model_name + and model_name.lower() == self.active_model_name.lower() + ): + model_name = self.active_model_name + # In-flight load: tear its subprocess down (shared loading-cancel logic; no + # worker command sent). + if self.cancel_load(model_name): return True if not self._ensure_subprocess_alive(): @@ -857,30 +1031,85 @@ class InferenceOrchestrator: self.active_model_name = None return True - try: - self._send_cmd( - { - "type": "unload", - "model_name": model_name, - } - ) - resp = self._wait_response("unloaded") - - # Update local state + # Nothing loaded under this name: don't unload a stale model. The worker falls + # back to unloading its *active* model when the name is absent, so a stale unload + # (lost a race to a concurrent load) would hit the wrong one. + if model_name != self.active_model_name and model_name not in self.models: self.models.pop(model_name, None) - if self.active_model_name == model_name: - self.active_model_name = None - - logger.info("Model '%s' unloaded from subprocess", model_name) return True - except Exception as exc: - logger.error("Error unloading model '%s': %s", model_name, exc) - # Clear local state anyway - self.models.pop(model_name, None) - if self.active_model_name == model_name: - self.active_model_name = None - return False + # The subprocess runs commands sequentially, so a bare unload queues behind a + # running generate (a 2-3 min hang). Cancel first (via the mp.Event the worker + # polls each token), then take _gen_lock as sole resp_queue reader (like GGUF). + self._unload_pending = True + # Cancelling only the running generation isn't enough: the worker clears + # cancel_event at each generate start, so a queued one would clear it and run the + # outgoing model to completion. drain_event, never cleared, makes any generate + # dequeued during the unload skip. + if self._drain_event is not None: + self._drain_event.set() + try: + self._cancel_generation() + acquired = self._gen_lock.acquire(timeout = _UNLOAD_GEN_LOCK_TIMEOUT) + if not acquired: + # Wedged worker: tear the subprocess down to free the GPU (next load respawns). + logger.warning( + "Unload: generation did not yield %.1fs after cancel; " + "shutting the inference subprocess down to free the model", + _UNLOAD_GEN_LOCK_TIMEOUT, + ) + self._shutdown_subprocess(timeout = 5) + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + + try: + # Stop the compare-mode dispatcher so it can't consume the "unloaded" reply + # off resp_queue before we do. A dispatched generation bypasses _gen_lock, so + # a wedged one slips past the acquire above; if the dispatcher is still active + # it owns resp_queue and the queued unload hangs _wait_response behind the + # stuck generate. Mirror the wedged locked path: tear the subprocess down. + if not self._wait_dispatcher_idle(): + logger.warning( + "Unload: compare-mode dispatcher still active after idle " + "wait; shutting the inference subprocess down to free the model" + ) + self._shutdown_subprocess(timeout = 5) + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + # Drop stale tokens so they can't be read as the unload reply. + self._drain_queue() + self._send_cmd( + { + "type": "unload", + "model_name": model_name, + } + ) + self._wait_response("unloaded") + + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + + logger.info("Model '%s' unloaded from subprocess", model_name) + return True + + except Exception as exc: + logger.error("Error unloading model '%s': %s", model_name, exc) + # Clear local state anyway + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return False + finally: + self._gen_lock.release() + finally: + self._unload_pending = False + if self._drain_event is not None: + self._drain_event.clear() def generate_chat_response( self, @@ -1068,6 +1297,7 @@ class InferenceOrchestrator: if not self.active_model_name: yield "Error: No active model" return + expected_model = self.active_model_name # Drain any prior compare-mode dispatcher so we can read resp_queue. self._wait_dispatcher_idle() @@ -1076,6 +1306,14 @@ class InferenceOrchestrator: # consume and drop each other's token events. Hold _gen_lock across the # cmd build + send + whole stream so we stay the sole resp_queue reader. with self._gen_lock: + # Recheck under the lock: an unload we raced may have cleared/swapped the model. + # _unload_pending resets after the lock releases, so it can read False by now; + # the active-model check catches that handoff and a reload that swapped models, + # so we never generate on the wrong one. + if self._unload_pending or self.active_model_name != expected_model: + # Won the lock handoff during a switch; don't start on the outgoing model. + yield "Error: model is being unloaded" + return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None cmd = self._build_generate_cmd( @@ -1143,53 +1381,62 @@ class InferenceOrchestrator: raise RuntimeError("Inference subprocess is not running") if not self.active_model_name: raise RuntimeError("No active model") + expected_model = self.active_model_name - request_id = str(uuid.uuid4()) + # Serialize under _gen_lock (sole resp_queue reader) and refuse to start on the + # outgoing model once an unload is pending, like the text and audio-input paths. + # Without this a concurrent /audio/generate could run TTS on a model being switched. + with self._gen_lock: + # Recheck under the lock (see _generate_inner): a raced unload/switch may have + # cleared or swapped the model while we waited. + if self._unload_pending or self.active_model_name != expected_model: + raise RuntimeError("model is being unloaded") - cmd = { - "type": "generate_audio", - "request_id": request_id, - "text": text, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "max_new_tokens": max_new_tokens, - "repetition_penalty": repetition_penalty, - } - if use_adapter is not None: - cmd["use_adapter"] = use_adapter + request_id = str(uuid.uuid4()) - self._send_cmd(cmd) + cmd = { + "type": "generate_audio", + "request_id": request_id, + "text": text, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + } + if use_adapter is not None: + cmd["use_adapter"] = use_adapter - # Wait for audio_done or audio_error - deadline = time.monotonic() + 120.0 - while time.monotonic() < deadline: - remaining = max(0.1, deadline - time.monotonic()) - resp = self._read_resp(timeout = min(remaining, 1.0)) + self._send_cmd(cmd) - if resp is None: - if not self._ensure_subprocess_alive(): - raise RuntimeError(self._subprocess_crash_message("audio generation")) - continue + deadline = time.monotonic() + 120.0 + while time.monotonic() < deadline: + remaining = max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) - rtype = resp.get("type", "") + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("audio generation")) + continue - if rtype == "audio_done": - wav_bytes = base64.b64decode(resp["wav_base64"]) - sample_rate = resp["sample_rate"] - return wav_bytes, sample_rate + rtype = resp.get("type", "") - if rtype == "audio_error": - raise RuntimeError(resp.get("error", "Audio generation failed")) + if rtype == "audio_done": + wav_bytes = base64.b64decode(resp["wav_base64"]) + sample_rate = resp["sample_rate"] + return wav_bytes, sample_rate - if rtype == "error": - raise RuntimeError(resp.get("error", "Unknown error")) + if rtype == "audio_error": + raise RuntimeError(resp.get("error", "Audio generation failed")) - if rtype == "status": - continue + if rtype == "error": + raise RuntimeError(resp.get("error", "Unknown error")) - raise RuntimeError("Timeout waiting for audio generation (120s)") + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for audio generation (120s)") def generate_whisper_response( self, @@ -1254,8 +1501,15 @@ class InferenceOrchestrator: if not self.active_model_name: yield "Error: No active model" return + expected_model = self.active_model_name with self._gen_lock: + # Recheck under the lock (see _generate_inner): a raced unload/switch may have + # cleared or swapped the model while we waited. + if self._unload_pending or self.active_model_name != expected_model: + # Won the lock handoff during a switch; don't start on the outgoing model. + yield "Error: model is being unloaded" + return request_id = str(uuid.uuid4()) # numpy array -> list for mp.Queue serialization diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 4e27183d88..615c5c5a0d 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -406,6 +406,32 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: ) +def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool: + """Skip a generate queued behind a cancelled one during an unload. + + The parent sets ``drain_event`` for the whole unload. Because the parent's + per-token ``cancel_event`` is cleared at the start of every generate, a cancel + set while this generate was still queued would otherwise be lost when it is + dequeued. If the drain is in effect, emit an immediate (empty) ``gen_done`` so + the parent's stream/mailbox drains fast and the switch stays fast, and report + the generate was skipped so the caller does not clear the cancel or run it. + """ + if drain_event is None or not drain_event.is_set(): + return False + request_id = cmd.get("request_id", "") + logger.info("Skipping generate for request %s: unload draining", request_id) + _send_response( + resp_queue, + { + "type": "gen_done", + "request_id": request_id, + "cancelled": True, + "stats": None, + }, + ) + return True + + def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: """Handle a generate command: stream tokens back via resp_queue. @@ -632,7 +658,14 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None: ) -def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None: +def run_inference_process( + *, + cmd_queue: Any, + resp_queue: Any, + cancel_event, + config: dict, + drain_event = None, +) -> None: """Subprocess entrypoint. Persistent — runs the command loop until shutdown. Args: @@ -640,6 +673,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf resp_queue: mp.Queue for sending responses to parent. cancel_event: mp.Event the parent sets to cancel generation. config: Initial configuration dict with model info. + drain_event: mp.Event the parent sets for the duration of an unload. Unlike + cancel_event (cleared at the start of every generate), it is never cleared + here, so a generate still queued behind a cancelled one is skipped rather + than run — the cancel survives the queue handoff. """ os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports @@ -715,7 +752,16 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf cmd_type = cmd.get("type", "") try: if cmd_type == "generate": + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue cancel_event.clear() + # Re-check the drain after clearing: the parent sets drain_event + # then cancel_event for an unload, so if that pair landed between + # the check above and this clear, the clear just erased the unload's + # cancel. Skip here so the outgoing model is not run to completion, + # which would stall the switch until the dispatcher idle-timeout. + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue _handle_generate(backend, cmd, resp_queue, cancel_event) elif cmd_type == "load": if backend.active_model_name: @@ -918,7 +964,16 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf try: if cmd_type == "generate": + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue cancel_event.clear() + # Re-check the drain after clearing: the parent sets drain_event then + # cancel_event for an unload, so if that pair landed between the check + # above and this clear, the clear just erased the unload's cancel. Skip + # here so the outgoing model is not run to completion, which would stall + # the switch until the dispatcher idle-timeout tears the subprocess down. + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue _handle_generate(backend, cmd, resp_queue, cancel_event) elif cmd_type == "load": diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 032a6e874a..e31c03f7f8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3394,12 +3394,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() - # Unload any active Unsloth model to free VRAM + # Unload any active Unsloth model to free VRAM (off the event loop: + # unload takes _gen_lock and can wait on an in-flight stream). if unsloth_backend.active_model_name: logger.info( f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" ) - unsloth_backend.unload_model(unsloth_backend.active_model_name) + await asyncio.to_thread( + 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 doesn't round-trip @@ -4063,28 +4066,76 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge # A deliberate unload means "stay unloaded": drop any idle reload stash so the # next /v1 request can't resurrect this model. The idle loop unloads via the # backend directly (not this route), so clearing here never fights keep-warm. - from core.inference.llama_keepwarm import note_model_unloaded + from core.inference.llama_keepwarm import inference_lifecycle_gate, note_model_unloaded try: - # Check if the GGUF backend has this model loaded or is loading it. - llama_backend = get_llama_cpp_backend() - if llama_backend.is_active and ( - llama_backend.model_identifier == request.model_path - or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) - or not llama_backend.is_loaded + # "Stop loading" (frontend cancelLoading -> /unload) must abort a still-loading + # model promptly. /load holds the lifecycle gate for the whole (multi-minute) load, + # so gating first would make the cancel wait it out. cancel_load only tears the + # loading subprocess down (no unload command), so it is safe off-gate. + backend = get_inference_backend() + loading = getattr(backend, "get_loading_model", lambda: None)() + if ( + loading is not None + and hasattr(backend, "cancel_load") + and (request.model_path == loading or request.model_path.lower() == loading.lower()) ): - # A manual unload is a deliberate user action: tear down now even if a - # request is mid-stream (only the automatic idle loop defers to it). - llama_backend.unload_model() + if await asyncio.to_thread(backend.cancel_load, request.model_path): + note_model_unloaded() + logger.info(f"Cancelled in-flight load: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) + + # Same "stop loading" fast path for a still-loading GGUF (llama-server spawned, + # health check not yet passed). A gated unload would wait out the multi-minute + # load; unload_model() sets the cancel_event load_model polls off its own lock and + # kills the child, sending no worker command, so it is safe off-gate like + # cancel_load. The gated GGUF branch below handles the already-loaded case. Gate on + # the loading model (identifier or native label): the single llama-server loads one + # GGUF at a time, so an unload for a different model must not cancel this load. + llama_backend = get_llama_cpp_backend() + if ( + llama_backend.is_active + and not llama_backend.is_loaded + and ( + llama_backend.model_identifier == request.model_path + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) + ) + ): + await asyncio.to_thread(llama_backend.unload_model) note_model_unloaded() - logger.info(f"Unloaded GGUF model: {request.model_path}") + logger.info(f"Cancelled in-flight GGUF load: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) - # Otherwise, unload from Unsloth backend - backend = get_inference_backend() - backend.unload_model(request.model_path) - note_model_unloaded() - logger.info(f"Unloaded model: {request.model_path}") - return UnloadResponse(status = "unloaded", model = request.model_path) + # Serialize with /load under the same lifecycle gate: the Unsloth unload now runs + # off the event loop (asyncio.to_thread), so without this a concurrent /load could + # swap in a fresh subprocess mid-unload and the unload command would land on the + # new worker. The gate makes load and unload exclusive. + async with inference_lifecycle_gate(): + # Check if the GGUF backend has this model loaded or is loading it. + llama_backend = get_llama_cpp_backend() + if llama_backend.is_active and ( + llama_backend.model_identifier == request.model_path + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) + or not llama_backend.is_loaded + ): + # A manual unload is a deliberate user action: tear down now even if a + # request is mid-stream (only the automatic idle loop defers to it). + llama_backend.unload_model() + note_model_unloaded() + logger.info(f"Unloaded GGUF model: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) + + # Unload from Unsloth backend off the event loop: unload takes _gen_lock, which + # a slow SSE stream paused between tokens still holds, so a sync call would block + # the loop that drives the stream's next token and the lock release. + backend = get_inference_backend() + await asyncio.to_thread(backend.unload_model, request.model_path) + note_model_unloaded() + logger.info(f"Unloaded model: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) except Exception as e: logger.error(f"Error unloading model: {e}", exc_info = True) diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py new file mode 100644 index 0000000000..e9d0f36fe2 --- /dev/null +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -0,0 +1,1430 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""unload_model cancels an in-flight generation instead of waiting it out. + +The sequential subprocess used to queue ``unload`` behind a running ``generate``, +hanging the UI. ``unload_model`` now cancels first (the mp.Event the worker checks +each token) and takes ``_gen_lock`` before the unload round-trip. +""" + +import threading +import time + +import pytest + +from core.inference import orchestrator as orch_mod +from core.inference.orchestrator import InferenceOrchestrator + + +def _bare_orchestrator(): + """An orchestrator without the real __init__ subprocess/network.""" + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._gen_lock = threading.Lock() + o._cancel_event = threading.Event() # stands in for the mp.Event + o._drain_event = threading.Event() # stands in for the unload-drain mp.Event + o._proc = object() # truthy so _ensure_subprocess_alive reports alive + o._cmd_queue = object() + o._resp_queue = object() + o._dispatcher_thread = None + o._unload_pending = False + o.active_model_name = "m" + o.models = {"m": {}} + o.loading_models = set() + return o + + +def test_unload_cancels_inflight_generation_then_unloads(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + # A generation holds _gen_lock and releases it only once cancelled. + o._gen_lock.acquire() + + def releaser(): + o._cancel_event.wait(timeout = 5) # released only after the cancel fires + o._gen_lock.release() + + t = threading.Thread(target = releaser) + t.start() + + start = time.monotonic() + ok = o.unload_model("m") + elapsed = time.monotonic() - start + t.join(timeout = 5) + + assert ok is True + assert o._cancel_event.is_set(), "generation must be cancelled before the unload" + assert {"type": "unload", "model_name": "m"} in sent + assert o.active_model_name is None + assert "m" not in o.models + # Waited on the released-after-cancel lock, not a full generation. + assert elapsed < 2.0 + + +def test_unload_no_active_generation_unloads_normally(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + ok = o.unload_model("m") + + assert ok is True + assert {"type": "unload", "model_name": "m"} in sent + assert o.active_model_name is None + # Lock released for the next caller. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_unload_falls_back_to_shutdown_when_generation_wont_yield(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send unload when wedged")) + + # A wedged worker never releases _gen_lock, even after the cancel. + o._gen_lock.acquire() + + ok = o.unload_model("m") + + assert ok is True + assert shutdown, "should tear the subprocess down to free the GPU" + assert o.active_model_name is None + + +def test_unload_tears_down_when_compare_dispatcher_wedged(monkeypatch): + # A wedged compare-mode generation bypasses _gen_lock, so the acquire guard + # misses it and _send_cmd/_wait_response would hang on resp_queue. Unload must + # instead tear the subprocess down, like the wedged locked-generation path. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_DISPATCH_IDLE_TIMEOUT", 0.2) + + # A live dispatcher whose mailbox never drains == a wedged compare-mode gen. + o._mailbox_lock = threading.Lock() + o._mailboxes = {"req-1": object()} + + class _AliveThread: + def is_alive(self): + return True + + o._dispatcher_thread = _AliveThread() + + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send unload with a wedged dispatcher") + ) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: pytest.fail( + "must not wait on resp_queue with a wedged dispatcher" + ), + ) + + # _gen_lock is free (compare mode never took it), so the acquire guard passes. + ok = o.unload_model("m") + + assert ok is True + assert shutdown, "should tear the subprocess down to free the GPU" + assert o.active_model_name is None + assert "m" not in o.models + + +def test_consume_token_stream_bails_when_subprocess_swapped(monkeypatch): + # After a wedged-worker teardown a fresh load swaps _proc/_resp_queue; the + # still-live generation thread must detect the swap and bail, not re-block on + # the new queue while holding _gen_lock. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_subprocess_crash_message", lambda ctx: "inference subprocess restarted" + ) + + def read_one(timeout): + o._proc = object() # simulate the reload swapping the subprocess + return None + + gen = o._consume_token_stream(read_one, lambda: None, crash_context = "generation") + msg = next(gen) + + assert "restarted" in msg + with pytest.raises(StopIteration): + next(gen) + + +def test_unload_pending_clears_after_unload(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: None) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + o.unload_model("m") + + # The flag must not leak past the unload, else every later generation bails. + assert o._unload_pending is False + + +def test_generation_bails_when_unload_pending(monkeypatch): + # Winning the _gen_lock handoff mid-switch must not start on the outgoing model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + o._unload_pending = True + + out = list(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + # It released (or never held) the lock, so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_dispatched_generation_bails_when_unload_pending(monkeypatch): + # Compare-mode bypasses _gen_lock, so it must early-out on a pending switch or + # it enqueues a generate on the outgoing model and delays the unload. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: pytest.fail("must not start a generation mid-switch") + ) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") + ) + o._unload_pending = True + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + + +def test_audio_input_generation_bails_when_unload_pending(monkeypatch): + # The audio path takes _gen_lock but must also skip the outgoing model mid-switch. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") + ) + o._unload_pending = True + + out = list(o._generate_audio_input_inner(audio_array = [0.0, 0.1])) + + assert any("unloaded" in chunk.lower() for chunk in out) + # Lock released so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_audio_response_bails_when_unload_pending(monkeypatch): + # TTS (generate_audio_response) is blocking, so it RAISES rather than starting on the + # outgoing model mid-switch; it takes _gen_lock and must release it either way. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send audio generate mid-switch") + ) + o._unload_pending = True + + with pytest.raises(RuntimeError, match = "unload"): + o.generate_audio_response("hello") + + # Lock released so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +# ---------------------------------------------------------------------------- +# Preserve unload cancels across the queue handoff (drain_event) — items #1/#4. +# ---------------------------------------------------------------------------- + + +def test_worker_drain_skip_emits_cancelled_gen_done_when_draining(): + # The worker clears cancel_event at the start of every generate, so a cancel set + # while a generate is still queued would be lost when it is dequeued. drain_event + # is the durable signal: while it is set the worker skips the generate (emitting an + # immediate gen_done so the stream/mailbox drains) instead of running it. + import queue as _queue + + from core.inference.worker import _drain_skip_generate + + drain = threading.Event() + rq: _queue.Queue = _queue.Queue() + cmd = {"type": "generate", "request_id": "r1"} + + # Not draining -> run normally (do not skip, emit nothing). + assert _drain_skip_generate(cmd, rq, drain) is False + assert rq.empty() + # Missing event (older worker) -> also runs normally. + assert _drain_skip_generate(cmd, rq, None) is False + assert rq.empty() + + # Draining -> skip and emit a cancelled gen_done for this request_id. + drain.set() + assert _drain_skip_generate(cmd, rq, drain) is True + resp = rq.get_nowait() + assert resp["type"] == "gen_done" + assert resp["request_id"] == "r1" + assert resp["cancelled"] is True + + +def test_worker_generate_branches_check_drain_before_clearing_cancel(): + # Both worker command loops (MLX fast-path + GPU) must consult the drain skip + # before clearing cancel_event and running, so a queued generate can't clear an + # unload-initiated cancel and run the outgoing model to completion. Each loop + # checks the drain twice -- once before the clear and once after -- so a + # drain+cancel pair that lands in the window between them is still caught. + import inspect + + from core.inference import worker + + src = inspect.getsource(worker.run_inference_process) + assert src.count("_drain_skip_generate(cmd, resp_queue, drain_event)") == 4 + + +def test_worker_generate_rechecks_drain_after_clearing_cancel(): + # The exact interleaving item #3 describes: the drain check reads unset, then the + # parent sets drain+cancel for an unload, then the worker clears cancel_event + # (erasing that cancel). A second drain check *after* the clear catches it and + # skips the generate instead of running the outgoing model to completion. + import queue as _queue + + from core.inference.worker import _drain_skip_generate + + drain = threading.Event() + cancel = threading.Event() + rq: _queue.Queue = _queue.Queue() + cmd = {"type": "generate", "request_id": "r1"} + + # 1. Pre-clear drain check: not draining yet -> run (no skip, no emit). + assert _drain_skip_generate(cmd, rq, drain) is False + assert rq.empty() + + # 2. Parent starts an unload: sets drain, then cancel (orchestrator order). + drain.set() + cancel.set() + + # 3. Worker clears cancel at the start of the generate -- erasing the cancel. + cancel.clear() + assert not cancel.is_set() + + # 4. Post-clear drain re-check catches the erased cancel and skips. + assert _drain_skip_generate(cmd, rq, drain) is True + resp = rq.get_nowait() + assert resp["type"] == "gen_done" and resp["cancelled"] is True + + +def test_unload_sets_drain_event_during_switch_and_clears_after(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + + seen = {} + + def record_send(cmd): + # drain_event must be set for the whole unload round-trip so any generate the + # worker dequeues in this window is skipped, not run. + seen["drain_set"] = o._drain_event.is_set() + + monkeypatch.setattr(o, "_send_cmd", record_send) + + assert o.unload_model("m") is True + assert seen.get("drain_set") is True + # Cleared on exit so a later generation (e.g. unloading a non-active model, or a + # reused subprocess) is not wrongly skipped. + assert o._drain_event.is_set() is False + + +def test_unload_clears_drain_event_even_on_wedged_teardown(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send when wedged")) + + # A wedged worker never releases _gen_lock; unload tears the subprocess down. The + # real teardown nulls _drain_event, so emulate that so the finally exercises its guard. + def fake_shutdown(timeout = 5): + o._drain_event = None + + monkeypatch.setattr(o, "_shutdown_subprocess", fake_shutdown) + o._gen_lock.acquire() + + assert o.unload_model("m") is True # must not raise in the drain_event clear + + +# ---------------------------------------------------------------------------- +# Recheck the active model after the lock wait — items #2/#3. +# ---------------------------------------------------------------------------- + + +def test_generation_rechecks_model_after_lock_wait(monkeypatch): + # A request passes the pre-lock active-model check, then blocks on _gen_lock while + # an unload clears/swaps the model. Even if _unload_pending was already reset (the + # unload's finally runs after the lock release), the under-lock active-model recheck + # must make it bail instead of sending a generate to the wrong/unloaded backend. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate on a swapped/unloaded model") + ) + + reached_lock = threading.Event() + # _wait_dispatcher_idle runs after the pre-lock check and before acquiring the lock; + # signalling here means the generator captured the model and is about to block. + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + + o.active_model_name = "m" + o._unload_pending = False + o._gen_lock.acquire() # stand in for an in-flight unload holding the lock + + out: list = [] + + def run(): + out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + + t = threading.Thread(target = run) + t.start() + assert reached_lock.wait(timeout = 5) + # Unload finished: model swapped, pending already cleared. Release the lock. + o.active_model_name = "other" + o._gen_lock.release() + t.join(timeout = 5) + + assert out and any("unloaded" in chunk.lower() for chunk in out) + + +def test_generation_rechecks_model_when_unloaded_to_none(monkeypatch): + # Same race, but the unload left no active model (a plain unload, not a switch). + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate after the model was unloaded") + ) + reached_lock = threading.Event() + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + + o.active_model_name = "m" + o._unload_pending = False + o._gen_lock.acquire() + + out: list = [] + t = threading.Thread( + target = lambda: out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + ) + t.start() + assert reached_lock.wait(timeout = 5) + o.active_model_name = None + o._gen_lock.release() + t.join(timeout = 5) + + assert out and any("unloaded" in chunk.lower() for chunk in out) + + +# ---------------------------------------------------------------------------- +# Don't unload a stale model name (worker's active-model fallback) — item #5. +# ---------------------------------------------------------------------------- + + +def test_unload_of_stale_name_does_not_touch_active_model(monkeypatch): + # If the named model isn't loaded (e.g. a concurrent load already swapped in a + # different one), unload must not send a command the worker would satisfy by + # unloading its *active* model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + ) + o.active_model_name = "current" + o.models = {"current": {}} + + assert o.unload_model("stale") is True + # The active model is left intact. + assert o.active_model_name == "current" + assert "current" in o.models + + +def test_unload_matches_active_model_case_insensitively(monkeypatch): + # active_model_name can differ in case from the raw model_path a client sends + # to /unload (the load path canonicalizes casing). The stale-name guard must + # match case-insensitively too; otherwise it no-ops the unload and leaves the + # model resident while reporting success. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + o.active_model_name = "unsloth/Qwen3-4B" + o.models = {"unsloth/Qwen3-4B": {}} + + # Client unloads with the casing it originally typed, before canonicalization. + assert o.unload_model("unsloth/qwen3-4b") is True + # The guard did not no-op: an unload for the canonical active model reached + # the worker (not the raw lowercase name, so the worker matches it directly). + assert {"type": "unload", "model_name": "unsloth/Qwen3-4B"} in sent + # Local state is cleared for the canonical name, not left stale. + assert o.active_model_name is None + assert o.models == {} + + +def test_unload_of_stale_name_still_no_ops_after_case_insensitive_match(monkeypatch): + # The case-insensitive match must only rescue the active model; a genuinely + # different model name (case-insensitively too) must still no-op so the + # worker's absent-name fallback can't tear down the active model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + ) + o.active_model_name = "unsloth/Qwen3-4B" + o.models = {"unsloth/Qwen3-4B": {}} + + assert o.unload_model("unsloth/Llama-3.1-8B") is True + assert o.active_model_name == "unsloth/Qwen3-4B" + assert "unsloth/Qwen3-4B" in o.models + + +def test_load_does_not_accumulate_stale_models_defeating_the_unload_guard(monkeypatch): + # A load always spawns a fresh subprocess holding only the new model, so + # self.models must mirror that instead of accumulating the previous model's name. + # Otherwise switching A -> B leaves 'A' in self.models, so a later unload('A') + # passes the "not in self.models" guard and the worker's absent-name fallback + # unloads the *active* model B. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + monkeypatch.setattr(orch_mod.time, "sleep", lambda *_a, **_k: None) + + def _load(name): + monkeypatch.setattr( + o, + "_wait_response", + lambda expected, timeout = 300.0: { + "type": "loaded", + "success": True, + "model_info": {"identifier": name, "display_name": name}, + }, + ) + assert o.load_model(types.SimpleNamespace(identifier = name, gguf_variant = None)) is True + + _load("modelA") + _load("modelB") # switch to B without unloading A first + + # self.models mirrors the single live model; the swapped-out name is gone. + assert o.active_model_name == "modelB" + assert set(o.models) == {"modelB"} + + # A stale unload of the swapped-out model must not reach the worker (whose + # absent-name fallback would unload the active model B). + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("stale unload reached the worker")) + assert o.unload_model("modelA") is True + assert o.active_model_name == "modelB" + assert "modelB" in o.models + + +def test_unload_route_serializes_with_loads_via_lifecycle_gate(monkeypatch): + # Item #5: /unload must hold the same lifecycle gate as /load so a concurrent load + # can't swap the backend subprocess/queues mid-unload. + import asyncio + + import routes.inference as inference_route + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + class _Llama: + is_active = False + is_loaded = False + model_identifier = None + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + + unloaded: list = [] + + class _Backend: + active_model_name = "m" + models = {"m": {}} + + def unload_model(self, name): + unloaded.append(name) + return True + + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: _Backend()) + + async def scenario(): + # Hold the real gate, exactly as an in-flight /load would. + assert kw._lifecycle_lock.acquire(blocking = False) + try: + task = asyncio.ensure_future( + inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + ) + # Yield to the loop repeatedly: the route must stay blocked on the gate. + for _ in range(10): + await asyncio.sleep(0.01) + assert unloaded == [], "unload ran while the lifecycle gate was held" + assert not task.done() + finally: + kw._lifecycle_lock.release() + resp = await task + assert resp.status == "unloaded" + assert unloaded == ["m"] + + asyncio.run(scenario()) + + +# ---------------------------------------------------------------------------- +# Cancel an in-flight load OFF the lifecycle gate (Stop-loading regression). +# /load holds the gate for the whole load, so a gated /unload could never +# interrupt it; cancel_load only tears the loading subprocess down. +# ---------------------------------------------------------------------------- + + +def test_cancel_load_terminates_loading_subprocess_and_sends_no_command(monkeypatch): + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + ) + + assert o.cancel_load("m") is True + assert shutdown, "must tear the loading subprocess down" + assert "m" not in o.loading_models + assert o.active_model_name is None + # A name that is not loading -> no-op, returns False so the caller takes the gate. + assert o.cancel_load("other") is False + + +def test_cancel_load_matches_loading_model_case_insensitively(monkeypatch): + o = _bare_orchestrator() + o.loading_models = {"unsloth/Qwen3-4B"} + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) + + assert o.cancel_load("unsloth/qwen3-4b") is True + assert o.loading_models == set() + + +def test_unload_model_cancels_a_loading_model_via_cancel_load(monkeypatch): + # unload_model still cancels an in-flight load (shared logic with cancel_load). + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send a command to cancel a load") + ) + + assert o.unload_model("m") is True + assert shutdown + assert "m" not in o.loading_models + + +def test_unload_route_cancels_in_flight_load_without_waiting_on_gate(monkeypatch): + # The regression: /unload wrapped its whole body in the lifecycle gate, so the + # Stop-loading button (cancelLoading -> /unload) could not interrupt a safetensors + # load that holds the gate for its full duration. The cancel must run off-gate. + import asyncio + + import routes.inference as inference_route + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + class _Llama: + is_active = False + is_loaded = False + model_identifier = None + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + + cancelled: list = [] + + class _Backend: + active_model_name = None + models: dict = {} + + def get_loading_model(self): + return "m" + + def cancel_load(self, name): + cancelled.append(name) + return True + + def unload_model(self, name): + pytest.fail("must not take the gated unload path for a still-loading model") + + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: _Backend()) + + async def scenario(): + # Hold the real gate, exactly as an in-flight /load would. + assert kw._lifecycle_lock.acquire(blocking = False) + try: + # Even with the gate held, the loading-cancel must go through. + resp = await inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + assert resp.status == "unloaded" + assert cancelled == ["m"] + finally: + kw._lifecycle_lock.release() + + asyncio.run(scenario()) + + +# ---------------------------------------------------------------------------- +# A dispatched (compare-mode) request that races an unload must not orphan its +# mailbox after _wait_dispatcher_idle stops the dispatcher. +# ---------------------------------------------------------------------------- + + +def test_dispatched_bails_when_unload_flips_before_mailbox_registration(monkeypatch): + # The request passes the pre-work _unload_pending check, then an unload sets + # _unload_pending and _wait_dispatcher_idle stops the dispatcher (mailboxes empty) + # before this request registers its mailbox. The recheck under _mailbox_lock must + # make it bail, or the worker's skipped-generate reply has nothing to route it and + # the compare stream hangs on an orphaned mailbox. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + # Flip the unload flag after the pre-work check (626) but before mailbox + # registration -- exactly the window _wait_dispatcher_idle exploits. + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +# ---------------------------------------------------------------------------- +# Dispatched path: bail when a cleared-pending unload swapped the model or +# tore the dispatcher down during the pre-registration window -- item #2. +# ---------------------------------------------------------------------------- + + +class _AliveDispatcher: + """Stand-in dispatcher thread that reports itself alive.""" + + def is_alive(self): + return True + + +def test_dispatched_bails_when_model_swapped_before_mailbox_registration(monkeypatch): + # The request passes the pre-work checks, then a full unload+reload completes + # (clearing _unload_pending) before this request registers its mailbox. The + # under-lock recheck must notice active_model_name changed and bail, instead of + # sending a generate that lands on the swapped-in model. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + # Swap the active model after the pre-work check but before registration, + # with _unload_pending already back to False (the unload finally ran). + def swap(*a, **k): + o.active_model_name = "other" + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", swap) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate on the swapped-in model") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration(monkeypatch): + # Same window, but the unload was a same-model reload so active_model_name is + # unchanged; the give-away is that the dispatcher was stopped. Registering a + # mailbox with no dispatcher to route the reply would hang the compare stream. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + def stop_dispatcher(*a, **k): + o._dispatcher_thread = None # unload's _stop_dispatcher cleared it + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", stop_dispatcher) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate with the dispatcher stopped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +def test_dispatched_happy_path_registers_and_sends(monkeypatch): + # Guard against a false bail: with the model unchanged and the dispatcher alive, + # the recheck must let the generate through (register a mailbox and send). + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + monkeypatch.setattr( + o, "_build_generate_cmd", lambda *a, **k: {"type": "generate", "request_id": "r1"} + ) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + + # Feed one gen_done so the consumer returns promptly. + def fake_consume(read_mailbox, drainer, **k): + mbox = o._mailboxes.get("r1") + if mbox is not None: + mbox.put({"type": "gen_done", "request_id": "r1"}) + yield "" + + monkeypatch.setattr(o, "_consume_token_stream", fake_consume) + + list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert sent, "happy path must send the generate command" + assert o._mailboxes == {}, "mailbox popped in finally" + + +# ---------------------------------------------------------------------------- +# load_model observes a cancel that discarded its loading marker -- item #4. +# ---------------------------------------------------------------------------- + + +def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch): + # Stop-loading during GPU placement discards the loading marker (cancel_load) with + # no child yet to kill. load_model must observe the removal and not spawn a worker + # that loads the model after /unload already reported it unloaded. + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = set() + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + monkeypatch.setattr( + o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn a worker after a cancel") + ) + + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + + # cancel_load discards the marker while we resolve GPU placement. + def cancel_during_gpu(gpu_ids, **k): + o.loading_models.discard("m") + return ([0], "sel") + + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", cancel_during_gpu) + + class _Cfg: + identifier = "m" + + ok = o.load_model(_Cfg()) + + assert ok is False + assert o.active_model_name is None + assert o.models == {} + + +def test_load_model_proceeds_when_not_cancelled(monkeypatch): + # Guard against a false abort: an uncancelled load keeps its marker and spawns. + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = set() + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + + spawned = [] + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: spawned.append(cfg)) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: {"success": True, "model_info": {"identifier": "m"}}, + ) + + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + + class _Cfg: + identifier = "m" + + ok = o.load_model(_Cfg()) + + assert ok is True + assert spawned, "uncancelled load must spawn a worker" + assert o.active_model_name == "m" + + +def test_load_model_aborts_when_cancelled_during_spawn(monkeypatch): + # Stop-loading can land AFTER the pre-spawn marker recheck but while + # _spawn_subprocess is still creating the queues/process, so cancel_load's + # _shutdown_subprocess finds _proc not yet alive and no-ops. load_model must + # recheck the marker once the child exists and tear the orphaned worker down, + # instead of waiting for "loaded" and publishing a model /unload already + # reported as unloaded (a live subprocess nothing later reaps). + import types + + from utils import transformers_version as tv + + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = {"m"} + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + + # The cancel lands during the spawn window: cancel_load already discarded the + # marker, but its teardown no-oped because _proc was not alive yet. + def spawn_then_cancel(cfg): + o.loading_models.discard("m") + + monkeypatch.setattr(o, "_spawn_subprocess", spawn_then_cancel) + + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: pytest.fail( + "must not wait for 'loaded' after a cancel during spawn" + ), + ) + + ok = o.load_model(types.SimpleNamespace(identifier = "m", gguf_variant = None)) + + assert ok is False + assert shutdown, "must tear the orphaned worker down" + assert o.active_model_name is None + assert o.models == {} + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# /unload cancels a still-loading GGUF off the lifecycle gate -- item #1. +# ---------------------------------------------------------------------------- + + +def test_unload_cancels_loading_gguf_off_gate(monkeypatch): + # A still-loading GGUF (is_active, not is_loaded) must be cancelled off the gate: + # /load holds the lifecycle gate for the whole load, so a gated unload would wait + # it out. Assert the gate is never entered and unload_model() runs. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = False + model_identifier = "gguf-model" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None # no Unsloth load in flight -> Unsloth fast path skipped + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-model") + resp = _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert getattr(resp, "status", None) == "unloaded" + assert llama.unloaded is True, "must cancel the loading GGUF via unload_model()" + assert gate_entered["v"] is False, "must handle the loading GGUF off the lifecycle gate" + + +def test_unload_loaded_gguf_still_uses_gate(monkeypatch): + # Guard: an already-loaded GGUF (is_loaded True) is NOT caught by the off-gate + # fast path; it goes through the gate as before. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = True + model_identifier = "gguf-model" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(ri, "is_registered_native_path_label", lambda a, b: False) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-model") + resp = _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert getattr(resp, "status", None) == "unloaded" + assert llama.unloaded is True + assert gate_entered["v"] is True, "loaded GGUF unload must still take the gate" + + +def test_unload_of_mismatched_loading_gguf_skips_off_gate_fast_path(monkeypatch): + # A still-loading GGUF X (is_active, not is_loaded) must NOT be torn down by the + # off-gate fast path when /unload names a DIFFERENT model Y. The single llama-server + # can only load one GGUF at a time, so this fast path is "stop loading THIS model"; + # without a target check it fires for any in-flight GGUF and would abort an unrelated + # load (e.g. a second tab unloading Y kills the load of X). A mismatched target must + # fall through to the lifecycle gate (where, in production, it waits out X's /load and + # then no-ops) instead of taking the off-gate teardown. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = False + model_identifier = "gguf-X" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None # no Unsloth load in flight -> Unsloth fast path skipped + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(ri, "is_registered_native_path_label", lambda a, b: False) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-Y") # different from the loading model X + _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert gate_entered["v"] is True, ( + "a mismatched-target unload must not use the off-gate GGUF fast path; " + "it would cancel the wrong in-flight load" + ) + + +# ---------------------------------------------------------------------------- +# cancel_load clears its loading marker BEFORE tearing the subprocess down, so a +# racing off-gate load_model observes the cancel during the shutdown window. +# ---------------------------------------------------------------------------- + + +def test_cancel_load_clears_marker_before_shutdown(monkeypatch): + # cancel_load runs off the lifecycle gate, concurrently with a load_model that + # rechecks the loading marker before each spawn to observe the cancel. + # _shutdown_subprocess can block (tearing a live child down / joining the compare + # dispatcher), so discarding the marker only AFTER it leaves a long window in which + # that load_model reads the marker still set, passes its pre-spawn recheck, and + # spawns + loads the model after /unload already reported it cancelled. The marker + # (and local state) must be cleared before the teardown. + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = "m" + o.models = {"m": {}} + + at_shutdown = {} + + def record_shutdown(timeout = 5): + at_shutdown["marker_present"] = "m" in o.loading_models + at_shutdown["active"] = o.active_model_name + at_shutdown["models"] = dict(o.models) + + monkeypatch.setattr(o, "_shutdown_subprocess", record_shutdown) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + ) + + assert o.cancel_load("m") is True + assert at_shutdown.get("marker_present") is False, ( + "the loading marker must be cleared before _shutdown_subprocess so a concurrent " + "load_model pre-spawn recheck observes the cancel during the shutdown window" + ) + assert at_shutdown.get("active") is None + assert at_shutdown.get("models") == {} + assert "m" not in o.loading_models + assert o.active_model_name is None + assert o.models == {} + + +def test_cancel_load_reclears_state_when_racing_load_repopulates_during_teardown(monkeypatch): + # cancel_load (off the lifecycle gate) can race a load_model whose worker already + # queued its successful "loaded" reply. cancel_load discards the loading marker and + # clears the local mirrors, then tears the subprocess down; but the still-running + # load_model thread can consume that "loaded" DURING the teardown window and repopulate + # active_model_name/models. _shutdown_subprocess nulls the queues but never touches those + # mirrors, so without a second clear /unload reports success while the backend keeps + # advertising a model whose worker was just killed. cancel_load must re-clear after the + # teardown so no phantom loaded model survives. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + o._proc = None # no prior subprocess -> load_model goes straight to the spawn loop + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + + parked = threading.Event() # load_model is parked in _wait_response("loaded") + release_loaded = threading.Event() # cancel_load lets the load consume "loaded" + load_done = threading.Event() + + def blocking_wait_response(expected, timeout = 300.0): + parked.set() + assert release_loaded.wait(timeout = 5) + return { + "type": "loaded", + "success": True, + "model_info": {"identifier": "m", "display_name": "m"}, + } + + monkeypatch.setattr(o, "_wait_response", blocking_wait_response) + + load_result: dict = {} + + def run_load(): + try: + load_result["ok"] = o.load_model( + types.SimpleNamespace(identifier = "m", gguf_variant = None) + ) + except Exception as exc: # noqa: BLE001 + load_result["exc"] = exc + finally: + load_done.set() + + loader = threading.Thread(target = run_load) + loader.start() + assert parked.wait(timeout = 5), "load_model must reach _wait_response" + + # The teardown IS the window in which the racing load repopulates the mirrors: the + # marker is already discarded here, so release the load and wait for it to finish + # repopulating, mirroring the 0.5s cancel-settle inside the real _shutdown_subprocess. + def racing_shutdown(timeout = 0.5): + release_loaded.set() + assert load_done.wait(timeout = 5), "the racing load must repopulate during teardown" + + monkeypatch.setattr(o, "_shutdown_subprocess", racing_shutdown) + + assert o.cancel_load("m") is True + loader.join(timeout = 5) + + # Fail-without: load_model set active_model_name/models during racing_shutdown and + # cancel_load left them set, so the backend advertises a model whose worker was killed. + assert o.active_model_name is None, "cancel_load must not leave a repopulated active model" + assert o.models == {}, "cancel_load must not leave a repopulated models mirror" + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# A dispatched (compare-mode) request that starts the dispatcher and then bails on +# a racing unload must stop the dispatcher it started, or that orphaned dispatcher +# steals the worker's "unloaded" reply and hangs unload_model on its 300s timeout. +# ---------------------------------------------------------------------------- + + +def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch): + # The request passes the pre-work _unload_pending check and starts the dispatcher + # (none was running), then an unload sets _unload_pending so the under-lock recheck + # bails. The just-started dispatcher, left running with no mailboxes, competes with + # unload_model()'s _wait_response for the worker's "unloaded" reply off the shared + # resp_queue and drops it as unroutable, hanging the unload until its 300s timeout. + # The bail must stop the dispatcher it started. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = None # none running -> this call starts it + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + + started = {"v": False} + stopped = {"v": False} + + def fake_start(): + started["v"] = True + o._dispatcher_thread = _AliveDispatcher() + + def fake_stop(): + stopped["v"] = True + o._dispatcher_thread = None + + monkeypatch.setattr(o, "_start_dispatcher", fake_start) + monkeypatch.setattr(o, "_stop_dispatcher", fake_stop) + + # An unload flips _unload_pending after the pre-work check but before registration. + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert started["v"], "this call started the dispatcher" + assert stopped["v"], "the bail must stop the dispatcher it started (no other mailboxes)" + assert o._mailboxes == {} + + +def test_dispatched_bail_keeps_dispatcher_with_other_active_mailbox(monkeypatch): + # Guard against over-stopping: if another compare request registered a mailbox on the + # dispatcher this call started, the bail must NOT stop it, or that request's token + # routing dies mid-stream. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: setattr(o, "_dispatcher_thread", _AliveDispatcher()) + ) + monkeypatch.setattr( + o, + "_stop_dispatcher", + lambda: pytest.fail("must not stop a dispatcher another compare request is using"), + ) + + # A concurrent compare request registers its mailbox, then an unload flips the flag. + def flip(*a, **k): + o._mailboxes["other"] = object() + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert set(o._mailboxes) == {"other"}, "the other request's mailbox is untouched" + + +def test_dispatched_bail_keeps_preexisting_dispatcher(monkeypatch): + # Guard: if the dispatcher was already running before this request (an earlier compare + # request started it), a bail must not stop it even with no mailboxes now -- this + # request did not start it and another may re-use it. Only the call that starts an + # otherwise-idle dispatcher during the race is responsible for stopping it. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() # already running + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + monkeypatch.setattr( + o, "_stop_dispatcher", lambda: pytest.fail("must not stop a pre-existing dispatcher") + ) + + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + + +# ---------------------------------------------------------------------------- +# load_model rechecks the loading marker AFTER _wait_response("loaded") and +# BEFORE publishing -- item #6. cancel_load's post-teardown re-clear only wipes a +# repopulation that lands during its shutdown; a publish that lands after +# cancel_load returns survives it, so the recheck must abort the publish itself. +# ---------------------------------------------------------------------------- + + +def test_load_model_aborts_publish_when_cancelled_after_wait_response(monkeypatch): + # cancel_load (off the lifecycle gate) discards the loading marker BEFORE its teardown + # and re-clears the mirrors AFTER it. A racing load_model can consume its worker's + # already-queued "loaded" reply and reach the publish block only AFTER cancel_load has + # fully returned -- so cancel_load's post-teardown re-clear cannot undo that publish. + # Without a marker recheck between _wait_response("loaded") and the publish, load_model + # advertises active_model_name/models for a model /unload already reported cancelled, + # over a subprocess cancel_load just killed. The recheck must observe the discarded + # marker and abort the publish. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + o._proc = None # no prior subprocess -> load_model goes straight to the spawn loop + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + # cancel_load tears the worker down; a no-op keeps the test off real subprocesses. + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) + + parked = threading.Event() # load_model reached _wait_response("loaded") + cancel_done = threading.Event() # cancel_load fully returned (marker discarded + re-clear) + load_done = threading.Event() + + def blocking_wait_response(expected, timeout = 300.0): + parked.set() + # Do not consume "loaded" until cancel_load has fully returned, so the publish + # would land AFTER cancel_load's post-teardown re-clear -- the window the + # re-clear alone cannot cover. + assert cancel_done.wait(timeout = 5) + return { + "type": "loaded", + "success": True, + "model_info": {"identifier": "m", "display_name": "m"}, + } + + monkeypatch.setattr(o, "_wait_response", blocking_wait_response) + + load_result: dict = {} + + def run_load(): + try: + load_result["ok"] = o.load_model( + types.SimpleNamespace(identifier = "m", gguf_variant = None) + ) + except Exception as exc: # noqa: BLE001 + load_result["exc"] = exc + finally: + load_done.set() + + loader = threading.Thread(target = run_load) + loader.start() + assert parked.wait(timeout = 5), "load_model must reach _wait_response" + + # cancel_load runs to completion while the load is parked: it discards the marker and + # re-clears the mirrors (post-teardown), then returns. Only then let the load consume + # "loaded" and attempt to publish. + assert o.cancel_load("m") is True + cancel_done.set() + + loader.join(timeout = 5) + assert load_done.is_set() + + # Fail-without: load_model published active_model_name/models for 'm' AFTER cancel_load + # returned, advertising a cancelled model over a killed subprocess. + assert load_result.get("ok") is False, "the cancelled load must not report success" + assert o.active_model_name is None, "must not publish a cancelled model's active name" + assert o.models == {}, "must not publish a cancelled model's mirror" + assert "m" not in o.loading_models From 46ab68306543e8e28d087f68b1cc0cb482e11a4a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 19:48:36 -0700 Subject: [PATCH 030/113] Studio: client-tool passthrough healing for safetensors and MLX (#6870) * Studio: client-tool passthrough healing for safetensors and MLX PR 6801 made response-side tool-call healing default-on for the client-tool passthrough, but only on the GGUF path: the passthrough branch in /v1/chat/completions is gated on using_gguf, and the safetensors section never reads payload.tools, so a client-tools request against a safetensors or MLX model silently dropped the tool schemas and returned prose with no tool_calls. Add the missing leg. When a non-GGUF model is loaded, the request declares client tools (or carries tool-role history), server-side tools are off, and the template supports tools, the route now: - renders the tools into the chat template for a single turn via the existing backend.generate_chat_response(..., tools=...) seam (worker templating already accepts role=tool and assistant.tool_calls messages, normalized with _openai_messages_for_passthrough); - non-streaming: promotes text-form calls with heal_openai_message, honors the opt-in nudge single retry (nudge_should_retry / nudge_messages), caps healed calls when parallel_tool_calls=false (covers the nudge retry too), and sets finish_reason=tool_calls with content null on a pure tool-call turn; - streaming: derives deltas from the worker's cumulative snapshots and feeds StreamToolCallHealer, emitting healed tool-call deltas and the correct finish chunk, guarded against repeated or shrinking snapshots. heal_gate semantics are identical to the GGUF passthrough: default on, auto_heal_tool_calls=false or UNSLOTH_DISABLE_TOOL_CALL_HEALING=1 relays verbatim, tool_choice narrows promotion, undeclared names stay text. MLX rides the same orchestrator seam, so both local backends gain the behavior. CompletionMessage.content becomes Optional so a promoted pure tool-call turn matches the OpenAI contract (content null when only tool_calls return). Adds tests/test_sf_client_tools_passthrough.py (22 cases: healing, gating, opt-outs, streaming deltas, tool-role history, dict-arguments history, forced tool_choice, parallel cap, usage, nudge on/off/double-failure, generator error hygiene, disconnect reset, empty output, MLX path). * Address review: tool_choice none, developer folding, retry fallback, monitor reply Four review follow-ups on the safetensors/MLX client-tool passthrough leg: - tool_choice="none" keeps the tool-history templating but no longer advertises the tools, so a forced final-answer turn is not prompted into emitting markup that the (correctly disabled) healer would relay as prose. Mirrors the GGUF passthrough where llama-server honors tool_choice itself. - OpenAI "developer" messages fold into a single leading system message via _set_or_prepend_system_message before templating; local templates reject the role and the fallback formatter drops it. - A nudge retry that fails or is cancelled after the original answer exists falls back to the first response instead of surfacing a 500, matching the GGUF nudge path. - The API monitor records the healed tool call summary instead of the raw markup on a promoted turn. Adds four regression tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: forced tool_choice templating, content-part flattening, stream monitor parity - A forced tool_choice function is now the only schema rendered into the local template, so the advertised tools and the healer allowlist can no longer disagree (llama-server enforces tool_choice itself on the GGUF path). - Content-part lists are flattened to their text parts before templating. Remote image URLs are not decodable locally, so such requests reached this path with part lists that raise inside apply_chat_template on text-only templates; the plain non-GGUF path has always flattened them. - The streaming monitor entry is now fed from the healed events the client actually receives, recording promoted calls as the [tool_calls] summary the non-streaming path records. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: gate passthrough on the engaged server path, deserialize templated arguments - The client-tools gate now keys on _sf_use_tools (whether the server-side tool path actually claimed the request) instead of the raw mcp_enabled flag: with an empty MCP registry or a CLI --disable-tools policy, a client that sets mcp_enabled while declaring its own tools fell through to plain generation with the tools silently dropped. The GGUF passthrough gate has no mcp_enabled clause either. - New _structured_tool_history_for_local_template deserializes assistant tool_calls[].function.arguments JSON strings into mappings for the templated copy only: spec-compliant clients send strings, but local chat templates iterate arguments as a mapping or raise on strings, which crashed or misrendered multi-turn tool history. The HTTP response and the GGUF wire shape keep strings. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments and docstrings in the client-tools passthrough * Report first-attempt usage when a nudge retry is discarded When nudge_should_retry fires but the retry produces no healable tool call (or raises), the first response is still delivered to the client. The retry's generate() had already overwritten stats_holder, so _monitor_usage recorded the unseen retry's token counts against the request instead of the first attempt that was actually returned. Capture the first attempt's stats before the retry and restore them on both the no-heal and exception paths so the monitor reports the usage of the response the caller received. * Do not promote buffered tool markup when a stream is cancelled The streaming client-tool heal path breaks out of the token loop when cancel_event is set (the registry "Stop" path), but then still fell through to healer.finalize(), which heals incomplete tool markup at EOF (allow_incomplete) and emits a tool_calls delta plus finish_reason=tool_calls. Because the Stop request only sets the event and leaves the SSE socket open, the client received that promoted call and executed a tool the user had just cancelled. The disconnect path already returns before finalize; guard finalize and the finish_reason on cancel_event too, so a cancelled stream ends with finish_reason=stop and no tool call. Adds a regression test driving a Stop mid-emission with buffered markup. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim comments in the client-tools passthrough * Trim client-tools passthrough comments further * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/models/inference.py | 3 +- studio/backend/routes/inference.py | 329 +++++++- .../tests/test_sf_client_tools_passthrough.py | 786 ++++++++++++++++++ 3 files changed, 1102 insertions(+), 16 deletions(-) create mode 100644 studio/backend/tests/test_sf_client_tools_passthrough.py diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 31c100dbec..1e7770a7f8 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1146,7 +1146,8 @@ class CompletionMessage(BaseModel): """The assistant's complete response message.""" role: Literal["assistant"] = "assistant" - content: str + # ``None`` on a pure tool-call turn (OpenAI content=null); string otherwise. + content: Optional[str] = None refusal: Optional[str] = None reasoning_content: Optional[str] = None tool_calls: Optional[list[dict]] = None diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e31c03f7f8..3f6f6a2a78 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -625,6 +625,66 @@ def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: ) +def _chat_tool_calls_chunk(completion_id, created, model_name, tool_calls) -> str: + """Delta chunk carrying OpenAI tool-call deltas (sibling of ``_chat_content_chunk``).""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(tool_calls = tool_calls), + finish_reason = None, + ) + + +def _sf_heal_events_to_sse( + events, + completion_id, + created, + model_name, + state, + parallel_tool_calls, + monitor_id = None, +): + """Serialize ``StreamToolCallHealer`` events into chat SSE lines. + + ``state["idx"]`` tracks the call index across ``feed``/``finalize``; + ``parallel_tool_calls is False`` caps promotion to one call (GGUF parity). + The monitor is fed from the same events the client receives, never the + healed-away markup.""" + lines = [] + for kind, value in events: + if kind == "text": + if value: + lines.append(_chat_content_chunk(completion_id, created, model_name, value)) + api_monitor.append_reply(monitor_id, value) + continue + if parallel_tool_calls is False and state["idx"] >= 1: + continue + lines.append( + _chat_tool_calls_chunk( + completion_id, + created, + model_name, + [ + { + "index": state["idx"], + "id": value["id"], + "type": "function", + "function": value["function"], + } + ], + ) + ) + _fn = value.get("function") or {} + api_monitor.append_reply( + monitor_id, + ("[tool_calls] " if state["idx"] == 0 else "; ") + + f"{_fn.get('name', '')}({_fn.get('arguments', '')})", + ) + state["idx"] += 1 + return lines + + def _rewrite_cmpl_id(raw: bytes) -> bytes: """Rewrite llama-server's chat-style ``chatcmpl-`` ids to the ``cmpl-`` prefix OpenAI's legacy /v1/completions use. Anchored on the ``"id":`` key @@ -7190,25 +7250,87 @@ async def openai_chat_completions( if payload.preserve_thinking is not None: gen_kwargs["preserve_thinking"] = payload.preserve_thinking + # ── Client-tool passthrough (safetensors + MLX) ────────────── + # Client tools (or tool-result history) without server-side tools: render + # tools into the template, generate one turn, heal text-form calls (#6801). + # supports_tools=False falls through to plain relay (GGUF gate parity). + _sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages) + # Gate on _sf_use_tools (did the server-side path claim the request?), not + # raw mcp_enabled: an empty MCP registry must not silently drop client tools. + _sf_client_tools = ( + not _effective_enable_tools(payload) + and not _sf_use_tools + and image is None + and not _sf_is_gptoss + and _sf_features.get("supports_tools", False) + and ((payload.tools and len(payload.tools) > 0) or _sf_has_tool_msgs) + ) + _sf_heal = ( + heal_gate(payload.auto_heal_tool_calls, payload.tools, payload.tool_choice) + if _sf_client_tools + else None + ) + if _sf_client_tools: + # Re-derive from payload.messages so tool_calls / role="tool" history + # survives templating; fold system/developer into one leading system + # message (templates reject "developer") and clear prompt to avoid a dup. + gen_kwargs["messages"] = _set_or_prepend_system_message( + _structured_tool_history_for_local_template( + _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)) + ), + system_prompt, + ) + gen_kwargs["system_prompt"] = "" + # tool_choice="none": keep history templating but advertise no tools + # (heal_gate is off, markup would relay as prose). A forced function + # narrows templating to that one schema. Both mirror the GGUF path, + # where llama-server honors tool_choice itself. + _sf_tc = payload.tool_choice + _sf_forced = None + if isinstance(_sf_tc, dict) and isinstance(_sf_tc.get("function"), dict): + _sf_forced = _sf_tc["function"].get("name") + if _sf_tc == "none": + gen_kwargs["tools"] = None + elif isinstance(_sf_forced, str): + gen_kwargs["tools"] = [ + t + for t in payload.tools or [] + if isinstance(t, dict) + and isinstance(t.get("function"), dict) + and t["function"].get("name") == _sf_forced + ] or None + else: + gen_kwargs["tools"] = payload.tools + # Request-scoped usage/timings receptacle (filled at gen_done). stats_holder: dict = {} if payload.use_adapter is not None: - def generate(): + def generate(messages_override = None): + kw = ( + gen_kwargs + if messages_override is None + else {**gen_kwargs, "messages": messages_override} + ) return backend.generate_with_adapter_control( use_adapter = payload.use_adapter, cancel_event = cancel_event, stats_holder = stats_holder, - **gen_kwargs, + **kw, ) else: - def generate(): + def generate(messages_override = None): + kw = ( + gen_kwargs + if messages_override is None + else {**gen_kwargs, "messages": messages_override} + ) return backend.generate_chat_response( cancel_event = cancel_event, stats_holder = stats_holder, - **gen_kwargs, + **kw, ) # ── Streaming response ──────────────────────────────────────── @@ -7224,6 +7346,11 @@ async def openai_chat_completions( try: yield _chat_role_chunk(completion_id, created, model_name) + # Client-tool passthrough: heal text-form calls on the fly + # (None => relay verbatim). + healer = StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None + heal_state = {"idx": 0} + prev_text = "" # Split prefilled into reasoning_content deltas (GGUF parity); single turn, serves MLX. reasoning_extractor = _new_sf_reasoning_extractor() @@ -7255,22 +7382,76 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue + # Split prefilled reasoning first (GGUF/MLX parity), + # then route only the visible text through the client-tool + # healer so tool markup inside a reasoning block is not promoted. reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) if reasoning_delta: yield _chat_reasoning_chunk( completion_id, created, model_name, reasoning_delta ) if visible_delta: - api_monitor.append_reply(monitor_id, visible_delta) - yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + if healer is None: + # Monitor mirrors the verbatim relay; with healing on, + # _sf_heal_events_to_sse records the healed events instead. + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk( + completion_id, created, model_name, visible_delta + ) + else: + for line in _sf_heal_events_to_sse( + healer.feed(visible_delta), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line final_reasoning, final_visible = reasoning_extractor.finish() if final_reasoning: yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning) if final_visible: - api_monitor.append_reply(monitor_id, final_visible) - yield _chat_content_chunk(completion_id, created, model_name, final_visible) - yield _chat_final_chunk(completion_id, created, model_name, "stop") + if healer is None: + api_monitor.append_reply(monitor_id, final_visible) + yield _chat_content_chunk(completion_id, created, model_name, final_visible) + else: + for line in _sf_heal_events_to_sse( + healer.feed(final_visible), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line + + # A cancelled stream must not promote buffered-but-incomplete + # markup: finalize()'s allow_incomplete heal would execute a tool + # the user just cancelled. Disconnect returns earlier; "Stop" only + # sets cancel_event, so guard on it here too. + _cancelled = cancel_event.is_set() + if healer is not None and not _cancelled: + for line in _sf_heal_events_to_sse( + healer.finalize(), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line + + _finish = ( + "tool_calls" + if (healer is not None and not _cancelled and healer.healed) + else "stop" + ) + yield _chat_final_chunk(completion_id, created, model_name, _finish) # Usage chunk (choices=[], usage set), same shape as the # GGUF path so the speed popover works for MLX too. # Request-scoped holder, so concurrent streams cannot @@ -7332,27 +7513,96 @@ async def openai_chat_completions( for token in generate(): full_text = token - # Split prefilled reasoning (GGUF parity); also covers MLX via the shared generate(). + # Split prefilled reasoning (GGUF parity); also covers MLX via + # the shared generate(). Client-tool healing then runs on the visible + # text so tool markup inside a reasoning block is never promoted. _reasoning_text, _visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _sf_parse_think, reasoning_prefilled = _sf_reasoning_prefilled, ) - _plain_msg_kwargs = {"content": _visible_text} + # Client-tool passthrough: promote text-form calls; opt-in single + # nudge retry on unparseable tool markup. + _msg = {"role": "assistant", "content": _visible_text} if _reasoning_text: - _plain_msg_kwargs["reasoning_content"] = _reasoning_text + _msg["reasoning_content"] = _reasoning_text + _finish = "stop" + if _sf_heal: + if heal_openai_message(_msg, _sf_heal, payload.tools): + _finish = "tool_calls" + elif nudge_enabled(payload.nudge_tool_calls): + _data = { + "choices": [{"message": {"role": "assistant", "content": _visible_text}}] + } + if nudge_should_retry(_data, _sf_heal, payload.tools): + # A failed retry must not 500 the request; keep the first + # response (GGUF nudge parity). The retry's generate() + # overwrites stats_holder, so save the first attempt's stats + # and restore them if the retry is discarded. + _first_stats = stats_holder.get("stats") + try: + retry_text = "" + for token in generate( + [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] + ): + retry_text = token + # Re-split reasoning on the retry so its visible text is + # what heals into a call (and reaches the monitor). + _retry_reasoning, _retry_visible = _extract_responses_reasoning( + retry_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + retry_msg = {"role": "assistant", "content": _retry_visible} + if _retry_reasoning: + retry_msg["reasoning_content"] = _retry_reasoning + if heal_openai_message(retry_msg, _sf_heal, payload.tools): + _visible_text, _msg, _finish = ( + _retry_visible, + retry_msg, + "tool_calls", + ) + else: + # Retry produced no healable call -> first response wins. + stats_holder["stats"] = _first_stats + except Exception as retry_exc: + logger.debug( + "Nudge retry failed; keeping first response: %s", retry_exc + ) + stats_holder["stats"] = _first_stats + # parallel_tool_calls=false: cap to one call (GGUF parity). + if payload.parallel_tool_calls is False: + _tcs = _msg.get("tool_calls") + if isinstance(_tcs, list) and len(_tcs) > 1: + _msg["tool_calls"] = _tcs[:1] + response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(**_plain_msg_kwargs), - finish_reason = "stop", + message = CompletionMessage( + content = _msg["content"], + reasoning_content = _msg.get("reasoning_content"), + tool_calls = _msg.get("tool_calls"), + ), + finish_reason = _finish, ) ], ) - api_monitor.set_reply(monitor_id, _visible_text) + _monitor_reply = _msg.get("content") or "" + if _finish == "tool_calls": + _tcs = _msg.get("tool_calls") or [] + _calls_text = "; ".join( + f"{(tc.get('function') or {}).get('name', '')}" + f"({(tc.get('function') or {}).get('arguments', '')})" + for tc in _tcs + ) + _monitor_reply = (_msg.get("content") or "") + ( + f"[tool_calls] {_calls_text}" if _calls_text else "" + ) + api_monitor.set_reply(monitor_id, _monitor_reply) _stats = stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) @@ -11189,6 +11439,55 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: return messages +def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict]: + """Flatten OpenAI content-part lists to plain strings. + + Local text templates take string content and raise on part lists (e.g. a + remote ``image_url`` that leaves ``image is None``): keep the text parts, + drop the rest, like the plain non-GGUF path. GGUF keeps the parts.""" + out = [] + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + text_parts = [ + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ] + msg = {**msg, "content": "\n".join(text_parts) if text_parts else ""} + out.append(msg) + return out + + +def _structured_tool_history_for_local_template(messages: list[dict]) -> list[dict]: + """Deserialize assistant ``tool_calls[].function.arguments`` JSON strings to + mappings for local templating. + + Clients send prior-turn arguments as JSON strings, but local templates take + mappings (some raise on strings). Only the internal messages copy is + rewritten; the HTTP response stays OpenAI-shaped and unparseable strings + are left untouched.""" + out = [] + for msg in messages: + tool_calls = msg.get("tool_calls") + if isinstance(tool_calls, list) and tool_calls: + new_calls = [] + for tc in tool_calls: + fn = tc.get("function") if isinstance(tc, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + try: + parsed = json.loads(args) + except ValueError: + parsed = None + if isinstance(parsed, dict): + tc = {**tc, "function": {**fn, "arguments": parsed}} + new_calls.append(tc) + msg = {**msg, "tool_calls": new_calls} + out.append(msg) + return out + + def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]: """Build llama-server messages for the standard GGUF chat path. diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py new file mode 100644 index 0000000000..01905b712c --- /dev/null +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -0,0 +1,786 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Client-tools passthrough healing for the safetensors/MLX backend. + +Parity for #6801: when a NON-GGUF model is loaded and the request declares its +own ``tools`` with server-side tools OFF, text-form tool calls are promoted back +into structured ``tool_calls`` (declared tools only) via the shared healer. MLX +rides the same orchestrator path, so a single scripted backend covers both. +""" + +import asyncio +import json +from types import SimpleNamespace + +from models.inference import ChatCompletionRequest, ChatMessage +from routes.inference import openai_chat_completions +from core.inference.api_monitor import ApiMonitor + + +LOOKUP_TOOL = { + "type": "function", + "function": { + "name": "lookup", + "description": "Look something up", + "parameters": { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + }, + }, +} +SEARCH_TOOL = { + "type": "function", + "function": { + "name": "search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, +} + +_CALL_XML = '{"name": "lookup", "arguments": {"q": "cats"}}' +_SEARCH_XML = '{"name": "search", "arguments": {"query": "dogs"}}' + + +class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + scope: dict = {} + + async def is_disconnected(self): + return False + + +class _ScriptedBackend: + """Non-GGUF backend: ``generate_chat_response`` replays scripted + CUMULATIVE snapshots. ``responder(messages, tools)`` returns the snapshot + list for one generation, so nudge tests can vary output across turns.""" + + active_model_name = "sf-model" + + def __init__( + self, + responder, + *, + stats = None, + ): + self.models = { + "sf-model": { + "chat_template_info": {"template": " chatml"}, + "context_length": 2048, + } + } + self._responder = responder + self._stats = stats + self.calls: list = [] + self.reset_count = 0 + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + snapshots = self._responder(messages, tools) + if stats_holder is not None and self._stats is not None: + stats_holder["stats"] = self._stats + for snap in snapshots: + yield snap + + def reset_generation_state(self): + self.reset_count += 1 + + +def _fixed(*snapshots): + """Responder that always replays the given cumulative snapshots.""" + return lambda messages, tools: list(snapshots) + + +def _llama_stub(): + return SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ) + + +def _install( + monkeypatch, + backend, + *, + supports_tools = True, +): + import routes.inference as inf + from state.tool_policy import reset_tool_policy + + reset_tool_policy() + monitor = ApiMonitor(max_entries = 8) + monkeypatch.setattr(inf, "api_monitor", monitor) + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _llama_stub()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: backend) + monkeypatch.setattr( + inf, + "_detect_safetensors_features", + lambda *a, **k: {"supports_tools": supports_tools}, + ) + return monitor + + +def _request(**kwargs): + base = dict(model = "default", messages = [ChatMessage(role = "user", content = "hi")]) + base.update(kwargs) + return ChatCompletionRequest(**base) + + +def _call(payload, monkeypatch, backend, **install_kwargs): + _install(monkeypatch, backend, **install_kwargs) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + return asyncio.run(_run()) + + +def _json_body(response): + return json.loads(response.body if hasattr(response, "body") else response.content) + + +def _collect_sse(response): + async def _run(): + return [c async for c in response.body_iterator] + + return asyncio.run(_run()) + + +def _sse_objects(chunks): + out = [] + for chunk in chunks: + if isinstance(chunk, bytes): + chunk = chunk.decode() + for line in str(chunk).splitlines(): + if line.startswith("data: "): + data = line.removeprefix("data: ") + if data != "[DONE]": + out.append(json.loads(data)) + return out + + +# ── Non-streaming ───────────────────────────────────────────────── + + +def test_xml_healed_to_tool_calls_non_streaming(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] is None + calls = choice["message"]["tool_calls"] + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "lookup" + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"} + # The client tools reached the generator (template injection). + assert backend.calls[0]["tools"] == [LOOKUP_TOOL] + + +def test_undeclared_call_stays_text(monkeypatch): + xml = '{"name": "other", "arguments": {}}' + backend = _ScriptedBackend(_fixed(xml)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == xml + + +def test_opt_out_relays_verbatim(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, auto_heal_tool_calls = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == _CALL_XML + + +def test_env_kill_switch_relays_verbatim(monkeypatch): + import core.inference.passthrough_healing as ph + + monkeypatch.setattr(ph, "_HEALING_DISABLED", True) + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == _CALL_XML + + +def test_no_tools_request_untouched(monkeypatch): + backend = _ScriptedBackend(_fixed("just a plain answer")) + payload = _request(stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + # No tools and no tool messages -> plain path, normal ChatCompletion. + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == "just a plain answer" + assert choice["message"].get("tool_calls") is None + + +def test_prose_around_call_retained(monkeypatch): + text = "Let me look:\n" + _CALL_XML + "\ndone" + backend = _ScriptedBackend(_fixed(text)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] == "Let me look:\n\ndone" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + + +def test_empty_output_is_valid_stop(monkeypatch): + backend = _ScriptedBackend(_fixed("")) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] in ("", None) + assert choice["message"].get("tool_calls") is None + + +def test_tool_role_follow_up_turn_preserves_history(monkeypatch): + backend = _ScriptedBackend(_fixed("The weather is sunny.")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "weather?"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "weather"}'}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "The weather is sunny." + # The tool history reached the generator intact (role=tool + assistant.tool_calls). + sent = backend.calls[0]["messages"] + roles = [m["role"] for m in sent] + assert "tool" in roles + assistant = next(m for m in sent if m["role"] == "assistant") + assert assistant.get("tool_calls") + + +def test_dict_arguments_history_does_not_crash(monkeypatch): + # Non-spec client: assistant tool_calls[].function.arguments as a dict. + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": {"q": "x"}}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "ok" + + +def test_forced_tool_choice_narrows_promotion(monkeypatch): + # tool_choice forces `search`; a `lookup` text call must NOT promote. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], + stream = False, + tool_choice = {"type": "function", "function": {"name": "search"}}, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + + +def test_parallel_cap_non_streaming(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) + payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False) + body = _json_body(_call(payload, monkeypatch, backend)) + calls = body["choices"][0]["message"]["tool_calls"] + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "lookup" + + +def test_usage_recorded_when_stats_present(monkeypatch): + stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} + backend = _ScriptedBackend(_fixed(_CALL_XML), stats = stats) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + [entry] = monitor.snapshot() + assert entry["prompt_tokens"] == 7 + assert entry["completion_tokens"] == 3 + + +# ── Nudge ───────────────────────────────────────────────────────── + + +def test_nudge_default_off_single_generation(monkeypatch): + # Signal present but unparseable; without opt-in, no retry. + truncated = '{"name": "lookup"' + backend = _ScriptedBackend(_fixed(truncated)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + _call(payload, monkeypatch, backend) + assert len(backend.calls) == 1 + + +def test_nudge_opt_in_retry_recovers(monkeypatch): + truncated = '{"name": "lookup"' + + def responder(messages, tools): + nudged = any( + "native tool-call format" in (m.get("content") or "") + for m in messages + if m.get("role") == "user" + ) + return [_CALL_XML] if nudged else [truncated] + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True) + body = _json_body(_call(payload, monkeypatch, backend)) + assert len(backend.calls) == 2 + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + + +def test_nudge_double_failure_relays_original(monkeypatch): + truncated = '{"name": "lookup"' + backend = _ScriptedBackend(_fixed(truncated)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True) + body = _json_body(_call(payload, monkeypatch, backend)) + assert len(backend.calls) == 2 # exactly one retry + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == truncated + + +# ── Streaming ───────────────────────────────────────────────────── + + +def test_streaming_heals_split_call_into_one_delta(monkeypatch): + # Cumulative snapshots that build the call across many increments. + pieces = ["{"name": "loo', '{"name": "lookup", "argum'] + cumulative = pieces + [_CALL_XML] + backend = _ScriptedBackend(_fixed(*cumulative)) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + assert tool_deltas[0]["function"]["name"] == "lookup" + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert finishes == ["tool_calls"] + + +def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch): + # A stream cancelled via the registry ("Stop") must NOT promote the + # buffered-but-unclosed tool markup at finalize, else it executes a tool + # the user just cancelled. Guarded on cancel_event at the finalize step. + import routes.inference as inf + + cancel_id = "cancel-me-6870" + # Balanced JSON but no closing -> healer HOLDS it until finalize. + held = '{"name": "lookup", "arguments": {"q": "cats"}}' + + class _CancelMidStream(_ScriptedBackend): + def __init__(self): + super().__init__(_fixed(held)) + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + yield held # healer holds the unclosed call + inf._cancel_by_cancel_id_or_stash(cancel_id) # user hits Stop before EOF + + backend = _CancelMidStream() + payload = _request(tools = [LOOKUP_TOOL], stream = True, cancel_id = cancel_id) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert tool_deltas == [] # no tool promoted after cancel + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert "tool_calls" not in finishes # ends with finish_reason=stop, not tool_calls + + +def test_streaming_no_tools_verbatim(monkeypatch): + backend = _ScriptedBackend(_fixed("hello ", "hello world")) + payload = _request(stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + text = "".join( + (o["choices"][0]["delta"].get("content") or "") + for o in objs + if o["choices"] and "delta" in o["choices"][0] + ) + assert text == "hello world" + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert finishes == ["stop"] + + +def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch): + # Repeated then shrunk cumulative snapshots must not double-heal. + backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + + +def test_streaming_parallel_cap(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) + payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + assert tool_deltas[0]["function"]["name"] == "lookup" + + +def test_streaming_generator_error_closes_cleanly(monkeypatch): + def responder(messages, tools): + raise RuntimeError("boom /secret/path") + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + chunks = _collect_sse(response) + joined = "".join(c.decode() if isinstance(c, bytes) else c for c in chunks) + assert "An internal error occurred" in joined + assert "secret/path" not in joined # CWE-209: no path leak + assert backend.reset_count >= 1 + + +def test_streaming_disconnect_resets_once(monkeypatch): + class _DisconnectRequest(_Request): + async def is_disconnected(self): + return True + + backend = _ScriptedBackend(_fixed("a", "ab", "abc")) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + _install(monkeypatch, backend) + + async def _run(): + resp = await openai_chat_completions( + payload, request = _DisconnectRequest(), current_subject = "u" + ) + return [c async for c in resp.body_iterator] + + asyncio.run(_run()) + assert backend.reset_count == 1 + + +def test_mlx_uses_same_path(monkeypatch): + # MLX and safetensors share get_inference_backend(); one scripted backend covers both. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["finish_reason"] == "tool_calls" + + +def test_tool_choice_none_does_not_advertise_tools(monkeypatch): + # tool_choice="none": no tools rendered into the template; history templating still applies. + backend = _ScriptedBackend(_fixed("plain answer")) + payload = _request(tools = [LOOKUP_TOOL], tool_choice = "none", stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "plain answer" + assert backend.calls[0]["tools"] is None + + +def test_developer_message_folded_into_system_prompt(monkeypatch): + # The "developer" role folds into one leading system message (local templates reject it). + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + messages = [ + ChatMessage(role = "developer", content = "always be terse"), + ChatMessage(role = "user", content = "hi"), + ], + tools = [LOOKUP_TOOL], + stream = False, + ) + _call(payload, monkeypatch, backend) + sent = backend.calls[0]["messages"] + assert sent[0]["role"] == "system" + assert "always be terse" in sent[0]["content"] + assert all(m.get("role") != "developer" for m in sent) + + +def test_failed_nudge_retry_keeps_original_response(monkeypatch): + # A raising retry must not 500; the first response is returned. + state = {"n": 0} + + def responder(messages, tools): + state["n"] += 1 + if state["n"] == 1: + return ['{"name":"lookup"'] # unhealable signal + raise RuntimeError("retry blew up") + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert state["n"] == 2 + assert body["choices"][0]["finish_reason"] == "stop" + assert body["choices"][0]["message"]["content"] == '{"name":"lookup"' + + +def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch): + # Double-failure nudge: the first response is delivered, but the retry's + # generate() overwrites stats_holder. The monitor must record the FIRST + # attempt's usage, not the discarded retry's. + first_stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} + retry_stats = {"usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198}} + + class _PerCallStatsBackend(_ScriptedBackend): + def __init__(self): + # Unhealable truncated markup on both attempts -> retry is discarded. + super().__init__(lambda m, t: ['{"name":"lookup"']) + self._stats_seq = [first_stats, retry_stats] + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + stats = self._stats_seq[min(len(self.calls) - 1, len(self._stats_seq) - 1)] + if stats_holder is not None: + stats_holder["stats"] = stats + for snap in self._responder(messages, tools): + yield snap + + backend = _PerCallStatsBackend() + payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + assert len(backend.calls) == 2 # first attempt + one discarded retry + [entry] = monitor.snapshot() + # The delivered response is the first attempt, so its usage must be reported. + assert entry["prompt_tokens"] == 7 + assert entry["completion_tokens"] == 3 + + +def test_monitor_records_healed_call_not_raw_xml(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + snap = monitor.snapshot(include_details = True) + replies = json.dumps(snap) + assert "" not in replies + assert "lookup" in replies + + +def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch): + # Monitor mirrors what the client received, never the healed-away raw markup. + backend = _ScriptedBackend( + _fixed("Sure. ", 'Sure. {"name": "loo', "Sure. " + _CALL_XML) + ) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + response = asyncio.run(_run()) + _collect_sse(response) + replies = json.dumps(monitor.snapshot(include_details = True)) + assert "" not in replies + assert "Sure. " in replies + assert "[tool_calls] lookup(" in replies + + +def test_forced_tool_choice_narrows_templated_tools(monkeypatch): + # A forced function is the only schema rendered into the template. + backend = _ScriptedBackend(_fixed(_SEARCH_XML)) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], + stream = False, + tool_choice = {"type": "function", "function": {"name": "search"}}, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + templated = backend.calls[0]["tools"] + assert [t["function"]["name"] for t in templated] == ["search"] + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "search" + + +def test_multimodal_content_parts_flattened_for_local_template(monkeypatch): + # Remote image URLs leave image=None, so content arrives as a part LIST: + # text parts are kept, the image part dropped. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request( + messages = [ + ChatMessage( + role = "user", + content = [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/cat.png"}, + }, + ], + ) + ], + tools = [LOOKUP_TOOL], + stream = False, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + templated = backend.calls[0]["messages"] + assert all(isinstance(m.get("content"), str) for m in templated) + assert any(m["content"] == "what is this?" for m in templated) + assert body["choices"][0]["finish_reason"] == "tool_calls" + + +def test_string_arguments_history_deserialized_for_template(monkeypatch): + # JSON-string tool_calls arguments become dicts in the templated copy; + # the HTTP response stays OpenAI-shaped. + backend = _ScriptedBackend(_fixed("done")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "weather?"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "weather"}'}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"), + ], + ) + _json_body(_call(payload, monkeypatch, backend)) + assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assert assistant["tool_calls"][0]["function"]["arguments"] == {"q": "weather"} + + +def test_unparseable_arguments_string_left_untouched(monkeypatch): + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "not json {"}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "ok" + assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assert assistant["tool_calls"][0]["function"]["arguments"] == "not json {" + + +def test_mcp_enabled_without_server_tools_uses_passthrough(monkeypatch): + # mcp_enabled=true with an empty registry must not silently drop the + # declared tools; the gate keys on the server-side path claiming the request. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, mcp_enabled = True) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + assert backend.calls[0]["tools"] == [LOOKUP_TOOL] From 35063716771f10b780ffb23eb0f938472eaa4291 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 21:57:56 -0700 Subject: [PATCH 031/113] Studio: keep the nudge wiring test collectable without the unsloth stack (#6924) test_nudge_tool_calls_wiring.py imported InferenceBackend from core.inference.inference, which pulls in unsloth (and thus unsloth_zoo) at module scope. The dependency-light backend CI matrix job does not install unsloth_zoo, so the import raised at collection time and aborted the whole job (831 tests never ran). Guard that one import and fold the safetensors InferenceBackend checks in only when the unsloth stack is importable; the orchestrator/llama_cpp/safetensors_agentic wiring is still asserted unconditionally, and local/full-stack runs keep the InferenceBackend coverage. --- .../tests/test_nudge_tool_calls_wiring.py | 29 +++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/studio/backend/tests/test_nudge_tool_calls_wiring.py b/studio/backend/tests/test_nudge_tool_calls_wiring.py index 2c27b220ba..e03fd0c7d7 100644 --- a/studio/backend/tests/test_nudge_tool_calls_wiring.py +++ b/studio/backend/tests/test_nudge_tool_calls_wiring.py @@ -23,11 +23,20 @@ Mechanism (verified here without loading a model): import inspect -from core.inference.inference import InferenceBackend from core.inference.llama_cpp import LlamaCppBackend from core.inference.orchestrator import InferenceOrchestrator from core.inference.safetensors_agentic import run_safetensors_tool_loop +try: + # core.inference.inference imports unsloth at module scope, which requires + # unsloth_zoo. The dependency-light backend CI matrix job does not install + # it, so the safetensors InferenceBackend is folded into the checks below + # only when the unsloth stack is importable (local runs / full CI); the + # other entry points are always checked. + from core.inference.inference import InferenceBackend +except ImportError: + InferenceBackend = None + def _params(fn): return inspect.signature(fn).parameters @@ -37,12 +46,14 @@ def test_shared_loop_accepts_nudge_flag(): assert "nudge_tool_calls" in _params(run_safetensors_tool_loop) -def test_all_three_backends_accept_the_flag(): - for method in ( - InferenceBackend.generate_chat_completion_with_tools, +def test_backends_accept_the_flag(): + methods = [ InferenceOrchestrator.generate_chat_completion_with_tools, LlamaCppBackend.generate_chat_completion_with_tools, - ): + ] + if InferenceBackend is not None: # safetensors path; needs the unsloth stack + methods.append(InferenceBackend.generate_chat_completion_with_tools) + for method in methods: assert "nudge_tool_calls" in _params(method), method.__qualname__ @@ -50,10 +61,10 @@ def test_delegating_backends_forward_the_flag_to_the_shared_loop(): # safetensors (in-process transformers) and MLX (parent-process orchestrator) # both delegate to run_safetensors_tool_loop; GGUF runs its own in-file loop # and consumes the flag directly (asserted separately by the gate test). - for method in ( - InferenceBackend.generate_chat_completion_with_tools, - InferenceOrchestrator.generate_chat_completion_with_tools, - ): + methods = [InferenceOrchestrator.generate_chat_completion_with_tools] + if InferenceBackend is not None: # safetensors path; needs the unsloth stack + methods.append(InferenceBackend.generate_chat_completion_with_tools) + for method in methods: src = inspect.getsource(method) assert "nudge_tool_calls = nudge_tool_calls" in src, method.__qualname__ From 9674e882c234479606ae98073fd7d734b97c173a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 22:09:41 -0700 Subject: [PATCH 032/113] Studio: serialize the compare-mode dispatcher lifecycle to fix a start race (#6922) * Studio: serialize the compare-mode dispatcher lifecycle to fix a start race _generate_dispatched (compare mode) bypasses _gen_lock so two concurrent compare requests can both reach _start_dispatcher. The check-then-spawn there had no lock, so both could observe no live dispatcher and each spawn one. The extra dispatcher is orphaned (self._dispatcher_thread tracks only the last) and during a later unload it can consume the 'unloaded' reply off _resp_queue before unload_model's _wait_response, hanging the unload on its timeout. Add _dispatcher_lifecycle_lock and take it around the whole body of both _start_dispatcher and _stop_dispatcher, so start/stop cannot interleave and the second concurrent starter sees the dispatcher alive and returns. _start_dispatcher now returns whether it actually spawned the thread, and _generate_dispatched derives dispatcher_preexisting from that atomic result instead of a separate unlocked is_alive() read. No call site holds _mailbox_lock when calling start/stop, so joining the dispatcher (which takes _mailbox_lock) under the new lock cannot deadlock; the lock order is always _gen_lock then _dispatcher_lifecycle_lock and is never inverted. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: refuse dispatcher start queued behind an unload's stop A compare request could pass the early _unload_pending check, then block in _start_dispatcher on _dispatcher_lifecycle_lock behind an unload's _stop_dispatcher. When the unload released the lock the start spawned a fresh dispatcher, which became the resp_queue reader and consumed the worker's unroutable 'unloaded' reply before unload_model's _wait_response saw it, hanging the unload for 300s. Gate _start_dispatcher on _unload_pending under the lifecycle lock, and set _unload_pending under the same lock ahead of the stop, so any start queued behind the stop observes the unload and refuses. Ordering stays _gen_lock -> _dispatcher_lifecycle_lock. Adds a regression test forcing the queued-behind-stop interleaving. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/orchestrator.py | 92 ++++++--- .../tests/test_orchestrator_unload_cancel.py | 192 ++++++++++++++++++ 2 files changed, 257 insertions(+), 27 deletions(-) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index e0e6cef6c9..394ff97eac 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -79,6 +79,12 @@ class InferenceOrchestrator: self._mailbox_lock = threading.Lock() self._dispatcher_thread: Optional[threading.Thread] = None self._dispatcher_stop = threading.Event() + # Serializes dispatcher start/stop. _generate_dispatched (compare mode) bypasses + # _gen_lock, so two concurrent compare requests can both reach _start_dispatcher; + # without this lock both could observe no live dispatcher and each spawn one, + # orphaning the extra thread (self._dispatcher_thread tracks only the last). The + # orphan later steals the "unloaded" reply off resp_queue and hangs unload_model. + self._dispatcher_lifecycle_lock = threading.Lock() # Local state mirrors (updated from subprocess responses) self.active_model_name: Optional[str] = None @@ -514,33 +520,56 @@ class InferenceOrchestrator: # Dispatcher — per-request mailbox routing for compare mode # ------------------------------------------------------------------ - def _start_dispatcher(self) -> None: + def _start_dispatcher(self) -> bool: """Start the dispatcher thread if not already running. The dispatcher reads the shared resp_queue and routes responses to per-request mailbox queues, letting multiple adapter-controlled (compare) requests be in-flight without holding _gen_lock. - """ - if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive(): - return - self._dispatcher_stop.clear() - self._dispatcher_thread = threading.Thread( - target = self._dispatcher_loop, - daemon = True, - name = "inference-dispatcher", - ) - self._dispatcher_thread.start() - logger.debug("Dispatcher thread started") + The whole check-then-spawn runs under _dispatcher_lifecycle_lock so + concurrent compare requests (which bypass _gen_lock) can't both observe + no live dispatcher and each spawn one. Returns True only for the caller + that actually started a new thread; False if one was already alive. + """ + with self._dispatcher_lifecycle_lock: + # Refuse to start while an unload is in progress. unload_model sets + # _unload_pending under this same lock before it stops the idle + # dispatcher, so a start queued behind that stop observes the unload + # here and bails. Without this a fresh dispatcher would be spawned + # after the stop, become the resp_queue reader, and consume the + # worker's "unloaded" reply (unroutable, so dropped) before + # unload_model's _wait_response sees it -- hanging the unload 300s. + if self._unload_pending: + return False + if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive(): + return False + + self._dispatcher_stop.clear() + self._dispatcher_thread = threading.Thread( + target = self._dispatcher_loop, + daemon = True, + name = "inference-dispatcher", + ) + self._dispatcher_thread.start() + logger.debug("Dispatcher thread started") + return True def _stop_dispatcher(self) -> None: - """Signal the dispatcher to stop and wait for it.""" - if self._dispatcher_thread is None: - return - self._dispatcher_stop.set() - self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT) - self._dispatcher_thread = None - logger.debug("Dispatcher thread stopped") + """Signal the dispatcher to stop and wait for it. + + Runs under _dispatcher_lifecycle_lock (paired with _start_dispatcher) so + a stop can't interleave with a concurrent start. Callers must NOT hold + _mailbox_lock here: this joins the dispatcher, and the dispatcher loop + takes _mailbox_lock, so holding it would deadlock the join. + """ + with self._dispatcher_lifecycle_lock: + if self._dispatcher_thread is None: + return + self._dispatcher_stop.set() + self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT) + self._dispatcher_thread = None + logger.debug("Dispatcher thread stopped") def _dispatcher_loop(self) -> None: """Background loop: read resp_queue → route to mailboxes by request_id.""" @@ -628,13 +657,14 @@ class InferenceOrchestrator: yield "Error: model is being unloaded" return - # Ensure the dispatcher runs. Track whether it was already running: if this call - # starts it and then bails on a racing unload, it must stop it again (see the - # unloading bail below). - dispatcher_preexisting = ( - self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() - ) - self._start_dispatcher() + # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under + # _dispatcher_lifecycle_lock and returns True only for the caller that actually spawned + # the thread, so at most one dispatcher ever exists even when two compare requests race + # here. Derive dispatcher_preexisting from that atomic result (not a separate unlocked + # is_alive() read): if THIS call started the dispatcher and then bails on a racing + # unload, it must stop it again (see the unloading bail below). + started = self._start_dispatcher() + dispatcher_preexisting = not started request_id = str(uuid.uuid4()) @@ -1041,7 +1071,15 @@ class InferenceOrchestrator: # The subprocess runs commands sequentially, so a bare unload queues behind a # running generate (a 2-3 min hang). Cancel first (via the mp.Event the worker # polls each token), then take _gen_lock as sole resp_queue reader (like GGUF). - self._unload_pending = True + # + # Set _unload_pending under _dispatcher_lifecycle_lock so it is ordered ahead of + # the dispatcher stop that _wait_dispatcher_idle runs under the same lock: a + # compare request's _start_dispatcher queued behind that stop then observes the + # unload and refuses to spawn a fresh dispatcher that would eat the "unloaded" + # reply off resp_queue. This is a standalone acquisition (no _gen_lock held yet), + # so it keeps the _gen_lock -> _dispatcher_lifecycle_lock order and can't deadlock. + with self._dispatcher_lifecycle_lock: + self._unload_pending = True # Cancelling only the running generation isn't enough: the worker clears # cancel_event at each generate start, so a queued one would clear it and run the # outgoing model to completion. drain_event, never cleared, makes any generate diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index e9d0f36fe2..fe3c6d5a0d 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -25,6 +25,8 @@ def _bare_orchestrator(): o._cmd_queue = object() o._resp_queue = object() o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() o._unload_pending = False o.active_model_name = "m" o.models = {"m": {}} @@ -1253,6 +1255,7 @@ def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch): def fake_start(): started["v"] = True o._dispatcher_thread = _AliveDispatcher() + return True # _start_dispatcher returns True for the caller that spawned it def fake_stop(): stopped["v"] = True @@ -1428,3 +1431,192 @@ def test_load_model_aborts_publish_when_cancelled_after_wait_response(monkeypatc assert o.active_model_name is None, "must not publish a cancelled model's active name" assert o.models == {}, "must not publish a cancelled model's mirror" assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# Concurrent compare-mode requests must not each spawn a dispatcher. Compare mode +# (_generate_dispatched) deliberately bypasses _gen_lock, so two requests can reach +# _start_dispatcher at once. Without _dispatcher_lifecycle_lock the check-then-spawn +# races: both observe no live dispatcher and each start one. The extra dispatcher is +# orphaned (self._dispatcher_thread tracks only the last) and later consumes the +# "unloaded" reply off the shared resp_queue before unload_model's _wait_response, +# hanging the unload on its 300s timeout. The lifecycle lock must serialize the +# check-then-spawn so exactly one dispatcher thread is ever created. +# ---------------------------------------------------------------------------- + + +def test_concurrent_start_dispatcher_spawns_exactly_one(): + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # real queue so the dispatcher loop blocks and stays alive + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + + n = 32 + # A barrier aligns every thread on the check-then-spawn window: without the lifecycle + # lock several would clear the "is a dispatcher alive?" check together and each spawn one. + barrier = threading.Barrier(n) + results: list = [] + results_lock = threading.Lock() + + def racer(): + barrier.wait() + started = o._start_dispatcher() + with results_lock: + results.append(started) + + threads = [threading.Thread(target = racer, name = f"racer-{i}") for i in range(n)] + for t in threads: + t.start() + for t in threads: + t.join(timeout = 5) + + try: + # _start_dispatcher returns True only for the caller that actually spawned a thread. + # Exactly one caller may win; every other must observe the dispatcher alive and bail. + assert results.count(True) == 1, f"expected exactly one spawn, got {results.count(True)}" + assert results.count(False) == n - 1 + # And exactly one live dispatcher thread exists -- no orphan racing resp_queue. + live = [ + t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive() + ] + assert len(live) == 1, f"expected one live dispatcher, found {len(live)}" + assert o._dispatcher_thread is live[0] + finally: + o._stop_dispatcher() + + # Stop joins and clears it; no dispatcher thread must survive. + assert o._dispatcher_thread is None + remaining = [ + t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive() + ] + assert remaining == [], "dispatcher must be stopped and joined" + + +# ---------------------------------------------------------------------------- +# A compare request whose _start_dispatcher is queued behind an unload's +# _stop_dispatcher must NOT spawn a fresh dispatcher. The idle-dispatcher stop +# and the queued start both serialize on _dispatcher_lifecycle_lock; if the +# queued start spawned a new dispatcher after the stop, it would become the +# resp_queue reader and consume unload_model's "unloaded" reply (unroutable, so +# dropped) before _wait_response saw it -- hanging the unload on its 300s +# timeout. unload_model sets _unload_pending under the SAME lifecycle lock ahead +# of the stop, so _start_dispatcher observes it and refuses. +# ---------------------------------------------------------------------------- + + +def test_start_dispatcher_refuses_while_unload_pending(): + # Direct unit guard: with an unload in progress (_unload_pending set under the + # lifecycle lock by unload_model), _start_dispatcher must refuse and spawn nothing, + # even though no dispatcher is currently running. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = True + + started = o._start_dispatcher() + + assert started is False, "must not start a dispatcher while an unload is pending" + assert o._dispatcher_thread is None, "no dispatcher thread may be created" + live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()] + assert live == [], "no dispatcher may exist to consume the unloaded reply" + + +def test_start_dispatcher_resumes_after_unload_clears(): + # Guard the other direction: once the unload finishes and clears _unload_pending, a + # later compare request must be able to start the dispatcher again (the gate must not + # wedge). Proves the refusal above is scoped to the unload, not permanent. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = False + + try: + assert ( + o._start_dispatcher() is True + ), "a fresh dispatcher must start once no unload is pending" + assert o._dispatcher_thread is not None and o._dispatcher_thread.is_alive() + finally: + o._stop_dispatcher() + + assert o._dispatcher_thread is None + + +def test_queued_start_behind_unload_stop_spawns_no_dispatcher(): + # Codex's exact ordering, forced deterministically: an unload holds + # _dispatcher_lifecycle_lock across its _stop_dispatcher (the idle dispatcher's join + # is gated by an event), while a compare request's _start_dispatcher is queued behind + # it on the same lock. When the stop releases the lock the queued start must observe + # _unload_pending (set under the lock ahead of the stop) and refuse: no fresh + # dispatcher may be left running to steal the "unloaded" reply. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = False + + start_queued = threading.Event() # release the stop's join once the start is queued behind it + join_may_finish = threading.Event() + + class _IdleDispatcher: + # Stand-in for the idle compare-mode dispatcher the unload stops. Its join blocks + # until we confirm the compare _start_dispatcher is queued behind the stop, so the + # stop provably holds _dispatcher_lifecycle_lock across that window. + def is_alive(self): + return True + + def join(self, timeout = None): + assert start_queued.wait(timeout = 5), "compare start must queue behind the stop" + assert join_may_finish.wait(timeout = 5) + + o._dispatcher_thread = _IdleDispatcher() + + def unload_side(): + # unload_model's sequence: set _unload_pending under the lifecycle lock, then stop + # the idle dispatcher (also under the lock, via _wait_dispatcher_idle). + with o._dispatcher_lifecycle_lock: + o._unload_pending = True + o._stop_dispatcher() + + started_result = {} + + def compare_side(): + started_result["v"] = o._start_dispatcher() + + u = threading.Thread(target = unload_side, name = "unload-side") + u.start() + # Let the unload set _unload_pending, enter _stop_dispatcher, and block in the gated join + # while holding the lifecycle lock. + time.sleep(0.2) + + c = threading.Thread(target = compare_side, name = "compare-side") + c.start() + # Let the compare _start_dispatcher block on the lifecycle lock (queued behind the stop). + time.sleep(0.2) + + start_queued.set() # the start is now queued behind the stop + join_may_finish.set() # let the stop's join complete and release the lock + + u.join(timeout = 5) + c.join(timeout = 5) + + assert started_result.get("v") is False, "the queued start must refuse while unloading" + assert o._dispatcher_thread is None, "the stop cleared it and the queued start spawned nothing" + live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()] + assert live == [], "no fresh dispatcher may be left to consume the unloaded reply" From 5608081c35d120343b2e196a292bc13c9a3a96af Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 22:24:47 -0700 Subject: [PATCH 033/113] Studio: apply presence_penalty on the safetensors and MLX inference paths (#6923) * Studio: apply presence_penalty on the safetensors and MLX inference paths The safetensors and MLX generate paths resolved the inference config and then dropped presence_penalty before generation, so the same model applied the configured value under GGUF and 0 under safetensors/MLX. Thread the already-resolved presence_penalty through the orchestrator command, worker gen_kwargs, and the safetensors/MLX generate calls, and apply it with a small logits processor (subtract once per distinct completion token, prompt excluded, presence not frequency, zero is a no-op, negatives raise). Backwards compatible: presence_penalty defaults to 0.0 (byte-identical output when unset) and the GGUF path is unchanged. Also forward min_p on the legacy /generate/stream route and add the missing min_p field to GenerateRequest. * Studio: bound presence_penalty generated ids to valid vocab range on both paths The presence-penalty logits processors index by generated token ids. The torch path filtered only the upper bound (seen < vocab_size), so a negative id would silently wrap to the wrong row; the MLX path had no bound at all, and MLX out-of-bounds indexing is documented undefined behavior (crash or memory corruption on Apple Silicon), unlike torch's harmless negative wrap. Bound generated ids to [0, vocab) consistently on both paths: - torch: seen[(seen >= 0) & (seen < vocab_size)] (zero-regression safety net; real completion tokens are always in range). - MLX: route out-of-range/negative ids to a discarded scratch slot via mx.where and a (vocab + 1)-wide scatter-assign mask, then subtract. MLX has no boolean-mask filtering (data-dependent output shape), so this keeps a fixed shape, stays on-device, and preserves once-per-distinct-token semantics without any torch/numpy dependency. Add torch tests for out-of-range and negative ids (only in-range distinct ids penalized, stray ids ignored, no wrong-index wrap) and a bound-documenting MLX test that runs on the arm64 macOS CI. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/inference.py | 26 ++ .../backend/core/inference/mlx_inference.py | 80 +++++- studio/backend/core/inference/orchestrator.py | 12 + .../core/inference/presence_penalty.py | 49 ++++ studio/backend/core/inference/worker.py | 1 + studio/backend/models/inference.py | 1 + studio/backend/routes/inference.py | 4 + studio/backend/tests/test_presence_penalty.py | 252 ++++++++++++++++++ 8 files changed, 419 insertions(+), 6 deletions(-) create mode 100644 studio/backend/core/inference/presence_penalty.py create mode 100644 studio/backend/tests/test_presence_penalty.py diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 064be30c06..167706f701 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -31,6 +31,7 @@ from core.inference.chat_eos import ( chat_eos_repair, resolve_chat_turn_end_eos_ids_using, ) +from core.inference.presence_penalty import _make_presence_penalty_processor from io import StringIO import structlog from loggers import get_logger @@ -832,6 +833,7 @@ class InferenceBackend: tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + presence_penalty: float = 0.0, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -865,6 +867,7 @@ class InferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) initial = list(messages) @@ -901,12 +904,14 @@ class InferenceBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate response for text or vision models (lock held by background thread). ``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking`` are forwarded into ``apply_chat_template`` so templates that understand them (Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls. + ``presence_penalty`` matches the GGUF sampling path (0 disables it). """ yield from self._generate_chat_response_inner( messages = messages, @@ -923,6 +928,7 @@ class InferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) def _generate_chat_response_inner( @@ -942,6 +948,7 @@ class InferenceBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Inner generation logic, called by generate_chat_response and generate_with_adapter_control. @@ -981,6 +988,7 @@ class InferenceBackend: max_new_tokens, repetition_penalty, cancel_event = cancel_event, + presence_penalty = presence_penalty, ) return else: @@ -1093,6 +1101,7 @@ class InferenceBackend: repetition_penalty, cancel_event = cancel_event, _adapter_state = _adapter_state, + presence_penalty = presence_penalty, ) def _generate_vision_response( @@ -1107,6 +1116,7 @@ class InferenceBackend: max_new_tokens, repetition_penalty, cancel_event = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Handle vision model generation with true token-by-token streaming.""" model_info = self.models[self.active_model_name] @@ -1196,6 +1206,14 @@ class InferenceBackend: top_k = top_k, min_p = min_p, ) + # Presence penalty (GGUF parity) for VLM chat. + _vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None + if _vision_input_ids is not None: + _pp = _make_presence_penalty_processor( + presence_penalty, int(_vision_input_ids.shape[1]) + ) + if _pp is not None: + generation_kwargs["logits_processor"] = _pp err: dict[str, str] = {} @@ -1424,11 +1442,13 @@ class InferenceBackend: repetition_penalty: float = 1.0, cancel_event = None, _adapter_state = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate a streaming text response (text models only). _adapter_state: if not None, the background thread toggles adapters before model.generate(), under _generation_lock. + ``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it). """ if not self.active_model_name: yield "Error: No active model" @@ -1489,6 +1509,12 @@ class InferenceBackend: if tokenizer.pad_token_id is None else tokenizer.pad_token_id, ) + # Presence penalty (GGUF parity); prompt_len excludes prompt tokens. + _pp = _make_presence_penalty_processor( + presence_penalty, int(inputs["input_ids"].shape[1]) + ) + if _pp is not None: + generation_kwargs["logits_processor"] = _pp if cancel_event is not None: from transformers.generation.stopping_criteria import ( StoppingCriteria, diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 45f46fef2f..62d268e15f 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -41,6 +41,50 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): } +def _make_mlx_presence_penalty_processor(penalty: float): + """Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path. + + generate_step calls processors as ``fn(tokens, logits)`` with ``tokens`` the + full running sequence; the first call is prompt-only, so latch that length + and penalize only after it. + """ + state = {"prompt_len": None} + + def _processor(tokens, logits): + if state["prompt_len"] is None: + # First call = prompt only; latch its length. + state["prompt_len"] = int(tokens.shape[0]) + return logits + generated = tokens[state["prompt_len"] :] + if generated.size == 0: + return logits + import mlx.core as mx + + vocab = logits.shape[-1] + # Bound generated ids to the valid range [0, vocab) before they index + # logits. MLX does no bounds checking and out-of-bounds indexing is + # documented undefined behavior (crash / memory corruption), unlike the + # torch path's harmless negative wrap -- so this bound is load-bearing + # here and matches the torch filter seen[(seen >= 0) & (seen < vocab)]. + # MLX has no boolean-mask filtering (data-dependent output shape is + # unsupported), so instead of compacting the id list we route every + # out-of-range or negative id to a scratch slot at index ``vocab`` that + # is dropped before the subtract. That scratch slot can never collide + # with a real token, so real ids (including id 0) are penalized exactly + # once and stray ids are ignored. + valid = (generated >= 0) & (generated < vocab) + safe = mx.where(valid, generated, vocab).astype(mx.int32) + # Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate + # ids are idempotent, so presence applies once per distinct token; the + # scratch column is discarded and the full-width subtract stays on-device. + mask = mx.zeros((vocab + 1,), dtype = logits.dtype) + mask[safe] = penalty + logits = logits - mask[:vocab] + return logits + + return _processor + + class MLXInferenceBackend: def __init__(self): self.models = {} @@ -282,6 +326,7 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ) -> Generator[str, None, None]: if self._model is None: raise RuntimeError("No model loaded") @@ -329,6 +374,7 @@ class MLXInferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) else: yield from self._generate_text( @@ -344,6 +390,7 @@ class MLXInferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) def _generate_text( @@ -361,6 +408,7 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ): from mlx_lm import stream_generate from mlx_lm.sample_utils import make_sampler, make_logits_processors @@ -407,15 +455,21 @@ class MLXInferenceBackend: min_p = float(min_p or 0.0), min_tokens_to_keep = 1, ) - # Only build a logits processor for a non-trivial repetition penalty. - logits_processors = None + # Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths). + logits_processors = [] if repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, 1.0, ): - logits_processors = make_logits_processors( - repetition_penalty = float(repetition_penalty), + logits_processors.extend( + make_logits_processors( + repetition_penalty = float(repetition_penalty), + ) ) + if presence_penalty: + logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + if not logits_processors: + logits_processors = None token_ids = [] logger.info( @@ -481,6 +535,7 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ): from mlx_vlm import stream_generate as vlm_stream @@ -528,10 +583,23 @@ class MLXInferenceBackend: top_k = int(top_k or 0), min_p = float(min_p or 0.0), ) - if repetition_penalty is not None and float(repetition_penalty) not in ( + _rep_active = repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, 1.0, - ): + ) + if presence_penalty: + # Presence needs a custom processor: pass the full list (repetition + + # presence) instead of the repetition_penalty shortcut so both apply once. + from mlx_lm.sample_utils import make_logits_processors + + _vlm_processors = [] + if _rep_active: + _vlm_processors.extend( + make_logits_processors(repetition_penalty = float(repetition_penalty)) + ) + _vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + vlm_kwargs["logits_processors"] = _vlm_processors + elif _rep_active: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) with self._generation_lock: diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 394ff97eac..19d2230278 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -428,6 +428,7 @@ class InferenceOrchestrator: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> dict: """Build the 'generate' command shared by the locked and dispatched paths.""" cmd = { @@ -442,6 +443,7 @@ class InferenceOrchestrator: "min_p": min_p, "max_new_tokens": max_new_tokens, "repetition_penalty": repetition_penalty, + "presence_penalty": presence_penalty, } # Only forward template kwargs the caller set, for older worker compat. if use_adapter is not None: @@ -631,6 +633,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Dispatched generation — sends command without holding _gen_lock. @@ -684,6 +687,7 @@ class InferenceOrchestrator: min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, @@ -1166,6 +1170,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate response, streaming tokens from subprocess. @@ -1175,6 +1180,8 @@ class InferenceOrchestrator: ``stats_holder``: caller-owned dict; on gen_done its "stats" key gets the worker's usage/timings. Request-scoped to avoid cross-stream reads. + + ``presence_penalty`` matches the GGUF sampling path (0 disables it). """ yield from self._generate_inner( messages = messages, @@ -1193,6 +1200,7 @@ class InferenceOrchestrator: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, stats_holder = stats_holder, + presence_penalty = presence_penalty, ) def generate_chat_completion_with_tools( @@ -1220,6 +1228,7 @@ class InferenceOrchestrator: bypass_permissions: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, **_unused, ): """Run the safetensors agentic tool loop in the parent process, @@ -1255,6 +1264,7 @@ class InferenceOrchestrator: preserve_thinking = preserve_thinking, # last turn wins, like the GGUF tool loop stats_holder = stats_holder, + presence_penalty = presence_penalty, ) if use_adapter is not None: yield from self.generate_with_adapter_control( @@ -1322,6 +1332,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Inner generation logic — sends command to subprocess, yields tokens. @@ -1365,6 +1376,7 @@ class InferenceOrchestrator: min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, diff --git a/studio/backend/core/inference/presence_penalty.py b/studio/backend/core/inference/presence_penalty.py new file mode 100644 index 0000000000..c73c513887 --- /dev/null +++ b/studio/backend/core/inference/presence_penalty.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Presence-penalty logits helpers for the safetensors/MLX inference paths. + +Kept in a dependency-light leaf module (torch + transformers only, no unsloth / +peft) so the pure logic can be imported and unit-tested without pulling in the +full inference backend. ``core.inference.inference`` re-exports these for the +runtime generate paths. +""" + +import torch + + +def apply_presence_penalty(input_ids, scores, penalty: float, prompt_len: int): + """OpenAI/llama.cpp presence penalty: subtract ``penalty`` once per distinct + completion token (positions >= prompt_len; prompt excluded, multiplicity + ignored, negatives raise). In place; zero is a no-op.""" + if not penalty: + return scores + vocab_size = scores.shape[-1] + for b in range(input_ids.shape[0]): + generated = input_ids[b, prompt_len:] + if generated.numel() == 0: + continue + seen = torch.unique(generated) + # Bound generated ids to the valid range [0, vocab_size). Real completion + # tokens are always in range, so this is a zero-regression safety net that + # drops any stray out-of-range or negative id before indexing (mirrors the + # MLX path's bound). Filtering both ends avoids indexing scores with a + # negative id (which would silently wrap to the wrong row). + seen = seen[(seen >= 0) & (seen < vocab_size)] + if seen.numel(): + scores[b, seen] = scores[b, seen] - penalty + return scores + + +def _make_presence_penalty_processor(penalty: float, prompt_len: int): + """``LogitsProcessorList`` for ``apply_presence_penalty``; ``None`` at zero penalty (generate call stays byte-identical).""" + if not penalty: + return None + from transformers import LogitsProcessor, LogitsProcessorList + + class _PresencePenaltyLogitsProcessor(LogitsProcessor): + @torch.no_grad() + def __call__(self, input_ids, scores): + return apply_presence_penalty(input_ids, scores, penalty, prompt_len) + + return LogitsProcessorList([_PresencePenaltyLogitsProcessor()]) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 615c5c5a0d..d4b102e422 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -457,6 +457,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: "min_p": cmd.get("min_p", 0.0), "max_new_tokens": cmd.get("max_new_tokens", 256), "repetition_penalty": cmd.get("repetition_penalty", 1.0), + "presence_penalty": cmd.get("presence_penalty", 0.0), "cancel_event": cancel_event, } diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 1e7770a7f8..0f27b695fe 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -177,6 +177,7 @@ class GenerateRequest(BaseModel): temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature") top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling") top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling") + min_p: float = Field(0.0, ge = 0.0, le = 1.0, description = "Min-p sampling") max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate") repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty") presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3f6f6a2a78..5332037e0d 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4353,8 +4353,10 @@ async def generate_stream( temperature = request.temperature, top_p = request.top_p, top_k = request.top_k, + min_p = request.min_p, max_new_tokens = request.max_new_tokens, repetition_penalty = request.repetition_penalty, + presence_penalty = request.presence_penalty, cancel_event = cancel_event, ) _DONE = object() @@ -6995,6 +6997,7 @@ async def openai_chat_completions( min_p = payload.min_p, max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, + presence_penalty = payload.presence_penalty, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, @@ -7240,6 +7243,7 @@ async def openai_chat_completions( min_p = payload.min_p, max_new_tokens = effective_max_tokens or 2048, repetition_penalty = payload.repetition_penalty, + presence_penalty = payload.presence_penalty, ) # Forward reasoning kwargs; the worker/template wrapper peels off any the # template doesn't accept. diff --git a/studio/backend/tests/test_presence_penalty.py b/studio/backend/tests/test_presence_penalty.py new file mode 100644 index 0000000000..030ddb6011 --- /dev/null +++ b/studio/backend/tests/test_presence_penalty.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""Presence-penalty parity between the GGUF path and the safetensors/MLX paths. + +The safetensors path historically dropped ``presence_penalty``, so the SAME model +looked worse served as safetensors. These tests pin the processor semantics +(subtract once per distinct completion token, prompt excluded, presence not +frequency, zero a no-op, negatives raise) plus a param-propagation regression +over route -> orchestrator cmd -> worker gen_kwargs. +""" + +import threading + +import pytest +import torch + +from core.inference.presence_penalty import ( + apply_presence_penalty, + _make_presence_penalty_processor, +) + + +def test_seen_token_gets_exactly_minus_penalty_unseen_unchanged(): + input_ids = torch.tensor([[0, 1, 3]]) # prompt [0, 1], completion [3] + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 2) + assert out[0, 3].item() == pytest.approx(-1.5) + for tok in (0, 1, 2, 4): + assert out[0, tok].item() == pytest.approx(0.0) + + +def test_multiplicity_ignored_presence_not_frequency(): + # Token 3 emitted three times -> still a single -penalty (presence, not freq). + input_ids = torch.tensor([[0, 3, 3, 3]]) + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 2.0, prompt_len = 1) + assert out[0, 3].item() == pytest.approx(-2.0) + + +def test_negative_penalty_raises_seen_logits(): + input_ids = torch.tensor([[0, 2]]) + scores = torch.zeros(1, 4) + out = apply_presence_penalty(input_ids, scores, penalty = -0.5, prompt_len = 1) + assert out[0, 2].item() == pytest.approx(0.5) + + +def test_prompt_tokens_excluded(): + # Token 7 is prompt-only (untouched); token 4 in the completion is penalized. + input_ids = torch.tensor([[7, 4, 4]]) + scores = torch.zeros(1, 8) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out[0, 7].item() == pytest.approx(0.0) + assert out[0, 4].item() == pytest.approx(-1.0) + + +def test_batch_rows_isolated(): + input_ids = torch.tensor([[0, 1], [0, 2]]) # row completions [1] and [2] + scores = torch.zeros(2, 4) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out[0, 1].item() == pytest.approx(-1.0) + assert out[0, 2].item() == pytest.approx(0.0) + assert out[1, 2].item() == pytest.approx(-1.0) + assert out[1, 1].item() == pytest.approx(0.0) + + +def test_zero_penalty_is_noop(): + input_ids = torch.tensor([[0, 1, 2]]) + scores = torch.randn(1, 5) + original = scores.clone() + out = apply_presence_penalty(input_ids, scores, penalty = 0.0, prompt_len = 1) + assert torch.equal(out, original) + + +def test_empty_completion_is_noop(): + # prompt_len covers the whole sequence -> nothing generated yet. + input_ids = torch.tensor([[0, 1, 2]]) + scores = torch.randn(1, 5) + original = scores.clone() + out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 3) + assert torch.equal(out, original) + + +def test_out_of_vocab_id_ignored(): + # A generated id >= vocab_size (defensive) must not index out of bounds. + input_ids = torch.tensor([[0, 9]]) + scores = torch.zeros(1, 5) # vocab 5, token 9 is out of range + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert torch.equal(out, torch.zeros(1, 5)) + + +def test_negative_generated_id_ignored(): + # A negative generated id (defensive) must be dropped, not wrap to scores[-1]. + input_ids = torch.tensor([[0, -1]]) + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + # Nothing penalized; in particular the last row (the numpy/torch wrap target + # for id -1) is untouched. + assert torch.equal(out, torch.zeros(1, 5)) + + +def test_mixed_oob_negative_and_valid_ids_only_in_range_penalized(): + # Completion mixes a valid id (1), an out-of-vocab id (9 >= vocab 5) and a + # negative id (-1). Only the in-range distinct id is penalized; OOB/negative + # ids are ignored with no crash and no wrong-index wrap. This fails under the + # old ``seen[seen < vocab_size]`` filter (id -1 wraps to the last row) and + # passes only with the both-ends bound. + input_ids = torch.tensor([[0, 1, 9, -1, 1]]) # prompt [0], completion [1, 9, -1, 1] + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + expected = torch.zeros(1, 5) + expected[0, 1] = -1.0 # once per distinct in-range id (multiplicity ignored) + assert torch.equal(out, expected) + assert out[0, 4].item() == pytest.approx(0.0) # id -1 did not wrap to the last row + + +def test_dtype_and_device_preserved(): + input_ids = torch.tensor([[0, 1]]) + scores = torch.zeros(1, 4, dtype = torch.float16) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out.dtype == torch.float16 + assert out.device == scores.device + + +def test_processor_none_when_zero(): + assert _make_presence_penalty_processor(0.0, prompt_len = 0) is None + + +def test_processor_applies_penalty(): + proc = _make_presence_penalty_processor(1.5, prompt_len = 2) + assert proc is not None + input_ids = torch.tensor([[0, 1, 3]]) + scores = torch.zeros(1, 5) + out = proc(input_ids, scores) + assert out[0, 3].item() == pytest.approx(-1.5) + + +def test_processor_composes_with_other_processors(): + # LogitsProcessorList must run our processor alongside a pre-existing one. + from transformers import LogitsProcessor, LogitsProcessorList + + class _AddToTokenZero(LogitsProcessor): + def __call__(self, input_ids, scores): + scores[:, 0] = scores[:, 0] + 100.0 + return scores + + presence = _make_presence_penalty_processor(1.0, prompt_len = 1) + combined = LogitsProcessorList([_AddToTokenZero(), *presence]) + input_ids = torch.tensor([[5, 2]]) # completion = [2] + scores = torch.zeros(1, 6) + out = combined(input_ids, scores) + assert out[0, 0].item() == pytest.approx(100.0) # other processor ran + assert out[0, 2].item() == pytest.approx(-1.0) # presence ran + + +def test_mlx_presence_penalty_callable(): + mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS") + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + proc = _make_mlx_presence_penalty_processor(1.5) + # First call = prompt only (latches prompt_len, penalizes nothing). + prompt = mx.array([10, 11]) + logits0 = mx.zeros((1, 20)) + out0 = proc(prompt, logits0) + assert float(out0[0, 10]) == pytest.approx(0.0) + # Second call: one completion token (5) appended -> penalized once. + seq = mx.array([10, 11, 5]) + logits1 = mx.zeros((1, 20)) + out1 = proc(seq, logits1) + assert float(out1[0, 5]) == pytest.approx(-1.5) + assert float(out1[0, 10]) == pytest.approx(0.0) # prompt token untouched + + +def test_mlx_presence_penalty_bounds_out_of_range_ids(): + # Documents (and, on Apple Silicon CI, enforces) the intended MLX bound: + # out-of-vocab and negative completion ids must be ignored. MLX does no + # bounds checking and OOB indexing is undefined behavior (crash / memory + # corruption), so the processor routes stray ids to a discarded scratch slot + # and penalizes only in-range distinct ids -- matching the torch filter + # seen[(seen >= 0) & (seen < vocab)]. Skips off arm64 macOS where MLX is absent. + mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS") + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + proc = _make_mlx_presence_penalty_processor(1.0) + proc(mx.array([10, 11]), mx.zeros((1, 8))) # first call latches prompt_len = 2 + # Completion appends a valid id (3), an out-of-vocab id (99 >= vocab 8) and a + # negative id (-1); only the in-range id is penalized and nothing crashes. + seq = mx.array([10, 11, 3, 99, -1]) + out = proc(seq, mx.zeros((1, 8))) + assert float(out[0, 3]) == pytest.approx(-1.0) + for tok in range(8): + if tok != 3: + assert float(out[0, tok]) == pytest.approx(0.0) + + +# Param propagation: route payload -> orchestrator cmd -> worker gen_kwargs +_SAMPLING = { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.05, + "repetition_penalty": 1.1, + "presence_penalty": 1.5, +} + + +def test_orchestrator_cmd_carries_all_sampling_params(): + from core.inference.orchestrator import InferenceOrchestrator + + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + cmd = o._build_generate_cmd( + "req1", + None, + messages = [{"role": "user", "content": "hi"}], + max_new_tokens = 128, + **_SAMPLING, + ) + for key, val in _SAMPLING.items(): + assert cmd[key] == val, f"{key} dropped/altered in orchestrator cmd" + + +def test_worker_forwards_all_sampling_params_to_backend(): + from core.inference.worker import _handle_generate + + class _RecordingBackend: + last_generation_stats = None + + def __init__(self): + self.received = None + + def generate_chat_response(self, **kwargs): + self.received = kwargs + return iter(()) # empty stream -> loop exits, gen_done is sent + + class _FakeQueue: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + cmd = { + "type": "generate", + "request_id": "r", + "messages": [{"role": "user", "content": "hi"}], + "max_new_tokens": 128, + **_SAMPLING, + } + backend = _RecordingBackend() + _handle_generate(backend, cmd, _FakeQueue(), threading.Event()) + + assert backend.received is not None + for key, val in _SAMPLING.items(): + assert backend.received[key] == val, f"{key} dropped/altered in worker gen_kwargs" From af93868760894958d3325775f3e2547d9bfa9c81 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Mon, 6 Jul 2026 23:38:29 -0700 Subject: [PATCH 034/113] Fix repeated base model downloads across checkpoint exports (#6896) * Fix repeated base model downloads across checkpoint exports (#6890) Pre-warm the HF hub cache with the 16bit base weights before merge_and_overwrite_lora runs. The merge fetches shards with hf_hub_download(local_dir=...), which never populates the hub cache, so temporary merge directories (GGUF checkpoint exports) forced a full re-download of the base model for every checkpoint. The first export now downloads once into the cache and later exports copy from it. Skips itself when already cached, offline, on Kaggle/Colab, for local or nf4/fp4 bases, non-downloading save methods, or low disk. Opt out with UNSLOTH_PREWARM_HUB_CACHE=0. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Show MB for small base models in the pre-warm download message * Harden pre-warm: getattr for model config, abspath for relative HF_HUB_CACHE - Read config._name_or_path via getattr so a model without a config skips cleanly instead of taking the outer error path. - abspath the cache probe so a relative HF_HUB_CACHE walks up to a real root rather than "", which would zero the free-space check and skip pre-warm. Both from PR review; each covered by a test that fails without the fix. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pre-warm the live-env HF cache so runtime redirects still hit (#6890) Resolve the hub cache the same way the merge does (unsloth_zoo _active_caches, live env) instead of huggingface_hub's import-time-frozen constants.HF_HUB_CACHE, and pass it as cache_dir to the cached probe, disk check and snapshot_download. Without this, a runtime HF_HOME/HF_HUB_CACHE redirect (unsloth_zoo redirect_hf_cache_if_readonly on a read-only default cache, or Studio) makes the pre-warm populate a different directory than the one the merge reads, so the cache-copy fast path misses and the base re-downloads on every export anyway. Adds 3 regression tests covering the cache_dir threading and the redirect case. * Apply ruff-format kwarg spacing to the pre-warm cache-dir changes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pre-warm the 16bit sibling for FP8 bases so their merged_16bit exports reuse the cache too For a merged_16bit export of an FP8 base with an existing 16bit sibling, the merge swaps to the sibling and downloads that (unsloth_zoo _resolve_fp8_16bit_sibling), so pre-warming the FP8 repo missed the cache and re-downloaded the sibling every export. Mirror the swap and pre-warm the sibling. No sibling still caches the FP8 repo for the in-place dequant path. Adds 2 regression tests. * Filter pre-warm shards through the safetensors index like the merge does Repos that ship a leftover shard set the index does not reference (e.g. granite-3.2) made the disk gate over-count and snapshot_download fetch shards the merge never reads. Mirror the merge: on the download path, keep only index-referenced shards. Runs after the already-cached check so the cached fast path stays network-free. Adds 2 tests. * Tighten pre-warm comments --------- Co-authored-by: Unsloth Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- .../test_prewarm_base_model_hub_cache.py | 459 ++++++++++++++++++ unsloth/save.py | 177 +++++++ 2 files changed, 636 insertions(+) create mode 100644 tests/saving/test_prewarm_base_model_hub_cache.py diff --git a/tests/saving/test_prewarm_base_model_hub_cache.py b/tests/saving/test_prewarm_base_model_hub_cache.py new file mode 100644 index 0000000000..cbb52863ba --- /dev/null +++ b/tests/saving/test_prewarm_base_model_hub_cache.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression tests for #6890: repeated base-model downloads across checkpoint exports. + +merge_and_overwrite_lora downloads missing 16-bit shards with hf_hub_download(local_dir), +which never populates the persistent HF hub cache; a temporary merge directory (Studio +GGUF exports delete it) means every checkpoint export re-downloads the full base model. +_prewarm_base_model_hub_cache snapshot-downloads the base into the hub cache first so +the zoo's cache-copy fast path is hit on later exports. + +unsloth.save cannot be imported on GPU-less hosts, so these tests extract the helper's +source via ast and exec it against fakes, mirroring the other GPU-free tests. +""" + +from __future__ import annotations + +import ast +import json +import os +import types +from pathlib import Path + +import pytest + + +_SAVE_PY = Path(__file__).resolve().parent.parent.parent / "unsloth" / "save.py" +_SOURCE = _SAVE_PY.read_text(encoding = "utf-8") + + +def _extract_function(name: str) -> str: + tree = ast.parse(_SOURCE) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return ast.get_source_segment(_SOURCE, node) + raise AssertionError(f"{name} not found in unsloth/save.py") + + +class _FakePeftModel: + def __init__(self, name_or_path = "unsloth/gemma-4-31b-it-bnb-4bit"): + self.config = types.SimpleNamespace(_name_or_path = name_or_path) + + +class _Recorder: + """Callable that records calls and returns/raises per configuration.""" + + def __init__( + self, + result = None, + exc = None, + results_fn = None, + ): + self.calls = [] + self.result = result + self.exc = exc + self.results_fn = results_fn + + def __call__(self, *args, **kwargs): + self.calls.append((args, kwargs)) + if self.exc is not None: + raise self.exc + if self.results_fn is not None: + return self.results_fn(*args, **kwargs) + return self.result + + +def _build_env( + monkeypatch, + tmp_path, + shards = None, + cached = False, + free_bytes = 10**15, + base_source = None, + kaggle = False, + colab = False, + hub_cache = None, + live_hub_cache = "__same__", + fp8_sibling = None, + sibling_source = None, + index_weight_map = None, +): + """Exec the extracted helper with stubbed collaborators; returns (fn, stubs).""" + shards = ( + shards + if shards is not None + else [ + ("model-00001-of-00002.safetensors", 30 * 1024**3), + ("model-00002-of-00002.safetensors", 29 * 1024**3), + ] + ) + if base_source is None: + base_source = ("unsloth/gemma-4-31b-it", False, None, False, None) + + class _FS: + def __init__(self, token = None): + pass + + def ls( + self, + repo, + detail = True, + ): + return [{"name": f"{repo}/{n}", "size": s} for n, s in shards] + + class _LocalMiss(Exception): + pass + + hf_hub_download = _Recorder(result = str(tmp_path / "cached")) + if not cached: + hf_hub_download.exc = _LocalMiss("not cached") + # Serve model.safetensors.index.json (the merge's shard filter) while shard cache + # probes still miss, so the index-filter path can be exercised without a network. + if index_weight_map is not None: + _idx_path = tmp_path / "model.safetensors.index.json" + _idx_path.write_text(json.dumps({"weight_map": index_weight_map})) + + def _hub_dl( + repo_id = None, + filename = None, + **kw, + ): + if filename == "model.safetensors.index.json": + return str(_idx_path) + raise _LocalMiss("not cached") + + hf_hub_download.exc = None + hf_hub_download.results_fn = _hub_dl + snapshot_download = _Recorder() + determine_base_model_source = _Recorder(result = base_source) + # For the FP8 -> 16bit sibling swap: return the sibling's (16bit) source when the + # helper re-resolves the sibling, else the original base source. + if fp8_sibling is not None: + _sib_src = sibling_source or (fp8_sibling, False, None, False, None) + determine_base_model_source.results_fn = ( + lambda name, token = None: _sib_src if name == fp8_sibling else base_source + ) + resolve_fp8_16bit_sibling = _Recorder(result = fp8_sibling) + + cache_dir = tmp_path / "hub_cache" + cache_dir.mkdir(exist_ok = True) + + hf_module = types.SimpleNamespace( + HfFileSystem = _FS, + hf_hub_download = hf_hub_download, + snapshot_download = snapshot_download, + constants = types.SimpleNamespace( + HF_HUB_CACHE = hub_cache if hub_cache is not None else str(cache_dir) + ), + ) + zoo_module = types.SimpleNamespace( + determine_base_model_source = determine_base_model_source, + _resolve_fp8_16bit_sibling = resolve_fp8_16bit_sibling, + ) + # Stub the live-env cache resolver the pre-warm uses (matches what the merge reads). + _live = ( + (hub_cache if hub_cache is not None else str(cache_dir)) + if live_hub_cache == "__same__" + else live_hub_cache + ) + hf_cache_module = types.SimpleNamespace(_active_caches = lambda: (None, _live, None)) + monkeypatch.setitem(__import__("sys").modules, "huggingface_hub", hf_module) + monkeypatch.setitem(__import__("sys").modules, "unsloth_zoo.saving_utils", zoo_module) + monkeypatch.setitem(__import__("sys").modules, "unsloth_zoo.hf_cache", hf_cache_module) + + fake_shutil = types.SimpleNamespace( + disk_usage = lambda path: types.SimpleNamespace(free = free_bytes) + ) + + prints = [] + namespace = { + "os": os, + "shutil": fake_shutil, + "PeftModel": _FakePeftModel, + "get_model_name": lambda name, load_in_4bit: name.removesuffix("-bnb-4bit"), + "IS_KAGGLE_ENVIRONMENT": kaggle, + "IS_COLAB_ENVIRONMENT": colab, + "print": lambda *a, **k: prints.append(" ".join(str(x) for x in a)), + } + exec( + compile(_extract_function("_prewarm_base_model_hub_cache"), str(_SAVE_PY), "exec"), + namespace, + ) + stubs = types.SimpleNamespace( + snapshot_download = snapshot_download, + hf_hub_download = hf_hub_download, + determine_base_model_source = determine_base_model_source, + resolve_fp8_16bit_sibling = resolve_fp8_16bit_sibling, + prints = prints, + ) + return namespace["_prewarm_base_model_hub_cache"], stubs + + +def test_downloads_base_into_hub_cache(monkeypatch, tmp_path): + monkeypatch.delenv("UNSLOTH_PREWARM_HUB_CACHE", raising = False) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(), save_method = "merged_16bit", token = "tok") + assert len(stubs.snapshot_download.calls) == 1 + _, kwargs = stubs.snapshot_download.calls[0] + assert kwargs["repo_id"] == "unsloth/gemma-4-31b-it" + # No local_dir: the whole point is populating the persistent cache. + assert "local_dir" not in kwargs + assert "model-00001-of-00002.safetensors" in kwargs["allow_patterns"] + assert "model.safetensors.index.json" in kwargs["allow_patterns"] + + +def test_skips_download_when_already_cached(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path, cached = True) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + # The cached check must not hit the network. + assert all(kwargs.get("local_files_only") for _, kwargs in stubs.hf_hub_download.calls) + + +def test_skips_when_disk_too_small_for_cache_copy(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path, free_bytes = 60 * 1024**3) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +@pytest.mark.parametrize("env_value", ["0", "false", "NO", "off"]) +def test_env_opt_out(monkeypatch, tmp_path, env_value): + monkeypatch.setenv("UNSLOTH_PREWARM_HUB_CACHE", env_value) + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +@pytest.mark.parametrize("var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) +def test_offline_skips(monkeypatch, tmp_path, var): + monkeypatch.setenv(var, "1") + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +def test_kaggle_and_colab_skip(monkeypatch, tmp_path): + for flag in ("kaggle", "colab"): + fn, stubs = _build_env(monkeypatch, tmp_path, **{flag: True}) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [], f"{flag} must skip pre-warm" + + +@pytest.mark.parametrize("save_method", ["merged_4bit", "forced_merged_4bit", "lora"]) +def test_non_downloading_save_methods_skip(monkeypatch, tmp_path, save_method): + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(), save_method = save_method) + assert stubs.snapshot_download.calls == [] + + +def test_local_base_model_skips(monkeypatch, tmp_path): + local_dir = tmp_path / "local_base" + local_dir.mkdir() + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(_FakePeftModel(name_or_path = str(local_dir)), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +def test_quantized_base_skips(monkeypatch, tmp_path): + fn, stubs = _build_env( + monkeypatch, + tmp_path, + base_source = ("unsloth/gemma-4-31b-it-bnb-4bit", False, None, True, "nf4"), + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls == [] + + +def test_non_peft_model_skips(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path) + fn(object(), save_method = "merged_16bit") + assert stubs.determine_base_model_source.calls == [] + assert stubs.snapshot_download.calls == [] + + +def test_consolidated_shard_excluded_when_proper_shards_exist(monkeypatch, tmp_path): + fn, stubs = _build_env( + monkeypatch, + tmp_path, + shards = [ + ("consolidated.safetensors", 14 * 1024**3), + ("model-00001-of-00001.safetensors", 14 * 1024**3), + ], + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + _, kwargs = stubs.snapshot_download.calls[0] + assert "consolidated.safetensors" not in kwargs["allow_patterns"] + assert "model-00001-of-00001.safetensors" in kwargs["allow_patterns"] + + +def test_consolidated_only_repo_is_kept(monkeypatch, tmp_path): + fn, stubs = _build_env( + monkeypatch, + tmp_path, + shards = [("consolidated.safetensors", 14 * 1024**3)], + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + _, kwargs = stubs.snapshot_download.calls[0] + assert "consolidated.safetensors" in kwargs["allow_patterns"] + + +def test_listing_failure_is_swallowed(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path) + stubs.determine_base_model_source.exc = RuntimeError("HF is down") + fn(_FakePeftModel(), save_method = "merged_16bit") # must not raise + assert stubs.snapshot_download.calls == [] + + +def test_gpt_oss_bf16_mxfp4_swap_skips(monkeypatch, tmp_path): + fn, stubs = _build_env(monkeypatch, tmp_path) + fn( + _FakePeftModel(name_or_path = "unsloth/gpt-oss-20b-BF16"), + save_method = "mxfp4", + ) + assert stubs.snapshot_download.calls == [] + + +def test_missing_config_skips_cleanly(monkeypatch, tmp_path): + # A model whose config is None must skip silently, not fall into the outer + # exception handler that prints a misleading "Could not pre-cache" warning. + fn, stubs = _build_env(monkeypatch, tmp_path) + model = _FakePeftModel() # a PeftModel instance so the isinstance guard passes + model.config = None + fn(model, save_method = "merged_16bit") + assert stubs.determine_base_model_source.calls == [] + assert stubs.snapshot_download.calls == [] + assert not any( + "Could not pre-cache" in p for p in stubs.prints + ), "missing config took the error path instead of a clean skip" + + +def test_relative_hub_cache_does_not_falsely_skip(monkeypatch, tmp_path): + # A relative HF_HUB_CACHE whose leaf does not exist yet must still resolve to a real + # root for the disk probe; without abspath the walk-up hits "" and pre-warm is skipped. + monkeypatch.chdir(tmp_path) + fn, stubs = _build_env(monkeypatch, tmp_path, hub_cache = "relcache/hub") + fn(_FakePeftModel(), save_method = "merged_16bit") + assert len(stubs.snapshot_download.calls) == 1, "relative cache path falsely skipped pre-warm" + + +def test_generic_save_calls_prewarm_before_merge(): + """unsloth_generic_save must pre-warm the cache before merge_and_overwrite_lora.""" + tree = ast.parse(_SOURCE) + fn = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "unsloth_generic_save" + ) + body_src = ast.get_source_segment(_SOURCE, fn) + prewarm_pos = body_src.find("_prewarm_base_model_hub_cache(") + merge_pos = body_src.find("merge_and_overwrite_lora(") + assert prewarm_pos != -1, "unsloth_generic_save no longer pre-warms the hub cache" + assert merge_pos != -1 + assert prewarm_pos < merge_pos, "pre-warm must run before the merge downloads shards" + + +def test_prewarm_downloads_into_live_env_cache(monkeypatch, tmp_path): + # Download must target the live-env cache (what the merge reads), via cache_dir. + fn, stubs = _build_env(monkeypatch, tmp_path, live_hub_cache = "/mnt/persistent/hf/hub") + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls[0][1]["cache_dir"] == "/mnt/persistent/hf/hub" + + +def test_prewarm_survives_runtime_cache_redirect(monkeypatch, tmp_path): + # Frozen constants (stale dir) vs the merge's runtime-redirected dir: the pre-warm + # must follow the redirect, else the cache-copy fast path misses and #6890 is unfixed. + fn, stubs = _build_env( + monkeypatch, + tmp_path, + hub_cache = "/read-only/original/hub", + live_hub_cache = "/writable/redirect/hub", + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.snapshot_download.calls[0][1]["cache_dir"] == "/writable/redirect/hub" + + +def test_cached_probe_uses_live_env_cache(monkeypatch, tmp_path): + # The already-cached fast path must probe the live-env cache dir too. + fn, stubs = _build_env( + monkeypatch, tmp_path, cached = True, live_hub_cache = "/mnt/persistent/hf/hub" + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + assert stubs.hf_hub_download.calls, "cached probe did not run" + assert all( + kw.get("cache_dir") == "/mnt/persistent/hf/hub" for _, kw in stubs.hf_hub_download.calls + ) + + +def test_fp8_base_prewarms_16bit_sibling_not_fp8_repo(monkeypatch, tmp_path): + # A merged_16bit export of an FP8 base with a 16bit sibling merges onto the sibling, + # so the pre-warm must cache the sibling (what the merge downloads), not the FP8 repo. + fn, stubs = _build_env( + monkeypatch, + tmp_path, + base_source = ("unsloth/Model-FP8", False, None, True, "fp8"), + fp8_sibling = "unsloth/Model", + ) + fn(_FakePeftModel(name_or_path = "unsloth/Model-FP8"), save_method = "merged_16bit") + assert stubs.resolve_fp8_16bit_sibling.calls, "sibling resolver was not consulted" + assert len(stubs.snapshot_download.calls) == 1 + assert stubs.snapshot_download.calls[0][1]["repo_id"] == "unsloth/Model" + + +def test_fp8_base_without_sibling_still_prewarms_fp8_repo(monkeypatch, tmp_path): + # No sibling: the merge dequants the FP8 base in place, so caching the FP8 repo helps. + fn, stubs = _build_env( + monkeypatch, + tmp_path, + base_source = ("unsloth/Model-FP8", False, None, True, "fp8"), + fp8_sibling = None, + ) + fn(_FakePeftModel(name_or_path = "unsloth/Model-FP8"), save_method = "merged_16bit") + assert len(stubs.snapshot_download.calls) == 1 + assert stubs.snapshot_download.calls[0][1]["repo_id"] == "unsloth/Model-FP8" + + +def test_prewarm_filters_shards_through_index(monkeypatch, tmp_path): + # A repo with a leftover shard not referenced by the index: the merge keeps only the + # indexed shards, so the pre-warm must too (else the disk gate over-counts and + # snapshot_download fetches the unused leftover). + fn, stubs = _build_env( + monkeypatch, + tmp_path, + shards = [ + ("model-00001-of-00002.safetensors", 10 * 1024**3), + ("model-00002-of-00002.safetensors", 10 * 1024**3), + ("leftover-00001-of-00001.safetensors", 10 * 1024**3), + ], + index_weight_map = { + "a.weight": "model-00001-of-00002.safetensors", + "b.weight": "model-00002-of-00002.safetensors", + }, + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + allow = stubs.snapshot_download.calls[0][1]["allow_patterns"] + assert "leftover-00001-of-00001.safetensors" not in allow + assert "model-00001-of-00002.safetensors" in allow + assert "model-00002-of-00002.safetensors" in allow + + +def test_prewarm_keeps_all_shards_when_index_matches(monkeypatch, tmp_path): + # No leftover: every listed shard is indexed, so none are dropped. + fn, stubs = _build_env( + monkeypatch, + tmp_path, + shards = [ + ("model-00001-of-00002.safetensors", 10 * 1024**3), + ("model-00002-of-00002.safetensors", 10 * 1024**3), + ], + index_weight_map = { + "a.weight": "model-00001-of-00002.safetensors", + "b.weight": "model-00002-of-00002.safetensors", + }, + ) + fn(_FakePeftModel(), save_method = "merged_16bit") + allow = stubs.snapshot_download.calls[0][1]["allow_patterns"] + assert "model-00001-of-00002.safetensors" in allow + assert "model-00002-of-00002.safetensors" in allow diff --git a/unsloth/save.py b/unsloth/save.py index 020c63a9e2..0b408a1889 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -3693,6 +3693,182 @@ from unsloth_zoo.llama_cpp import ( ) +def _prewarm_base_model_hub_cache( + model, + save_method = "merged_16bit", + token = None, +): + """Download the 16-bit base weights into the persistent HF hub cache before the merge. + + merge_and_overwrite_lora fetches missing shards with hf_hub_download(local_dir = ...), + which never populates the hub cache. When the merge directory is temporary (GGUF + checkpoint exports delete it after conversion), every export re-downloads the full + base model (#6890). Pre-warming the cache makes the first export download once and + later exports copy from the cache. Best-effort: any failure or skip falls back to + the streaming download. Disable with UNSLOTH_PREWARM_HUB_CACHE=0. + """ + _false = ("0", "false", "no", "off") + if os.environ.get("UNSLOTH_PREWARM_HUB_CACHE", "1").strip().lower() in _false: + return + if IS_KAGGLE_ENVIRONMENT or IS_COLAB_ENVIRONMENT: + return + _true = ("1", "true", "yes", "on") + if ( + os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _true + or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _true + ): + return + # Only the 16bit / mxfp4 merges download the base model; merged_4bit and lora do not. + if save_method not in ("merged_16bit", "mxfp4"): + return + if not isinstance(model, PeftModel): + return + + try: + # getattr so a model without a config / _name_or_path skips instead of raising. + name_or_path = getattr(getattr(model, "config", None), "_name_or_path", None) + if not name_or_path: + return + try: + model_name = get_model_name(name_or_path, load_in_4bit = False) + except Exception: + model_name = name_or_path + if not model_name or os.path.isdir(model_name): + return # local checkpoints are copied, never downloaded + + # The merge may swap a gpt-oss "-BF16" repo for its MXFP4 variant, so skip it. + if save_method == "mxfp4" and model_name.endswith("-BF16"): + return + + from unsloth_zoo.saving_utils import determine_base_model_source + + model_name, is_local_path, _, base_is_quantized, quant_type = determine_base_model_source( + model_name, token + ) + if not model_name or is_local_path: + return + # Mirror the merge: an FP8 base with a 16bit sibling merges onto the sibling, so + # pre-warm the sibling (what the merge downloads), not the FP8 repo (#6890). + if base_is_quantized and quant_type == "fp8" and save_method == "merged_16bit": + try: + from unsloth_zoo.saving_utils import _resolve_fp8_16bit_sibling + sibling = _resolve_fp8_16bit_sibling(model_name, token) + except Exception: + sibling = None + if sibling: + model_name, is_local_path, _, base_is_quantized, quant_type = ( + determine_base_model_source(sibling, token) + ) + if not model_name or is_local_path: + return + if base_is_quantized and quant_type in ("nf4", "fp4"): + return # the 16bit merge refuses these bases; nothing worth caching + + from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download + + # Resolve the cache from the live env like the merge, not huggingface_hub's frozen + # constants: a runtime cache redirect (read-only default, Studio) would else miss (#6890). + try: + from unsloth_zoo.hf_cache import _active_caches + _hub_cache = _active_caches()[1] + hub_cache_dir = str(_hub_cache) if _hub_cache is not None else None + except Exception: + hub_cache_dir = None + + # Mirror the zoo's shard listing (drop consolidated.safetensors when proper + # shards coexist) so the cached set is a superset of what the merge looks up. + shard_names = [] + total_size_in_bytes = 0 + for x in HfFileSystem(token = token).ls(model_name, detail = True): + if x["name"].endswith(".safetensors"): + shard_names.append((os.path.split(x["name"])[-1], int(x.get("size") or 0))) + if any(name != "consolidated.safetensors" for name, _ in shard_names): + shard_names = [x for x in shard_names if x[0] != "consolidated.safetensors"] + if not shard_names: + return + + try: + for filename, _ in shard_names: + hf_hub_download( + repo_id = model_name, + filename = filename, + cache_dir = hub_cache_dir, + local_files_only = True, + token = token, + ) + return # already fully cached + except Exception: + pass + + # Mirror the merge's index filter (download path only): some repos ship leftover shards + # the index omits; keep only indexed ones, else the disk gate over-counts and we fetch + # unused shards. + if len(shard_names) > 1: + try: + import json as _json + + _idx = hf_hub_download( + repo_id = model_name, + filename = "model.safetensors.index.json", + cache_dir = hub_cache_dir, + token = token, + ) + with open(_idx, encoding = "utf-8") as _f: + _indexed = { + os.path.split(v)[-1] for v in _json.load(_f).get("weight_map", {}).values() + } + if _indexed and not {n for n, _ in shard_names}.issubset(_indexed): + _kept = [x for x in shard_names if x[0] in _indexed] + if _kept: + shard_names = _kept + except Exception: + pass + total_size_in_bytes = sum(size for _, size in shard_names) + + # The cache copy is extra disk on top of the merge working copy; need room for both. + from huggingface_hub import constants as _hf_constants + + # abspath so a relative HF_HUB_CACHE walks up to an existing root, not "". + cache_probe = os.path.abspath( + os.path.expanduser(str(hub_cache_dir or _hf_constants.HF_HUB_CACHE)) + ) + while cache_probe and not os.path.exists(cache_probe): + parent = os.path.dirname(cache_probe) + if parent == cache_probe: + break + cache_probe = parent + free_space = shutil.disk_usage(cache_probe).free if os.path.exists(cache_probe) else 0 + if free_space < 2 * total_size_in_bytes: + print( + f"Unsloth: Not enough free disk to keep `{model_name}` in the Hugging Face " + f"cache (need ~{round(2 * total_size_in_bytes / 1024**3, 1)}GB free, have " + f"{round(free_space / 1024**3, 1)}GB). Downloading straight to the merge " + f"directory instead; the next export will re-download it." + ) + return + + if total_size_in_bytes >= 0.1 * 1024**3: + size_str = f"{round(total_size_in_bytes / 1024**3, 1)}GB" + else: + size_str = f"{max(1, round(total_size_in_bytes / 1024**2))}MB" + print( + f"Unsloth: Downloading `{model_name}` into the Hugging Face cache so future " + f"exports skip the {size_str} download..." + ) + snapshot_download( + repo_id = model_name, + allow_patterns = [name for name, _ in shard_names] + + ["model.safetensors.index.json", "tokenizer.model"], + cache_dir = hub_cache_dir, + token = token, + ) + except Exception as e: + print( + f"Unsloth: Could not pre-cache the base model weights ({e}). " + f"Falling back to downloading into the merge directory." + ) + + @torch.inference_mode def save_to_gguf_generic( model, @@ -3888,6 +4064,7 @@ def unsloth_generic_save( print(f"Unsloth: Model saved successfully to '{save_directory}'") else: + _prewarm_base_model_hub_cache(model, save_method = save_method, token = token) merge_and_overwrite_lora( get_model_name, model = model, From 08226c2475e81a2c46336787f60d20a47e47ccea Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 23:48:23 -0700 Subject: [PATCH 035/113] Studio: fix torch CUDA undefined-symbol errors from a conflicting LD_LIBRARY_PATH (#6905) * Studio: re-exec to prepend torch's bundled CUDA libs to LD_LIBRARY_PATH On Linux the dynamic linker reads LD_LIBRARY_PATH before the RUNPATH baked into torch's .so files, so a pre-existing LD_LIBRARY_PATH pointing at a system CUDA (conda, a Docker base image, /usr/local/cuda-*/lib64) shadows torch's bundled nvidia/*/lib libraries and causes undefined-symbol errors when the Studio backend imports torch. Detect torch's lib dirs without importing torch, prepend them to LD_LIBRARY_PATH, and re-exec once (LD_LIBRARY_PATH is only read at process start). Linux-only, sentinel-guarded against re-exec loops, and called only from run.py's __main__ so library/embedder imports (e.g. Colab's `from run import run_server`) are never re-exec'd. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/run.py | 79 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/studio/backend/run.py b/studio/backend/run.py index ccd0113972..2cc6c4a93e 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -12,6 +12,79 @@ import time from pathlib import Path from typing import Optional + +def _fix_torch_cuda_ld_path(): + """Prepend torch's bundled CUDA libs to LD_LIBRARY_PATH. + + PyTorch wheels ship their own CUDA runtime (libcudart, libcublas, ...) in + ``site-packages/nvidia/*/lib``. On Linux the dynamic linker reads + LD_LIBRARY_PATH before the RUNPATH baked into torch's .so files, so a + pre-existing LD_LIBRARY_PATH pointing at a different system CUDA (e.g. + /usr/local/cuda-13/lib64 from conda or a Docker base image) shadows torch's + libs and triggers "undefined symbol" errors when torch is imported. Detect + torch's lib dirs (without importing torch) and prepend them. Returns True if + LD_LIBRARY_PATH was changed. + """ + if sys.platform != "linux": + return False + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + if not ld_path: + return False + try: + import importlib.util + + spec = importlib.util.find_spec("torch") + if not spec or not spec.origin: + return False + torch_dir = os.path.dirname(spec.origin) + site_pkgs = os.path.dirname(torch_dir) + nvidia_dir = os.path.join(site_pkgs, "nvidia") + + lib_dirs = [] + torch_lib = os.path.join(torch_dir, "lib") + if os.path.isdir(torch_lib): + lib_dirs.append(torch_lib) + if os.path.isdir(nvidia_dir): + for sub in sorted(os.listdir(nvidia_dir)): + lib = os.path.join(nvidia_dir, sub, "lib") + if os.path.isdir(lib): + lib_dirs.append(lib) + if not lib_dirs: + return False + + existing = ld_path.split(":") + if existing[: len(lib_dirs)] == lib_dirs: + return False # already at the front, nothing to do + + torch_set = set(lib_dirs) + cleaned = [p for p in existing if p not in torch_set] + os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + cleaned) + return True + except Exception: + return False + + +_LD_FIXED_SENTINEL = "_UNSLOTH_STUDIO_LD_FIXED" + + +def _maybe_reexec_for_cuda_ld_path(): + """Re-exec once so the dynamic linker sees the corrected LD_LIBRARY_PATH. + + LD_LIBRARY_PATH is read at process start, so editing os.environ in-process + cannot fix the running interpreter; a single re-exec is required. Call only + from a true entry point (the ``if __name__ == "__main__"`` block), never at + import time, because os.execv replaces the whole process (an embedder such + as Colab that does ``from run import run_server`` must not be re-exec'd). + """ + if _LD_FIXED_SENTINEL in os.environ: + return + if not _fix_torch_cuda_ld_path(): + return + os.environ[_LD_FIXED_SENTINEL] = "1" + argv = getattr(sys, "orig_argv", None) or [sys.executable, *sys.argv] + os.execv(sys.executable, argv) + + # Suppress C-level dependency warnings globally (e.g. SwigPyPacked). os.environ["PYTHONWARNINGS"] = "ignore" @@ -1457,6 +1530,12 @@ def _build_arg_parser(): # For direct execution (also invoked by CLI via os.execvp / subprocess). if __name__ == "__main__": + # Correct a conflicting system CUDA on LD_LIBRARY_PATH before torch is + # imported (below, via run_server). Re-execs once on Linux so the dynamic + # linker uses torch's bundled CUDA libs; no-op on other platforms, when + # LD_LIBRARY_PATH is unset or already correct, or after the single re-exec. + _maybe_reexec_for_cuda_ld_path() + import signal import traceback From 69f8e0b228200ea8e748b5c3a118eef71f441168 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 00:06:48 -0700 Subject: [PATCH 036/113] Clear stale yolo approval state on no-launch reruns (#6868) * Clear stale yolo approval state on no-launch reruns The no-launch session config dir is deliberately reused across runs, but the config writers only ever added the --yolo auto-approval settings and never removed them. After one --yolo --no-launch run, every later run without --yolo kept OpenClaw's tools.exec security=full/ask=off policy plus exec-approvals.json, and OpenCode's permission allow block, so tool execution stayed silently pre-approved. Non-yolo runs now reset that state: OpenClaw drops the exec policy keys and the yolo defaults in exec-approvals.json (approvals OpenClaw itself recorded are kept; the file is removed when only the yolo payload is left), and OpenCode drops the permission block. Launch mode is untouched since it already uses an ephemeral temp dir. * Strip only yolo-written values on non-yolo cleanup Match each field against the exact value the yolo path writes before removing it, so a stricter exec policy, approvals defaults set by the user or the OpenClaw UI, and deny/ask OpenCode permission entries all survive a plain no-launch rerun. An unparseable exec-approvals.json is left in place, matching how an unparseable config is handled. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Write a prompting policy on non-yolo instead of deleting to a permissive default OpenClaw and OpenCode both treat an omitted policy as permissive: OpenClaw's effective exec policy for an unset tools.exec is security=full/ask=off on the gateway host, and OpenCode defaults an unset permission to allow. So clearing the yolo values on a non-yolo run did not restore prompting, it fell back to those permissive defaults and left tool execution auto-approved. A non-yolo run now writes an explicit prompting policy: OpenClaw gets security=allowlist/ask=on-miss (verified to prompt even with the approvals file removed, since the stricter of config and approvals wins), and OpenCode gets edit/bash/webfetch=ask. Only a permissive/yolo value is tightened; a stricter deny (or an ask the user set) is preserved, and the yolo approvals defaults are still stripped. The file-edit CI path opts opencode/openclaw into --yolo, since those agents now prompt by default and the headless test needs auto-approval. * Respect existing exec mode, sandbox/node host, and global permission rules on non-yolo reset The non-yolo reset for openclaw/opencode assumed an omitted policy was the permissive yolo default and rewrote it, which corrupted or weakened stricter setups it should have preserved: - OpenClaw tools.exec.mode is the normalized policy knob and cannot be combined with explicit security/ask (OpenClaw rejects the whole config), so writing security+ask alongside a mode:deny/ask policy both broke the config and relaxed it. Leave a mode-based policy untouched. - host=sandbox defaults to security=deny and host=node routes to a paired node; neither is written by --yolo (which only writes host=gateway). Treating the missing security as full and popping host broadened those into gateway/auto exec. Only rewrite a gateway-routed permissive policy, and never pop a non-gateway host. - OpenCode permission can be a string ("deny") or a {"*": ...} catch-all. The old code dropped a string form and overrode a catch-all by writing per-tool ask, weakening a stricter user rule. Now a string is left in place, a catch-all governs absent tools, and only an effective allow is tightened. - The non-yolo ask policy only lived in OPENCODE_CONFIG, which loads below project opencode.json, so a project config allowing edit/bash/webfetch still auto-approved. Carry the ask policy in OPENCODE_CONFIG_CONTENT (above project config) too, symmetric to how yolo carries its allow. Also harden the openclaw path against a malformed non-dict tools value. Adds tests for mode/sandbox/node hosts, string and catch-all permissions, and the inline ask policy over a project config. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope non-yolo resets to the exact yolo fingerprint and preserve granular denies OpenClaw: reset only the exact host=gateway/security=full/ask=off policy --yolo writes, so an omitted or host=auto/sandbox/node policy (which can resolve to a sandbox security=deny default) is no longer broadened to allowlist/on-miss, and a deliberate tools.exec.mode is left alone (OpenClaw never migrates our security/ask write into a mode). OpenCode: carry a granular object or a deny inline verbatim so a per-tool user rule is not collapsed to a blanket ask, but floor any object that grants allow anywhere to the string ask (which fully replaces a project object) so no inline allow pattern can leak through into a silent auto-approve on a non-yolo session. * Stop overriding project config on non-yolo; require full approvals fingerprint The non-yolo OpenCode reset carried a session permission in OPENCODE_CONFIG_CONTENT, which outranks the project opencode.json we cannot read. That inline override could not correctly reflect the project: it weakened a project deny to a prompt, mishandled global string rules, leaked through a granular object's permissive default when no catch-all was present, collapsed an object with an allow (losing its deny), and missed per-agent permissions. All of these stem from forcing a value over an unknown project config. A non-yolo run now only undoes what --yolo wrote: it flips our own explicit per-tool allow back to ask in our config file and carries no permission inline, so the project's own permissions are honored as written. Clearing our persisted yolo state is the actual fix; --yolo still carries its allow inline so it works over a project config. OpenClaw approvals cleanup now strips the yolo defaults only when the full fingerprint (security=full, ask=off, askFallback=full) is present, so a mixed user policy that merely shares askFallback=full (whose omitted default is deny) is kept intact. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/scripts/agent-guides-drive.sh | 12 +- unsloth_cli/commands/start.py | 102 +++++++-- unsloth_cli/tests/test_start.py | 293 +++++++++++++++++++++++++- 3 files changed, 389 insertions(+), 18 deletions(-) diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 9b85b20177..d430d2c172 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -154,7 +154,12 @@ raw_env() { # $1 = var name -> value (one shlex-quote layer stripped) # writers as a side effect (it writes each agent's relocated session config). parse_connect() { local raw="$LOGS_DIR/connect-${AGENT}.txt" - if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + # CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their + # config (which now prompts by default), so the file-edit test opts into auto-approval + # here, the same intent as claude/codex's per-call bypass flags. + local yolo=() + [ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo) + if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then cat_redacted "$raw" guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero" fi @@ -394,7 +399,10 @@ case "$MODE" in T2='Run hello.py with python and show me the exact output.' # The start.py recipe writers + crosscheck must see the repo; run them - # from the repo root BEFORE cd-ing into the scratch work dir. + # from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw + # gate tool approval through their config (prompting by default), so file-edit + # opts them into auto-approval to run edits/commands headlessly. + case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac parse_connect crosscheck_contract # File-edit needs real tools, so we cannot zero them as in connection. diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index b188180188..1895125b11 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -983,7 +983,9 @@ def _session_config(agent: str, launch: bool): else: # Never wipe this dir: a previously printed recipe may still be running # an agent whose sessions/state live here, and every config writer - # merges idempotently into an existing home anyway. + # merges idempotently into an existing home anyway. Writers must also + # reset any state a previous run's flags left behind (--yolo especially), + # since files here outlive the invocation that wrote them. path = _agents_config_root() / agent path.mkdir(parents = True, exist_ok = True, mode = 0o700) yield path @@ -1043,6 +1045,62 @@ def write_openclaw_config( {"version": 1, "defaults": {"security": "full", "ask": "off", "askFallback": "full"}}, ) typer.echo(f"Updated {approvals}") + else: + # The no-launch config dir is reused across runs, so a previous --yolo run may + # have left auto-approval state behind. OpenClaw treats an omitted exec policy as + # security=full, ask=off on the gateway host, so deleting the keys would keep + # auto-approval on: a non-yolo run must WRITE a prompting policy. Only a + # permissive/yolo policy is replaced; a stricter one set by hand survives. + tools = config.get("tools") + exec_policy = tools.get("exec") if isinstance(tools, dict) else None + exec_policy = exec_policy if isinstance(exec_policy, dict) else {} + # Match ONLY the exact fingerprint --yolo writes (host=gateway, security=full, + # ask=off, all explicit, no mode); anything else is left untouched. host=auto or an + # omitted host resolves to security=deny under an active sandbox, so treating those + # as the permissive gateway default would broaden a fresh sandboxed config from + # deny to allowlist. host=node and host=sandbox are user-set (--yolo only writes + # gateway). tools.exec.mode is OpenClaw's normalized knob (it cannot be combined + # with security/ask, and OpenClaw never rewrites our security/ask write into it), + # so a mode is always a deliberate user policy; never clobber it. + permissive = ( + "mode" not in exec_policy + and exec_policy.get("host") == "gateway" + and exec_policy.get("security") == "full" + and exec_policy.get("ask") == "off" + ) + if permissive: + exec_policy = _subdict(_subdict(config, "tools"), "exec") + exec_policy.pop("host", None) # routing only; defaults to the gateway host + exec_policy["security"] = "allowlist" # only allowlisted commands skip approval + exec_policy["ask"] = "on-miss" # prompt on every non-allowlisted command + # Drop the yolo defaults from the host approvals file (a stricter default set by + # the user or OpenClaw is kept). With a prompting tools.exec the stricter of the + # two layers wins, so an omitted approvals default still prompts. + approvals = path.parent / "exec-approvals.json" + if approvals.exists(): + state = _read_json_object(approvals) + if state is not None: + defaults = state.get("defaults") + # Strip the defaults only when they are exactly the yolo fingerprint; a + # user-managed mixed policy that merely shares a field (e.g. askFallback=full, + # whose omitted default is deny) must be kept intact. + yolo_defaults = (("security", "full"), ("ask", "off"), ("askFallback", "full")) + is_yolo = isinstance(defaults, dict) and all( + defaults.get(k) == v for k, v in yolo_defaults + ) + if is_yolo: + for k, _ in yolo_defaults: + del defaults[k] + if not defaults: + del state["defaults"] + if set(state) <= {"version"}: + # Nothing left but our own yolo payload: remove it. + approvals.unlink() + typer.echo(f"Removed {approvals}") + else: + # Keep approvals OpenClaw itself recorded; only the yolo defaults go. + _write_private_json(approvals, state) + typer.echo(f"Updated {approvals}") if json.dumps(config, sort_keys = True) != before: _write_private_json(path, config) typer.echo(f"Updated {path}") @@ -1054,7 +1112,7 @@ def write_opencode_config( model: dict, path: Path, yolo: bool = False, -) -> None: +) -> dict: config = _read_json_object(path) if config is None: typer.echo( @@ -1062,7 +1120,7 @@ def write_opencode_config( "yourself, or move the file aside and re-run.", err = True, ) - return + return {} before = json.dumps(config, sort_keys = True) config.setdefault("$schema", "https://opencode.ai/config.json") model_entry = {"name": model["id"]} @@ -1087,13 +1145,32 @@ def write_opencode_config( compaction = _subdict(config, "compaction") compaction["auto"] = True compaction["reserved"] = max(1, window // 10) + tools = ("edit", "bash", "webfetch") if yolo: # OpenCode has no --yolo flag; auto-approve is the config `permission` block - # (singular). Allow the prompting tools so tool calls don't block on the TUI. - config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + # (singular). Allow the prompting tools so tool calls don't block on the TUI. This + # rides inline (OPENCODE_CONFIG_CONTENT) so --yolo works even over a project config. + session_permission = {t: "allow" for t in tools} + config["permission"] = dict(session_permission) + else: + # Undo only what --yolo wrote: our yolo sets an explicit per-tool "allow" for these + # three tools, so flip exactly those explicit allows back to "ask". A "deny"/"ask", + # a granular object, a string, or a "*" catch-all is the user's own rule and is left + # untouched. We do NOT carry a permission inline for a non-yolo session: since + # OPENCODE_CONFIG_CONTENT outranks the project opencode.json we cannot read, any + # value forced there would override the user's project rules (weakening a project + # deny, or auto-approving through a granular object's permissive default). Clearing + # our own persisted yolo state is the fix; the project's own permissions are honored. + session_permission: dict = {} + permission = config.get("permission") + if isinstance(permission, dict): + for tool in tools: + if permission.get(tool) == "allow": + permission[tool] = "ask" if json.dumps(config, sort_keys = True) != before: _write_private_json(path, config) typer.echo(f"Updated {path}") + return session_permission def write_hermes_config(base: str, model: dict, path: Path) -> None: @@ -1379,14 +1456,15 @@ def opencode( # OPENCODE_CONFIG is an overlay (loaded between the user's global and project # configs), so this adds the Unsloth provider/model for the session without # changing the user's default model. Key lives in the config, not the env. - write_opencode_config(base, key, entry, config_path, yolo = yolo) - # A project's own opencode.json outranks OPENCODE_CONFIG, so the session model - # pin (and --yolo permissions) would silently lose to a repo config. Carry the - # settings that must win in OPENCODE_CONFIG_CONTENT, which outranks project - # config; the API key stays in the private file, never in the printed env. + session_permission = write_opencode_config(base, key, entry, config_path, yolo = yolo) + # A project's own opencode.json outranks OPENCODE_CONFIG, so the session model pin + # would silently lose to a repo config. Carry it in OPENCODE_CONFIG_CONTENT, which + # outranks project config; the API key stays in the private file, never the env. + # Only --yolo carries a permission here (its allow must win over a project config); + # a non-yolo session returns no permission, so the project's own rules are honored. inline_config: dict = {"model": f"unsloth/{entry['id']}"} - if yolo: - inline_config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + if session_permission: + inline_config["permission"] = session_permission env = { "OPENCODE_CONFIG": str(config_path), "OPENCODE_CONFIG_CONTENT": json.dumps(inline_config), diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index a6a092a17c..affc24626e 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -446,7 +446,10 @@ def test_opencode_inline_config_beats_project_config(fake_studio): assert "sk-unsloth" not in content_line # key stays in the private file -def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio): +def test_opencode_inline_config_omits_permission_without_yolo(fake_studio): + # A non-yolo session carries no permission inline. OPENCODE_CONFIG_CONTENT outranks the + # project opencode.json we cannot read, so forcing any value there would override the + # user's project rules; clearing our own config is the fix, and the inline pins the model. result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert result.exit_code == 0, result.output content_line = next( @@ -455,7 +458,8 @@ def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio): inline = json.loads( shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] ) - assert inline == {"model": f"unsloth/{MODEL['id']}"} + assert inline["model"] == f"unsloth/{MODEL['id']}" + assert "permission" not in inline def test_https_loopback_never_auto_serves(fake_studio, monkeypatch): @@ -1676,9 +1680,31 @@ def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path): result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert result.exit_code == 0, result.output config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + # A non-yolo run on a fresh config writes no permission block; it only flips a prior + # --yolo run's explicit allow back to ask (see the yolo-then-plain test below). assert "permission" not in config +def test_no_yolo_opencode_flips_prior_yolo_allow_to_ask(fake_studio, tmp_path): + # The core reset: a --yolo run wrote explicit per-tool allow; a later non-yolo run + # must flip exactly those back to ask so nothing stays auto-approved. + yolo = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) + assert yolo.exit_code == 0, yolo.output + config_path = tmp_path / "agents" / "opencode" / "opencode.json" + assert json.loads(config_path.read_text())["permission"] == { + "edit": "allow", + "bash": "allow", + "webfetch": "allow", + } + plain = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert plain.exit_code == 0, plain.output + assert json.loads(config_path.read_text())["permission"] == { + "edit": "ask", + "bash": "ask", + "webfetch": "ask", + } + + def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path): result = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"]) assert result.exit_code == 0, result.output @@ -1691,12 +1717,17 @@ def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path): assert approvals["defaults"] == {"security": "full", "ask": "off", "askFallback": "full"} -def test_no_yolo_openclaw_has_no_exec_policy(fake_studio, tmp_path): +def test_no_yolo_openclaw_leaves_fresh_config_untouched(fake_studio, tmp_path): + # A fresh non-yolo run only undoes state a prior --yolo wrote; with no yolo + # fingerprint present it must not synthesize an exec policy. An omitted policy can + # resolve to a sandbox default of security=deny, so writing allowlist here would + # BROADEN it. The reset is scoped to the exact yolo write, verified by the + # yolo-then-plain round trip below. result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) assert result.exit_code == 0, result.output state = tmp_path / "agents" / "openclaw" config = json.loads((state / "openclaw.json").read_text()) - assert "exec" not in config.get("tools", {}) # no auto-approve policy without --yolo + assert "exec" not in config.get("tools", {}) assert not (state / "exec-approvals.json").exists() @@ -1719,6 +1750,260 @@ def test_write_openclaw_config_yolo_unit(tmp_path): } +def test_no_launch_rerun_clears_stale_opencode_yolo_permissions(fake_studio, tmp_path): + # The no-launch config dir is reused across runs, so a --yolo run persists its + # auto-approve settings; a later run without --yolo must strip them, not leave + # tool execution silently pre-approved. + yolo = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) + assert yolo.exit_code == 0, yolo.output + config_path = tmp_path / "agents" / "opencode" / "opencode.json" + assert "permission" in json.loads(config_path.read_text()) + plain = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert plain.exit_code == 0, plain.output + config = json.loads(config_path.read_text()) + # The yolo allow policy is replaced by a prompting one, not deleted (which would + # revert to OpenCode's permissive "allow" default). + assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"} + # The session provider survives the cleanup. + assert "unsloth" in config["provider"] + + +def test_no_launch_rerun_clears_stale_openclaw_yolo_state(fake_studio, tmp_path): + yolo = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"]) + assert yolo.exit_code == 0, yolo.output + state = tmp_path / "agents" / "openclaw" + assert (state / "exec-approvals.json").exists() + plain = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert plain.exit_code == 0, plain.output + config = json.loads((state / "openclaw.json").read_text()) + # The yolo policy is replaced by a prompting one, not deleted (which would revert + # to OpenClaw's permissive default), and the yolo approvals file is gone. + assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"} + assert not (state / "exec-approvals.json").exists() + # The session provider survives the cleanup. + assert "unsloth" in config["models"]["providers"] + + +def test_write_openclaw_config_yolo_then_plain_unit(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + # A plain rerun replaces the yolo policy with a prompting one (deleting it would + # fall back to OpenClaw's permissive default) and removes the yolo approvals file. + assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"} + assert not (path.parent / "exec-approvals.json").exists() + + +def test_write_opencode_config_yolo_then_plain_unit(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + # A plain rerun replaces the yolo allow policy with a prompting one. + assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"} + + +def test_openclaw_non_yolo_keeps_runtime_approvals(tmp_path): + # OpenClaw records its own entries in exec-approvals.json (OPENCLAW_STATE_DIR is + # this dir); the non-yolo reset drops only the yolo defaults, not those. + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + approvals = path.parent / "exec-approvals.json" + state = json.loads(approvals.read_text()) + state["agents"] = {"main": {"allowlist": ["git status"]}} + approvals.write_text(json.dumps(state)) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + remaining = json.loads(approvals.read_text()) + assert "defaults" not in remaining + assert remaining["agents"] == {"main": {"allowlist": ["git status"]}} + + +def test_openclaw_non_yolo_keeps_mixed_approval_defaults(tmp_path): + # A mixed user-managed defaults block that only shares a field with the yolo payload + # (here askFallback=full, whose omitted default is deny) is not stale yolo state, so a + # non-yolo run leaves it intact rather than stripping the shared field. + path = tmp_path / "openclaw.json" + approvals = path.parent / "exec-approvals.json" + mixed = { + "version": 1, + "defaults": {"security": "allowlist", "ask": "on-miss", "askFallback": "full"}, + } + approvals.write_text(json.dumps(mixed)) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert json.loads(approvals.read_text()) == mixed + + +def test_openclaw_non_yolo_leaves_partial_policy_untouched(tmp_path): + # A policy that lacks the full yolo fingerprint (here no host and no security) is not + # our --yolo write, so a non-yolo run leaves it as-is rather than assuming ask=off + # means permissive: an omitted host/security can resolve to a sandbox deny default. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"timeout": 30, "ask": "off"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"timeout": 30, "ask": "off"} + + +def test_openclaw_non_yolo_leaves_no_permissive_values(tmp_path): + # The whole point of the reset: after a yolo run, a plain run must leave neither the + # config nor the approvals file at OpenClaw's permissive (security=full, ask=off) + # default, or exec still auto-approves. + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + exec_policy = json.loads(path.read_text())["tools"]["exec"] + assert exec_policy.get("security") != "full" + assert exec_policy.get("ask") != "off" + assert not (path.parent / "exec-approvals.json").exists() + + +def test_openclaw_non_yolo_preserves_stricter_exec_policy(tmp_path): + # A policy that doesn't carry the yolo values (for example stricter security or + # prompting turned on) was not written by --yolo and must survive a plain run. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"security": "deny", "ask": "on"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"security": "deny", "ask": "on"} + + +def test_openclaw_non_yolo_preserves_stricter_approval_defaults(tmp_path): + # exec-approvals.json defaults that don't match the yolo payload (stricter + # settings from the user or the OpenClaw UI) are kept, and the file stays. + path = tmp_path / "openclaw.json" + approvals = path.parent / "exec-approvals.json" + approvals.write_text( + json.dumps({"version": 1, "defaults": {"security": "allowlist", "ask": "on"}}) + ) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + state = json.loads(approvals.read_text()) + assert state["defaults"] == {"security": "allowlist", "ask": "on"} + + +def test_openclaw_non_yolo_leaves_unparseable_approvals(tmp_path): + # An unreadable approvals file is left in place rather than deleted, matching + # how an unparseable config is handled. + path = tmp_path / "openclaw.json" + approvals = path.parent / "exec-approvals.json" + approvals.write_text("{not json") + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert approvals.read_text() == "{not json" + + +def test_opencode_non_yolo_flips_only_explicit_allow(tmp_path): + # Only a tool explicitly set to "allow" (what --yolo writes) is flipped to "ask". A + # deny/ask a user set is kept, and an absent tool is not added. + path = tmp_path / "opencode.json" + path.write_text(json.dumps({"permission": {"edit": "allow", "bash": "deny", "read": "ask"}})) + session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["permission"] == {"edit": "ask", "bash": "deny", "read": "ask"} + assert session == {} # a non-yolo session carries no permission inline + + +def test_opencode_non_yolo_leaves_string_permission(tmp_path): + # A global string rule ("deny") is a user-managed catch-all; leave it untouched and + # carry no inline override. + path = tmp_path / "opencode.json" + path.write_text(json.dumps({"permission": "deny"})) + session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert json.loads(path.read_text())["permission"] == "deny" + assert session == {} + + +def test_opencode_non_yolo_leaves_catch_all_and_flips_explicit_allow(tmp_path): + # A "*" catch-all is the user's own rule, never something --yolo writes (yolo sets + # explicit per-tool allow), so it is left intact; an explicit per-tool "allow" is still + # flipped to "ask", but an absent tool inheriting the catch-all is not touched. + path = tmp_path / "opencode.json" + path.write_text(json.dumps({"permission": {"*": "allow", "bash": "allow"}})) + session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert json.loads(path.read_text())["permission"] == {"*": "allow", "bash": "ask"} + assert session == {} + + +def test_opencode_non_yolo_leaves_granular_object(tmp_path): + # A granular object value is a user rule (yolo only ever writes a plain "allow" string), + # so it is left in the file verbatim and never carried inline. + path = tmp_path / "opencode.json" + obj = {"read *": "deny", "git *": "ask"} + path.write_text(json.dumps({"permission": {"bash": dict(obj)}})) + session = start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + assert json.loads(path.read_text())["permission"]["bash"] == obj + assert session == {} + + +def test_openclaw_non_yolo_leaves_mode_policy(tmp_path): + # tools.exec.mode is OpenClaw's normalized knob and cannot be combined with explicit + # security/ask (the config is rejected), so a mode-based policy must be left as-is. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"mode": "deny"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"mode": "deny"} + + +def test_openclaw_non_yolo_preserves_sandbox_host(tmp_path): + # host=sandbox defaults to security=deny (stricter than the gateway "full" default), + # so a non-yolo run must not treat the missing security as permissive nor pop host + # (which would broaden routing to the gateway). + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"host": "sandbox"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"host": "sandbox"} + + +def test_openclaw_non_yolo_preserves_node_host(tmp_path): + # host=node routes to a paired node and is only ever set by the user (--yolo writes + # host=gateway), so a non-yolo run must not pop it and reroute to the gateway. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"host": "node"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"host": "node"} + + +def test_openclaw_non_yolo_preserves_auto_host_permissive(tmp_path): + # host=auto (or omitted) with security=full/ask=off is NOT the --yolo write: under an + # active sandbox, auto resolves to security=deny. --yolo only ever writes host=gateway, + # so the reset must not treat auto/None as the permissive gateway default and broaden a + # sandboxed deny to allowlist. + for exec_policy in ( + {"host": "auto", "security": "full", "ask": "off"}, + {"security": "full", "ask": "off"}, + ): + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": dict(exec_policy)}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == exec_policy + + +def test_openclaw_non_yolo_resets_only_gateway_yolo_fingerprint(tmp_path): + # The reset fires on exactly the host=gateway + security=full + ask=off write --yolo + # makes, and nothing else. + path = tmp_path / "openclaw.json" + path.write_text( + json.dumps({"tools": {"exec": {"host": "gateway", "security": "full", "ask": "off"}}}) + ) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"security": "allowlist", "ask": "on-miss"} + + +def test_openclaw_non_yolo_preserves_full_mode(tmp_path): + # OpenClaw never normalizes our security=full/ask=off yolo write into mode:"full" + # (verified against the binary: doctor --fix and config get leave security/ask as-is), + # so a mode:"full" is always a deliberate user policy, not stale yolo state; leave it. + path = tmp_path / "openclaw.json" + path.write_text(json.dumps({"tools": {"exec": {"mode": "full"}}})) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"mode": "full"} + + def test_yolo_command_flags_unmapped_agent_is_empty(): # Config-based agents (and any typo) must yield no flag, not a KeyError. assert start._yolo_command_flags("opencode", True) == [] From 296cacb5a176a31fce91b4ef90ed1a044dfc4fa5 Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Tue, 7 Jul 2026 04:29:37 -0500 Subject: [PATCH 037/113] ROCm-on-WSL: support discrete Radeon (RDNA 3/4) in WSL, not just Strix Halo (#6915) * WSL ROCm: generalize ROCm-on-WSL bootstrap from Strix-only to any RDNA arch install_rocm_wsl_strixhalo.sh hardcoded gfx1151, so its verify step died on discrete Radeon cards even though the ROCm + librocdxg setup is arch-agnostic. Auto-detect the GPU arch from rocminfo (override via UNSLOTH_WSL_GFX), verify any GPU agent enumerates over DXG, and map the arch to AMD's per-arch wheel family for the optional smoke test (injecting librocdxg into torch/lib so torch's bundled ROCr finds the DXG bridge). Verified on gfx1200 (Radeon RX 9060 XT) in WSL2 + Ubuntu 24.04 -- torch.cuda now enumerates the GPU. * WSL ROCm: trigger the ROCm-on-WSL bootstrap for discrete Radeon GPUs too _maybe_bootstrap_rocm_wsl only fired for Strix APUs (matched via /proc/cpuinfo, which discrete cards don't appear in). Add _wsl_amd_gpu_name() -- queries the Windows host via WMI -- and broaden the trigger gate plus the 'already-usable ROCm' rocminfo check from gfx1151-only to any real GPU agent (gfxNNNN, excluding the gfx11-generic fallback ISA). The generalized bootstrap then auto-detects the arch. Enables 'curl install.sh | sh' to set up ROCm-on-WSL on discrete Radeon RX 7000/9000 in WSL2 + Ubuntu 24.04, not just Strix Halo/Point. * WSL ROCm: address review -- filter generic ISA in bootstrap, bound the host GPU query - install_rocm_wsl_strixhalo.sh: exclude the gfx11-generic fallback ISA in arch detection (grep -v generic), matching install.sh's rocminfo check, so a generic agent listed before the real one can't be picked as the arch. - install.sh: wrap the powershell.exe Win32_VideoController query in _run_bounded (10s timeout) so an unstable WSL-interop / busy host can't hang the installer. * WSL ROCm: harden arch-detect + librocdxg copy under set -eo pipefail (review) - _detected_gfx: append '|| true' so a no-GPU rocminfo (empty pipeline, non-zero under pipefail) doesn't abort the assignment before the '[ -z ]' branch prints the diagnostic + die message. - smoke-test librocdxg copy: gate on '[ -d "$_tlib" ]' instead of '[ -n ]' so a non-directory value can't make cp rename librocdxg to 'lib'. * WSL ROCm: address Codex review (gfx000, 24.04 reroute for discrete, test locator) - Exclude gfx000 (the CPU agent) from the WSL 'usable ROCm' check and the bootstrap arch-detect: match gfx[1-9] (nonzero arch), so a partial ROCm install that only reports the CPU ISA no longer short-circuits the librocdxg setup. (P2) - Reuse the Ubuntu-24.04 reroute for discrete Radeon: broaden _maybe_reroute_strixhalo_to_2404's gate with the same _wsl_amd_gpu_name (WMI) fallback, so a discrete card on 26.04 reroutes to a 24.04 distro like Strix does instead of falling to CPU. Moved _wsl_amd_gpu_name above the reroute and made it self-contained + 10s-bounded (it runs before _run_bounded is defined). (P2) - Update TestInstallShDropinPersistence to locate the gate by its unique '!/generic/' clause now that the gfx1151 literal is gone. (P1) * Condense ROCm-on-WSL comments in install.sh and bootstrap helper * Guard WSL reroute from NVIDIA hybrid hosts and fix GFX-override pipefail check * Honor CUDA_VISIBLE_DEVICES-hidden NVIDIA in the WSL reroute guard * Reuse _has_usable_nvidia_gpu in the WSL reroute guard --------- Co-authored-by: Daniel Han --- install.sh | 162 +++++++++++++--------- scripts/install_rocm_wsl_strixhalo.sh | 57 +++++--- tests/studio/install/test_rocm_support.py | 48 ++++++- 3 files changed, 182 insertions(+), 85 deletions(-) diff --git a/install.sh b/install.sh index 81c50bc899..14fbba478d 100755 --- a/install.sh +++ b/install.sh @@ -1483,6 +1483,81 @@ elif [ "$OS" = "macos" ]; then fi tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none" +# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in +# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded +# to 10s. Defined here so the reroute below can use it before _run_bounded exists. +_WSL_AMD_GPU_NAME_CACHE="" +_wsl_amd_gpu_name() { + if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then + [ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1 + printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0 + fi + command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; } + _wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name" + if command -v timeout >/dev/null 2>&1; then + _wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')" + else + _wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')" + fi + if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi + _WSL_AMD_GPU_NAME_CACHE="-"; return 1 +} + +# ── Bounded command runner ── +# Runs a command under a 10s timeout when the `timeout` binary is available, +# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during +# driver init or after a reset) from hanging the installer: a timed-out probe +# exits nonzero and is treated exactly like a failed probe. No-op semantics on +# hosts without `timeout` (e.g. macOS) or when the probe is healthy. +_run_bounded() { + if command -v timeout >/dev/null 2>&1; then + timeout 10 "$@" + else + "$@" + fi +} + +# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every +# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to +# the AMD card). Unset means all devices visible. nvidia-smi ignores this env +# var, so the probes below cannot see the distinction on their own. +_cvd_hides_nvidia() { + [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 + _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') + [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] +} + +# ── NVIDIA usable-GPU helper ── +# Returns 0 (true) if an NVIDIA GPU is present and usable. +# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, +# which the NVIDIA driver populates on Linux regardless of nvidia-smi state +# -- handles PATH gaps, subprocess timeouts, and driver init races that +# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. +# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches +# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. +_has_usable_nvidia_gpu() { + if _cvd_hides_nvidia; then + return 1 + fi + _nvsmi="" + if command -v nvidia-smi >/dev/null 2>&1; then + _nvsmi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _nvsmi="/usr/bin/nvidia-smi" + fi + if [ -n "$_nvsmi" ]; then + if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + return 0 + fi + fi + # Fallback: NVIDIA driver exposes one subdir per GPU under this path. + if [ -d /proc/driver/nvidia/gpus ] && \ + [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then + return 0 + fi + return 1 +} + # Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04) # with a 24.04 distro present, re-run the install there and stop; else fall through # to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before @@ -1493,7 +1568,15 @@ _maybe_reroute_strixhalo_to_2404() { [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 [ -e /dev/dxg ] || return 0 - grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + # A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on + # this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors + # CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps. + if _has_usable_nvidia_gpu; then return 0; fi + # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + && ! _wsl_amd_gpu_name >/dev/null 2>&1; then + return 0 + fi # Already ROCm-on-WSL? leave a working GPU alone, whatever the version. if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then return 0 @@ -1959,61 +2042,6 @@ _has_amd_rocm_gpu() { return 1 } -# ── Bounded command runner ── -# Runs a command under a 10s timeout when the `timeout` binary is available, -# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during -# driver init or after a reset) from hanging the installer: a timed-out probe -# exits nonzero and is treated exactly like a failed probe. No-op semantics on -# hosts without `timeout` (e.g. macOS) or when the probe is healthy. -_run_bounded() { - if command -v timeout >/dev/null 2>&1; then - timeout 10 "$@" - else - "$@" - fi -} - -# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every -# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to -# the AMD card). Unset means all devices visible. nvidia-smi ignores this env -# var, so the probes below cannot see the distinction on their own. -_cvd_hides_nvidia() { - [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 - _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') - [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] -} - -# ── NVIDIA usable-GPU helper ── -# Returns 0 (true) if an NVIDIA GPU is present and usable. -# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, -# which the NVIDIA driver populates on Linux regardless of nvidia-smi state -# -- handles PATH gaps, subprocess timeouts, and driver init races that -# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. -# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches -# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. -_has_usable_nvidia_gpu() { - if _cvd_hides_nvidia; then - return 1 - fi - _nvsmi="" - if command -v nvidia-smi >/dev/null 2>&1; then - _nvsmi="nvidia-smi" - elif [ -x "/usr/bin/nvidia-smi" ]; then - _nvsmi="/usr/bin/nvidia-smi" - fi - if [ -n "$_nvsmi" ]; then - if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then - return 0 - fi - fi - # Fallback: NVIDIA driver exposes one subdir per GPU under this path. - if [ -d /proc/driver/nvidia/gpus ] && \ - [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then - return 0 - fi - return 1 -} - # ── Detect GPU and choose PyTorch index URL ── # Mirrors Get-TorchIndexUrl in install.ps1. # On CPU-only machines this returns the cpu index, avoiding the solver @@ -2327,19 +2355,19 @@ _persist_rocm_wsl_dropin() { fi } +# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it. _maybe_bootstrap_rocm_wsl() { [ "${OS:-}" = "wsl" ] || return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 # Leave any already-usable GPU completely alone (NVIDIA, or working ROCm). if _has_usable_nvidia_gpu; then return 0; fi - # "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the - # generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and - # would skip this bootstrap while the real GPU is still unusable. awk consumes - # all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail. + # Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000, + # the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so + # rocminfo isn't SIGPIPE'd like `grep -q` under pipefail. _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then # rocminfo may work only via the transient env _ensure_rocm_probe_env # just set, which dies with the installer. Persist the drop-in so login # shells (Studio, llama.cpp) inherit it -- else a reinstall over an @@ -2349,9 +2377,12 @@ _maybe_bootstrap_rocm_wsl() { fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). [ -e /dev/dxg ] || return 0 - # Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match - # the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S"). - grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + # Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also + # ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + && ! _wsl_amd_gpu_name >/dev/null 2>&1; then + return 0 + fi command -v bash >/dev/null 2>&1 || return 0 # Fast path: already configured (librocdxg present) but launched from a @@ -2369,7 +2400,8 @@ _maybe_bootstrap_rocm_wsl() { fi echo "" - substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN" + _rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU" + substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN" substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU." substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)" diff --git a/scripts/install_rocm_wsl_strixhalo.sh b/scripts/install_rocm_wsl_strixhalo.sh index 5ef9ee386a..aa560fc432 100644 --- a/scripts/install_rocm_wsl_strixhalo.sh +++ b/scripts/install_rocm_wsl_strixhalo.sh @@ -3,13 +3,14 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # # ────────────────────────────────────────────────────────────────────────────── -# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151) +# Enable ROCm-on-WSL for AMD GPUs (Strix Halo/Point APUs AND discrete Radeon RX +# 7000/9000). Verified on gfx1151 (Radeon 8060S) and gfx1200 (Radeon RX 9060 XT). # ────────────────────────────────────────────────────────────────────────────── -# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime -# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG -# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04 -# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via -# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies). +# install.sh routes the detected arch to the right ROCm wheels once a runtime exists; +# what it does NOT do is install AMD's ROCm userspace + the WSL DXG bridge (librocdxg). +# This helper does that Linux-side prerequisite on Ubuntu 24.04 WSL2, invoked by +# install.sh when it sees an AMD GPU via /dev/dxg but no ROCm yet. Arch-agnostic: the +# arch is auto-detected from rocminfo (override UNSLOTH_WSL_GFX=gfx1200). Idempotent. # # Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with # production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once @@ -34,10 +35,12 @@ set -euo pipefail # ── Tunables (override via env) ────────────────────────────────────────────── ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install -GFX="gfx1151" +# GPU arch: empty = auto-detect from rocminfo after install (override UNSLOTH_WSL_GFX=gfx1200). +# The ROCm + librocdxg setup is arch-agnostic; only verify + the smoke test need the arch. +GFX="${UNSLOTH_WSL_GFX:-}" LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build -# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test. -TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/" +# AMD's wheel index for the (optional) smoke test; resolved after arch detection. +TORCH_INDEX="" # Optional torch smoke test (throwaway venv). OFF by default: install.sh installs # torch itself into the real venv right after, so a duplicate download is wasteful. SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}" @@ -220,12 +223,12 @@ $SUDO ldconfig say "Persisting ROCm-on-WSL environment" _envfile="/etc/profile.d/unsloth-rocm-wsl.sh" $SUDO tee "$_envfile" >/dev/null <>> Unsloth ROCm-on-WSL (gfx1151) >>> +# >>> Unsloth ROCm-on-WSL >>> export HSA_ENABLE_DXG_DETECTION=1 export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 export PATH="${ROCM_DIR}/bin:\${PATH}" export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}" -# <<< Unsloth ROCm-on-WSL (gfx1151) <<< +# <<< Unsloth ROCm-on-WSL <<< EOF # also drop into ~/.bashrc for interactive shells if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then @@ -237,32 +240,50 @@ export PATH="${ROCM_DIR}/bin:${PATH}" export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}" # ── Step 5: verify the runtime enumerates the GPU ──────────────────────────── -say "Verifying rocminfo sees ${GFX}" +say "Verifying rocminfo enumerates the GPU over DXG" # Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs # rocminfo on first match, which under `set -o pipefail` turns a successful match -# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a -# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass. +# into a pipeline failure. _rocminfo_out="$(rocminfo 2>/dev/null || true)" -if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then +# GPU agents advertise an ISA "Name: gfxNNNN". Match gfx[1-9] (excludes gfx000, the CPU +# agent), drop the "gfx*-generic" fallback ISA, and take the first real GPU arch. +_detected_gfx="$(printf '%s\n' "$_rocminfo_out" | grep -E 'Name:[[:space:]]*gfx[1-9]' | grep -v 'generic' | grep -oE 'gfx[1-9][0-9a-z]*' | head -1 || true)" +if [ -z "$_detected_gfx" ]; then printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true - die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." + die "rocminfo did not enumerate any GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." fi +# Honour a caller-pinned arch (sanity-check via a consuming grep, not grep -q: under +# pipefail -q would SIGPIPE printf on large output and misreport the arch); else adopt. +if [ -n "$GFX" ] && ! printf '%s\n' "$_rocminfo_out" | grep -E "Name:[[:space:]]*${GFX}([^0-9]|$)" >/dev/null; then + die "rocminfo enumerated '${_detected_gfx}' but not the requested UNSLOTH_WSL_GFX='${GFX}'." +fi +GFX="${GFX:-$_detected_gfx}" # Display-only summary: best-effort (|| true) so head's early pipe-close under # `set -o pipefail` can't fail the bootstrap after verification already passed. printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true note "ROCm-on-WSL runtime is live for ${GFX}." -# ── Step 6 (optional): torch smoke test from the gfx1151 index ─────────────── +# ── Step 6 (optional): torch smoke test from AMD's per-arch wheel index ─────── if [ "$SMOKE_TEST" = "1" ]; then say "Smoke-testing PyTorch on ${GFX} (throwaway venv)" + # Map the detected arch to AMD's repo.amd.com wheel family index. + case "$GFX" in + gfx1200|gfx1201) _fam="gfx120X-all" ;; + gfx1100|gfx1101|gfx1102|gfx1103) _fam="gfx110X-all" ;; + *) _fam="$GFX" ;; # gfx1150/gfx1151/gfx90a: own index + esac + TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${_fam}/" _venv="${HOME}/.unsloth/rocm-smoketest" rm -rf "$_venv"; python3 -m venv "$_venv" "$_venv/bin/pip" install --quiet --upgrade pip - # gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py + # AMD arch index is primary (torch + triton); PyPI only an extra for pure-py # deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch. "$_venv/bin/pip" install --index-url "$TORCH_INDEX" \ --extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \ die "torch install from ${TORCH_INDEX} failed." + # WSL: torch's bundled ROCr must load the DXG bridge -- drop librocdxg into torch/lib. + _tlib="$("$_venv/bin/python" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))' 2>/dev/null || true)" + [ -d "$_tlib" ] && cp -f "${ROCM_DIR}"/lib/librocdxg.so* "$_tlib"/ 2>/dev/null || true "$_venv/bin/python" - <<'PY' import torch ok = torch.cuda.is_available() diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index c8f2053946..bfc8132683 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -3426,8 +3426,9 @@ class TestInstallShDropinPersistence: def test_gate5_early_return_persists_dropin(self): """The rocminfo-already-works early return must call the persist helper before returning.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") - # The persist call must precede `return 0` at the rocminfo gfx1151 gate. - gate = source.find("Name:[[:space:]]*gfx1151") + # The persist call must precede `return 0` at the rocminfo GPU-agent gate + # (uniquely identified by the `!/generic/` clause the other probes lack). + gate = source.find("Name:[[:space:]]*gfx[1-9]/ && !/generic/") assert gate != -1 window = source[gate : gate + 900] assert "_persist_rocm_wsl_dropin" in window @@ -3441,6 +3442,49 @@ class TestInstallShDropinPersistence: assert "profile.d/unsloth-rocm-wsl.sh" in body +_STRIXHALO_WSL_PATH = PACKAGE_ROOT / "scripts" / "install_rocm_wsl_strixhalo.sh" + + +class TestWslRerouteNvidiaGuard: + """_maybe_reroute_strixhalo_to_2404 must skip the AMD reroute on hybrid AMD+NVIDIA hosts by + reusing _has_usable_nvidia_gpu (CUDA_VISIBLE_DEVICES-aware + /proc/driver/nvidia fallback), + which must be defined before the reroute's call site so it is actually available.""" + + def test_reroute_calls_nvidia_helper_before_amd_signal(self): + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + start = source.find("_maybe_reroute_strixhalo_to_2404()") + assert start != -1 + body = source[start : start + 1200] + nv = body.find("_has_usable_nvidia_gpu") + wmi = body.find("_wsl_amd_gpu_name") + assert nv != -1, "reroute must consult _has_usable_nvidia_gpu before deciding to reroute" + assert wmi != -1 + # The NVIDIA guard must precede the AMD/WMI signal and return early. + assert nv < wmi + assert body.find("return 0", nv) < wmi + + def test_nvidia_helper_and_deps_defined_before_reroute_callsite(self): + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + call = source.find("\n_maybe_reroute_strixhalo_to_2404 || true") + assert call != -1 + for fn in ("_run_bounded() {", "_cvd_hides_nvidia() {", "_has_usable_nvidia_gpu() {"): + idx = source.find(fn) + assert idx != -1 and idx < call, f"{fn} must be defined before the reroute call" + + +class TestStrixhaloGfxOverridePipefail: + """The UNSLOTH_WSL_GFX override check must use a consuming grep, not grep -q: under + `set -o pipefail` an early -q exit SIGPIPEs printf and misreports the arch on large output.""" + + def test_gfx_override_uses_consuming_grep(self): + source = _STRIXHALO_WSL_PATH.read_text(encoding = "utf-8") + idx = source.find('grep -E "Name:[[:space:]]*${GFX}') + assert idx != -1, "GFX override must use a consuming grep -E (not grep -q)" + line = source[idx : source.find("\n", idx)] + assert ">/dev/null" in line + assert 'grep -qE "Name:[[:space:]]*${GFX}' not in source + + class TestLlamaCppRuntimeWslOrdering: """The serve-time launcher mirrors binary_env: system HIP before the bundle dir on WSL.""" From bdb958e052eca6a47c17410558562f00481f71b8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:16:57 -0700 Subject: [PATCH 038/113] Guard RoPE scaling against the transformers v5 buffer blank; honor extended RoPE factor (#6925) * Guard RoPE scaling against the transformers v5 buffer blank; honor extended factor Add a family-agnostic guard that builds each rotary from a scaled config, blanks its non-persistent buffers (what transformers v5 does on load), runs loader._fix_rope_inv_freq, and asserts every buffer is restored to its scaled value (llama3 and longrope). This catches the whole bug class, not just the one call site, and is validated to fail on the pre-fix repair. Also make LlamaExtendedRotaryEmbedding read the llama3 factor from the config instead of hardcoding 8 (wrong for Llama-3.2, factor 32), falling back to the Llama-3.1 defaults when built without a config. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pass config into extended rotary codegen; skip v5 round-trip on transformers 4.x - patch_llama_rope_scaling now builds the llama3 extended rotary with config=self.config so it reads the real factor (32 for Llama-3.2) instead of falling back to 8; the template already references self.config. - test_v5_blank_repair_roundtrip now skips when loader._NEEDS_ROPE_FIX is False, since _fix_rope_inv_freq is a no-op on transformers 4.x and cannot restore the blanked buffers there. * Raise stream deadlock-guard timeouts from 0.2s to 5.0s in passthrough tests These asyncio.wait_for guards bound test setup and cross-task event signaling that complete near-instantly on success; the 0.2s budget is a latency assertion in disguise and times out under CI scheduling load (seen on the 3.11 matrix leg while 3.10/3.12/3.13 pass the same commit). 5.0s matches the timeout used elsewhere in the suite and still fails fast on a real hang. No test relies on the guard expiring. * Extended rotary reads rope_parameters as well as rope_scaling transformers v5 stores llama3 scaling under config.rope_parameters and exposes rope_scaling only as a back-compat property. Reading that property works on 5.0-5.13 (verified: factor resolves to 32 for Llama-3.2), but a future release may drop the shim, after which the subclass path would fall back to factor 8. Read either field so the factor survives the rename. Adds test_extended_rotary_reads_rope_parameters_v5 (fails on the old single-field read: rope_parameters-only config resolves to 8, not 32). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../tests/test_openai_tool_passthrough.py | 18 +-- tests/utils/test_rope_scaling_drift.py | 138 ++++++++++++++++++ unsloth/models/_utils.py | 1 + unsloth/models/llama.py | 17 ++- 4 files changed, 160 insertions(+), 14 deletions(-) diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index ccbd78e2b1..05e017ba7f 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1900,7 +1900,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2044,7 +2044,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) gate.set() @@ -2107,7 +2107,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2190,7 +2190,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2252,7 +2252,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) assert cancel_id in inf_mod._CANCEL_REGISTRY @@ -2323,13 +2323,13 @@ class TestApiMonitorProviderAndCompletionStreams: monitor_id = monitor_id, ) ) - await asyncio.wait_for(entered.wait(), timeout = 0.2) + await asyncio.wait_for(entered.wait(), timeout = 5.0) assert cancel_id in inf_mod._CANCEL_REGISTRY task.cancel() with pytest.raises(asyncio.CancelledError): await task - await asyncio.wait_for(cancelled.wait(), timeout = 0.2) + await asyncio.wait_for(cancelled.wait(), timeout = 5.0) assert cancel_id not in inf_mod._CANCEL_REGISTRY asyncio.run(_run()) @@ -2389,13 +2389,13 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) assert cancel_id in inf_mod._CANCEL_REGISTRY gate.set() - await asyncio.wait_for(returned.wait(), timeout = 0.2) + await asyncio.wait_for(returned.wait(), timeout = 5.0) await asyncio.sleep(0) await response._unstarted_cleanup() assert upstream_response.is_closed diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index 98f7e2db62..7a738e236c 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -257,6 +257,63 @@ def test_recompute_helper_scales_on_cpu(): ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled." +def test_extended_rotary_reads_config_factor(): + # LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8 + # (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405). + from types import SimpleNamespace + + from unsloth.models.llama import LlamaExtendedRotaryEmbedding + + rot = object.__new__(LlamaExtendedRotaryEmbedding) + rot.base = ROPE_THETA + rot.dim = HEAD_DIM + rot._unsloth_rope_config = SimpleNamespace( + rope_scaling = { + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + } + ) + vanilla = _vanilla_inv_freq() + scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1) + ratio = float(vanilla[-1]) / float(scaled[-1]) + assert abs(ratio - 32.0) < 1e-3, ( + f"LlamaExtendedRotaryEmbedding ignored config factor 32 (ratio {ratio}); the " + "low-frequency band must be divided by the config factor (issue #2405)." + ) + + +def test_extended_rotary_reads_rope_parameters_v5(): + # transformers v5 stores scaling under rope_parameters (rope_scaling is a + # back-compat shim that may be removed); the factor must still be read. + from types import SimpleNamespace + + from unsloth.models.llama import LlamaExtendedRotaryEmbedding + + rot = object.__new__(LlamaExtendedRotaryEmbedding) + rot.base = ROPE_THETA + rot.dim = HEAD_DIM + rot._unsloth_rope_config = SimpleNamespace( + rope_scaling = None, + rope_parameters = { + "rope_type": "llama3", + "factor": 32.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + }, + ) + vanilla = _vanilla_inv_freq() + scaled = rot._apply_inv_freq_scaling(vanilla).reshape(-1) + ratio = float(vanilla[-1]) / float(scaled[-1]) + assert abs(ratio - 32.0) < 1e-3, ( + f"Extended rotary ignored rope_parameters factor 32 (ratio {ratio}); v5 " + "keeps the factor under rope_parameters, not rope_scaling." + ) + + def _cos_at_position(rot, position): """cos row at one position, built like _set_cos_sin_cache but CPU-only.""" inv_freq = rot.inv_freq.float().cpu() @@ -324,6 +381,87 @@ def test_extended_cache_keeps_scaling_after_growth(): ) +def _blank_nonpersistent_buffers(module): + """Mimic transformers v5 meta-load: overwrite non-persistent buffers with garbage.""" + for name, buf in list(module.named_buffers()): + leaf = module + *parents, attr = name.split(".") + for part in parents: + leaf = getattr(leaf, part) + if attr in getattr(leaf, "_non_persistent_buffers_set", set()): + setattr(leaf, attr, torch.rand_like(buf)) + + +def _build_llama3_rotary(): + from unsloth.models import llama as llama_mod + config = _make_config(LLAMA3_ROPE_SCALING) + return llama_mod.LlamaRotaryEmbedding(config = config), config + + +def _build_longrope_rotary(): + from types import SimpleNamespace + + from unsloth.models import llama as llama_mod + + short_factor, long_factor = [1.05] * 48, [1.3] * 48 + rot = llama_mod.LongRopeRotaryEmbedding( + dim = 96, + max_position_embeddings = 131072, + original_max_position_embeddings = 4096, + base = ROPE_THETA, + short_factor = short_factor, + long_factor = long_factor, + ) + config = SimpleNamespace( + rope_scaling = { + "rope_type": "longrope", + "short_factor": short_factor, + "long_factor": long_factor, + "original_max_position_embeddings": 4096, + } + ) + return rot, config + + +@requires_cuda +@pytest.mark.parametrize( + "build", [_build_llama3_rotary, _build_longrope_rotary], ids = ["llama3", "longrope"] +) +def test_v5_blank_repair_roundtrip(build): + # Build scaled -> blank non-persistent buffers (what transformers v5 does on + # load) -> run the repair -> every buffer must return to its scaled value. + # Family-agnostic: encodes no scaling math, so it guards any rotary that + # keeps scaling in a buffer (issue #2405 / PR #6907). + from unsloth.models import loader + + # The repair only runs on transformers v5 (it is what blanks the buffers); + # on v4 _fix_rope_inv_freq is a no-op, so the round-trip cannot restore. + if not loader._NEEDS_ROPE_FIX: + pytest.skip("transformers < 5 does not blank rope buffers; repair is a no-op") + + rot, config = build() + snapshot = {name: buf.detach().clone() for name, buf in rot.named_buffers()} + assert snapshot, "rotary registers no buffers; nothing to guard" + + _blank_nonpersistent_buffers(rot) + assert any( + not torch.equal(rot.get_buffer(name), snapshot[name]) for name in snapshot + ), "blanking changed no buffer; the round-trip would be vacuous" + + wrapper = torch.nn.Module() + wrapper.add_module("rotary_emb", rot) + wrapper.config = config + loader._fix_rope_inv_freq(wrapper) + + for name in snapshot: + assert torch.allclose( + rot.get_buffer(name).cpu(), snapshot[name].cpu(), rtol = 1e-4, atol = 1e-6 + ), ( + f"{name} was not restored to its scaled value by loader._fix_rope_inv_freq " + "after the transformers v5 buffer blank (issue #2405 / PR #6907)." + ) + + def test_object_style_rope_scaling_does_not_crash(): # Object-style rope_scaling must be normalized, not .get()'d directly. from dataclasses import dataclass diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 169b610988..1aa2c6e820 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -2834,6 +2834,7 @@ def patch_llama_rope_scaling( dim = self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, + config=self.config, ) elif scaling_type == "longrope": self.rotary_emb = {longrope_rope_function}( diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c25a031b82..a1da099758 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1930,11 +1930,18 @@ class LlamaExtendedRotaryEmbedding(LlamaRotaryEmbedding): # From https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py#L41 def _apply_inv_freq_scaling(self, freqs: torch.Tensor): - # Values obtained from grid search - scale_factor = 8 - low_freq_factor = 1 - high_freq_factor = 4 - old_context_len = 8192 # original llama3 length + # llama3 factors from config; Llama-3.1 defaults when built without one + # (legacy codegen path). Hardcoding 8 is wrong for e.g. Llama-3.2 (32). + # v5 renames rope_scaling -> rope_parameters; read either so the factor + # survives even if the rope_scaling back-compat shim is dropped. + config = getattr(self, "_unsloth_rope_config", None) + rope_scaling = _rope_scaling_as_dict( + getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {} + ) + scale_factor = rope_scaling.get("factor", 8) + low_freq_factor = rope_scaling.get("low_freq_factor", 1) + high_freq_factor = rope_scaling.get("high_freq_factor", 4) + old_context_len = rope_scaling.get("original_max_position_embeddings", 8192) low_freq_wavelen = old_context_len / low_freq_factor high_freq_wavelen = old_context_len / high_freq_factor From 414503745e71b52c37117d7a96894ea46fc13ec5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:30:21 -0700 Subject: [PATCH 039/113] Run the malware gate on the RAG embedding model before it loads (#6887) * Run the malware gate on the RAG embedding model before it loads Setting the RAG embedding model through PUT /api/settings/embedding-model persisted an arbitrary repo and later handed it straight to SentenceTransformer, which deserializes pickle weights. Unlike the normal model-load paths, this route never ran evaluate_file_security, and force skipped verification entirely, so a repo Hugging Face flags as unsafe (or any repo under force) could be downloaded and loaded in the backend process without a scan. Run the malware/pickle scan at both ends: the settings endpoint now scans before persisting and returns 409 on a flagged repo even under force (force still only skips the is-embedding-model type check for offline or local repos), and the embedder scans again at the load sink so a name that arrives via env or default is covered too. Local paths and unreachable scans fail open inside evaluate_file_security, and the sink never bricks the embedder on a gate error. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Thread the load token into the embedding scan and hard-fail on a block The load-sink scan ran without a token, so evaluate_file_security (which passes token=False when none is given) could not reach a gated or private repo and failed open for exactly the model SentenceTransformer would still load. Resolve the loader's own token (HF_TOKEN env or the cached login) and pass it to the sink scan, and fall back to it in the settings endpoint when the request omits one. The sink previously raised a plain RuntimeError, which the llama-server fallback in encode() and _build_st_backend_or_fallback() swallowed as a routine ST failure, silently switching backends instead of blocking. Raise a distinct UnsafeEmbeddingModelError that both fallback paths re-raise, so a flagged model hard-fails. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scan sentence-transformers module dirs and scope the embedding pickle gate to the ST backend Extend the RAG embedding malware gate so a poisoned pickle under a SentenceTransformer module dir (for example 0_Transformer/pytorch_model.bin) blocks. Those dirs are read from the repo's modules.json and passed as load roots to evaluate_file_security at both the settings endpoint and the load sink, so such a pickle is treated as root-level there instead of an unreferenced nested shard that was previously allowed. Scope the ST pickle scan to the sentence-transformers backend. On the llama-server backend the embedder loads GGUF files (inert) from the -GGUF companion repo, never the ST repo's pickle, so a custom ST repo with a flagged pickle and a clean GGUF companion is no longer rejected. The existing GGUF availability checks already cover that path. Return 403 for the hard security block instead of 409. The settings UI routes every 409 into the forceable save-anyway flow, but this block cannot be bypassed by force, so it now uses a distinct status the client treats as non-forceable. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Base the embedding pickle scan on the actual backend, not just the resolver _llama_backend_active only consulted the auto resolver, so on a GPU box where auto resolves to sentence-transformers but the process already fell back to the llama-server backend at runtime (a torch or CUDA load/encode failure), it returned False and the settings endpoint hard-blocked a save whose ST pickle is flagged even though the process loads only inert GGUF. Add active_backend_is_llama, which reflects the actual built backend (True when the cached backend is a LlamaServerBackend, including a runtime fallback) and otherwise defers to the resolver as a fresh process would, and delegate _llama_backend_active to it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Report the cached embedding backend verbatim, not the resolver active_backend_is_llama() fell through to the config resolver whenever a backend was already built but was not llama-server, so a live sentence-transformers backend could report llama=True once the resolver picked llama (GPU heuristic or a runtime config change) and wrongly skip its pickle scan. Once a backend exists, return isinstance(backend, LlamaServerBackend) directly; only defer to the resolver before any backend is built. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/rag/embeddings.py | 117 ++++++ studio/backend/routes/settings.py | 78 +++- .../test_embedding_model_security_gate.py | 365 ++++++++++++++++++ .../tests/test_security_gate_consistency.py | 13 + .../features/settings/api/embedding-model.ts | 9 + .../features/settings/tabs/general-tab.tsx | 6 +- 6 files changed, 573 insertions(+), 15 deletions(-) create mode 100644 studio/backend/tests/test_embedding_model_security_gate.py diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 345b4dd853..47d26209b4 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -63,6 +63,87 @@ def _install_torchao_stub_once() -> None: install_torchao_windows_rocm_stub() +class UnsafeEmbeddingModelError(RuntimeError): + """Raised when the embedding model repo is flagged unsafe. A distinct type so the + llama-server fallback paths re-raise it instead of masking a security block as a + routine ST failure.""" + + +def _ambient_hf_token() -> str | None: + """The HF token the loader itself would use (HF_TOKEN env or the cached login), so + the scan can reach a gated/private repo instead of failing open. None if unavailable.""" + try: + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: + """The module directories a SentenceTransformer load reads weights from, taken from + the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``). + ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the + security scan: a flagged pickle directly under one must block. Returns () on any + failure (no modules.json, offline, malformed) so the guard never bricks the embedder. + """ + try: + import json + + from utils.paths import is_local_path + + if is_local_path(name): + from pathlib import Path + from utils.paths import normalize_path + + path = Path(normalize_path(name)).expanduser() / "modules.json" + if not path.is_file(): + return () + data = json.loads(path.read_text()) + else: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + try: + local = hf_hub_download(name, "modules.json", token = token or None) + except EntryNotFoundError: + return () + data = json.loads(open(local).read()) + subdirs = [] + for module in data or (): + sub = str((module or {}).get("path", "")).strip().strip("/") + if sub: + subdirs.append(sub) + return tuple(dict.fromkeys(subdirs)) + except Exception: + return () + + +def _guard_model_security(name: str) -> None: + """Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside + SentenceTransformer regardless of trust_remote_code. Defense in depth behind the + /settings gate (a name can also arrive via env/default); local paths and unreachable + scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error. + """ + try: + from utils.security import evaluate_file_security, security_load_subdirs + + token = _ambient_hf_token() + # Union the audio-model load roots with the ST module dirs so a flagged pickle + # directly under a Transformer module dir (0_Transformer/) blocks instead of + # passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token))) + ) + blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked + except Exception: + return + if blocked: + raise UnsafeEmbeddingModelError( + f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security " + "scan; refusing to load. Set a different RAG embedding model." + ) + + def _get(model_name: str | None = None): """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 for a ~1.5x speedup at negligible accuracy loss.""" @@ -75,6 +156,7 @@ def _get(model_name: str | None = None): device = _device() logger.info("loading embedding model %s on %s", name, device) + _guard_model_security(name) _model = SentenceTransformer( name, device = device, model_kwargs = {"torch_dtype": "float16"} ) @@ -159,6 +241,8 @@ class _SentenceTransformersBackend: ): try: return _st_encode(texts, model_name = model_name, normalize = normalize) + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure # ST loaded but this encode blew up; swap the process to the llama-server # embedder (so later encodes stay in one space) and retry. @@ -222,6 +306,8 @@ def _build_st_backend_or_fallback(): try: backend.warm(model_name = None) return backend + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure fallback = _try_make_llama_backend() if fallback is None: @@ -290,6 +376,37 @@ def _reset_backend() -> None: _backend_key = None +def active_backend_is_llama() -> bool: + """True when this process actually embeds via the llama-server (GGUF) backend. + + Reflects the ACTUAL built backend once one exists: an ``auto`` install that + resolves to sentence-transformers but then falls back to llama-server at + runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or + ``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so + callers gating on the ST pickle must see llama here. Before any backend is + built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw + key) exactly as a fresh process would. Never raises: a backend probe must not + block saving a model.""" + try: + with _backend_lock: + backend = _backend + if backend is not None: + # A backend exists: report what it ACTUALLY is. A concrete + # sentence-transformers backend must return False even if the + # resolver would now pick llama, so its pickle stays gated. If the + # llama import fails we cannot be llama, so fall to the safe False. + try: + from .embed_llama_server import LlamaServerBackend + except Exception: # noqa: BLE001 - llama plumbing import must never block + return False + return isinstance(backend, LlamaServerBackend) + raw = (config.EMBED_BACKEND or "auto").strip().lower() + key = _resolve_auto() if raw in _AUTO_ALIASES else raw + return key in _LLAMA_ALIASES + except Exception: # noqa: BLE001 - a backend probe must never block saving + return False + + def warm(model_name: str | None = None) -> None: """Eagerly load the embedder so the first real request isn't slow.""" _get_backend().warm(model_name = model_name) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 862bce8be8..914699f540 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -260,17 +260,29 @@ def _embedding_model_response() -> EmbeddingModelResponse: ) -def _llama_backend_active() -> bool: - """True when this install embeds via the llama-server (GGUF) backend.""" - from core.rag import config as rag_config - from core.rag import embeddings - +def _ambient_hf_token() -> Optional[str]: + """The HF token the loader would use (HF_TOKEN env or the cached login), so a gated + repo is scanned rather than failing open. None if unavailable.""" try: - raw = (rag_config.EMBED_BACKEND or "auto").strip().lower() - key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _llama_backend_active() -> bool: + """True when this install actually embeds via the llama-server (GGUF) backend. + + Delegates to the embeddings module so a runtime fallback from + sentence-transformers to llama-server (after a torch/CUDA load or encode + failure) is honored: in that state the process loads only inert GGUF, so the + ST pickle gate below must not hard-block a repo whose GGUF companion is clean. + Before any backend is built this still reflects the resolver.""" + from core.rag import embeddings + try: + return embeddings.active_backend_is_llama() except Exception: # noqa: BLE001 - backend probe must never block saving return False - return key in embeddings._LLAMA_ALIASES def _resolves_as_local_gguf(model: str) -> bool: @@ -357,6 +369,8 @@ def update_embedding_model( """Set the RAG embedding model. Unless ``force`` is set, the repo is verified to be an embedding model via HF metadata; an unverifiable model (wrong type, typo, gated repo, or no network) returns 409 so the UI can offer "save anyway". + A repo flagged unsafe by HF's security scan returns 403 instead: a hard block + that ``force`` cannot bypass, so the UI must not offer "save anyway". Documents indexed under the previous model must be re-uploaded.""" from utils.models import is_embedding_model @@ -370,15 +384,51 @@ def update_embedding_model( event = "settings.update_embedding_model_failed", log = logger, ) from exc + hf_token = (payload.hf_token or "").strip() or None # The env/default model needs no verification; saving it is a no-op override. # A local GGUF on the llama-server backend is accepted as-is: it is exactly # what the backend loads, and HF metadata cannot verify a local path. - if ( - model != default_embedding_model() - and not payload.force - and not (_llama_backend_active() and _resolves_as_local_gguf(model)) - ): - hf_token = (payload.hf_token or "").strip() or None + is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model) + # The pickle gate only matters for the sentence-transformers backend, which is what + # deserializes pickles. On the llama-server backend the embedder loads GGUF files + # (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would + # wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability + # checks below cover that path instead. + scan_st_pickle = ( + model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() + ) + if scan_st_pickle: + # Malware/pickle gate before we persist a repo the embedder later loads with + # SentenceTransformer. Runs even under force (force only skips the is-embedding + # type check for offline/local repos HF cannot verify); local paths and + # unreachable scans fail open inside evaluate_file_security. + from utils.security import evaluate_file_security, security_load_subdirs + from core.rag.embeddings import _st_module_subdirs + + # Fall back to the loader's own token so a gated/private repo is actually scanned + # (a token-less scan fails open for exactly the repo that would still load). + scan_token = hf_token or _ambient_hf_token() + # Include the ST module dirs (0_Transformer/) so a flagged pickle directly under + # one blocks instead of passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys( + ( + *security_load_subdirs(model, scan_token), + *_st_module_subdirs(model, scan_token), + ) + ) + ) + if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked: + # 403, not 409: the client routes every 409 into the forceable "save anyway" + # flow, but this block is a hard, non-forceable security refusal. + raise HTTPException( + status_code = 403, + detail = ( + f"{model!r} is flagged as unsafe by Hugging Face's security scan and " + "cannot be used as the embedding model." + ), + ) + if model != default_embedding_model() and not payload.force and not is_local_gguf: from core.rag import config as rag_config # A GGUF-named repo on the llama-server backend is loaded from its .gguf diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py new file mode 100644 index 0000000000..940b35d7ba --- /dev/null +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The RAG embedding model must pass the malware/pickle gate before it is persisted or +loaded. A flagged repo (or any repo saved with force) previously reached +SentenceTransformer unscanned, bypassing the normal model-load protections.""" + +from pathlib import Path +import sys +import types as _types + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import routes.settings as settings + + +class _Decision: + def __init__(self, blocked): + self.blocked = blocked + + +def _security_stub(blocked): + mod = _types.ModuleType("utils.security") + mod.evaluate_file_security = lambda *a, **k: _Decision(blocked) + mod.security_load_subdirs = lambda *a, **k: () + return mod + + +@pytest.fixture +def client(monkeypatch): + # The settings scan unions in the ST module dirs read from modules.json; keep it + # offline and deterministic for the endpoint tests that use this fixture. + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + return TestClient(app, raise_server_exceptions = False), saved + + +def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put( + "/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True} + ) + # 403, not the forceable 409, so the client does not offer "save anyway". + assert r.status_code == 403 + assert "model" not in saved # force must not persist a flagged repo + + +def test_flagged_repo_is_blocked_without_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert r.status_code == 403 + assert "model" not in saved + + +def test_hard_block_uses_non_forceable_status(client, monkeypatch): + # The forceable verification path uses 409; the hard security block must be distinct + # (403) so the frontend never routes it into the "save anyway" force flow. + c, _saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert blocked.status_code == 403 + + # A verification failure (not-an-embedding-model) stays forceable at 409. + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False) + import utils.models as _models + + monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False) + unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"}) + assert unverified.status_code == 409 + + +def test_llama_backend_skips_the_st_pickle_scan(monkeypatch): + # On the llama-server backend the embedder loads GGUF (inert), not the ST repo's + # pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: True) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + # force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped. + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama path + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch): + # auto resolves to sentence-transformers (GPU present) but the embedder fell back to + # llama-server at runtime (torch/CUDA load or encode failure), so the process now loads + # only inert GGUF. The real _llama_backend_active() must reflect that cached fallback, + # so a flagged ST repo with a clean GGUF companion must not be hard-blocked here. + import core.rag.embeddings as embeddings + from core.rag.embed_llama_server import LlamaServerBackend + + # Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even + # though the auto resolver would still say sentence-transformers. + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + # Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the + # real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored. + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch): + # active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers + # to the resolver (fresh-process behavior) when none has been built yet. + import core.rag.embeddings as embeddings + import core.rag.config as rag_config + from core.rag.embed_llama_server import LlamaServerBackend + + # A cached llama backend wins even when auto would resolve to sentence-transformers. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto") + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + assert embeddings.active_backend_is_llama() is True + + # A cached ST backend reports False even when the resolver now picks llama, so its + # pickle stays gated (the cached backend, not the resolver, is what actually embeds). + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend()) + assert embeddings.active_backend_is_llama() is False + + # No cached backend -> the resolver decides, unchanged from before. + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", None) + assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers + + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + assert embeddings.active_backend_is_llama() is True # auto -> llama-server + + # An explicit (non-auto) key is honored verbatim without a cached backend. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server") + assert embeddings.active_backend_is_llama() is True + + +def test_settings_scan_scopes_module_subdirs(monkeypatch): + # The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a + # pickle directly under one blocks; assert those subdirs reach evaluate_file_security. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + import core.rag.embeddings as embeddings + + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda *a, **k: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True} + ) + assert r.status_code == 200 + assert "0_Transformer" in seen["subdirs"] + + +def test_clean_repo_saves_under_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True}) + assert r.status_code == 200 + assert saved.get("model") == "acme/clean-embed" + + +def test_load_sink_refuses_flagged_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + import core.rag.embeddings as embeddings + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._guard_model_security("attacker/malicious-embed") + + +def test_load_sink_allows_clean_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import core.rag.embeddings as embeddings + embeddings._guard_model_security("acme/clean-embed") # no raise + + +def test_sink_threads_ambient_token_into_scan(monkeypatch): + # A gated repo set via env/default has no request token; the guard must feed the + # loader's own token to the scan, or it fails open for the repo that still loads. + seen = {} + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = ( + lambda name, token = None: seen.setdefault("subdirs_token", token) or () + ) + mod.evaluate_file_security = lambda *a, **k: seen.setdefault( + "scan_token", k.get("hf_token") + ) or _Decision(False) + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient") + embeddings._guard_model_security("acme/gated-embed") + assert seen["scan_token"] == "hf_ambient" + assert seen["subdirs_token"] == "hf_ambient" + + +def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch): + # A flagged pickle directly under a Transformer module dir (0_Transformer/) must + # reach the scan as a load root; assert the guard unions the module dirs into + # load_subdirs so evaluate_file_security treats such a pickle as root-level. + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda name, token = None: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None) + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + embeddings._guard_model_security("acme/embed-with-module-dir") + assert "0_Transformer" in seen["subdirs"] + + +def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch): + # The helper must parse each module's non-empty "path" from a local repo's + # modules.json and drop the root-level ("") Transformer entry. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps( + [ + {"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."}, + {"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."}, + {"idx": 2, "name": "2", "path": "", "type": "..."}, + ] + ) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None) + assert subdirs == ("0_Transformer", "1_Pooling") + + +def test_st_module_subdirs_swallows_errors(monkeypatch): + # Any failure (no modules.json, offline, malformed) returns () so the guard never + # bricks the embedder. + import huggingface_hub + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise RuntimeError("offline") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom) + assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == () + + +def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch): + # The ST encode fallback must re-raise a security block, not swap to llama-server. + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise embeddings.UnsafeEmbeddingModelError("flagged") + + monkeypatch.setattr(embeddings, "_st_encode", _boom) + monkeypatch.setattr( + embeddings, + "_switch_to_llama_fallback", + lambda err: pytest.fail("security block must not fall back to llama-server"), + ) + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._SentenceTransformersBackend().encode(["hi"]) diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py index db66df8a30..b5f1069f12 100644 --- a/studio/backend/tests/test_security_gate_consistency.py +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base(): if runs_gate and not resolves_base: offenders.append(f"{rel} runs a load gate but never resolves the LoRA base") assert not offenders, "\n".join(offenders) + + +def test_rag_embedding_path_runs_the_malware_gate(): + """The RAG embedding model is set through /settings and later loaded by + SentenceTransformer, which deserializes pickles; both sites must run the malware gate + or a flagged repo loads unscanned (bypassing the normal model-load protections).""" + offenders = [] + for rel in ("routes/settings.py", "core/rag/embeddings.py"): + if "evaluate_file_security(" not in (_BACKEND / rel).read_text(): + offenders.append( + f"{rel} loads/persists an embedding model without evaluate_file_security" + ) + assert not offenders, "\n".join(offenders) diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts index 8b6bc7ee7f..9a61142f73 100644 --- a/studio/frontend/src/features/settings/api/embedding-model.ts +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -23,6 +23,10 @@ type ApiEmbeddingModelSettings = { * (wrong type, gated repo, or offline). Retry with force to save anyway. */ export class EmbeddingModelVerificationError extends Error {} +/** 403 from the backend: the repo is flagged unsafe by Hugging Face's security scan. + * A hard block; force cannot bypass it, so it must not enter the "save anyway" flow. */ +export class EmbeddingModelBlockedError extends Error {} + function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { return { embeddingModel: settings.embedding_model, @@ -56,6 +60,11 @@ export async function updateEmbeddingModelSettings( force: options?.force ?? false, }), }); + if (res.status === 403) { + throw new EmbeddingModelBlockedError( + await readFastApiError(res, "This model is blocked by a security scan"), + ); + } if (res.status === 409) { throw new EmbeddingModelVerificationError( await readFastApiError(res, "Could not verify the embedding model"), diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 7670aae5fa..8fd70cc1b7 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -42,6 +42,7 @@ import { updatePreviewSharing, } from "../api/preview-sharing"; import { + EmbeddingModelBlockedError, type EmbeddingModelSettings, EmbeddingModelVerificationError, loadEmbeddingModelSettings, @@ -410,7 +411,10 @@ export function GeneralTab() { description: t("settings.general.rag.reindexWarning"), }); } catch (error) { - if (error instanceof EmbeddingModelVerificationError) { + // A hard security block cannot be forced; keep the "save anyway" action hidden. + if (error instanceof EmbeddingModelBlockedError) { + setEmbeddingModelNeedsForce(false); + } else if (error instanceof EmbeddingModelVerificationError) { setEmbeddingModelNeedsForce(true); } setEmbeddingModelError( From d79495dc96cf7c6ab0d60585e232fa381549c14e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 04:41:37 -0700 Subject: [PATCH 040/113] Add RDNA 2/3/4 ROCm routing tests via a CPU-only torch spoof (#6935) * Add RDNA 2/3/4 ROCm routing tests via a CPU-only torch spoof Introduces tests/_zoo_rocm_spoof.py, the ROCm sibling of _zoo_aggressive_cuda_spoof.py: it reuses the CUDA spoof's torch.cuda no-op machinery and overlays an AMD Radeon identity (torch.version.hip, gcnArchName, capability) for any RDNA 2/3/4 gfx target, so hip code paths run on CPU-only CI with no AMD hardware. tests/studio/install/test_rocm_rdna_routing.py then asserts unsloth_zoo routes every RDNA arch (gfx1030/1031/1032/1034, gfx1100/1101/1102, gfx1150/1151, gfx1200/1201) correctly: device_type resolves to hip, llama.cpp target resolves to (rocm, gfx), and the per-family ROCm bundle suffix (gfx103X/gfx110X/gfx120X, or self for gfx1150/1151) is picked. The torch-facing checks run in a subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached at import) resolves from a clean process; the pure gfx-family mapping runs in-process. Guarded by importorskip so it runs where torch and unsloth_zoo are installed (the Repo tests CPU job) and skips elsewhere. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/_zoo_rocm_spoof.py | 84 +++++++++++++++++++ .../studio/install/test_rocm_rdna_routing.py | 84 +++++++++++++++++++ 2 files changed, 168 insertions(+) create mode 100644 tests/_zoo_rocm_spoof.py create mode 100644 tests/studio/install/test_rocm_rdna_routing.py diff --git a/tests/_zoo_rocm_spoof.py b/tests/_zoo_rocm_spoof.py new file mode 100644 index 0000000000..050191e9d1 --- /dev/null +++ b/tests/_zoo_rocm_spoof.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""ROCm/RDNA spoof: present torch as an AMD Radeon (RDNA 2/3/4) card on a +GPU-less host, so hip paths (device_type -> "hip", llama.cpp ROCm bundle) are +testable in CPU-only CI with no AMD hardware. The ROCm sibling of +_zoo_aggressive_cuda_spoof.py: it reuses that spoof's torch.cuda no-op machinery +and overlays the AMD identity (torch.version.hip, gcnArchName, Radeon name). +Apply BEFORE importing unsloth/unsloth_zoo, since DEVICE_TYPE is cached there. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys + +# gfx -> (marketing name, (capability major, minor), torch.version.hip). hip is +# the ROCm build torch was made against (RDNA2/3 ship 6.x; gfx1102/115x/RDNA4 7.2). +_PROFILES: dict[str, tuple[str, tuple[int, int], str]] = { + "gfx1030": ("AMD Radeon RX 6900 XT", (10, 3), "6.4.43483"), # RDNA2 + "gfx1031": ("AMD Radeon RX 6700 XT", (10, 3), "6.4.43483"), + "gfx1032": ("AMD Radeon RX 6600", (10, 3), "6.4.43483"), + "gfx1034": ("AMD Radeon RX 6400", (10, 3), "6.4.43483"), + "gfx1100": ("AMD Radeon RX 7900 XTX", (11, 0), "6.4.43483"), # RDNA3 + "gfx1101": ("AMD Radeon RX 7800 XT", (11, 0), "6.4.43483"), + "gfx1102": ("AMD Radeon RX 7600", (11, 0), "7.2.1"), + "gfx1150": ("AMD Radeon 890M", (11, 5), "7.2.1"), # RDNA3.5 APU + "gfx1151": ("AMD Radeon 8060S", (11, 5), "7.2.1"), + "gfx1200": ("AMD Radeon RX 9060 XT", (12, 0), "7.2.1"), # RDNA4 + "gfx1201": ("AMD Radeon RX 9070 XT", (12, 0), "7.2.1"), +} + + +def _cuda_spoof(): + """Load the sibling CUDA spoof by path (robust to sys.path), so we reuse its + torch.cuda machinery instead of duplicating it.""" + if "_zoo_aggressive_cuda_spoof" in sys.modules: + return sys.modules["_zoo_aggressive_cuda_spoof"] + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_zoo_aggressive_cuda_spoof.py") + spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + sys.modules["_zoo_aggressive_cuda_spoof"] = mod + return mod + + +def apply(gfx: str = "gfx1100", device_count: int = 1) -> None: + """Present torch as `gfx`. Re-callable to switch arch (identity is overlaid; + the underlying no-op machinery is applied once).""" + import torch + + if gfx not in _PROFILES: + raise KeyError(f"Unknown gfx {gfx!r}; known: {', '.join(_PROFILES)}") + name, cap, hip = _PROFILES[gfx] + + _cuda_spoof().apply() # is_available/device_count/streams/rng/amp/... + + # Overlay the AMD identity on top of the (NVIDIA-shaped) CUDA spoof. + torch.version.hip = hip + torch.version.cuda = None + torch.cuda.device_count = lambda: device_count + torch.cuda.get_device_name = lambda *a, **k: name + torch.cuda.get_device_capability = lambda *a, **k: cap + torch.cuda.get_arch_list = lambda: [gfx] + + class _Props: + pass + + _p = _Props() + _p.name = name + _p.gcnArchName = f"{gfx}:sramecc-:xnack-" # ROCm advertises feature flags + _p.major, _p.minor = cap + _p.total_memory = 16 * 1024**3 + _p.multi_processor_count = 40 + _p.warp_size = 32 # RDNA wavefront (CDNA is 64) + _p.is_integrated = gfx in ("gfx1150", "gfx1151") + _p.is_multi_gpu_board = False + torch.cuda.get_device_properties = lambda *a, **k: _p + + +if __name__ == "__main__": + apply() + import torch + print("ROCm spoof applied:", torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName) diff --git a/tests/studio/install/test_rocm_rdna_routing.py b/tests/studio/install/test_rocm_rdna_routing.py new file mode 100644 index 0000000000..b4aeafb7e4 --- /dev/null +++ b/tests/studio/install/test_rocm_rdna_routing.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""RDNA 2/3/4 routing, validated on CPU-only CI with no AMD hardware. + +tests/_zoo_rocm_spoof.py presents torch as each Radeon gfx arch, then we assert +unsloth_zoo routes it: device_type -> "hip", llama.cpp target -> ("rocm", gfx), +and the per-family ROCm bundle suffix. The torch-facing checks run in a +subprocess so the spoof never leaks into sibling tests and DEVICE_TYPE (cached +at import) resolves from a clean process. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("unsloth_zoo") + +_TESTS_DIR = Path(__file__).resolve().parents[2] # tests/ + +# gfx -> (expected llama.cpp target, expected ROCm bundle family). +_ARCHES = { + "gfx1030": (("rocm", "gfx1030"), "gfx103X"), # RDNA2 + "gfx1031": (("rocm", "gfx1031"), "gfx103X"), + "gfx1032": (("rocm", "gfx1032"), "gfx103X"), + "gfx1034": (("rocm", "gfx1034"), "gfx103X"), + "gfx1100": (("rocm", "gfx1100"), "gfx110X"), # RDNA3 + "gfx1101": (("rocm", "gfx1101"), "gfx110X"), + "gfx1102": (("rocm", "gfx1102"), "gfx110X"), + "gfx1150": (("rocm", "gfx1150"), "gfx1150"), # RDNA3.5 APU (self-family) + "gfx1151": (("rocm", "gfx1151"), "gfx1151"), + "gfx1200": (("rocm", "gfx1200"), "gfx120X"), # RDNA4 + "gfx1201": (("rocm", "gfx1201"), "gfx120X"), +} + +# Child: spoof each arch, then record device_type once (fresh import) and the +# live llama.cpp target per arch. Emits one JSON line the parent parses. +_CHILD = """ +import json, sys +sys.path.insert(0, {tests!r}) +import _zoo_rocm_spoof as spoof +arches = {arches!r} +spoof.apply(arches[0]) +from unsloth_zoo.device_type import get_device_type, is_hip +device_type = [get_device_type(), is_hip()] +from unsloth_zoo import llama_cpp as lc +targets = {{}} +for gfx in arches: + spoof.apply(gfx) + targets[gfx] = list(lc._detect_gpu_target()) +print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}})) +""" + + +@pytest.fixture(scope = "module") +def routed(): + code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES)) + proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True) + line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None) + assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}" + return json.loads(line[len("RESULT ") :]) + + +@pytest.mark.parametrize("gfx", list(_ARCHES)) +def test_detect_gpu_target(routed, gfx): + # RDNA card is routed to its ROCm gfx target (drives the llama.cpp bundle). + assert tuple(routed["targets"][gfx]) == _ARCHES[gfx][0] + + +def test_device_type_is_hip(routed): + # An RDNA card must resolve the compute device_type to "hip". + assert routed["device_type"] == ["hip", True] + + +@pytest.mark.parametrize("gfx", list(_ARCHES)) +def test_rocm_gfx_family(gfx): + # Pure mapping (no torch): each gfx picks the right per-family ROCm bundle. + from unsloth_zoo import llama_cpp as lc + assert lc._rocm_gfx_family(gfx) == _ARCHES[gfx][1] From 59977f95c318c1ba81b53c5b50d4a0ef9c342fea Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 05:49:24 -0700 Subject: [PATCH 041/113] GRPO: default router_aux_loss_coef to 0 on TRL >= 1.7.0 (#6938) TRL 1.7.0 enables the MoE router load-balancing aux loss by default (router_aux_loss_coef = 0.001). Unsloth's optimized GRPO forward does not compute it, so default the coefficient to 0, matching pre-1.7.0 behaviour. Users can still opt in with router_aux_loss_coef > 0. No-op on TRL < 1.7.0. --- unsloth/models/rl.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 62ef9e916a..b5cadf2dea 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1370,6 +1370,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): # [TODO] See https://fengyao.notion.site/off-policy-rl # https://github.com/huggingface/trl/pull/3867 (August 7th) "vllm_importance_sampling_correction": False, + # TRL >= 1.7.0 enables the MoE router aux loss by default (0.001); the optimized + # GRPO forward does not compute it, so default off. Opt in via router_aux_loss_coef > 0. + "router_aux_loss_coef": 0.0, } for k, v in replacements.items(): x = f"{k}( = [^,\n]{{1,}})?,\n" From 411c4d1e50362c0fe6c27d4eea0c29171f79a27a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 06:13:43 -0700 Subject: [PATCH 042/113] Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning (#6908) * Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning Adds unsloth/DeepSeek-V4-Flash-GGUF as a default selectable model with the recommended decoding defaults (temperature 1.0, top_p 1.0 from the official generation_config.json) and its three tier reasoning control. The high/max ladder is surfaced for deepseek-v4 model ids and flows through the existing enable_thinking_effort reasoning style via chat_template_kwargs, so no frontend changes are needed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio DeepSeek-V4: segment-scope high, enable thinking for lone effort, render tests Match deepseek-v4 on whole repo-name segments so a future deepseek-v40 or deepseek40 cannot false-match the synthetic 'high'. In _request_reasoning_kwargs, emit enable_thinking when a named effort level is sent without it, so the newly exposed High mode renders thinking-on over the API (the UI already sent it explicitly). Add a none/high/max render-path test file (jinja behind importorskip) with a lone-high regression. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../assets/configs/inference_defaults.json | 9 +- studio/backend/core/inference/defaults.py | 2 + studio/backend/core/inference/llama_cpp.py | 18 +- .../tests/test_deepseek_v4_thinking_effort.py | 181 ++++++++++++++++++ .../test_safetensors_capability_advertise.py | 38 ++++ 5 files changed, 245 insertions(+), 3 deletions(-) create mode 100644 studio/backend/tests/test_deepseek_v4_thinking_effort.py diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index 1c7a409bc1..0633f80bbc 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -235,6 +235,13 @@ "min_p": 0.1, "repetition_penalty": 1.0 }, + "deepseek-v4": { + "temperature": 1.0, + "top_p": 1.0, + "top_k": -1, + "min_p": 0.0, + "repetition_penalty": 1.0 + }, "deepseek-r1": { "temperature": 0.6, "top_p": 0.95, @@ -394,7 +401,7 @@ "phi-4", "phi-3", "mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral", "devstral", "pixtral", - "deepseek-r1", "deepseek-v3", "deepseek-ocr", + "deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr", "glm-5", "glm-4", "nemotron", "minimax-m2.7", "minimax-m2.5", "minimax", diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py index b64605e16f..a1d03c03e0 100644 --- a/studio/backend/core/inference/defaults.py +++ b/studio/backend/core/inference/defaults.py @@ -8,6 +8,7 @@ import utils.hardware.hardware as hw DEFAULT_MODELS_GGUF = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", @@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [ DEFAULT_MODELS_STANDARD = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 5e6287f528..8b40f5fccd 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -686,6 +686,16 @@ def detect_reasoning_flags( else [] ) if effort_levels: + # DeepSeek-V4's encoder accepts reasoning_effort {'high', 'max'} but its + # template only branches on 'max', so the literal scan misses 'high'. Add it + # (matched on whole repo-name segments, so 'deepseek-v40' won't false-match) + # to expose the full none/high/max ladder instead of none/max. + segments = re.split(r"[-_.]", (model_identifier or "").lower().split("/")[-1]) + is_dsv4 = "deepseek4" in segments or any( + a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:]) + ) + if is_dsv4 and "high" not in effort_levels: + effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index) # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort # level among a discrete set (e.g. 'high' | 'max'). Distinct from # gpt-oss (reasoning_effort only, no on/off gate) and Qwen @@ -1741,9 +1751,13 @@ class LlamaCppBackend: # 'low' effort the way gpt-oss does (those models genuinely # cannot disable). thinking_off = enable_thinking is False or reasoning_effort == "none" - if enable_thinking is not None or reasoning_effort == "none": + # A named effort level implies thinking on, so emit enable_thinking + # even if the caller sent only reasoning_effort (else the template + # defaults it off and the requested level never renders). + effort_on = reasoning_effort in self._reasoning_effort_levels + if enable_thinking is not None or reasoning_effort == "none" or effort_on: kwargs["enable_thinking"] = not thinking_off - if not thinking_off and reasoning_effort in self._reasoning_effort_levels: + if not thinking_off and effort_on: kwargs["reasoning_effort"] = reasoning_effort elif self._reasoning_style == "reasoning_effort": if reasoning_effort in ("none", "low", "medium", "high"): diff --git a/studio/backend/tests/test_deepseek_v4_thinking_effort.py b/studio/backend/tests/test_deepseek_v4_thinking_effort.py new file mode 100644 index 0000000000..19808ad0d7 --- /dev/null +++ b/studio/backend/tests/test_deepseek_v4_thinking_effort.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""DeepSeek-V4-Flash reasoning toggle: None / High / Max. + +The GGUF template gates thinking with ``enable_thinking`` and only branches +``reasoning_effort`` on ``'max'`` (an escalation layered over plain thinking). +Detection used to return the single level ``['max']``, so the UI collapsed to +None / Max and the plain-thinking tier was unreachable. Detection now surfaces +``'high'`` as that plain tier, giving None / High / Max. These tests pin the +classifier, the GLM-style parity case, and the full request-kwargs -> rendered +prompt path for each state (the model itself is too large to load here). +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + + +# Faithful slice of the DeepSeek-V4-Flash GGUF template: the enable_thinking +# gate, the sole ``reasoning_effort == 'max'`` escalation, and the plain-think +# fallback. Any non-'max' effort renders as ordinary thinking. +DEEPSEEK_V4_TEMPLATE = """ +{%- if not thinking is defined -%} + {%- if enable_thinking is defined -%} + {%- set thinking = enable_thinking -%} + {%- else -%} + {%- set thinking = false -%} + {%- endif -%} +{%- endif -%} +{%- if not reasoning_effort is defined -%} + {%- set reasoning_effort = none -%} +{%- endif -%} +{{- bos_token -}} +{%- if thinking and reasoning_effort == 'max' -%} + {{- 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\\n' -}} +{%- endif -%} +{%- for message in messages -%} + {{- '<|User|>' + (message['content'] or '') -}} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- '<|Assistant|>' -}} + {%- if thinking -%}{{- '' -}}{%- else -%}{{- '' -}}{%- endif -%} +{%- endif -%} +""" + + +# GLM-5.2-style: branches on two effort literals, so 'high' already exists as +# the sub-'max' tier and detection must leave the pair untouched. +GLM_STYLE_TEMPLATE = """ +{%- if enable_thinking -%} + {%- if reasoning_effort == 'high' -%}{{- 'H' -}} + {%- elif reasoning_effort == 'max' -%}{{- 'M' -}} + {%- endif -%} +{%- endif -%} +""" + + +# A ['max']-only template under a non-deepseek id: the synthetic 'high' is scoped +# to deepseek-v4, so this must stay ['max'] (no phantom 'high'). +NON_DEEPSEEK_MAX_ONLY_TEMPLATE = DEEPSEEK_V4_TEMPLATE + + +# A template whose sole effort literal is a sub-'max' level: the guard targets +# only the ['max']-alone case, so a lone 'high' stays a singleton. +HIGH_ONLY_TEMPLATE = """ +{%- if enable_thinking and reasoning_effort == 'high' -%}{{- 'H' -}}{%- endif -%} +""" + + +def _render(template: str, **kwargs) -> str: + jinja2 = pytest.importorskip("jinja2") + env = jinja2.Environment() + tmpl = env.from_string(template) + return tmpl.render(bos_token = "", add_generation_prompt = True, **kwargs) + + +# -- Classifier ------------------------------------------------------- + + +def test_deepseek_v4_surfaces_high_as_plain_tier(): + """Sole 'max' escalation expands to ['high', 'max'] so None/High/Max show.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_glm_style_two_level_template_unchanged(): + """A template that already names a sub-'max' tier is left as-is.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(GLM_STYLE_TEMPLATE, "unsloth/GLM-5.2") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_synthetic_high_scoped_to_deepseek_v4(): + """The same ['max']-only template under a non-deepseek id keeps ['max'].""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_effort_levels"] == ["max"] + + +def test_guard_does_not_fire_for_sub_max_singleton(): + """The expansion targets only ['max']; a lone 'high' stays a singleton.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(HIGH_ONLY_TEMPLATE, "custom/high-only") + assert flags["reasoning_effort_levels"] == ["high"] + + +# -- Request kwargs -> rendered prompt, for each state ---------------- + + +def _kwargs_for(flags: dict, enable_thinking, reasoning_effort): + """Drive the real backend method with a shim carrying the detected flags.""" + from core.inference.llama_cpp import LlamaCppBackend + + shim = SimpleNamespace( + _supports_reasoning = flags["supports_reasoning"], + _reasoning_always_on = flags["reasoning_always_on"], + _reasoning_style = flags["reasoning_style"], + _reasoning_effort_levels = flags["reasoning_effort_levels"], + _supports_preserve_thinking = flags["supports_preserve_thinking"], + ) + build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim) + return build(enable_thinking, reasoning_effort, None) or {} + + +def _flags(): + from core.inference.llama_cpp import detect_reasoning_flags + return detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + + +def test_none_state_renders_non_thinking(): + """UI 'None' -> enable_thinking=false -> closed , no preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None) + assert kwargs == {"enable_thinking": False} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_high_state_renders_plain_thinking(): + """UI 'High' -> et=true, effort=high -> open , no max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_max_state_injects_max_preamble(): + """UI 'Max' -> et=true, effort=max -> open plus the max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" in out + + +def test_high_effort_alone_enables_thinking(): + """API caller sending only reasoning_effort='high' (no enable_thinking) still + gets thinking on, so the newly exposed High mode renders correctly.""" + kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 9fd1535f22..0ed670ac01 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -48,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }} """ +# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort +# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders +# identically to thinking-on-without-the-preamble), so the literal scan alone +# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to +# expose the encoder's full none/high/max ladder. +DEEPSEEK_V4_TEMPLATE = ( + "{%- if not thinking is defined %}" + "{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}" + "{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n" + "{%- if thinking and reasoning_effort == 'max' %}" + "{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n" + "{%- for message in messages %}{{- message.content }}{%- endfor %}" +) + + PLAIN_TEMPLATE = """ {%- for message in messages %} {{- message.role + ': ' + message.content + '\\n' }} @@ -90,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false(): assert flags["reasoning_style"] == "enable_thinking" +def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max(): + """DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble. + Classified as the hybrid style with the full none/high/max ladder even + though the template only branches on 'max'.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + assert flags["reasoning_always_on"] is False + + +def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected(): + """The 'high' injection is scoped to deepseek-v4: a different model whose + template only branches on 'max' keeps ['max'] (no phantom 'high').""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["max"] + + def test_detect_safetensors_features_passes_template_through_to_classifier(): """Route wrapper forwards a real template to the inner classifier.""" from routes.inference import _detect_safetensors_features From 10d8f985a270cd4cacf2518e9e895874da201f3e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 06:19:25 -0700 Subject: [PATCH 043/113] Versioning --- pyproject.toml | 6 +++--- unsloth/models/_utils.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 844ead2454..80b3d757e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "wheel>=0.42.0", "packaging", "numpy", @@ -94,7 +94,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "torchvision", "unsloth[triton]", ] @@ -579,7 +579,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.1", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1aa2c6e820..1c75f8ce66 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.6.9" +__version__ = "2026.7.1" __all__ = [ "SUPPORTS_BFLOAT16", From ba450b437eb34ee476df7ffb61bb25db365c7353 Mon Sep 17 00:00:00 2001 From: Etherll <61019402+Etherll@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:26:00 +0300 Subject: [PATCH 044/113] Studio: add assistant response details panel (#6842) * Studio: add assistant response details panel * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Hide model badge by default, show on hover/focus Wrap MessageResponseModelBadge in a span with hidden/group-hover visibility classes to reduce visual clutter. The badge now only displays when hovering or focusing on the assistant message, improving the UI presentation. Updated corresponding tests to verify the new CSS classes. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../message-response-details-sheet.tsx | 483 ++++++++++++++++++ .../src/components/assistant-ui/reasoning.tsx | 10 +- .../src/components/assistant-ui/thread.tsx | 118 +++-- .../src/features/chat/api/chat-adapter.ts | 92 +++- studio/frontend/src/features/chat/index.ts | 7 +- .../chat/stores/chat-preferences-store.ts | 7 + .../src/features/settings/tabs/chat-tab.tsx | 15 + .../test_chat_response_details_ui_contract.py | 93 ++++ 8 files changed, 776 insertions(+), 49 deletions(-) create mode 100644 studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx create mode 100644 tests/studio/test_chat_response_details_ui_contract.py diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx new file mode 100644 index 0000000000..823696693a --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -0,0 +1,483 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { + Sheet, + SheetContent, + SheetDescription, + SheetHeader, + SheetTitle, +} from "@/components/ui/sheet"; +import { + customProviderDisplayName, + parseExternalModelId, + useChatPreferencesStore, + useChatRuntimeStore, + useExternalProvidersStore, +} from "@/features/chat"; +import { cn } from "@/lib/utils"; +import { FileDatabaseIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useMessage, useMessageTiming } from "@assistant-ui/react"; +import type { FC, ReactNode } from "react"; + +type ResponseDetailsMetadata = { + modelId?: string; + modelLabel?: string; + responseModelId?: string; + providerId?: string; + providerName?: string; + providerType?: string; + startedAt?: number; + finishedAt?: number; + durationMs?: number; + sessionId?: string | null; + cancelId?: string; + toolCalls?: string[]; + tools?: Record; +}; + +type ContextUsageMetadata = { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + cachedTokens?: number; + cacheWriteTokens?: number; + modelId?: string; +}; + +type MessageCustomMetadata = { + responseDetails?: ResponseDetailsMetadata; + contextUsage?: ContextUsageMetadata; + serverTimings?: Record; + reasoningDuration?: number; +}; + +function asNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) + ? value + : undefined; +} + +function formatNumber(value: number | undefined): string | null { + return value == null ? null : value.toLocaleString(); +} + +function formatMs(value: number | undefined): string | null { + if (value == null) return null; + if (value < 1000) return `${Math.round(value)}ms`; + return `${(value / 1000).toFixed(2)}s`; +} + +function formatRate(value: number | undefined): string | null { + if (value == null) return null; + return `${value.toFixed(1)} tok/s`; +} + +function formatDate(value: Date | number | string | undefined): string | null { + if (value == null) return null; + const date = value instanceof Date ? value : new Date(value); + if (Number.isNaN(date.getTime())) return null; + return new Intl.DateTimeFormat(undefined, { + dateStyle: "medium", + timeStyle: "medium", + }).format(date); +} + +const TOOL_CATEGORY_LABELS: Record = { + search: "Search", + fetch: "Fetch", + code: "Code", + images: "Images", + mcp: "MCP", + docs: "Docs", + artifacts: "Canvas", +}; + +const TOOL_CALL_LABELS: Record = { + web_search: "Search", + web_fetch: "Fetch", + code_execution: "Code", + python: "Python", + terminal: "Terminal", + image_generation: "Images", + search_knowledge_base: "Docs", + render_html: "Canvas", +}; + +function uniqueValues(values: string[]): string[] { + return Array.from(new Set(values)); +} + +function toolCategoryFromCall(toolName: string): string | null { + const normalized = toolName.toLowerCase(); + if (normalized === "web_search") return "search"; + if (normalized === "web_fetch") return "fetch"; + if ( + normalized === "code_execution" || + normalized === "python" || + normalized === "terminal" + ) { + return "code"; + } + if (normalized === "image_generation") return "images"; + if (normalized === "search_knowledge_base") return "docs"; + if (normalized === "render_html") return "artifacts"; + if (normalized.startsWith("mcp__")) return "mcp"; + return null; +} + +function formatToolCallName(toolName: string): string { + const normalized = toolName.toLowerCase(); + if (TOOL_CALL_LABELS[normalized]) return TOOL_CALL_LABELS[normalized]; + if (normalized.startsWith("mcp__")) return `MCP: ${toolName.slice(5)}`; + return toolName + .replace(/[_-]+/g, " ") + .replace(/\b\w/g, (letter) => letter.toUpperCase()); +} + +function toolCallsFromContent(content: unknown): string[] { + if (!Array.isArray(content)) return []; + return uniqueValues( + content + .map((part) => + part && typeof part === "object" && "type" in part + ? (part as { type?: unknown; toolName?: unknown }) + : null, + ) + .filter( + (part): part is { type: "tool-call"; toolName: string } => + part?.type === "tool-call" && + typeof part.toolName === "string" && + part.toolName.length > 0, + ) + .map((part) => part.toolName), + ); +} + +function enabledTools( + tools: Record | undefined, + toolCalls: string[], +): string | null { + if (!tools && toolCalls.length === 0) return null; + const activeKeys = new Set(); + for (const key of Object.keys(TOOL_CATEGORY_LABELS)) { + if (tools?.[key] === true) activeKeys.add(key); + } + for (const toolName of toolCalls) { + const key = toolCategoryFromCall(toolName); + if (key) activeKeys.add(key); + } + const active = Object.keys(TOOL_CATEGORY_LABELS) + .filter((key) => activeKeys.has(key)) + .map((key) => TOOL_CATEGORY_LABELS[key]); + return active.length > 0 ? active.join(", ") : "None"; +} + +function calledTools(toolCalls: string[]): string | null { + if (toolCalls.length === 0) return null; + return uniqueValues(toolCalls.map(formatToolCallName)).join(", "); +} + +function DetailSection({ + title, + children, +}: { + title: string; + children: ReactNode; +}) { + return ( +
+

{title}

+
{children}
+
+ ); +} + +function DetailRow({ + label, + value, + mono = false, +}: { + label: string; + value: ReactNode | null | undefined; + mono?: boolean; +}) { + if (value == null || value === "") return null; + return ( +
+ {label} + + {value} + +
+ ); +} + +function useResponseModelDisplay() { + const message = useMessage(); + const models = useChatRuntimeStore((s) => s.models); + const providers = useExternalProvidersStore((s) => s.providers); + + const custom = ( + message.metadata as Record | undefined + )?.custom as MessageCustomMetadata | undefined; + const responseDetails = custom?.responseDetails; + const usage = custom?.contextUsage; + const serverTimings = custom?.serverTimings; + + const recordedModelId = + responseDetails?.responseModelId ?? + responseDetails?.modelId ?? + usage?.modelId; + const parsedExternal = parseExternalModelId(recordedModelId); + const provider = parsedExternal + ? providers.find((candidate) => candidate.id === parsedExternal.providerId) + : null; + const modelSummary = models.find( + (candidate) => candidate.id === recordedModelId, + ); + const modelLabel = + responseDetails?.modelLabel ?? + responseDetails?.responseModelId ?? + parsedExternal?.modelId ?? + modelSummary?.name ?? + recordedModelId ?? + "Not recorded"; + const providerLabel = + responseDetails?.providerName ?? + provider?.name ?? + (responseDetails?.providerType + ? customProviderDisplayName(responseDetails.providerType) + : parsedExternal + ? customProviderDisplayName(provider?.providerType) + : recordedModelId + ? "Local model" + : null); + + return { + message, + custom, + responseDetails, + usage, + serverTimings, + modelLabel, + providerLabel, + }; +} + +export const MessageResponseModelBadge: FC<{ className?: string }> = ({ + className, +}) => { + const showResponseModel = useChatPreferencesStore( + (state) => state.showResponseModel, + ); + const { modelLabel, providerLabel } = useResponseModelDisplay(); + + if (!showResponseModel || modelLabel === "Not recorded") { + return null; + } + + return ( + + {modelLabel} + + ); +}; + +export const MessageResponseDetailsSheet: FC<{ + open: boolean; + onOpenChange: (open: boolean) => void; +}> = ({ open, onOpenChange }) => { + const timing = useMessageTiming(); + const { + message, + responseDetails, + usage, + serverTimings, + modelLabel, + providerLabel, + } = useResponseModelDisplay(); + const promptTokens = + usage?.promptTokens ?? asNumber(serverTimings?.prompt_n); + const completionTokens = + usage?.completionTokens ?? + timing?.tokenCount ?? + asNumber(serverTimings?.predicted_n); + const totalTokens = + usage?.totalTokens ?? + (promptTokens != null && completionTokens != null + ? promptTokens + completionTokens + : undefined); + const totalTime = + responseDetails?.durationMs ?? timing?.totalStreamTime ?? undefined; + const summaryLabel = + modelLabel === "Not recorded" ? "Model not recorded" : `Used ${modelLabel}`; + const messageToolCalls = toolCallsFromContent(message.content); + const toolCalls = + responseDetails?.toolCalls && responseDetails.toolCalls.length > 0 + ? responseDetails.toolCalls + : messageToolCalls; + + return ( + + + + + + Response details + + + Timing, model, token, and tool details for this response. + + + +
+
+

+ {summaryLabel} +

+ {providerLabel ? ( +

+ {providerLabel} +

+ ) : null} +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
+ + + ); +}; diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index c5b53577cb..96d21d6fe7 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -6,6 +6,7 @@ /* eslint-disable react-refresh/only-export-components */ import { MarkdownText } from "@/components/assistant-ui/markdown-text"; +import { MessageResponseModelBadge } from "@/components/assistant-ui/message-response-details-sheet"; import { Collapsible, CollapsibleContent, @@ -390,14 +391,17 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ onOpenChange={handleOpenChange} variant={variant} > -
+
-
+ + + +
{isOpen && !isReasoningStreaming && ( )} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 7fe98b701f..09551cd413 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -12,6 +12,10 @@ import { import { downloadImagePart } from "@/components/assistant-ui/image"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts"; +import { + MessageResponseDetailsSheet, + MessageResponseModelBadge, +} from "@/components/assistant-ui/message-response-details-sheet"; import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources"; @@ -3564,6 +3568,9 @@ const AssistantMessage: FC = () => { const aui = useAui(); const messageId = useAuiState(({ message }) => message.id); const messageContent = useAuiState(({ message }) => message.content); + const hasReasoningParts = useAuiState(({ message }) => + message.parts.some((part) => part.type === "reasoning"), + ); const incognito = useChatRuntimeStore((s) => s.incognito); // Use global store for editing state to ensure a single source of truth @@ -3620,7 +3627,7 @@ const AssistantMessage: FC = () => { return (
@@ -3649,6 +3656,11 @@ const AssistantMessage: FC = () => {
) : ( <> + {!hasReasoningParts ? ( +
+ +
+ ) : null} @@ -3893,58 +3905,76 @@ const EditAssistantMessageButton: FC = () => { const AssistantActionBar: FC = () => { const { forkMessage, forkDisabled } = useForkMessageAction(); + const [detailsOpen, setDetailsOpen] = useState(false); return ( - - - - - - - - - - - - - - + <> + + + + + + - - e.preventDefault()} - className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" - > - void forkMessage()} - className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + + + + + + + + + + e.preventDefault()} + className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" > - - Fork in new chat - - - + void forkMessage()} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + > + + Fork in new chat + + + + + Export as Markdown + + + setDetailsOpen(true)} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground" + > - Export as Markdown + See response details - - - - - + + + + + + ); }; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index df266fd749..2ca6ceddf7 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -140,6 +140,32 @@ interface ServerTimings { diffusion_steps_per_second?: number; } +interface ResponseDetailsMetadata { + modelId: string; + modelLabel: string; + responseModelId: string; + providerId?: string; + providerName: string; + providerType: string; + startedAt: number; + finishedAt: number; + durationMs: number; + sessionId?: string; + cancelId: string; + toolCalls: string[]; + tools: { + search: boolean; + fetch: boolean; + code: boolean; + images: boolean; + mcp: boolean; + docs: boolean; + artifacts: boolean; + confirmToolCalls: boolean; + bypassPermissions: boolean; + }; +} + type RunMessages = Parameters[0]["messages"]; type RunMessage = RunMessages[number]; @@ -1769,6 +1795,9 @@ export function createOpenAIStreamAdapter( (provider) => provider.id === externalSelection.providerId, ) : null; + const selectedModelSummary = runtime.models.find( + (model) => model.id === params.checkpoint, + ); const externalApiKey = externalProvider ? getExternalProviderApiKey(externalProvider.id).trim() : ""; @@ -2151,6 +2180,7 @@ export function createOpenAIStreamAdapter( let waitingFirstChunk = true; let firstTokenSettled = false; const streamStartTime = Date.now(); + let responseModelId = externalSelection?.modelId ?? params.checkpoint; let firstTokenTime: number | undefined; let totalChunks = 0; let resolveFirstToken: (() => void) | null = null; @@ -2372,6 +2402,59 @@ export function createOpenAIStreamAdapter( const externalBackendProviderType = toExternalBackendProviderType( externalProvider?.providerType, ); + const buildResponseDetails = ( + finishedAt: number, + ): ResponseDetailsMetadata => ({ + modelId: params.checkpoint, + modelLabel: + (isExternalRequest || responseModelId !== params.checkpoint + ? responseModelId + : selectedModelSummary?.name || responseModelId) || + params.checkpoint || + "Unknown model", + responseModelId: + responseModelId || + externalSelection?.modelId || + params.checkpoint, + ...(externalProvider?.id ? { providerId: externalProvider.id } : {}), + providerName: + externalProvider?.name ?? + (isExternalRequest ? "External provider" : "Local model"), + providerType: externalProvider?.providerType ?? "local", + startedAt: streamStartTime, + finishedAt, + durationMs: finishedAt - streamStartTime, + ...(sandboxSessionId ? { sessionId: sandboxSessionId } : {}), + cancelId, + toolCalls: Array.from( + new Set( + toolCallParts + .map((part) => part.toolName) + .filter( + (toolName): toolName is string => + typeof toolName === "string" && toolName.length > 0, + ), + ), + ), + tools: { + search: + webSearchEnabledForThisTurn || + (!isExternalRequest && supportsTools && toolsEnabled), + fetch: webFetchEnabledForThisTurn, + code: + codeExecEnabledForThisTurn || + (!isExternalRequest && supportsTools && codeToolsEnabled), + images: imageGenerationEnabledForThisTurn, + mcp: !isExternalRequest && supportsTools && mcpEnabledForChat, + docs: + !isExternalRequest && + supportsTools && + (ragEnabled || projectRagEnabled), + artifacts: renderHtmlToolEnabledForThisTurn, + confirmToolCalls, + bypassPermissions, + }, + }); const externalCapabilities = getProviderCapabilities( externalProvider?.providerType, ); @@ -2768,6 +2851,11 @@ export function createOpenAIStreamAdapter( const stream = streamChatCompletions(requestPayload, abortSignal); for await (const chunk of stream) { + const chunkModel = (chunk as { model?: unknown }).model; + if (typeof chunkModel === "string" && chunkModel.length > 0) { + responseModelId = chunkModel; + } + // Handle tool status events const toolStatusText = ( chunk as unknown as { _toolStatus?: string } @@ -3435,11 +3523,12 @@ export function createOpenAIStreamAdapter( }); } + const finishedAt = Date.now(); const finalTiming = buildTiming( streamStartTime, totalChunks, serverPromptEvalTime ?? firstTokenTime, - Date.now() - streamStartTime, + finishedAt - streamStartTime, finalTokenCount, toolCallParts.length, finalTokPerSec, @@ -3475,6 +3564,7 @@ export function createOpenAIStreamAdapter( modelId: params.checkpoint, } : undefined, + responseDetails: buildResponseDetails(finishedAt), timing: finalTiming, }, }, diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 75749b0040..7cd9611c71 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -26,7 +26,12 @@ export { type PlusMenuItemId, } from "./stores/plus-menu-prefs-store"; export { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; -export { isExternalModelId } from "./external-providers"; +export { + customProviderDisplayName, + isExternalModelId, + parseExternalModelId, +} from "./external-providers"; +export { useExternalProvidersStore } from "./stores/external-providers-store"; export { ChatSearchDialog } from "./components/chat-search-dialog"; export { setTrainingCompareHandoff } from "./lib/training-compare-handoff"; export type { ProjectRecord } from "./types"; diff --git a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts index f019b7dc3d..5b5011a78f 100644 --- a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts @@ -7,11 +7,14 @@ import { persist } from "zustand/middleware"; // Client-side chat UI prefs kept in localStorage, not the chat DB. // confirmDeleteChats: when off, deleting a chat skips the confirm dialog. // showModelDisclaimer: when off, hide the "LLMs can make mistakes" footer note. +// showResponseModel: when on, assistant responses show the producing model. export interface ChatPreferencesState { confirmDeleteChats: boolean; setConfirmDeleteChats: (value: boolean) => void; showModelDisclaimer: boolean; setShowModelDisclaimer: (value: boolean) => void; + showResponseModel: boolean; + setShowResponseModel: (value: boolean) => void; } export const useChatPreferencesStore = create()( @@ -23,6 +26,9 @@ export const useChatPreferencesStore = create()( showModelDisclaimer: true, setShowModelDisclaimer: (showModelDisclaimer) => set({ showModelDisclaimer }), + showResponseModel: false, + setShowResponseModel: (showResponseModel) => + set({ showResponseModel }), }), { name: "unsloth_chat_preferences", @@ -32,6 +38,7 @@ export const useChatPreferencesStore = create()( ...current, confirmDeleteChats: saved?.confirmDeleteChats ?? true, showModelDisclaimer: saved?.showModelDisclaimer ?? true, + showResponseModel: saved?.showResponseModel ?? false, }; }, }, diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx index b26b09e022..9519898b21 100644 --- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx @@ -213,6 +213,12 @@ export function ChatTab() { const setShowModelDisclaimer = useChatPreferencesStore( (state) => state.setShowModelDisclaimer, ); + const showResponseModel = useChatPreferencesStore( + (state) => state.showResponseModel, + ); + const setShowResponseModel = useChatPreferencesStore( + (state) => state.setShowResponseModel, + ); useEffect(() => { void countAllChats().then(setCount); @@ -412,6 +418,15 @@ export function ChatTab() { onCheckedChange={setShowModelDisclaimer} /> + + + diff --git a/tests/studio/test_chat_response_details_ui_contract.py b/tests/studio/test_chat_response_details_ui_contract.py new file mode 100644 index 0000000000..04301a0de6 --- /dev/null +++ b/tests/studio/test_chat_response_details_ui_contract.py @@ -0,0 +1,93 @@ +"""Static contract for the chat response-details action and metadata.""" + +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" +DETAILS_TSX = ( + REPO / "studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx" +) +REASONING_TSX = REPO / "studio/frontend/src/components/assistant-ui/reasoning.tsx" +ADAPTER_TS = REPO / "studio/frontend/src/features/chat/api/chat-adapter.ts" +CHAT_PREFS_TS = REPO / "studio/frontend/src/features/chat/stores/chat-preferences-store.ts" +CHAT_TAB_TSX = REPO / "studio/frontend/src/features/settings/tabs/chat-tab.tsx" + + +def test_assistant_more_menu_exposes_response_details_action(): + src = THREAD_TSX.read_text() + assert "MessageResponseDetailsSheet" in src + assert "See response details" in src + assert "setDetailsOpen(true)" in src + + +def test_response_details_sheet_uses_unsloth_sheet_and_key_sections(): + src = DETAILS_TSX.read_text() + assert "SheetContent" in src + assert "Response details" in src + assert "MessageResponseModelBadge" in src + assert "showResponseModel" in src + assert "ChipIcon" not in src + assert "s.params.checkpoint" not in src + assert "Not recorded" in src + assert "min-w-0 break-words font-heading" in src + assert "toolCallsFromContent(message.content)" in src + assert 'label="Called"' in src + for section in ["Response", "Tokens", "Timing", "Tools"]: + assert f'title="{section}"' in src + for field in ["Model", "Provider", "Total", "Cache hits", "Enabled", "Called"]: + assert f'label="{field}"' in src + + +def test_response_model_chip_is_user_configurable_and_rendered_in_metadata_rows(): + prefs_src = CHAT_PREFS_TS.read_text() + chat_tab_src = CHAT_TAB_TSX.read_text() + thread_src = THREAD_TSX.read_text() + reasoning_src = REASONING_TSX.read_text() + + assert "showResponseModel: boolean" in prefs_src + assert "showResponseModel: false" in prefs_src + assert "showResponseModel: saved?.showResponseModel ?? false" in prefs_src + assert "Show response model" in chat_tab_src + assert "setShowResponseModel" in chat_tab_src + assert "aui-response-model-badge inline-flex min-h-5" in DETAILS_TSX.read_text() + assert "leading-5" in DETAILS_TSX.read_text() + assert "group-hover/assistant-message:opacity-100" in DETAILS_TSX.read_text() + assert "MessageResponseModelBadge" in thread_src + assert "hasReasoningParts" in thread_src + assert "group/assistant-message aui-assistant-message-root" in thread_src + assert "pointer-events-none relative h-0" in thread_src + assert "MessageResponseModelBadge" in reasoning_src + assert 'className="min-w-0 flex-none"' in reasoning_src + assert "hidden min-w-0 max-w-[12rem]" in reasoning_src + assert "group-hover/assistant-message:inline-flex" in reasoning_src + + +def test_response_details_metadata_is_persisted_without_backend_schema_change(): + src = ADAPTER_TS.read_text() + assert "interface ResponseDetailsMetadata" in src + assert "buildResponseDetails" in src + assert "responseDetails: buildResponseDetails(finishedAt)" in src + assert "toolCalls: Array.from(" in src + assert "!isExternalRequest && supportsTools && toolsEnabled" in src + assert "!isExternalRequest && supportsTools && codeToolsEnabled" in src + assert re.search(r"selectedModelSummary\?\.name\s*\|\|\s*responseModelId", src) + assert "providerName" in src + assert "cancelId" in src + metadata_block = src[ + src.find("interface ResponseDetailsMetadata") : src.find("type RunMessages") + ] + builder_block = src[ + src.find("const buildResponseDetails") : src.find("const externalCapabilities") + ] + for forbidden in [ + "encrypted_api_key", + "externalApiKey", + "apiKey", + "providerKey", + "secret", + ]: + assert forbidden not in metadata_block + assert forbidden not in builder_block From 8efcc17f476c21dd6f5534ce719b014d7e6a1e4c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 07:20:31 -0700 Subject: [PATCH 045/113] Studio: account for DeepSeek-V4 compute buffer in context auto-fit (#6940) * Studio: account for DeepSeek-V4 compute buffer in context auto-fit DeepSeek-V4-Flash's lightning indexer plus compressed sparse attention reserve a large context-scaling compute buffer that _compute_buffer_ctx_bytes did not model (the KQ-mask and dequant-scratch rates both miss it, even with an f16 cache). Measured on UD-Q4_K_XL at ub 512 it is about 65.5 GiB at 1M context, which the mask estimate puts near 1.5 GiB, so the auto-fit kept the full 1M train context and llama-server OOM'd allocating the ~70 GB buffer, then spilled to CPU (~4 tok/s). Add a deepseek4-gated flat plus per-token term so the fit caps the context (about 256k on a B200) and the model stays fully on GPU. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 14 +++++ studio/backend/tests/test_compute_buffer.py | 61 +++++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8b40f5fccd..11a4ebb3ec 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -3143,6 +3143,12 @@ class LlamaCppBackend: _CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch) _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) _CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) + # DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large + # context-scaling compute buffer the rates above miss (present even with an f16 + # cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without + # it auto-fit commits the full 1M train context, OOMs the reserve, and spills to CPU. + _DSV4_CTX_COMPUTE_FLAT_BYTES = 2 * 1024**3 # ctx-independent indexer scratch + _DSV4_CTX_COMPUTE_BYTES_PER_TOK = 72000 # per token at ub=512 (~72 GiB at 1M) def _estimate_compute_buffer_bytes( self, @@ -3192,6 +3198,14 @@ class LlamaCppBackend: if n_embd <= 0 or n_ctx <= 0: return 0 ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + if getattr(self, "_architecture", None) == "deepseek4": + # DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires + # for any KV type -- the indexer scratch is present even with an f16 cache. + ub_scale = ub / self._DEFAULT_N_UBATCH + return int( + self._DSV4_CTX_COMPUTE_FLAT_BYTES + + self._DSV4_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale + ) if _kv_bytes_per_elem(cache_type_kv) < 2.0: # Quantized cache: the dequant scratch dominates and scales with n_embd. # MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 5d14c5c5bd..8408f8203d 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -65,12 +65,14 @@ def _backend( vocab = 248320, embd = 5120, mla = None, + arch = None, ): """Backend with just the dims the compute-buffer estimate reads.""" b = LlamaCppBackend.__new__(LlamaCppBackend) b._vocab_size = vocab b._embedding_length = embd b._key_length_mla = mla # non-None -> MLA (compressed attention) + b._architecture = arch # GGUF general.architecture (e.g. 'deepseek4') return b @@ -290,3 +292,62 @@ class TestContextBufferMLA: b = _backend(embd = 6144, mla = 256) est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB assert est <= 4141 * 1.7 + + +class TestContextBufferDSV4: + """DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention + compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache). + Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit + must see this so it does not commit the full 1M train context and OOM (spilling + to CPU at ~4 tok/s).""" + + _MEASURED_1M_GIB = 65.5 # 70353790464 B compute-graph reserve that OOM'd at 1M ctx + GIB = 1024**3 + + def test_covers_measured_1m_buffer(self): + b = _backend(embd = 4096, arch = "deepseek4") + gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB + assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}" + + def test_not_wildly_over_at_1m(self): + # Within ~1.3x of measured so the fit still grants a large (~256k) context. + b = _backend(embd = 4096, arch = "deepseek4") + gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB + assert gib <= self._MEASURED_1M_GIB * 1.3 + + def test_fires_for_f16_cache(self): + # The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must + # reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx. + dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) + other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) + assert dsv4 > 40 * other + + def test_cache_type_independent(self): + # Indexer scratch is present for an f16 and a quantized cache alike. + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + + def test_flat_floor_at_small_ctx(self): + # ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB). + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0 + + def test_scales_with_context_and_ubatch(self): + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536) + assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes( + 131072, n_ubatch = 256 + ) + + def test_non_dsv4_unchanged(self): + # Regression guard: a non-deepseek4 model keeps the mask-only f16 rate. + b = _backend(embd = 4096, arch = "llama") + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000 + expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY + assert per_tok == pytest.approx(expected, rel = 1e-6) From 37075c542258e87bb556f6ed7496ea88a818c3b6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 07:49:59 -0700 Subject: [PATCH 046/113] Bump install.sh / install.ps1 pin to unsloth>=2026.7.1 (#6943) Co-authored-by: danielhanchen --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 8c667df079..9114f80af9 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2155,7 +2155,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2169,7 +2169,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2235,7 +2235,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2247,7 +2247,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2275,7 +2275,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index 14fbba478d..796d80e401 100755 --- a/install.sh +++ b/install.sh @@ -2706,7 +2706,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2721,7 +2721,7 @@ if [ "$_MIGRATED" = true ]; then # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2925,7 +2925,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2943,7 +2943,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2975,7 +2975,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." From 07ecdb34c092cb0107f74dbf117616e320ec6ff2 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Tue, 7 Jul 2026 22:24:32 +0530 Subject: [PATCH 047/113] Sort chat recents by last activity (#6844) * show chat by by last activity * Update chat thread updated_at logic and enhance sidebar chat item handling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/routes/chat_history.py | 4 +- studio/backend/storage/studio_db.py | 77 +++++++++- .../tests/test_chat_history_storage.py | 134 ++++++++++++++++++ .../frontend/src/components/app-sidebar.tsx | 16 ++- .../chat/hooks/use-chat-sidebar-items.ts | 26 +++- studio/frontend/src/features/chat/types.ts | 1 + .../chat/utils/chat-history-storage.ts | 5 +- studio/frontend/src/i18n/locales/en.ts | 1 + studio/frontend/src/i18n/locales/zh-CN.ts | 1 + 9 files changed, 252 insertions(+), 13 deletions(-) diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 963d584303..7a27a58a52 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -56,6 +56,7 @@ class ChatThread(BaseModel): projectId: Optional[str] = None archived: bool = False createdAt: int + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None forkedFromThreadId: Optional[str] = None @@ -70,6 +71,7 @@ class ChatThreadPatch(BaseModel): projectId: Optional[str] = None archived: Optional[bool] = None createdAt: Optional[int] = None + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None @@ -252,7 +254,7 @@ async def patch_thread( current_subject: str = Depends(get_current_subject), ): patch = payload.model_dump(exclude_unset = True) - for field in ("title", "modelType", "modelId", "archived", "createdAt"): + for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"): if field in patch and patch[field] is None: raise HTTPException(status_code = 400, detail = f"{field} cannot be null") if patch.get("projectId") and get_chat_project(patch["projectId"]) is None: diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 41a9adcc29..87aa50ee26 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -240,6 +240,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: project_id TEXT, archived INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, + updated_at INTEGER, openai_code_exec_container_id TEXT, anthropic_code_exec_container_id TEXT, forked_from_thread_id TEXT, @@ -261,6 +262,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT") if "forked_from_message_id" not in chat_thread_cols: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT") + if "updated_at" not in chat_thread_cols: + conn.execute("ALTER TABLE chat_threads ADD COLUMN updated_at INTEGER") + # Floor at created_at: forked threads copy older ancestor messages, + # so the fork's creation time must win over the branch message times. + conn.execute( + """ + UPDATE chat_threads SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + """ + ) conn.execute( """ CREATE TABLE IF NOT EXISTS chat_messages ( @@ -992,6 +1011,9 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict: "projectId": data.get("project_id") or None, "archived": bool(data["archived"]), "createdAt": data["created_at"], + "updatedAt": data.get("updated_at") + if data.get("updated_at") is not None + else data["created_at"], "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"), "anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"), "forkedFromThreadId": data.get("forked_from_thread_id"), @@ -1039,8 +1061,8 @@ def upsert_chat_thread(thread: dict) -> dict: conn.execute( """ INSERT INTO chat_threads - (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + (id, title, model_type, model_id, pair_id, project_id, archived, created_at, updated_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, model_type = excluded.model_type, @@ -1049,6 +1071,7 @@ def upsert_chat_thread(thread: dict) -> dict: project_id = excluded.project_id, archived = excluded.archived, created_at = excluded.created_at, + updated_at = COALESCE(excluded.updated_at, chat_threads.updated_at), openai_code_exec_container_id = excluded.openai_code_exec_container_id, anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id, forked_from_thread_id = excluded.forked_from_thread_id, @@ -1063,6 +1086,7 @@ def upsert_chat_thread(thread: dict) -> dict: thread.get("projectId"), 1 if thread.get("archived") else 0, int(thread["createdAt"]), + int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None, thread.get("openaiCodeExecContainerId"), thread.get("anthropicCodeExecContainerId"), thread.get("forkedFromThreadId"), @@ -1084,6 +1108,7 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]: "projectId": ("project_id", patch.get("projectId")), "archived": ("archived", 1 if patch.get("archived") else 0), "createdAt": ("created_at", patch.get("createdAt")), + "updatedAt": ("updated_at", patch.get("updatedAt")), "openaiCodeExecContainerId": ( "openai_code_exec_container_id", patch.get("openaiCodeExecContainerId"), @@ -1155,7 +1180,8 @@ def list_chat_threads( conn = get_connection() try: rows = conn.execute( - f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC", + f"SELECT * FROM chat_threads {where} " + "ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC", values, ).fetchall() return [_chat_thread_from_row(row) for row in rows] @@ -1394,6 +1420,44 @@ def _raise_if_chat_message_thread_conflicts( ) +def _bump_chat_thread_updated_at( + conn: sqlite3.Connection, thread_id: str, message_created_at: int +) -> None: + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX(COALESCE(updated_at, created_at), ?) + WHERE id = ? + """, + (message_created_at, thread_id), + ) + + +def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) -> None: + """Set updated_at from the remaining messages, floored at created_at. + + Unlike the ratchet-only bump, this can lower updated_at -- needed after + pruning, which may delete the thread's newest message. + """ + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + WHERE id = ? + """, + (thread_id,), + ) + + def upsert_chat_message(message: dict) -> dict: conn = get_connection() try: @@ -1432,6 +1496,7 @@ def upsert_chat_message(message: dict) -> dict: int(message["createdAt"]), ), ) + _bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"])) conn.commit() return message except Exception: @@ -1484,6 +1549,12 @@ def sync_chat_messages( for m in messages ], ) + if prune_missing: + _recompute_chat_thread_updated_at(conn, thread_id) + elif messages: + _bump_chat_thread_updated_at( + conn, thread_id, max(int(m["createdAt"]) for m in messages) + ) conn.commit() return list_chat_messages(thread_id) except ChatMessageConflictError: diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index aa19df15fe..0239410734 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -4,6 +4,7 @@ import os import platform import shutil +import sqlite3 import threading import uuid from pathlib import Path @@ -11,6 +12,7 @@ from pathlib import Path import pytest from storage import studio_db +from utils.paths import studio_db_path def _reset_studio_db( @@ -108,6 +110,138 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch): assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}] +def test_chat_thread_updated_at_bumps_on_message_writes(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + assert thread["updatedAt"] == thread["createdAt"] + + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.upsert_chat_message(_message("msg-0", 1_600_000_000_000, "old")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-2", 1_700_000_001_000, "newer")], + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + +def test_chat_thread_updated_at_recomputed_when_pruning(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + studio_db.sync_chat_messages( + "thread-1", + [ + _message("msg-1", 1_700_000_000_500, "older"), + _message("msg-2", 1_700_000_001_000, "newest"), + ], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + # Pruning the newest message must lower updated_at to the remaining one. + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-1", 1_700_000_000_500, "older")], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + # Pruning every message falls back to created_at. + studio_db.sync_chat_messages("thread-1", [], prune_missing = True) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == thread["createdAt"] + + +def test_chat_thread_updated_at_survives_thread_resave(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + + studio_db.upsert_chat_thread(_thread()) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + +def test_list_chat_threads_orders_by_last_activity(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + older = _thread("thread-old") + older["createdAt"] = 1_700_000_000_000 + newer = _thread("thread-new") + newer["createdAt"] = 1_700_000_100_000 + studio_db.upsert_chat_thread(older) + studio_db.upsert_chat_thread(newer) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-new", "thread-old"] + + studio_db.upsert_chat_message( + _message("msg-1", 1_700_000_200_000, "hi", thread_id = "thread-old") + ) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-old", "thread-new"] + + +def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + db_path = studio_db_path() + db_path.parent.mkdir(parents = True, exist_ok = True) + conn = sqlite3.connect(str(db_path)) + try: + conn.execute( + """ + CREATE TABLE chat_threads ( + id TEXT NOT NULL PRIMARY KEY, + title TEXT NOT NULL, + model_type TEXT NOT NULL, + model_id TEXT, + pair_id TEXT, + archived INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE chat_messages ( + id TEXT NOT NULL PRIMARY KEY, + thread_id TEXT NOT NULL, + parent_id TEXT, + role TEXT NOT NULL, + content_json TEXT NOT NULL, + attachments_json TEXT, + metadata_json TEXT, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-with-msgs", "Old", "base", 1_700_000_000_000), + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-empty", "Empty", "base", 1_700_000_050_000), + ) + # Fork-like thread: copied ancestor messages predate the thread itself. + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-fork", "Fork", "base", 1_700_000_100_000), + ) + conn.executemany( + "INSERT INTO chat_messages (id, thread_id, role, content_json, created_at) VALUES (?, ?, ?, ?, ?)", + [ + ("m1", "thread-with-msgs", "user", "[]", 1_700_000_001_000), + ("m2", "thread-with-msgs", "assistant", "[]", 1_700_000_002_000), + ("m3", "thread-fork", "user", "[]", 1_700_000_001_000), + ], + ) + conn.commit() + finally: + conn.close() + + assert studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000 + assert studio_db.get_chat_thread("thread-empty")["updatedAt"] == 1_700_000_050_000 + assert studio_db.get_chat_thread("thread-fork")["updatedAt"] == 1_700_000_100_000 + + def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) project = studio_db.upsert_chat_project(_project()) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index fb1a1fc9c7..2dd0d02515 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -359,7 +359,11 @@ export function AppSidebar() { const activeProjectId = isChatRoute ? ((search.project as string | undefined) ?? null) : null; - const { items: allChatItems } = useChatSidebarItems({ + const { + items: allChatItems, + archivedItems: archivedChatItems, + loaded: chatItemsLoaded, + } = useChatSidebarItems({ enabled: !isStudioRoute, requireMessages: false, }); @@ -1306,6 +1310,16 @@ export function AppSidebar() { renderChatSidebarItem(item, "recent"), )} + {/* "No chats yet" only when there is truly no history: + project-scoped and archived threads leave Recents empty + but still count as existing chats. */} + {chatItemsLoaded && + allChatItems.length === 0 && + archivedChatItems.length === 0 && ( +

+ {t("shell.navigation.noChatsYet")} +

+ )} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts index 55d7777e43..0a0df1139b 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts @@ -27,16 +27,21 @@ export interface SidebarItem { id: string; title: string; createdAt: number; + updatedAt: number; isFork?: boolean; projectId?: string | null; } +function lastActivityAt(thread: ThreadRecord): number { + return thread.updatedAt ?? thread.createdAt; +} + export function groupThreads( threads: ThreadRecord[], archived = false, ): SidebarItem[] { const items: SidebarItem[] = []; - const seenPairs = new Set(); + const pairItems = new Map(); for (const t of threads) { // Coerce archived to a boolean before comparing. Legacy threads (from the @@ -48,30 +53,35 @@ export function groupThreads( continue; } if (t.pairId) { - if (seenPairs.has(t.pairId)) { + const existing = pairItems.get(t.pairId); + if (existing) { + existing.updatedAt = Math.max(existing.updatedAt, lastActivityAt(t)); continue; } - seenPairs.add(t.pairId); - items.push({ + const item: SidebarItem = { type: "compare", id: t.pairId, title: t.title, createdAt: t.createdAt, + updatedAt: lastActivityAt(t), projectId: t.projectId ?? null, - }); + }; + pairItems.set(t.pairId, item); + items.push(item); } else if (!t.pairId) { items.push({ type: "single", id: t.id, title: t.title, createdAt: t.createdAt, + updatedAt: lastActivityAt(t), isFork: Boolean(t.forkedFromThreadId), projectId: t.projectId ?? null, }); } } - return items.sort((a, b) => b.createdAt - a.createdAt); + return items.sort((a, b) => b.updatedAt - a.updatedAt); } // Streaming fires CHAT_HISTORY_UPDATED_EVENT per chunk. Debounce so each quiet @@ -84,6 +94,7 @@ export function useChatSidebarItems(options?: { requireMessages?: boolean; }) { const [allThreads, setAllThreads] = useState([]); + const [loaded, setLoaded] = useState(false); const enabled = options?.enabled ?? true; const requireMessages = options?.requireMessages ?? true; @@ -111,6 +122,7 @@ export function useChatSidebarItems(options?: { // were in flight, or if the effect was torn down. if (cancelled || seq !== requestSeq) return; setAllThreads(threads); + setLoaded(true); } catch (error) { if (isExpectedBackgroundChatStorageError(error)) { return; @@ -144,7 +156,7 @@ export function useChatSidebarItems(options?: { const archivedItems = groupThreads(allThreads ?? [], true); const canCompare = useChatRuntimeStore((s) => Boolean(s.params.checkpoint)); - return { items, archivedItems, canCompare }; + return { items, archivedItems, canCompare, loaded }; } function cancelIfRunning(threadId: string): void { diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 3fe69ccc26..111510d925 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -36,6 +36,7 @@ export interface ThreadRecord { projectId?: string | null; archived: boolean; createdAt: number; + updatedAt?: number; /** * OpenAI shell tool container id from a prior response. When set, the * next turn reuses it via `environment.type="container_reference"` so diff --git a/studio/frontend/src/features/chat/utils/chat-history-storage.ts b/studio/frontend/src/features/chat/utils/chat-history-storage.ts index 4294887df0..00df3657b1 100644 --- a/studio/frontend/src/features/chat/utils/chat-history-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-history-storage.ts @@ -590,7 +590,10 @@ export async function listStoredChatThreads( } return Array.from(byId.values()) .filter((thread) => matchesThreadListArgs(thread, args)) - .sort((a, b) => b.createdAt - a.createdAt); + .sort( + (a, b) => + (b.updatedAt ?? b.createdAt) - (a.updatedAt ?? a.createdAt), + ); } export async function listStoredChatThreadsWithMessages( diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 3d73ad4343..b67fd5ca1d 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -41,6 +41,7 @@ export const en = { recipes: "Recipes", export: "Export", recents: "Recents", + noChatsYet: "No chats yet", settings: "Settings", api: "API", lightMode: "Light Mode", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index 5fd31dc7cb..f6dc265fc7 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -41,6 +41,7 @@ export const zhCN = { recipes: "配方", export: "导出", recents: "最近", + noChatsYet: "暂无对话", settings: "设置", api: "API", lightMode: "浅色模式", From 93c9d6d0dd0bfd48d5766ac952c3a880a47b3727 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 7 Jul 2026 15:13:53 -0300 Subject: [PATCH 048/113] Studio: render \[ \] and \( \) LaTeX delimiters in chat (#6914) --- studio/frontend/src/lib/latex.ts | 169 ++++++++++++++++++++++++++++--- 1 file changed, 155 insertions(+), 14 deletions(-) diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 1a4ebdace8..86a9634048 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -1,8 +1,13 @@ // Adapted from LibreChat's latex.ts // https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts // -// Escapes currency dollar signs so they are not misinterpreted as LaTeX math -// delimiters when singleDollarTextMath is enabled. +// Two jobs, in order: +// 1. Convert LaTeX bracket delimiters (`\[...\]`, `\(...\)`) into the dollar +// forms remark-math understands (`$$...$$`, `$...$`). remark-math only +// tokenizes dollar delimiters, so models that emit `\[...\]` / `\(...\)` +// would otherwise render as literal text. +// 2. Escape currency dollar signs so they are not misinterpreted as LaTeX +// math delimiters when singleDollarTextMath is enabled. /** * Matches a single $ followed by a number pattern (currency), e.g.: @@ -15,14 +20,14 @@ const CURRENCY_REGEX = /(? { const regions: Array<[number, number]> = []; - // Fenced code blocks: ```...``` - const fencedRe = /```[\s\S]*?```/g; + // Fenced code blocks: ```...``` and ~~~...~~~ (both are code in GFM) + const fencedRe = /```[\s\S]*?```|~~~[\s\S]*?~~~/g; let match: RegExpExecArray | null; while ((match = fencedRe.exec(content)) !== null) { regions.push([match.index, match.index + match[0].length]); @@ -51,9 +56,38 @@ function findCodeBlockRegions(content: string): Array<[number, number]> { } /** - * Binary search to check if a position falls inside any code region. + * Match an inline link/image `[text](DEST)`, capturing the destination as group 1 + * with the `d` flag so its span is read straight from `match.indices` (the text + * can contain an escaped `\](`, so a string search for the separator is unsafe). + * The text disallows unescaped `]`; the destination allows escapes and one level + * of balanced parens. */ -function isInCodeBlock( +const LINK_DEST_RE = + /!?\[(?:\\.|[^\]\\])*?\]\(((?:\\.|[^()\\]|\([^()]*\))*)\)/gd; + +/** + * Find the destination spans of inline links/images, so a `\(...\)` written with + * escaped parens inside a URL isn't rewritten as math (which would break the + * link). Only the destination is returned, not the link text, so math in the + * visible text still converts. Sorted, non-overlapping (matches are disjoint). + */ +function findLinkDestinationRegions(content: string): Array<[number, number]> { + if (!content.includes("](")) return []; + const regions: Array<[number, number]> = []; + let match: RegExpExecArray | null; + LINK_DEST_RE.lastIndex = 0; + while ((match = LINK_DEST_RE.exec(content)) !== null) { + // `indices` is present (the `d` flag); group 1 spans the destination. + regions.push(match.indices![1]); + } + return regions; +} + +/** + * Binary search to check if a position falls inside any region. Regions must be + * sorted by start and non-overlapping. + */ +function isInRegion( position: number, regions: Array<[number, number]>, ): boolean { @@ -174,9 +208,109 @@ function hasInlineMathCloser(content: string, offset: number): boolean { } /** - * Preprocess a markdown string to escape currency dollar signs so they are not - * parsed as LaTeX math delimiters. + * Matches a `\[...\]` (display) or `\(...\)` (inline) LaTeX span. Non-greedy so + * the first closer wins; dotall so display spans can wrap lines. `(? block `$$...$$` and `\(...\)` -> inline `$...$` so + * remark-math can tokenize them. Bodies are trimmed: remark-math won't open an + * inline span on `$ ` (a `$` followed by whitespace), and display fences must + * sit on their own line to render as a centered block (not inline math), so + * `\[...\]` becomes `\n$$\n...\n$$\n`. * + * Spans inside code blocks/spans are left intact (a code sample showing `\(x\)` + * must not be rewritten). + * + * A space is inserted between a converted span and a following `$` so their + * delimiters can't fuse (`\(a\)\(b\)` -> `$a$$b$` would mis-tokenize into one + * broken span). A preceding currency (`$5\(x\)`) is instead broken later by the + * currency escape pass. + * + * Returns the rewritten text and the `[start, end)` ranges (in the rewritten + * string) of every span it produced, so the currency pass can skip them. + */ +function convertLatexDelimiters(content: string): { + text: string; + mathRegions: Array<[number, number]>; +} { + if (!content.includes("\\[") && !content.includes("\\(")) { + return { text: content, mathRegions: [] }; + } + + const codeRegions = findCodeBlockRegions(content); + const linkRegions = findLinkDestinationRegions(content); + const inSkipZone = (pos: number) => + isInRegion(pos, codeRegions) || isInRegion(pos, linkRegions); + // Pushed in ascending, non-overlapping order (offset only grows), so this + // stays valid for isInRegion's binary search without a sort. + const mathRegions: Array<[number, number]> = []; + // Accumulate into an array, not a string: reading the last char off a growing + // `+=` accumulator flattens its rope every append (O(n^2) over many spans, on + // the per-frame streaming path), so track the tail char and length instead. + const parts: string[] = []; + let offset = 0; + let lastChar = ""; + let last = 0; + // Append a chunk, separating a trailing `$` from a leading `$` so two spans + // can't fuse. Returns where the chunk landed (after any inserted space). + const append = (chunk: string): number => { + if (!chunk) return offset; + if (lastChar === "$" && chunk.startsWith("$")) { + parts.push(" "); + offset += 1; + } + const start = offset; + parts.push(chunk); + offset += chunk.length; + lastChar = chunk[chunk.length - 1]; + return start; + }; + let match: RegExpExecArray | null; + LATEX_DELIM_RE.lastIndex = 0; + while ((match = LATEX_DELIM_RE.exec(content)) !== null) { + const matchEnd = match.index + match[0].length; + // Skip if either delimiter is inside code or a link destination: an opener + // outside such a zone must not consume a closer inside one and rewrite + // across the boundary. Resume right after this opener (not past the whole + // match) so a valid span that this match spanned across (a stray code `\(` + // paired with a real closer) is still found on the next pass, not swallowed. + if (inSkipZone(match.index) || inSkipZone(matchEnd - 1)) { + LATEX_DELIM_RE.lastIndex = match.index + 1; + continue; + } + const isDisplay = match[1] !== undefined; + const body = (isDisplay ? match[1] : match[2]).trim(); + // Leave an empty span (`\(\)`) literal; a bare `$$` would open a stray + // display block that swallows following text. + if (!body) { + continue; + } + append(content.slice(last, match.index)); + const wrapped = isDisplay ? `\n$$\n${body}\n$$\n` : `$${body}$`; + const start = append(wrapped); + mathRegions.push([start, offset]); + last = matchEnd; + } + append(content.slice(last)); + return { text: parts.join(""), mathRegions }; +} + +/** + * Preprocess a markdown string so LaTeX renders: convert bracket delimiters to + * dollar forms, then escape currency dollar signs so they are not parsed as + * math delimiters. + * + * - `\[E = mc^2\]` becomes a `$$` display block on its own lines (display math) + * - `\(\alpha\)` becomes `$\alpha$` (inline math) + * - `\(x\)` in a code span is untouched * - `$5` alone becomes `\$5` (currency, not math) * - `$\alpha$` is untouched (real LaTeX) * - `$30^\circ$` is untouched (LaTeX whose body starts with a digit) @@ -185,15 +319,22 @@ function hasInlineMathCloser(content: string, offset: number): boolean { * - Currency inside code blocks/spans is untouched */ export function preprocessLaTeX(content: string): string { - if (!content.includes("$")) return content; + const { text, mathRegions } = convertLatexDelimiters(content); - const codeRegions = findCodeBlockRegions(content); + if (!text.includes("$")) return text; - return content.replace(CURRENCY_REGEX, (match, offset) => { - if (isInCodeBlock(offset, codeRegions)) { + const codeRegions = findCodeBlockRegions(text); + + return text.replace(CURRENCY_REGEX, (match, offset) => { + if (isInRegion(offset, codeRegions)) { return match; } - if (hasInlineMathCloser(content, offset)) { + // Skip the spans we just created from `\(...\)` so a numeric body like + // `$5$` isn't re-escaped back to literal `\$5$`. + if (isInRegion(offset, mathRegions)) { + return match; + } + if (hasInlineMathCloser(text, offset)) { return match; } return "\\" + match; From 304b8eca7ae8a7bd743850a248308931c062dab9 Mon Sep 17 00:00:00 2001 From: Ayushman <139611211+InfoSage05@users.noreply.github.com> Date: Tue, 7 Jul 2026 23:58:18 +0530 Subject: [PATCH 049/113] fix: match qwen3-thinking double-newline in train_on_responses_only response pattern (#6926) * fix: match qwen3-thinking chat template double-newline in response pattern The Qwen3-thinking chat template generates `\n\n` (double newline) after the think tag, but `train_on_responses_only` was looking for `\n` (single newline). `\n\n` is token 271 while `\n` is token 198 -- different tokens, so the pattern match in `train_on_responses_only` fails, masking ALL tokens and dropping 100% of training samples. Update the response pattern from `\n` to `\n\n` to match what the actual qwen3-thinking template generates. Fixes #6919 * fix qwen3 thinking response marker --------- Co-authored-by: Ayushman Paul Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> --- studio/backend/utils/datasets/model_mappings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 463d26a692..9d2c983aed 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -487,7 +487,7 @@ TEMPLATE_TO_RESPONSES_MAPPER = { }, "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", From a9db53e189f2c23586bfc4a6f472448afc4807ef Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 7 Jul 2026 19:50:40 -0300 Subject: [PATCH 050/113] Studio: stream reasoning tokens in the tool-loop generator (fixes DeepSeek thinking not streaming with a pill on) (#6947) --- .../core/inference/anthropic_compat.py | 28 +++ studio/backend/core/inference/llama_cpp.py | 54 +++- .../backend/tests/test_anthropic_messages.py | 43 ++++ .../backend/tests/test_llama_cpp_tool_loop.py | 231 +++++++++++++++++- 4 files changed, 338 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 7b572a28ff..3c7a4cb182 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -258,6 +258,10 @@ class AnthropicStreamEmitter: self._open_tool_use_id: Optional[str] = None self._open_tool_args_sent: bool = False self._prev_text: str = "" + # Net minus in the text emitted to the client. Tracked + # from emitted deltas (not _prev_text, which a final bare shrink clobbers) + # so an unclosed reasoning-only block can be balanced before close. + self._open_think_tags: int = 0 self._usage: dict = {} def start( @@ -317,6 +321,7 @@ class AnthropicStreamEmitter: """Close any open block and emit message_delta + message_stop.""" events = [] if self._text_block_open or self._open_tool_call_id is not None: + events.extend(self._close_open_think()) events.append(self._close_block()) self._open_tool_call_id = None self._open_tool_use_id = None @@ -344,12 +349,33 @@ class AnthropicStreamEmitter: ) return events + def _close_open_think(self) -> list[str]: + """Emit a ```` delta when the streamed text left a ```` + open. This emitter diffs cumulative snapshots and drops the generator's + final bare shrink, so a reasoning-only reply would otherwise end on an + unclosed tag. Mirrors the chat route's reasoning extractor, which closes + the block on finish; balances the block before it is closed.""" + if not self._text_block_open or self._open_think_tags <= 0: + return [] + self._open_think_tags = 0 + return [ + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": ""}, + }, + ) + ] + def _handle_content(self, event: dict) -> list[str]: cumulative = event.get("text", "") new_text = cumulative[len(self._prev_text) :] self._prev_text = cumulative if not new_text: return [] + self._open_think_tags += new_text.count("") - new_text.count("") if not self._text_block_open: events = self._open_text_block() else: @@ -374,6 +400,7 @@ class AnthropicStreamEmitter: events = [] if self._text_block_open: + events.extend(self._close_open_think()) events.append(self._close_block()) # Defensive: close a stale open tool_use block before starting another. elif self._open_tool_call_id is not None: @@ -452,6 +479,7 @@ class AnthropicStreamEmitter: events.extend(self._open_text_block()) # Reset text tracking for the next synthesis turn self._prev_text = "" + self._open_think_tags = 0 return events def _open_text_block(self) -> list[str]: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 11a4ebb3ec..757467a008 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8603,13 +8603,31 @@ class LlamaCppBackend: } def _flush_reasoning_and_buffer(): - """Append buffered reasoning (as a block) then the held + """Close a live-streamed block (or emit the buffered reasoning + as one block if it never streamed), then append the held content_buffer to the cumulative display text.""" - nonlocal cumulative_display - if reasoning_accum: + nonlocal cumulative_display, in_thinking + if in_thinking: + cumulative_display += "" + in_thinking = False + elif reasoning_accum: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _close_streamed_think() -> bool: + """Close a live-streamed before a tool call drains, so + consumers without a reasoning extractor (Anthropic) get a balanced + block. Returns True when the caller should yield the result.""" + nonlocal cumulative_display, in_thinking, _last_emitted + if not in_thinking: + return False + cumulative_display += "" + in_thinking = False + if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output: + _last_emitted = cumulative_display + return True + return False + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" probe = strip_llama3_leading_sentinels(text.lstrip()) @@ -8797,6 +8815,10 @@ class LlamaCppBackend: # the structured tool call. has_structured_tc = True detect_state = _S_DRAINING + # Close the reasoning prefix before the tool card + # (mirrors the is_match path). + if _close_streamed_think(): + yield {"type": "content", "text": cumulative_display} for tc_d in tc_deltas: idx = tc_d.get("index", 0) if idx not in tool_calls_acc: @@ -8882,17 +8904,17 @@ class LlamaCppBackend: continue # ── Reasoning tokens ── - # Yield only in STREAMING. In BUFFERING and - # DRAINING, accumulate silently so we don't - # corrupt the consumer's prev_text tracker - # (routes/inference.py never resets it - # between tool iterations). + # Stream live except while DRAINING: reasoning is + # orthogonal to tool detection (content_buffer + # only), and the route resets prev_text on + # tool_start, so the block stays a + # monotonic prefix like the no-tool path. reasoning = delta.get("reasoning_content", "") if reasoning: if _reasoning_started_at is None: _reasoning_started_at = time.monotonic() reasoning_accum += reasoning - if detect_state == _S_STREAMING: + if detect_state != _S_DRAINING: if not in_thinking: cumulative_display += "" in_thinking = True @@ -9020,9 +9042,15 @@ class LlamaCppBackend: _hold_buffer = True if _drain_silently: - # No visible prefix -- the buffered text IS - # the call; drain without yielding it. + # The buffered content IS the call; drain it + # without yielding. A live prefix is + # separate from it -- close that. detect_state = _S_DRAINING + if _close_streamed_think(): + yield { + "type": "content", + "text": cumulative_display, + } elif is_match: # Tool signal -- flush any visible # prefix before DRAINING so the @@ -9115,7 +9143,9 @@ class LlamaCppBackend: ), } elif reasoning_accum and not has_content_tokens: - # Reasoning-only reply: show it as plain text. + # Reasoning-only reply: show it as the main response, + # not a thinking block (mirrors the no-tool path; the + # route's extractor closes the streamed ). if _reasoning_started_at is not None and not _reasoning_summary_emitted: _reasoning_summary_emitted = True yield _reasoning_summary_event(_reasoning_started_at) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 170b456eac..0c6550a3bb 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +def _emitter_client_text(events: list[str]) -> str: + """Concatenate the text_delta payloads an SSE event list carries.""" + text = "" + for line in events: + for raw in line.split("\n"): + raw = raw.strip() + if not raw.startswith("data: "): + continue + data = json.loads(raw[len("data: ") :]) + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + text += delta.get("text", "") + return text + + +def test_anthropic_emitter_closes_reasoning_only_think_block(): + # A reasoning-only reply streams X live then shrinks to bare X at EOF. + # This emitter diffs cumulative snapshots and drops the shrink, so without a + # closing pass the client text would end on an unclosed . finish() + # must balance it. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "The capital"}) + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + # The generator's final bare-text shrink (dropped by the cumulative diff). + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "The capital of France is Paris." + + +def test_anthropic_emitter_does_not_double_close_balanced_think(): + # A reasoning-then-answer reply already closes its own ; the balancer + # must not append a second one. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "Thinking."}) + events += emitter.feed({"type": "content", "text": "Thinking.Answer."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "Thinking.Answer." + + def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index fb1b0e52b7..afac1f5249 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -221,7 +221,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" -def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): +def test_streamed_reasoning_answer_emits_backend_summary(monkeypatch): stream = [ _sse({"reasoning_content": "I am thinking."}), _sse({"reasoning_content": " Still thinking."}), @@ -240,17 +240,236 @@ def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): ) ) + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streams live during BUFFERING instead of arriving as one block: + # each reasoning delta is emitted immediately, wrapped in . + assert content_texts[0] == "I am thinking." + assert content_texts[1] == "I am thinking. Still thinking." + # The final event closes the block and appends the answer. + assert content_texts[-1] == "I am thinking. Still thinking.Final answer." + summary_index = next( i for i, event in enumerate(events) if event["type"] == "reasoning_summary" ) - content_index = next(i for i, event in enumerate(events) if event["type"] == "content") - assert summary_index < content_index + final_content_index = max(i for i, event in enumerate(events) if event["type"] == "content") + assert summary_index < final_content_index assert events[summary_index]["duration_ms"] == 62000 - assert ( - events[content_index]["text"] - == "I am thinking. Still thinking.Final answer." + + +def test_reasoning_streams_incrementally_with_tools(monkeypatch): + # Regression (DeepSeek "thinking doesn't stream"): with a tool/pill active the + # tool-loop generator must stream reasoning token-by-token like the no-tool + # path, not accumulate it and dump one buffered block. + stream = [ + _sse({"reasoning_content": "Step one."}), + _sse({"reasoning_content": " Step two."}), + _sse({"reasoning_content": " Step three."}), + _sse({"content": "Done."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "think then answer"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) ) + reasoning_stage = [ + e["text"] + for e in events + if e["type"] == "content" + and e["text"].startswith("") + and "" not in e["text"] + ] + # One live emission per reasoning delta -- not a single dump. + assert reasoning_stage == [ + "Step one.", + "Step one. Step two.", + "Step one. Step two. Step three.", + ] + final = [e["text"] for e in events if e["type"] == "content"][-1] + assert final == "Step one. Step two. Step three.Done." + + +def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch): + # A reasoning-only turn (whole answer in reasoning_content, no content, no + # tool) with a tool active streams the reasoning live, then resolves to the + # bare reasoning text -- identical to the no-tool generate_chat_completion + # path -- so the non-streaming drain still returns it as `content`, not an + # empty answer. + stream = [ + _sse({"reasoning_content": "The capital of France is Paris."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 5.0, 5.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "just think"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streamed live during BUFFERING (the fix). + assert content_texts[0] == "The capital of France is Paris." + # Resolves to bare reasoning, matching the no-tool sibling. + assert content_texts[-1] == "The capital of France is Paris." + + +def test_reasoning_before_structured_tool_closes_think_block(monkeypatch): + # Regression: reasoning streamed live during BUFFERING must be closed with + # before a structured tool_call drains, so consumers without a + # reasoning extractor (Anthropic /v1/messages) never receive an unclosed + # . Mirrors the is_match (XML tool signal) path. + tool_stream = [ + _sse({"reasoning_content": "Let me search."}), + *_structured_tool_call("web_search", {"query": "weather"}, "call_1"), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + # Reasoning streamed live, then closed before the tool -- balanced block. + assert content_before_tool[0] == "Let me search." + assert content_before_tool[-1] == "Let me search." + + +def _replay_route_reasoning_extractor(cumulatives: list[str]) -> tuple[str, str]: + """Replay the route's cumulative suffix-diff + reasoning extractor (the + shared core of routes/inference.py gguf_stream_chunks and the tool-loop + consumer) over content snapshots. Returns (visible, reasoning).""" + from routes.inference import _ResponsesReasoningExtractor + + extractor = _ResponsesReasoningExtractor(parse_think_markers = True) + prev_text = "" + visible: list[str] = [] + reasoning: list[str] = [] + for cumulative in cumulatives: + new_text = cumulative[len(prev_text) :] + prev_text = cumulative + if not new_text: + continue + reasoning_delta, visible_delta = extractor.feed(new_text) + if reasoning_delta: + reasoning.append(reasoning_delta) + if visible_delta: + visible.append(visible_delta) + final_reasoning, final_visible = extractor.finish() + if final_reasoning: + reasoning.append(final_reasoning) + if final_visible: + visible.append(final_visible) + return "".join(visible), "".join(reasoning) + + +def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch): + # Parity contract: a reasoning-only reply must reach the client identically + # whether tools are on or off. Both generators stream live then + # resolve to the bare reasoning text; the route's suffix-diff + extractor + # must therefore produce the same (visible, reasoning) split for both. + stream = [ + _sse({"reasoning_content": "The capital"}), + _sse({"reasoning_content": " of France is Paris."}), + _done(), + ] + + tool_backend = _make_backend(monkeypatch, [list(stream)], []) + _patch_monotonic(monkeypatch, [1.0, 2.0, 2.0]) + tool_cumulatives = [ + e["text"] + for e in tool_backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "capital of France?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + if e.get("type") == "content" + ] + + no_tool_backend = _make_backend(monkeypatch, [list(stream)], []) + no_tool_cumulatives = [ + y + for y in no_tool_backend.generate_chat_completion( + messages = [{"role": "user", "content": "capital of France?"}], + ) + if isinstance(y, str) + ] + + # Both paths stream the reasoning live with the same leading shape. (Raw + # yield lists aren't compared verbatim: the tool path emits a pre-existing + # duplicate trailing event that the route's suffix-diff dedupes.) + assert tool_cumulatives[:3] == no_tool_cumulatives[:3] + # The contract that matters: identical route-level output. + tool_out = _replay_route_reasoning_extractor(tool_cumulatives) + no_tool_out = _replay_route_reasoning_extractor(no_tool_cumulatives) + assert tool_out == no_tool_out + # Pin the shared contract so a change to either path shows up here. + _visible, reasoning = tool_out + assert reasoning == "The capital of France is Paris." + + +def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch): + # _drain_silently sibling of the structured-tool close: a bare-JSON tool call + # with a live reasoning prefix must also close before draining, and + # must never leak the drained call text as content. + tool_stream = [ + _sse({"reasoning_content": "Searching now."}), + _sse({"content": '{"name":"web_search","arguments":{"query":"weather"}}'}), + _done(), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + assert content_before_tool[0] == "Searching now." + assert content_before_tool[-1] == "Searching now." + # The bare-JSON call text was drained, never surfaced as content. + assert not any('"name"' in t for t in content_before_tool) + def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch): tool_stream = [ From 01b8085dc2e8fae5d99ee7d236d58b706988773c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 17:10:01 -0700 Subject: [PATCH 051/113] Create ossf.yml (#6952) --- .github/workflows/ossf.yml | 78 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 .github/workflows/ossf.yml diff --git a/.github/workflows/ossf.yml b/.github/workflows/ossf.yml new file mode 100644 index 0000000000..f9a270540f --- /dev/null +++ b/.github/workflows/ossf.yml @@ -0,0 +1,78 @@ +# This workflow uses actions that are not certified by GitHub. They are provided +# by a third-party and are governed by separate terms of service, privacy +# policy, and support documentation. + +name: Scorecard supply-chain security +on: + # For Branch-Protection check. Only the default branch is supported. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection + branch_protection_rule: + # To guarantee Maintained check is occasionally updated. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained + schedule: + - cron: '21 20 * * 0' + push: + branches: [ "main" ] + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + # `publish_results: true` only works when run from the default branch. conditional can be removed if disabled. + if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request' + permissions: + # Needed to upload the results to code-scanning dashboard. + security-events: write + # Needed to publish results and get a badge (see publish_results below). + id-token: write + # Uncomment the permissions below if installing in a private repository. + # contents: read + # actions: read + + steps: + - name: "Checkout code" + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: "Run analysis" + uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1 + with: + results_file: results.sarif + results_format: sarif + # (Optional) "write" PAT token. Uncomment the `repo_token` line below if: + # - you want to enable the Branch-Protection check on a *public* repository, or + # - you are installing Scorecard on a *private* repository + # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional. + # repo_token: ${{ secrets.SCORECARD_TOKEN }} + + # Public repositories: + # - Publish results to OpenSSF REST API for easy access by consumers + # - Allows the repository to include the Scorecard badge. + # - See https://github.com/ossf/scorecard-action#publishing-results. + # For private repositories: + # - `publish_results` will always be set to `false`, regardless + # of the value entered here. + publish_results: true + + # (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore + # file_mode: git + + # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF + # format to the repository Actions tab. + - name: "Upload artifact" + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: SARIF file + path: results.sarif + retention-days: 5 + + # Upload the results to GitHub's code scanning dashboard (optional). + # Commenting out will disable upload of results to your repo's Code Scanning dashboard + - name: "Upload to code-scanning" + uses: github/codeql-action/upload-sarif@v3 + with: + sarif_file: results.sarif From 49d1fb38633b2e5b640f7034a3e69734a16396ac Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Wed, 8 Jul 2026 03:08:07 +0200 Subject: [PATCH 052/113] Speed up Studio startup path (#6899) * Speed up Studio startup path * Studio: recheck managed binary executability on preflight cache hit and ignore stale unauthenticated platform fetches Preflight: a matching capability cache fingerprint no longer skips the runnability check when the managed binary's executable bit was cleared (size and mtime unchanged, since chmod bumps ctime not mtime). The cache fast path now confirms the binary is still executable, otherwise it falls back to the CLI help probe so preflight reports Stale and can repair, instead of returning Ready and failing later at backend start. Adds a regression test. Frontend: now that first render is no longer gated on fetchDeviceType, the initial unauthenticated health call can resolve after an authenticated platform fetch. Guard the store so a late unauthenticated or failed non-forced response cannot overwrite an already authoritative device type, tunnel URL, or secure flag. Forced refreshes and the first unauthenticated load are unaffected. * Studio: use access(X_OK) for the preflight cache executability guard A mode bitmask treats any execute bit as launchable, but the executable bits can be set only for another owner or group, or be denied by an ACL, so the current user could still hit PermissionDenied at launch and the cached fast path would wrongly return Ready. access(X_OK) checks real executability for the calling user, so an ownership or permission change correctly falls back to the CLI help probe and the Stale repair path. * Studio: ignore any stale non-forced platform fetch once authoritative Extend the platform store guard so a non-forced health response never overwrites an already authoritative result, not only unauthenticated ones. With a saved token the post-render non-forced request can be authenticated but older than a later forced refresh that already picked up the tunnel URL and secure flag; if that earlier request resolves last it would null those fields. Now any non-forced response is dropped once the store holds a server-reported platform. Forced refreshes and the first authoritative write are unaffected. * Studio: run the managed CLI help probe before trusting the preflight cache Restore running the managed CLI help probe before returning Ready from the desktop capability cache, so a managed install whose venv interpreter or a runtime dependency is broken (while path, size, mtime, and markers are unchanged) is reported Stale for repair rather than proceeding to a backend start that cannot spawn. The capability cache still skips the heavier desktop-capabilities probe on a hit, so a warm cache runs one probe instead of two. Removes the executable-access shortcut, which the help probe now subsumes. --------- Co-authored-by: Daniel Han --- studio/backend/core/inference/orchestrator.py | 8 +- ...t_inference_default_models_non_blocking.py | 42 ++++++ .../frontend/src/components/app-sidebar.tsx | 45 ++++-- studio/frontend/src/config/env.ts | 26 +++- .../frontend/src/features/chat/chat-page.tsx | 43 +++++- .../chat/hooks/use-chat-model-runtime.ts | 12 +- .../src/features/chat/runtime-provider.tsx | 12 +- studio/frontend/src/main.tsx | 16 +-- studio/src-tauri/src/preflight.rs | 132 +++++++++++++++++- studio/src-tauri/src/preflight/managed.rs | 16 +++ 10 files changed, 308 insertions(+), 44 deletions(-) create mode 100644 studio/backend/tests/test_inference_default_models_non_blocking.py diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 19d2230278..cf5d24c367 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -108,13 +108,11 @@ class InferenceOrchestrator: @property def default_models(self) -> list[str]: - # Wait up to 5s for background HF fetch - self._top_models_ready.wait(timeout = 5) top_gguf = self._top_gguf_cache or [] top_hub = self._top_hub_cache or [] - # Curated static defaults first, then HF download-ranked to backfill. - # Send extras so the frontend keeps 4 per category after removing - # downloaded ones. + # Never wait for the remote Hugging Face ranking during startup. Chat's + # first /api/models/list needs curated defaults immediately; the + # background fetch backfills extra choices on later calls. result: list[str] = [] seen: set[str] = set() for m in self._static_models + top_gguf + top_hub: diff --git a/studio/backend/tests/test_inference_default_models_non_blocking.py b/studio/backend/tests/test_inference_default_models_non_blocking.py new file mode 100644 index 0000000000..83a8e7bbfb --- /dev/null +++ b/studio/backend/tests/test_inference_default_models_non_blocking.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Default Chat model metadata must not block on remote Hugging Face discovery.""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.orchestrator import InferenceOrchestrator # noqa: E402 + + +def test_default_models_returns_static_defaults_before_top_fetch(monkeypatch): + sleep_seconds = 2.0 + + def _slow_fetch(self: InferenceOrchestrator) -> None: + time.sleep(sleep_seconds) + self._top_gguf_cache = ["unsloth/slow-GGUF"] + self._top_models_ready.set() + + monkeypatch.setattr(InferenceOrchestrator, "_fetch_top_models", _slow_fetch) + + orchestrator = InferenceOrchestrator() + started = time.monotonic() + defaults = orchestrator.default_models + elapsed = time.monotonic() - started + + assert elapsed < 0.5, f"default_models blocked for {elapsed:.2f}s" + assert defaults == orchestrator._static_models + assert "unsloth/slow-GGUF" not in defaults + + deadline = time.monotonic() + sleep_seconds + 5 + while not orchestrator._top_models_ready.is_set() and time.monotonic() < deadline: + time.sleep(0.05) + + assert "unsloth/slow-GGUF" in orchestrator.default_models diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 2dd0d02515..f59a952b3a 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -81,11 +81,6 @@ import { TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; -import { - exportConversationRawJsonl, - exportConversationCsv, - exportConversationShareGPT, -} from "@/features/chat/prompt-storage/prompt-storage-dialog"; import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, @@ -174,6 +169,36 @@ const TestTubeOutlineIcon = TestTube01Icon.slice( 3, ) as typeof TestTube01Icon; + +type ConversationExportFormat = "raw-jsonl" | "csv" | "sharegpt-jsonl"; + +const CHAT_EXPORT_OPTIONS: Array<{ + label: string; + format: ConversationExportFormat; +}> = [ + { label: "Raw JSONL", format: "raw-jsonl" }, + { label: "CSV", format: "csv" }, + { label: "ShareGPT JSONL", format: "sharegpt-jsonl" }, +]; + +async function exportConversationByFormat( + threadId: string, + format: ConversationExportFormat, +): Promise { + const exports = await import( + "@/features/chat/prompt-storage/prompt-storage-dialog" + ); + switch (format) { + case "raw-jsonl": + return exports.exportConversationRawJsonl(threadId); + case "csv": + return exports.exportConversationCsv(threadId); + case "sharegpt-jsonl": + return exports.exportConversationShareGPT(threadId); + } +} + + function runStatusDotClass(status: TrainingRunSummary["status"]): string { switch (status) { case "running": @@ -899,11 +924,7 @@ export function AppSidebar() { Export - {[ - { label: "Raw JSONL", fn: exportConversationRawJsonl }, - { label: "CSV", fn: exportConversationCsv }, - { label: "ShareGPT JSONL", fn: exportConversationShareGPT }, - ].map(({ label, fn }) => ( + {CHAT_EXPORT_OPTIONS.map(({ label, format }) => ( { @@ -911,7 +932,9 @@ export function AppSidebar() { const ids = item.type === "single" ? [item.id] : (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id); - await Promise.all(ids.map((id) => fn(id))); + await Promise.all( + ids.map((id) => exportConversationByFormat(id, format)), + ); } catch { toast.error("Export failed."); } diff --git a/studio/frontend/src/config/env.ts b/studio/frontend/src/config/env.ts index def1b9bad9..63cdd03141 100644 --- a/studio/frontend/src/config/env.ts +++ b/studio/frontend/src/config/env.ts @@ -54,6 +54,17 @@ export const usePlatformStore = create()((_, get) => ({ isChatOnly: () => get().chatOnly, })); +// Once an authoritative (server-reported) platform has been fetched, a +// non-forced response must not overwrite it. The post-render fetchDeviceType() +// in main.tsx runs before auth is ready and can resolve after the authed +// root-route/provider fetches; such a late write would reset deviceType, +// cloudflareUrl/serverUrl/secure, and fetched, whether it is a browser fallback +// (unauthenticated) or an earlier authenticated request that landed after a +// later forced refresh. Forced refreshes are explicit re-reads, so they still write. +function shouldKeepAuthoritativePlatform(force?: boolean): boolean { + return !force && usePlatformStore.getState().fetched; +} + // `force` re-reads /api/health even if cached, to pick up a late-arriving tunnel URL. export async function fetchDeviceType(options?: { force?: boolean; @@ -81,6 +92,15 @@ export async function fetchDeviceType(options?: { server_url?: string | null; secure?: boolean; }; + // Once the store holds an authoritative (server-reported) platform, a + // non-forced response must not overwrite it. It may be an unauthenticated + // fallback, or an earlier authenticated request that resolved after a + // later forced refresh already picked up device_type and the tunnel + // fields; writing either would reset device type or null the tunnel + // fields. Forced refreshes are explicit re-reads, so they still write. + if (shouldKeepAuthoritativePlatform(options?.force)) { + return usePlatformStore.getState().deviceType; + } const deviceType = data.device_type ?? detectLocalPlatform(); const chatOnly = data.chat_only ?? false; const chatOnlyReason = data.chat_only_reason ?? null; @@ -101,7 +121,11 @@ export async function fetchDeviceType(options?: { } catch { // Backend not ready: use client-side detection so chat-only guard works // on initial load (important for macOS). Keep fetched=false so a later - // call retries against the backend. + // call retries against the backend. But a late non-forced failure must not + // wipe an authoritative platform that already resolved. + if (shouldKeepAuthoritativePlatform(options?.force)) { + return usePlatformStore.getState().deviceType; + } const deviceType = detectLocalPlatform(); const chatOnly = deviceType === "mac"; usePlatformStore.setState({ deviceType, chatOnly, fetched: false }); diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index cd7cfc77fc..b155eff780 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -35,7 +35,6 @@ import { useNativeModelDrop, useNativePathLeasesSupported, } from "@/features/native-intents"; -import { ProjectSourcesPanel } from "@/features/rag/components/project-sources-panel"; import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { isTauri } from "@/lib/api-base"; import { toast } from "@/lib/toast"; @@ -51,7 +50,9 @@ import { Tooltip as TooltipPrimitive } from "radix-ui"; import { type CSSProperties, type ReactElement, + lazy, memo, + Suspense, useCallback, useEffect, useMemo, @@ -134,6 +135,13 @@ import { } from "./utils/chat-history-storage"; import { isAssistantLocalThreadId } from "./utils/thread-ids"; + +const ProjectSourcesPanel = lazy(() => + import("@/features/rag/components/project-sources-panel").then((module) => ({ + default: module.ProjectSourcesPanel, + })), +); + type LoraCandidate = { id: string; baseModel: string; @@ -1018,7 +1026,15 @@ function ProjectLanding({
{projectTab === "sources" ? ( - + + Loading sources… +
+ } + > + + ) : (
{items.map((item) => { @@ -2246,12 +2262,29 @@ export function ChatPage({ return [...fromLoras, ...localModels]; }, [lorasFromStore, localModels]); - useEffect(() => { - if (getTrainingCompareHandoff()) return; - void refresh(); + const inventoryRefreshStartedRef = useRef(false); + const refreshDeferredModelInventories = useCallback(() => { + inventoryRefreshStartedRef.current = true; + void refresh({ includeLoras: true }); refreshLocalModels(); }, [refresh, refreshLocalModels]); + useEffect(() => { + if (getTrainingCompareHandoff()) return; + void refresh({ includeLoras: false }); + const timeoutId = window.setTimeout(() => { + if (!inventoryRefreshStartedRef.current) { + refreshDeferredModelInventories(); + } + }, 1200); + return () => window.clearTimeout(timeoutId); + }, [refresh, refreshDeferredModelInventories]); + + useEffect(() => { + if (!active || !modelSelectorOpen) return; + refreshDeferredModelInventories(); + }, [active, modelSelectorOpen, refreshDeferredModelInventories]); + useEffect(() => { // ChatPage no longer remounts on navigation, so re-check the handoff whenever // we return to /chat (e.g. from the training progress "compare in chat" action). diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e23d1b0b33..798c1658f9 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -312,14 +312,18 @@ export function useChatModelRuntime() { [], ); - const refresh = useCallback(async (options?: { signal?: AbortSignal }) => { + const refresh = useCallback(async (options?: { + signal?: AbortSignal; + includeLoras?: boolean; + }) => { const signal = options?.signal; + const includeLoras = options?.includeLoras ?? true; setModelsError(null); try { const [listRes, statusRes, lorasRes] = await Promise.all([ listModels(), getInferenceStatus(), - listLoras(), + includeLoras ? listLoras() : Promise.resolve(null), ]); // Cancellation can land while the requests above are in flight. Bail @@ -327,7 +331,9 @@ export function useChatModelRuntime() { if (signal?.aborted) return; setModels(listRes.models.map(toChatModelSummary)); - setLoras(lorasRes.loras.map(toLoraSummary)); + if (lorasRes) { + setLoras(lorasRes.loras.map(toLoraSummary)); + } const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const isExternalSelectionActive = isExternalModelId(selectedCheckpoint); diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 67980b94c6..b545695f9e 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -22,7 +22,6 @@ import { unstable_useRemoteThreadListRuntime as useRemoteThreadListRuntime, } from "@assistant-ui/react"; import { createAssistantStream } from "assistant-stream"; -import mammoth from "mammoth"; import { type ReactElement, type ReactNode, @@ -33,7 +32,6 @@ import { useMemo, useRef, } from "react"; -import { extractText, getDocumentProxy } from "unpdf"; import { toast } from "sonner"; import { StudioWebSpeechDictationAdapter } from "./adapters/studio-web-speech-dictation-adapter"; import { @@ -181,7 +179,10 @@ class PDFAttachmentAdapter implements AttachmentAdapter { } async send(attachment: PendingAttachment): Promise { - const buffer = new Uint8Array(await attachment.file.arrayBuffer()); + const [{ extractText, getDocumentProxy }, buffer] = await Promise.all([ + import("unpdf"), + attachment.file.arrayBuffer().then((bytes) => new Uint8Array(bytes)), + ]); const pdf = await getDocumentProxy(buffer); const { text } = await extractText(pdf, { mergePages: true }); return { @@ -298,7 +299,10 @@ class DocxAttachmentAdapter implements AttachmentAdapter { } async send(attachment: PendingAttachment): Promise { - const arrayBuffer = await attachment.file.arrayBuffer(); + const [{ default: mammoth }, arrayBuffer] = await Promise.all([ + import("mammoth"), + attachment.file.arrayBuffer(), + ]); const { value } = await mammoth.extractRawText({ arrayBuffer }); return { id: attachment.id, diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx index 0922e764bb..d0ddf2fc6e 100644 --- a/studio/frontend/src/main.tsx +++ b/studio/frontend/src/main.tsx @@ -5,8 +5,8 @@ import { StrictMode } from "react"; import { createRoot } from "react-dom/client"; import "./index.css"; -import { fetchDeviceType } from "./config/env"; import { App } from "./app/app"; +import { fetchDeviceType } from "./config/env"; import { initializeLocale } from "./i18n"; const globalCrypto = globalThis.crypto as Crypto | undefined; @@ -36,10 +36,10 @@ if (!rootElement) { initializeLocale(); -fetchDeviceType().then(() => { - createRoot(rootElement).render( - - - , - ); -}); +createRoot(rootElement).render( + + + , +); + +fetchDeviceType().catch(() => undefined); diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 5a48d26632..7ef5244754 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -498,19 +498,68 @@ mod tests { } #[cfg(unix)] - fn remove_managed_capability_cache() { - let _ = std::fs::remove_file( - dirs::home_dir() + static MANAGED_CAPABILITY_CACHE_TEST_LOCK: std::sync::LazyLock> = + std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); + + #[cfg(unix)] + struct ManagedCapabilityCacheHome { + path: PathBuf, + previous: Option, + } + + #[cfg(unix)] + impl ManagedCapabilityCacheHome { + fn new(test_name: &str) -> Self { + use std::time::{SystemTime, UNIX_EPOCH}; + + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) .unwrap() - .join(".unsloth") - .join("studio") - .join("desktop_capability_cache.json"), - ); + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "unsloth-preflight-cache-{test_name}-{}-{nanos}", + std::process::id() + )); + std::fs::create_dir_all(&path).unwrap(); + let previous = std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME"); + std::env::set_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME", &path); + Self { path, previous } + } + } + + #[cfg(unix)] + impl Drop for ManagedCapabilityCacheHome { + fn drop(&mut self) { + if let Some(previous) = &self.previous { + std::env::set_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME", previous); + } else { + std::env::remove_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME"); + } + let _ = std::fs::remove_dir_all(&self.path); + } + } + + #[cfg(unix)] + fn managed_capability_cache_path_for_test() -> PathBuf { + std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME") + .map(PathBuf::from) + .or_else(dirs::home_dir) + .unwrap() + .join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json") + } + + #[cfg(unix)] + fn remove_managed_capability_cache() { + let _ = std::fs::remove_file(managed_capability_cache_path_for_test()); } #[cfg(unix)] #[tokio::test] async fn managed_cli_capability_probe_classifies_core_cases() { + let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; + let _cache_home = ManagedCapabilityCacheHome::new("core-cases"); remove_managed_capability_cache(); for (name, script, stale_reason) in [ @@ -567,6 +616,75 @@ exit 1 } } + #[cfg(unix)] + #[tokio::test] + async fn managed_cli_capability_help_probe_runs_before_cache() { + use std::fs; + + let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; + let _cache_home = ManagedCapabilityCacheHome::new("cache-hit"); + + remove_managed_capability_cache(); + // `-h` always succeeds unless `modeh` exists; the desktop-capabilities + // probe always succeeds unless `modecap` exists. Toggling those lets us + // prove the ordering: -h runs on every probe (even a cache hit), while + // the heavier capability probe is skipped once the cache is warm. + let fake = fake_cli( + "cap-cache-hit", + r#"#!/bin/sh +log="$0.calls" +modeh="$0.modeh" +modecap="$0.modecap" +printf '%s\n' "$*" >> "$log" +if [ "$1" = "-h" ]; then + if [ -f "$modeh" ]; then exit 42; fi + exit 0 +fi +if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then + if [ -f "$modecap" ]; then exit 42; fi + printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"version":"2026.5.3"}' + exit 0 +fi +exit 1 +"#, + ); + let bin = fake.bin.clone(); + let calls = bin.with_extension("calls"); + let modeh = bin.with_extension("modeh"); + let modecap = bin.with_extension("modecap"); + + // Cold probe: runs -h and the capability probe, then caches the result. + assert!(matches!( + probe_managed_bin(bin.clone()).await, + ManagedProbe::Ready { .. } + )); + let first_calls = fs::read_to_string(&calls).unwrap(); + assert!(first_calls.contains("-h")); + assert!(first_calls.contains("studio desktop-capabilities --json")); + + // Cache hit: -h still runs, but the capability probe is skipped (breaking + // it via `modecap` proves it is not invoked). + fs::write(&modecap, "broken").unwrap(); + fs::write(&calls, "").unwrap(); + assert!(matches!( + probe_managed_bin(bin.clone()).await, + ManagedProbe::Ready { .. } + )); + assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n"); + + // A non-launchable CLI is caught by the -h probe even with a warm cache: + // preflight reports Stale (for repair) and never trusts the cache. + fs::write(&modeh, "broken").unwrap(); + fs::write(&calls, "").unwrap(); + assert!(matches!( + probe_managed_bin(bin).await, + ManagedProbe::Stale { .. } + )); + assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n"); + + remove_managed_capability_cache(); + } + const EXPECTED_ROOT_ID: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; const OTHER_ROOT_ID: &str = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 57b8365ec5..0d20f271c5 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -188,6 +188,16 @@ fn managed_bin_fingerprint(bin: &Path) -> Option { } fn capability_cache_path() -> Option { + #[cfg(test)] + if let Some(home) = std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME") { + return Some( + PathBuf::from(home) + .join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json"), + ); + } + dirs::home_dir().map(|home| { home.join(".unsloth") .join("studio") @@ -400,6 +410,12 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool { pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { let started = Instant::now(); + // Always verify the managed CLI actually launches before trusting the cache. + // A matching capability fingerprint does not prove the binary can still run: + // its venv interpreter or a runtime dependency can be broken while the + // path/size/mtime/markers are unchanged, so the -h probe runs first and a + // non-launchable install is reported Stale for repair. The capability cache + // below still skips the heavier desktop-capabilities probe on a hit. if !run_cli_probe(&bin, &["-h"]).await { info!( "Managed preflight: cli unusable for {:?} in {}ms", From e7e6a0fb475963e9747358e84cb0a22994c1a54a Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 8 Jul 2026 01:50:35 -0700 Subject: [PATCH 053/113] Polish assistant message actions menu (#6962) * Polish assistant message actions menu Use the circle question mark (HelpCircleIcon) for the "See response details" action instead of the file-database icon, and lowercase the "Export as markdown" label. * Align response details sheet icon --- .../assistant-ui/message-response-details-sheet.tsx | 4 ++-- studio/frontend/src/components/assistant-ui/thread.tsx | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx index 823696693a..331a06a4c6 100644 --- a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -18,7 +18,7 @@ import { useExternalProvidersStore, } from "@/features/chat"; import { cn } from "@/lib/utils"; -import { FileDatabaseIcon } from "@hugeicons/core-free-icons"; +import { HelpCircleIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useMessage, useMessageTiming } from "@assistant-ui/react"; import type { FC, ReactNode } from "react"; @@ -341,7 +341,7 @@ export const MessageResponseDetailsSheet: FC<{ diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 09551cd413..d987092c48 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -127,6 +127,7 @@ import { FileDatabaseIcon, Folder01Icon, FolderAddIcon, + HelpCircleIcon, Image03Icon, McpServerIcon, PencilRulerIcon, @@ -3952,7 +3953,7 @@ const AssistantActionBar: FC = () => { strokeWidth={1.75} className="size-icon" /> - Export as Markdown + Export as markdown { className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground" > From 7f9964f21ed540fdf1ecc1549947bfa5353e6127 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 8 Jul 2026 01:51:34 -0700 Subject: [PATCH 054/113] Move New badge to System settings tab (#6963) * Move New badge to System settings tab Show the "New" badge on the System tab and drop it from Connections. * Stabilize refresh revocation UI test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../src/features/settings/settings-dialog.tsx | 2 +- tests/studio/playwright_chat_ui.py | 58 ++++++++++++++++--- 2 files changed, 50 insertions(+), 10 deletions(-) diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 0d201edc0d..d8a0d45d3f 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -56,6 +56,7 @@ const TABS: TabDef[] = [ id: "resources", labelKey: "settings.tabs.resources", icon: CpuIcon, + badgeKey: "common.new", }, { id: "chat", @@ -72,7 +73,6 @@ const TABS: TabDef[] = [ id: "connections", labelKey: "settings.tabs.connections", icon: CloudIcon, - badgeKey: "common.new", }, { id: "about", labelKey: "settings.tabs.about", icon: HelpCircleIcon }, ]; diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index a53534acc0..71297f9043 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -380,6 +380,18 @@ with sync_playwright() as p: fail(f"/api/auth/refresh wedged: {refresh_resp['error']!r}") refresh = refresh_resp.get("body") or {} token = (refresh or {}).get("access_token") + next_refresh_token = (refresh or {}).get("refresh_token") + if token and next_refresh_token: + robust_evaluate( + page, + """([accessToken, refreshToken]) => { + localStorage.setItem('unsloth_auth_token', accessToken); + localStorage.setItem('unsloth_auth_refresh_token', refreshToken); + }""", + [token, next_refresh_token], + ) + elif token: + fail("/api/auth/refresh returned access_token but no refresh_token") if not token: fail("could not obtain auth token after change-password") @@ -1169,6 +1181,13 @@ with sync_playwright() as p: fail(f"curl login returned no access_token: {login_body!r}") info("CLI obtained an access token") + browser_refresh_token = robust_evaluate( + page, + "() => localStorage.getItem('unsloth_auth_refresh_token')", + ) + if not browser_refresh_token: + fail("browser refresh token missing before CLI rotation") + change_proc = subprocess.run( [ "curl", @@ -1203,18 +1222,39 @@ with sync_playwright() as p: # /change-password revoked refresh tokens server-side (auth.py), so # the browser's /api/auth/refresh must now fail. - refresh_after = evaluate_fetch( - page, - f"{BASE}/api/auth/refresh", - method = "POST", - timeout_ms = FETCH_TIMEOUT_MS, + refresh_proc = subprocess.run( + [ + "curl", + "-sS", + "-o", + os.devnull, + "-w", + "%{http_code}", + "-X", + "POST", + f"{BASE}/api/auth/refresh", + "-H", + "Content-Type: application/json", + "-d", + json.dumps({"refresh_token": browser_refresh_token}), + ], + capture_output = True, + text = True, + timeout = 15, ) - if refresh_after.get("error"): - fail(f"/api/auth/refresh wedged: {refresh_after['error']!r}") - if refresh_after["status"] == 200: + if refresh_proc.returncode != 0: + fail( + f"curl refresh-token check failed: rc={refresh_proc.returncode} " + f"stderr={refresh_proc.stderr!r} stdout={refresh_proc.stdout!r}" + ) + try: + refresh_status = int(refresh_proc.stdout.strip()) + except ValueError: + fail(f"curl refresh-token check returned invalid status: " f"{refresh_proc.stdout!r}") + if refresh_status == 200: fail(f"/api/auth/refresh should fail after CLI rotation; got 200") info( - f"OK browser /api/auth/refresh now {refresh_after['status']} " + f"OK browser /api/auth/refresh now {refresh_status} " "(refresh token revoked) -- old studio session can no longer renew" ) From 393d7e9c2b48b928126008399b0c9443af8a9ec7 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:22:36 +0100 Subject: [PATCH 055/113] Fix opencode Unsloth provider selection (#6906) * fix: force Unsloth provider selection for opencode * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * opencode: pin the model without clobbering the user's disabled providers The session overlay wrote disabled_providers unconditionally and the inline OPENCODE_CONFIG_CONTENT set disabled_providers to an empty list. Since that inline layer outranks the user's global and project config and opencode replaces the array rather than merging it, every provider the user had disabled was silently re-enabled for the session. Only strip 'unsloth' from an existing disable list, and drop disabled_providers from the inline config. Also insert --model only on a bare launch: it is a global flag for the TUI, so placing it before a passthrough subcommand (serve/run) breaks arg parsing; a subcommand takes the model from the pinned config instead. Parse the printed OPENCODE_CONFIG_CONTENT with shlex.split in the test so it round-trips under POSIX shell quoting. * Re-enable a globally disabled opencode unsloth provider for the session A fresh OPENCODE_CONFIG overlay omits disabled_providers, and opencode replaces that array across config layers only when a higher layer sets the key, so a user's global disabled_providers of ['unsloth', ...] survived the merge and left the session provider disabled even though the overlay defines provider.unsloth and pins the model. Consult the user's global opencode config (XDG_CONFIG_HOME/opencode, or %APPDATA%/opencode on Windows) when the overlay has no list of its own, and when the effective list disables unsloth write it back to the overlay minus unsloth. The provider loads while the user's other disabled providers stay disabled. Best-effort read: a missing or unparseable global config is a no-op. * Override opencode disabled_providers in the inline layer; keep model flag for TUI flags Re-enabling a disabled unsloth provider now rides in the inline OPENCODE_CONFIG_CONTENT layer instead of the session overlay. The overlay sits below a project opencode.json, which could re-disable the provider; the inline layer outranks both global and project configs and is recomputed each run, so no-launch reruns never reuse a stale generated list. The effective disabled list is read from the project config if the repo sets one, else the global config, across config.json/opencode.json/opencode.jsonc (JSONC tolerated), and written back minus unsloth only when unsloth is disabled. Also keep the pinned --model when the opencode passthrough starts with a top-level TUI flag such as --dir or --continue; only a real subcommand (serve/run/...) takes the model from config, so a leading '-' now still gets --model injected. * Discover the opencode project config by walking up from the cwd opencode finds a project config by searching ancestor directories, not just the cwd. Walk from the cwd up to the filesystem root and use the nearest directory that sets disabled_providers, so running unsloth start opencode from a subdirectory of a repo whose root config disables unsloth still gets the inline override. * Only inject opencode --model on a bare launch; rely on the inline model pin Injecting --model whenever the passthrough started with a flag could place it before a subcommand (e.g. opencode --print-logs serve), which opencode can misparse. --model is unnecessary for any passthrough because the inline OPENCODE_CONFIG_CONTENT pins the model in the highest-priority layer, so the session model is forced without the flag. Restrict --model to the bare launch and pass any other invocation through untouched. * Register the session provider under a dedicated OpenCode id Selecting the Unsloth model reliably required the wrapper to re-enable a user-disabled unsloth provider, which meant reconstructing OpenCode's full disabled_providers resolution (global, OPENCODE_CONFIG overlay, project config discovered via --dir or an ancestor walk, .opencode directories, OPENCODE_CONFIG_DIR, config.json/opencode.json/opencode.jsonc precedence, and {env:} variable substitution) and overriding it in the inline layer. That is unbounded and cannot be kept correct. Register the session provider under a dedicated id (unsloth-studio) instead. A user's disabled_providers list would never target it, so the session model is always selectable and the overlay no longer reads or writes disabled_providers at all: the user's own disables, in whatever config layer, are left exactly as they are. This removes the JSONC parser, the config-directory scan, and the ancestor/global resolution helpers, and the tests that exercised them. * Scope the opencode session to the Studio provider opencode filters every provider, including a config-defined custom one, through its enabled_providers allowlist and disabled_providers denylist, and pinning the model does not bypass that gate (a filtered provider resolves to a not-found error). The provider arrays are also replaced, not merged, across config layers. So a user with an enabled_providers allowlist that omits the session provider would still have the Studio model filtered out. Set enabled_providers to just the session provider and clear disabled_providers in the inline OPENCODE_CONFIG_CONTENT overlay (the highest-priority layer, which replaces these arrays). This guarantees the Studio model loads regardless of the user's provider filters, without reading or reconstructing their multi-layer config. It is session-only: the overlay lives in the env for this launch and never touches the user's config files, so their normal opencode is unchanged and only this session is limited to the Studio provider. Also drop the redundant --model on --no-launch so the printed command stays append-safe for drivers that append a subcommand (the inline pin forces the model), and parse both POSIX and PowerShell no-launch output in the opencode tests so they are not shell-specific. * Pin opencode small_model to the session provider The session allowlists only the Studio provider, but opencode's separate small_model (used for lightweight tasks) could still point at another provider from the user or project config; under the allowlist that provider is filtered, so the lightweight task would resolve a not-found error mid-session even with the main model pinned. Pin small_model to the session model in the same inline overlay so every model use stays on the enabled provider. The session serves one model, so it is the only valid target, and this stays session-only like the rest of the overlay. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Wasim Yousef Said --- unsloth_cli/commands/start.py | 54 +++++++++++-- unsloth_cli/tests/test_start.py | 132 ++++++++++++++++++++++++++------ 2 files changed, 158 insertions(+), 28 deletions(-) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 1895125b11..d959d20c83 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -50,6 +50,12 @@ _HERMES_PROVIDER = "unsloth" # windows and scales the compaction threshold back down to the real window. _HERMES_MIN_CONTEXT = 65536 _PI_PROVIDER = "unsloth" +# OpenCode selects a model by "/" and honors a user +# disabled_providers list. Register the session provider under a dedicated id a +# user's disable list would never target, so the model is always selectable +# without the wrapper having to reconstruct (and override) OpenCode's full, +# multi-layer disabled_providers resolution. +_OPENCODE_PROVIDER = "unsloth-studio" _PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" _PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} _CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") @@ -1116,13 +1122,17 @@ def write_opencode_config( config = _read_json_object(path) if config is None: typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", + f"Warning: couldn't parse {path} — add an '{_OPENCODE_PROVIDER}' provider " + "there yourself, or move the file aside and re-run.", err = True, ) return {} before = json.dumps(config, sort_keys = True) config.setdefault("$schema", "https://opencode.ai/config.json") + # The session provider is registered under a dedicated id (_OPENCODE_PROVIDER) + # that a user's disabled_providers list would never target, so it is always + # selectable without this overlay having to reconstruct or override OpenCode's + # disabled_providers resolution. model_entry = {"name": model["id"]} window = model.get("context_length") or model.get("max_context_length") if window: @@ -1131,14 +1141,14 @@ def write_opencode_config( # disables OpenCode's auto-compaction; declare the real window (and a sane # output cap) so it compacts instead of overflowing the server. model_entry["limit"] = {"context": window, "output": min(window // 4, 8192)} - _subdict(config, "provider")["unsloth"] = { + _subdict(config, "provider")[_OPENCODE_PROVIDER] = { "npm": "@ai-sdk/openai-compatible", "name": "Unsloth Studio", "options": {"baseURL": f"{base}/v1", "apiKey": key}, "models": {model["id"]: model_entry}, } # OpenCode selects a model by "/". - config["model"] = f"unsloth/{model['id']}" + config["model"] = f"{_OPENCODE_PROVIDER}/{model['id']}" if window: # Compact with ~10% headroom (near 90% full). The fixed 20k-token default # buffer over-compacts, or never settles, on a small local context. @@ -1450,7 +1460,20 @@ def opencode( serve = serve, launch = launch, ) - command = ["opencode", *ctx.args] + opencode_model = f"{_OPENCODE_PROVIDER}/{entry['id']}" + # The inline OPENCODE_CONFIG_CONTENT below pins the model in the highest-priority + # layer, so the session model is forced without a --model flag. Only add --model for + # an interactive bare launch (a convenience so the TUI opens on our model). It is + # omitted for passthrough (inserting it before a subcommand can be misparsed) and for + # --no-launch, where the printed command is consumed by drivers that append a + # subcommand such as `run `; a leading --model would land before that + # subcommand and break it. Those paths rely on the inline pin instead. + if ctx.args: + command = ["opencode", *ctx.args] + elif launch: + command = ["opencode", "--model", opencode_model] + else: + command = ["opencode"] with _session_config("opencode", launch) as cfg: config_path = cfg / "opencode.json" # OPENCODE_CONFIG is an overlay (loaded between the user's global and project @@ -1462,7 +1485,26 @@ def opencode( # outranks project config; the API key stays in the private file, never the env. # Only --yolo carries a permission here (its allow must win over a project config); # a non-yolo session returns no permission, so the project's own rules are honored. - inline_config: dict = {"model": f"unsloth/{entry['id']}"} + # opencode filters every provider (a config-defined custom one included) through + # its enabled_providers allowlist and disabled_providers denylist, and a model pin + # does not bypass that gate -- a filtered provider resolves to ModelNotFoundError. + # To guarantee the session model loads without reading or modifying the user's real + # config, scope THIS session to our provider alone: allowlist _OPENCODE_PROVIDER and + # clear the denylist. These arrays are replaced (not merged) by higher layers, so + # setting them in the highest-priority inline overlay neutralizes any user allowlist + # or denylist for the launch. It is session-only: it lives in OPENCODE_CONFIG_CONTENT + # for this invocation and never touches the user's config files, so their normal + # `opencode` is unchanged; only this session is limited to the Studio provider. + # small_model is opencode's separate model for lightweight tasks; pin it to the + # session model too, or a user/project small_model on another (now filtered) + # provider would resolve a not-found error mid-session. The session serves one + # model, so the session model is the only valid target here anyway. + inline_config: dict = { + "model": opencode_model, + "small_model": opencode_model, + "enabled_providers": [_OPENCODE_PROVIDER], + "disabled_providers": [], + } if session_permission: inline_config["permission"] = session_permission env = { diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index affc24626e..58888010b3 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -435,15 +435,10 @@ def test_opencode_inline_config_beats_project_config(fake_studio): # permissions) ride in OPENCODE_CONFIG_CONTENT, which outranks project config. result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "--yolo"]) assert result.exit_code == 0, result.output - content_line = next( - ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") - ) - inline = json.loads( - shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] - ) - assert inline["model"] == f"unsloth/{MODEL['id']}" + inline = _opencode_inline_config(result.output) + assert inline["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" assert inline["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} - assert "sk-unsloth" not in content_line # key stays in the private file + assert "sk-unsloth" not in result.output # key stays in the private file, not the env def test_opencode_inline_config_omits_permission_without_yolo(fake_studio): @@ -452,13 +447,8 @@ def test_opencode_inline_config_omits_permission_without_yolo(fake_studio): # user's project rules; clearing our own config is the fix, and the inline pins the model. result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert result.exit_code == 0, result.output - content_line = next( - ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") - ) - inline = json.loads( - shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] - ) - assert inline["model"] == f"unsloth/{MODEL['id']}" + inline = _opencode_inline_config(result.output) + assert inline["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" assert "permission" not in inline @@ -1376,14 +1366,17 @@ def test_write_opencode_config_fresh(tmp_path): path = tmp_path / "opencode.json" start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) config = json.loads(path.read_text()) - provider = config["provider"]["unsloth"] + provider = config["provider"][start._OPENCODE_PROVIDER] assert provider["npm"] == "@ai-sdk/openai-compatible" assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} # Context limit must be declared, or OpenCode treats it as 0 and disables compaction. assert provider["models"] == { MODEL["id"]: {"name": MODEL["id"], "limit": {"context": 131072, "output": 8192}} } - assert config["model"] == f"unsloth/{MODEL['id']}" + assert config["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + # The overlay never writes disabled_providers; the dedicated provider id is one a + # user's disable list would not target, so nothing needs re-enabling. + assert "disabled_providers" not in config # Compaction buffer scaled to ~10% of the window (compact near 90%). assert config["compaction"] == {"auto": True, "reserved": 131072 // 10} @@ -1391,18 +1384,98 @@ def test_write_opencode_config_fresh(tmp_path): def test_write_opencode_config_preserves_and_idempotent(tmp_path): path = tmp_path / "opencode.json" path.write_text( - json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) + json.dumps( + { + "theme": "tokyonight", + "disabled_providers": ["ollama", "unsloth"], + "provider": {"anthropic": {"name": "Anthropic"}}, + } + ) ) start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) config = json.loads(path.read_text()) assert config["theme"] == "tokyonight" + # The overlay no longer edits disabled_providers; re-enabling unsloth is done in + # the inline layer, so an existing list here is preserved untouched. + assert config["disabled_providers"] == ["ollama", "unsloth"] assert config["provider"]["anthropic"]["name"] == "Anthropic" - assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" + assert config["provider"][start._OPENCODE_PROVIDER]["options"]["baseURL"] == f"{BASE}/v1" before = path.read_text() start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) assert path.read_text() == before +def test_write_opencode_config_keeps_foreign_disabled_providers(tmp_path): + # A user who disabled other providers (but not unsloth) must keep them disabled: + # the overlay must not rewrite disabled_providers, or those providers get silently + # re-enabled for the session. + path = tmp_path / "opencode.json" + path.write_text(json.dumps({"disabled_providers": ["openai", "gemini"]})) + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["disabled_providers"] == ["openai", "gemini"] + + +def _opencode_inline_config(output: str) -> dict: + # --no-launch prints OPENCODE_CONFIG_CONTENT as a POSIX `export NAME=` + # line on Unix/WSL and a PowerShell `$env:NAME = ""` line on native Windows; + # parse whichever the host emitted so the opencode tests are shell-agnostic. + name = "OPENCODE_CONFIG_CONTENT" + for raw in output.splitlines(): + line = raw.strip() + if line.startswith(f"export {name}="): + return json.loads(shlex.split(line.removeprefix(f"export {name}="))[0]) + prefix = f'$env:{name} = "' + if line.startswith(prefix) and line.endswith('"'): + escaped = line[len(prefix) : -1] + # Reverse _print_env's PowerShell escaping (backtick is the escape char). + value = escaped.replace("`$", "$").replace('`"', '"').replace("``", "`") + return json.loads(value) + raise AssertionError(f"{name} not found in:\n{output}") + + +def test_opencode_inline_scopes_session_to_studio_provider(fake_studio): + # opencode filters even config-defined providers through enabled/disabled_providers, + # and a model pin does not bypass that gate. The inline overlay (session-only, highest + # layer, arrays replace) allowlists our provider and clears the denylist so the Studio + # model always loads regardless of the user's config, without reading or editing it. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + inline = _opencode_inline_config(result.output) + assert inline["enabled_providers"] == [start._OPENCODE_PROVIDER] + assert inline["disabled_providers"] == [] + assert inline["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + # small_model stays on the enabled provider too, so lightweight tasks do not resolve a + # filtered provider mid-session. + assert inline["small_model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + + +def test_opencode_passthrough_flags_omit_model_flag(fake_studio): + # Any passthrough (top-level flags that may precede a subcommand, or a subcommand) + # is left untouched; --model is not injected. The model is pinned by the inline + # OPENCODE_CONFIG_CONTENT (highest layer) instead, so it is still forced. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "--dir", "repo"]) + assert result.exit_code == 0, result.output + command = _launch_command(result.output) + assert command == ["opencode", "--dir", "repo"] + assert "--model" not in command + assert ( + _opencode_inline_config(result.output)["model"] + == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + ) + + +def test_opencode_passthrough_subcommand_omits_model_flag(fake_studio): + # A passthrough subcommand (e.g. `serve`) takes the model from the pinned config; + # inserting --model before it would break opencode's arg parsing. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "serve"]) + assert result.exit_code == 0, result.output + command = _launch_command(result.output) + assert command[0] == "opencode" + assert command[1] == "serve" + assert "--model" not in command + + def test_connect_opencode_no_launch(fake_studio, tmp_path): result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert result.exit_code == 0, result.output @@ -1410,9 +1483,24 @@ def test_connect_opencode_no_launch(fake_studio, tmp_path): config_path = tmp_path / "agents" / "opencode" / "opencode.json" # OPENCODE_CONFIG overlay points at the session file, not the user's global config. _assert_env_set(result.output, "OPENCODE_CONFIG", str(config_path)) + inline_config = _opencode_inline_config(result.output) config = json.loads(config_path.read_text()) - assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["model"] == f"unsloth/{MODEL['id']}" + provider = config["provider"][start._OPENCODE_PROVIDER] + assert provider["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + # The session config file (a throwaway overlay, not the user's real config) does not + # carry provider filters; the session scoping rides in the inline env layer only. + assert "disabled_providers" not in config + assert "enabled_providers" not in config + assert inline_config == { + "model": f"{start._OPENCODE_PROVIDER}/{MODEL['id']}", + "small_model": f"{start._OPENCODE_PROVIDER}/{MODEL['id']}", + "enabled_providers": [start._OPENCODE_PROVIDER], + "disabled_providers": [], + } + # --no-launch prints an append-safe base command (no --model before a subcommand a + # driver may append); the model is forced by the inline pin above. + assert _launch_command(result.output) == ["opencode"] assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) @@ -1765,7 +1853,7 @@ def test_no_launch_rerun_clears_stale_opencode_yolo_permissions(fake_studio, tmp # revert to OpenCode's permissive "allow" default). assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"} # The session provider survives the cleanup. - assert "unsloth" in config["provider"] + assert start._OPENCODE_PROVIDER in config["provider"] def test_no_launch_rerun_clears_stale_openclaw_yolo_state(fake_studio, tmp_path): From baacbd025d9ce92034cd3a41ad0d1ceacabf8202 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:23:25 +0100 Subject: [PATCH 056/113] Fix Hermes install hint on Windows (#6903) * fix: use Windows Hermes installer from unsloth start * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Skip the Hermes setup wizard during unattended start-install unsloth start hermes auto-installs Hermes and then writes its own session-scoped Hermes config. The install commands, as written, drop into the installer's interactive setup wizard (hermes setup), which prompts for global API keys and model choice and points the user at a different global provider than the one Unsloth just configured, blocking the launch. Pass the installer's skip flag on both platforms: the PowerShell scriptblock form with -SkipSetup, and bash -s -- --skip-setup for the piped POSIX installer. * Refresh PATH from the registry after a Windows agent install A Windows installer persists the agent's directory to the User/Machine PATH in the registry and updates only its own process, so the current process keeps a stale PATH until it restarts (the installers print 'restart your terminal'). The post-install shutil.which then misses the just-installed agent and unsloth start fails with 'installed but isn't on PATH yet', forcing a re-run in a new shell. Merge the registry PATH hives back into the process before re-resolving so a freshly installed agent launches in the same invocation. No-op off Windows and on any read error; only ever augments PATH. * Fix/adjust PATH refresh for PR #6903 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- unsloth_cli/commands/start.py | 73 +++++++++++++++++++++++-- unsloth_cli/tests/test_start.py | 96 +++++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+), 4 deletions(-) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index d959d20c83..01014de4d3 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -44,6 +44,19 @@ _CODEX_PROFILE = "unsloth_api" _CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" _HERMES_ENV_KEY = "UNSLOTH_API_KEY" _HERMES_PROVIDER = "unsloth" +# Skip the installer's interactive setup wizard: `unsloth start hermes` runs +# this hint unattended and then writes its own session-scoped Hermes config, so +# the wizard's global API-key/model prompts would block the launch and point the +# user at a different (global) provider than the one Unsloth just configured. +# Both installers expose a skip flag: `-SkipSetup` (PowerShell) and +# `--skip-setup` (POSIX; passed to the piped script via `bash -s --`). +_HERMES_WINDOWS_INSTALL_HINT = ( + "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup" +) +_HERMES_POSIX_INSTALL_HINT = ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash -s -- --skip-setup" +) # Hermes refuses to initialize when the model window is under 64,000 tokens; its # error message points at the model.context_length / auxiliary.compression # overrides in config.yaml. write_hermes_config claims this value for smaller @@ -138,6 +151,10 @@ def _yolo_command_flags(agent: str, yolo: bool) -> list: return _YOLO_COMMAND_FLAGS.get(agent, []) if yolo else [] +def _hermes_install_hint() -> str: + return _HERMES_WINDOWS_INSTALL_HINT if os.name == "nt" else _HERMES_POSIX_INSTALL_HINT + + class LoadOptions(NamedTuple): """Model-load knobs forwarded to /api/inference/load when --model triggers a load.""" @@ -848,6 +865,54 @@ def _print_env( typer.echo(" ".join((*inline, shlex.join(command)))) +def _refresh_windows_path() -> None: + # Merge Windows registry PATH hives after the current process PATH so a + # freshly installed agent is visible without changing existing precedence. + if os.name != "nt": + return + try: + import winreg + except Exception: + return + + entries = [] + seen = set() + + def add_path(value: str) -> bool: + added = False + for entry in str(value).split(os.pathsep): + entry = entry.strip() + if not entry: + continue + key = os.path.normcase(entry).casefold() + if key in seen: + continue + seen.add(key) + entries.append(entry) + added = True + return added + + add_path(os.environ.get("PATH", "")) + added_registry = False + hives = ( + (winreg.HKEY_CURRENT_USER, "Environment"), + ( + winreg.HKEY_LOCAL_MACHINE, + r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment", + ), + ) + for root, sub in hives: + try: + with winreg.OpenKey(root, sub) as key: + value, _ = winreg.QueryValueEx(key, "Path") + except OSError: + continue + if value: + added_registry = add_path(os.path.expandvars(str(value))) or added_registry + if added_registry: + os.environ["PATH"] = os.pathsep.join(entries) + + def _install_agent(name: str, install_hint: str) -> Optional[str]: # Missing agent under --launch: offer to run its documented install command, then # re-resolve it on PATH. Consent-based (we never auto-run a remote install script @@ -866,6 +931,9 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]: install_command = ["/bin/sh", "-c", install_hint] if subprocess.run(install_command).returncode != 0: _fail(f"Install command failed. Run it yourself, then re-run: {install_hint}") + # The installer just wrote PATH to the registry (Windows); pull it into this + # process so the freshly installed agent resolves without a shell restart. + _refresh_windows_path() executable = shutil.which(name) if executable is None: _fail( @@ -1536,10 +1604,7 @@ def hermes( launch = launch, ) command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] - install_hint = ( - "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" - "/main/scripts/install.sh | bash" - ) + install_hint = _hermes_install_hint() with _session_config("hermes", launch) as home: # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session. diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 58888010b3..bd964d5e54 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -92,6 +92,7 @@ def test_claude_flags_detected_when_version_not_first_token(monkeypatch): def test_install_agent_prompts_then_installs(monkeypatch): # TTY + yes: run the documented install command, then re-resolve the now-present binary. + monkeypatch.setattr(start.os, "name", "posix") monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True) ran = [] @@ -108,6 +109,101 @@ def test_install_agent_prompts_then_installs(monkeypatch): assert ran == [["/bin/sh", "-c", "npm install -g @openai/codex"]] +def test_install_agent_uses_powershell_on_windows(monkeypatch): + monkeypatch.setattr(start.os, "name", "nt") + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True) + ran = [] + monkeypatch.setattr( + start.subprocess, + "run", + lambda command, *a, **k: ran.append(command) or SimpleNamespace(returncode = 0), + ) + monkeypatch.setattr(start.shutil, "which", lambda _: r"C:\Users\samle\bin\hermes.exe") + + install_hint = "& ([scriptblock]::Create((irm https://x/install.ps1))) -SkipSetup" + executable = start._install_agent("hermes", install_hint) + + assert executable == r"C:\Users\samle\bin\hermes.exe" + assert ran == [["powershell", "-NoProfile", "-Command", install_hint]] + + +def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch): + monkeypatch.setattr(start.os, "name", "nt") + + # Scriptblock form so `-SkipSetup` reaches the installer and the interactive + # setup wizard is skipped during the unattended `unsloth start hermes` run. + assert start._hermes_install_hint() == ( + "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1)))" + " -SkipSetup" + ) + + +def test_hermes_install_hint_is_bash_on_posix(monkeypatch): + monkeypatch.setattr(start.os, "name", "posix") + + # `bash -s -- --skip-setup` forwards the skip flag to the piped installer. + assert start._hermes_install_hint() == ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash -s -- --skip-setup" + ) + + +def test_refresh_windows_path_noop_off_windows(monkeypatch): + monkeypatch.setattr(start.os, "name", "posix") + before = os.environ.get("PATH", "") + monkeypatch.setenv("PATH", before) + start._refresh_windows_path() + assert os.environ.get("PATH", "") == before + + +def test_refresh_windows_path_merges_registry_hives(monkeypatch): + # Fake Windows registry PATH values written after this process started. + hkcu, hklm = object(), object() + reg = { + (hkcu, "Environment"): r"C:\existing;C:\Users\me\hermes\bin", + ( + hklm, + r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment", + ): r"C:\Windows\System32", + } + + class _Key: + def __init__(self, value): + self._value = value + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def open_key(root, sub): + if (root, sub) in reg: + return _Key(reg[(root, sub)]) + raise OSError("missing hive") + + fake_winreg = SimpleNamespace( + HKEY_CURRENT_USER = hkcu, + HKEY_LOCAL_MACHINE = hklm, + OpenKey = open_key, + QueryValueEx = lambda key, name: (key._value, 1), + ) + monkeypatch.setattr(start.os, "name", "nt") + monkeypatch.setattr(start.os, "pathsep", ";") + monkeypatch.setitem(sys.modules, "winreg", fake_winreg) + monkeypatch.setenv("PATH", r"C:\custom;C:\existing") + + start._refresh_windows_path() + + assert os.environ["PATH"].split(";") == [ + r"C:\custom", + r"C:\existing", + r"C:\Users\me\hermes\bin", + r"C:\Windows\System32", + ] + + def test_install_agent_declined_returns_none(monkeypatch): # TTY + no: never runs anything; caller falls back to the print-hint failure. monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) From a113f893ea767a701b589cd321408f718389a2b8 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Wed, 8 Jul 2026 06:30:37 -0300 Subject: [PATCH 057/113] Studio: heal DiffusionGemma tool calls into structured tool_calls (#6851) * Studio: heal DiffusionGemma tool calls into structured tool_calls * Fall back to supports_tools for backends without the passthrough capability * Route DiffusionGemma client tools through passthrough when enable_tools is on * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop orphaned strip_tool_call_markup import after syncing with main * Tighten supports_tool_passthrough comment * Re-run CI on current main --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 8 +++++++- studio/backend/routes/inference.py | 19 ++++++++++--------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 757467a008..b4f8fe1ca6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -67,7 +67,6 @@ from core.tool_healing import ( _strip_bracket_tag_calls, apply_tool_strip_patterns, strip_outside_think, - strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback @@ -1781,6 +1780,13 @@ class LlamaCppBackend: return False return self._supports_tools + @property + def supports_tool_passthrough(self) -> bool: + # supports_tools is forced off for DiffusionGemma (its agentic loop drops the + # per-step canvas frames), but client passthrough skips that loop, so it uses + # the real _supports_tools. + return self._supports_tools + @property def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5332037e0d..ce755ae1fb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6053,14 +6053,13 @@ async def openai_chat_completions( # free-form sampling. Guided decoding does not require ``supports_tools`` -- # the grammar machinery is independent of tool-call parsing. _has_response_format = _extract_response_format(payload) is not None - _tools_passthrough = llama_backend.supports_tools and ( - (payload.tools and len(payload.tools) > 0) or _has_tool_messages - ) - if ( - using_gguf - and not _effective_enable_tools(payload) - and (_tools_passthrough or _has_response_format) - ): + _tools_passthrough = getattr( + llama_backend, "supports_tool_passthrough", llama_backend.supports_tools + ) and ((payload.tools and len(payload.tools) > 0) or _has_tool_messages) + # DiffusionGemma keeps supports_tools off, so the server-side tool loop can't + # claim the request; fall through to client passthrough, matching /v1/messages. + _server_tool_loop = _effective_enable_tools(payload) and llama_backend.supports_tools + if using_gguf and not _server_tool_loop and (_tools_passthrough or _has_response_format): if _wants_multiple_choices(payload): raise _reject_unsupported_n("GGUF tool or response_format passthrough") if payload.audio_base64: @@ -10222,7 +10221,9 @@ async def anthropic_messages( and not _has_image ) client_tools = ( - not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools + not server_tools + and len(openai_client_tools) > 0 + and getattr(llama_backend, "supports_tool_passthrough", llama_backend.supports_tools) ) # Anthropic tool_choice.disable_parallel_tool_use caps the response to a From df6b5a57d97ab206dd85572e9760a36ef779622a Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:32:06 +0100 Subject: [PATCH 058/113] Fix case-variant model matching and GGUF cache reuse in unsloth start (#6900) * fix: handle case-variant GGUF cache hits for unsloth start * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gguf cache: keep split shards co-located and isolate cache tests properly When a cached main shard was reused from an older snapshot, the extra shards were resolved independently and could come from a different snapshot dir (or a fresh download into the current ref), leaving llama.cpp unable to load a multi-shard GGUF whose pieces are split across directories. Only reuse a cached main shard when every sibling shard sits in the same snapshot; otherwise fetch the whole set together so they stay co-located. Also patch huggingface_hub.constants.HF_HUB_CACHE (not just the HF_HUB_CACHE env var) in the two cache tests that seeded a temp cache: the snapshot lookup reads the module constant, so the env-only override let the real cache leak in and skip an asserted download. * Do not let a companion-only cache snapshot shadow real GGUF variants When listing GGUF variants from the local HF cache, a newer snapshot may contain only a companion file (for example a vision projector fetched on demand) while the actual quant files live in an older snapshot. The prior scan returned the first snapshot whose vision flag was set, yielding an empty variant list and hiding the real quants. Keep scanning older snapshots for actual variants and carry the vision flag across snapshots. Also record the disk-space fallback variant's size in expected_sizes so the later cache-reuse probe can size-verify the fallback main shard instead of only checking for its existence. * Propagate cached repo casing to companions and preflight split co-location Two fixes to the case-variant GGUF cache reuse: - Resolve the requested repo id to its cached canonical casing once in load_model, up front, and pass it to the main GGUF and its companions (mmproj / MTP drafter). Previously only _download_gguf resolved the casing internally, so a case-variant request loaded the main file from the canonical cache dir while the companions kept the requested casing and missed the cached vision projector / drafter offline. Extracted the resolution into a shared _resolve_repo_id_casing helper. - Apply the split-shard co-location check in the disk-space preflight. When a split GGUF's shards are cached across different snapshots the whole set is refetched later, so counting them as cached made the preflight read 0 bytes to download, skip the smaller-variant fallback, and then fail the full download on a low-disk machine. * Reuse a co-located split GGUF snapshot and fix split fallback size probe - When reusing a cached split GGUF, scan snapshots for one that holds the whole set co-located instead of taking the newest snapshot's first shard. A newer snapshot with only the first shard no longer shadows an older complete snapshot, so an already-cached split model is reused rather than refetched (which would fail offline). - The disk-space fallback records its size in expected_sizes only for a single-file fallback. _find_smallest_fitting_variant returns the whole variant size, so using it as the first shard's expected size rejected a valid cached first shard of a split fallback and forced a re-download. * Scan for a complete split snapshot in the preflight; require a loaded catalog hit - The disk-space preflight now uses the same co-located snapshot scan as the download path (_cached_colocated_split_main) instead of the newest-snapshot probe, so a newer snapshot holding only the first shard no longer masks an older complete one and trips the smaller-variant fallback for a fully cached split model. - _resolve_model only attaches to a /v1/models entry that is actually loaded (loaded != False). /v1/models also lists cached-but-unloaded catalog entries, and matching one by case skipped /api/inference/load and left the agent pointed at a model that is not resident. * Restrict cross-snapshot GGUF cache reuse to offline Reusing a same-name blob from an older or case-variant snapshot bypasses the Hub revision/etag check, so a repo that updates a GGUF in place could serve stale weights online. Gate the cross-snapshot and case-variant reuse (both the disk-space preflight accounting and the download path) on HF_HUB_OFFLINE. Online, hf_hub_download fetches the current revision and resumes a partial download, so the reuse is unnecessary there; offline it remains the resilience fallback. Marked the two reuse regression tests as the offline scenarios they represent and added an online test asserting a fresh fetch. * Harden offline cache reuse and hub-id detection Three follow-ups on the case-variant GGUF cache path: - Honor every truthy HF_HUB_OFFLINE spelling (1/true/yes/on), not just "1", when gating the cross-snapshot and case-variant cache reuse. With HF_HUB_OFFLINE=true the Hub calls are already offline, so the reuse must trigger or the cached GGUF fails to load; route both the preflight accounting and the download path through the same offline parse the rest of the backend uses. - Resolve mmproj/MTP companions from the actual cached snapshot when offline. resolve_cached_repo_id_case can keep a partial lower-case spelling when any dir exists under the requested casing, so an hf_hub_download on that casing misses the canonical companion; scan every case-variant snapshot and return the cached path. - Restrict the case-insensitive model-id match to syntactically valid hub ids (a single namespace/name over the HF charset). A server-side relative path such as models/Llama/Foo.gguf is no longer treated as a hub id, so it cannot casefold-match a differently cased path on a case-sensitive filesystem. This is host independent, unlike the local-existence probe which cannot see a server path. * Only casefold-match model ids against a loopback Studio A two-segment string like Models/Foo is indistinguishable from a hub id, and the local Path.exists() probe in _is_hub_model_id cannot see a path that exists only on a remote Studio host. So against a remote server, casefolding could attach to a distinct server-side path (Models/Foo vs models/foo) on a case-sensitive filesystem. Gate the case-insensitive match on is_loopback_url(base): only a local Studio, where the existence probe is authoritative, casefolds. For a remote Studio the match is exact and a case-mismatched request falls through to /api/inference/load, whose already-loaded dedup resolves it correctly. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Wasim Yousef Said --- studio/backend/core/inference/llama_cpp.py | 239 +++++++++++- .../tests/test_offline_gguf_cache_fallback.py | 352 +++++++++++++++++- studio/backend/utils/models/model_config.py | 54 ++- unsloth_cli/commands/start.py | 80 +++- unsloth_cli/tests/test_start.py | 183 +++++++++ 5 files changed, 874 insertions(+), 34 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b4f8fe1ca6..b07bd33076 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -377,6 +377,19 @@ def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool: return True if result[0] is None else result[0] +def _hf_env_offline() -> bool: + """True when an HF offline env var is set to any truthy value (1/true/yes/on). + + Mirrors utils.models.model_config._env_offline so a user-set HF_HUB_OFFLINE=true + (not just "1") still routes through the local-cache reuse path below. + """ + try: + from utils.models.model_config import _env_offline + return _env_offline() + except Exception: + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} + + @contextlib.contextmanager def _hf_offline_if_dns_dead(): """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; @@ -837,6 +850,112 @@ def _gguf_snapshot_files(snapshot: Path) -> list[str]: ] +def _cached_hf_snapshot_file( + repo_id: str, + filename: str, + *, + expected_size: Optional[int] = None, +) -> Optional[str]: + """Return a cached snapshot file even when HF's current-ref probe misses it.""" + if not filename: + return None + parts = [part for part in filename.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + candidate = snap.joinpath(*parts) + if not candidate.is_file(): + continue + if expected_size: + try: + if candidate.stat().st_size < expected_size: + continue + except OSError: + continue + return str(candidate) + except Exception as e: + logger.debug("Snapshot cache lookup failed for %s/%s: %s", repo_id, filename, e) + return None + + +def _snapshot_has_all_shards( + main_path: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> bool: + """True when every shard sits beside ``main_path`` in the same cache snapshot. + + llama.cpp loads a split GGUF by resolving its siblings from the main shard's + directory, so a cached main shard is only safe to reuse when the rest of the + set is co-located; otherwise the caller must fetch the whole set together. + """ + root = Path(main_path) + for _ in [part for part in main_filename.replace("\\", "/").split("/") if part]: + root = root.parent + for shard in shards: + parts = [part for part in shard.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return False + sibling = root.joinpath(*parts) + try: + if not sibling.is_file(): + return False + expected = expected_sizes.get(shard) + if expected and sibling.stat().st_size < expected: + return False + except OSError: + return False + return True + + +def _resolve_repo_id_casing(hf_repo: str) -> str: + """Map a requested repo id to its cached canonical casing, or return it unchanged. + + A case-variant request (for example a lowercased id) resolves to the + canonical-cased cache directory so the main GGUF and its companions + (mmproj / MTP drafter) all read the same cache entry. Returns ``hf_repo`` + unchanged when resolution is unavailable or errors. + """ + try: + from utils.paths import resolve_cached_repo_id_case + return resolve_cached_repo_id_case(hf_repo) + except Exception: + return hf_repo + + +def _cached_colocated_split_main( + repo_id: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> Optional[str]: + """Main-shard path from a cache snapshot that also holds every sibling shard. + + A newer snapshot may hold only the first shard while an older snapshot has the + complete split set. ``_cached_hf_snapshot_file`` would return that newer partial + main and the co-location check would then force a refetch, so scan snapshots for + one where the whole set is present and return that main path instead. None when + no snapshot holds the full set. + """ + main_parts = [part for part in main_filename.replace("\\", "/").split("/") if part] + if not main_parts or any(part in (".", "..") for part in main_parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + main_path = snap.joinpath(*main_parts) + if not main_path.is_file(): + continue + expected_main = expected_sizes.get(main_filename) + try: + if expected_main and main_path.stat().st_size < expected_main: + continue + except OSError: + continue + if _snapshot_has_all_shards(str(main_path), main_filename, shards, expected_sizes): + return str(main_path) + except Exception as e: + logger.debug("Co-located split snapshot lookup failed for %s: %s", repo_id, e) + return None + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -3992,6 +4111,15 @@ class LlamaCppBackend: "Install it with: pip install huggingface_hub" ) + resolved_hf_repo = _resolve_repo_id_casing(hf_repo) + if resolved_hf_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved_hf_repo, + hf_repo, + ) + hf_repo = resolved_hf_repo + # Resolve the filename from the variant gguf_filename = None gguf_extra_shards: list[str] = [] @@ -4037,10 +4165,12 @@ class LlamaCppBackend: # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards + expected_sizes: dict[str, int] = {} try: from huggingface_hub import get_paths_info, try_to_load_from_cache path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token)) + expected_sizes = {p.path: p.size for p in path_infos if p.size} total_bytes = sum((p.size or 0) for p in path_infos) # Subtract bytes already in the HF cache so we only preflight @@ -4049,7 +4179,26 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - if not force: + # Cross-snapshot / case-variant cache reuse is offline-only (see the download + # path below); online, hf_hub_download fetches the current revision and + # resumes partials, so an old snapshot must not be counted as cached here or + # the preflight would under-count the download and skip the disk fallback. + offline = _hf_env_offline() + # A split GGUF whose shards are not co-located in a single snapshot is + # refetched as a whole set later, so it must not be counted as cached here. + split_needs_refetch = False + if offline and not force and gguf_extra_shards: + # Scan all snapshots for one that holds the whole set co-located, so a + # newer snapshot with only the first shard does not mask an older + # complete one and needlessly trip the disk fallback. + if ( + _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + is None + ): + split_needs_refetch = True + if not force and not split_needs_refetch: for p in path_infos: if not p.size: continue @@ -4057,6 +4206,15 @@ class LlamaCppBackend: cached_path = try_to_load_from_cache(hf_repo, p.path) except Exception: cached_path = None + if ( + not (isinstance(cached_path, str) and os.path.exists(cached_path)) + and offline + ): + cached_path = _cached_hf_snapshot_file( + hf_repo, + p.path, + expected_size = p.size, + ) if isinstance(cached_path, str) and os.path.exists(cached_path): try: on_disk = os.path.getsize(cached_path) @@ -4119,6 +4277,13 @@ class LlamaCppBackend: ) else: gguf_extra_shards = [] + # Record the fallback's size so the later cache-reuse probe can + # size-verify it; only for a single-file fallback, since + # _find_smallest_fitting_variant returns the whole-variant size + # and using that as the first shard's expected size would reject + # a valid cached first shard of a split fallback. + if not gguf_extra_shards: + expected_sizes[fallback_file] = fallback_size else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -4138,25 +4303,45 @@ class LlamaCppBackend: raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - local_path = hf_hub_download_with_xet_fallback( - hf_repo, - gguf_filename, - hf_token, - cancel_event = cancel_event, - on_status = lambda m: logger.info(m), - force_download = force, - ) - for shard in gguf_extra_shards: - if cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download_with_xet_fallback( + local_path = None + # Reuse a cached copy from another snapshot / case-variant repo dir only when + # offline. Online, fall through to hf_hub_download so its revision/etag check + # fetches the current file (and resumes a partial) instead of serving a stale + # same-name blob from an older revision. + if not force and _hf_env_offline(): + if gguf_extra_shards: + # A split GGUF must load every shard from one snapshot; reuse only a + # snapshot that holds the whole set co-located, scanning past a newer + # snapshot that has just the first shard while an older one is complete. + local_path = _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + else: + local_path = _cached_hf_snapshot_file( + hf_repo, + gguf_filename, + expected_size = expected_sizes.get(gguf_filename), + ) + if local_path is None: + local_path = hf_hub_download_with_xet_fallback( hf_repo, - shard, + gguf_filename, hf_token, cancel_event = cancel_event, + on_status = lambda m: logger.info(m), force_download = force, ) + for shard in gguf_extra_shards: + if cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Resolving GGUF shard: {shard}") + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = cancel_event, + force_download = force, + ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -4234,6 +4419,17 @@ class LlamaCppBackend: if target is None or cancel_event.is_set(): return None + # Offline, resolve the companion straight from the cache snapshot that + # holds it. resolve_cached_repo_id_case can return a partial lower-case + # spelling when any dir exists under the requested casing, so calling + # hf_hub_download with hf_repo would miss the canonical file and silently + # drop the companion. _cached_hf_snapshot_file scans every case variant. + if _hf_env_offline(): + cached = _cached_hf_snapshot_file(hf_repo, target) + if cached: + logger.info("Resolved %s from local HF cache: %s", label, cached) + return cached + try: logger.info(f"Downloading {label}: {hf_repo}/{target}") # Same policy; companions are best-effort (caller below swallows failures to None). @@ -5012,6 +5208,19 @@ class LlamaCppBackend: # dead; cleanup runs even on exception so a transient hiccup # can't quarantine future loads. if hf_repo: + # Resolve the requested repo id to its cached canonical casing once, + # up front, so the main GGUF and its companions (mmproj / MTP drafter) + # all resolve from the same cache entry. Otherwise a case-variant + # request resolves the main file from the canonical cache dir while the + # companions keep the requested casing and miss the cached files. + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): model_path = self._download_gguf( hf_repo = hf_repo, diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 4499881c4d..e24e2ca451 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -79,9 +79,11 @@ from huggingface_hub import constants as hf_constants from core.inference.llama_cpp import ( LlamaCppBackend, + _cached_colocated_split_main, _gguf_files_for_variant, _hf_offline_if_dns_dead, _probe_dns_dead, + _resolve_repo_id_casing, ) from utils.models.model_config import ( _detect_gguf_from_hf_cache, @@ -217,7 +219,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch( "huggingface_hub.list_repo_files", @@ -239,6 +241,214 @@ class TestGgufVariantFileResolution: assert downloaded == ["tinyllamas/stories260K.gguf"] assert out == "/fake/ggml-org/models/tinyllamas/stories260K.gguf" + def test_download_reuses_older_snapshot_when_current_ref_snapshot_is_partial( + self, monkeypatch, hf_cache + ): + # Cross-snapshot reuse is an offline-resilience path: online, hf_hub_download + # resumes the partial current-ref download and revalidates the revision instead + # of serving an older snapshot's same-name blob. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache( + hf_cache, + repo, + {"model-UD-Q4_K_XL.gguf": 4}, + snapshot_sha = "a" * 40, + ) + _build_cache( + hf_cache, + repo, + {"mtp-model.gguf": 1}, + snapshot_sha = "b" * 40, + ) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf", "mtp-model.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_reuses_cached_gguf_when_lowercase_partial_cache_shadows_it( + self, monkeypatch, hf_cache + ): + # Case-variant cross-dir reuse is offline-only; online the canonical repo id + # resolves up front and hf_hub_download fetches the current revision. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + gguf_file = "gemma-4-E2B-it-UD-Q4_K_XL.gguf" + snap = _build_cache( + hf_cache, + canonical_repo, + {gguf_file: 4}, + snapshot_sha = "a" * 40, + ) + lower_snap = _build_cache( + hf_cache, + requested_repo, + {"mtp-gemma-4-E2B-it.gguf": 1}, + snapshot_sha = "b" * 40, + ) + os.utime(lower_snap, (2000, 2000)) + os.utime(snap, (1000, 1000)) + seen_repos: list[str] = [] + + def fake_list_repo_files(repo_id, token = None): + seen_repos.append(repo_id) + return [gguf_file] + + def fake_get_paths_info( + repo_id, + paths, + token = None, + ): + seen_repos.append(repo_id) + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fake_cache(repo_id, filename, *args, **kwargs): + seen_repos.append(repo_id) + return str(snap / filename) if repo_id == canonical_repo else None + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", fake_cache), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = requested_repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(snap / gguf_file) + assert seen_repos + + def test_download_online_does_not_reuse_old_snapshot(self, monkeypatch, hf_cache): + # Online, an older same-name snapshot must not be served (it may be a stale + # revision); hf_hub_download is called so the current revision is fetched and + # its etag revalidated. + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + downloaded.append(filename) + return f"/fresh/{filename}" + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert downloaded == ["model-UD-Q4_K_XL.gguf"] + assert out == "/fresh/model-UD-Q4_K_XL.gguf" + + def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache): + # HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline + # cache reuse must trigger for those too, otherwise the earlier Hub calls run + # offline while this branch still attempts hf_hub_download and the cached GGUF + # cannot load. + monkeypatch.setenv("HF_HUB_OFFLINE", "true") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_companion_resolves_from_case_variant_snapshot_offline( + self, monkeypatch, hf_cache + ): + # Offline, resolve_cached_repo_id_case can keep a partial lower-case spelling, + # so the companion (mmproj) must resolve from whichever case-variant snapshot + # actually holds it rather than being dropped by an hf_hub_download on the + # wrong casing. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + snap = _build_cache(hf_cache, canonical_repo, {"mmproj-F16.gguf": 4}, snapshot_sha = "a" * 40) + # A partial lower-case dir exists so casing resolution keeps the requested spelling. + _build_cache(hf_cache, requested_repo, {"config.json": 1}, snapshot_sha = "b" * 40) + + _offline_exc = type("OfflineModeIsEnabled", (Exception,), {}) + + def fake_list_repo_files(repo_id, token = None): + raise _offline_exc("offline") + + def fail_download(*_args, **_kwargs): + raise AssertionError("should resolve the companion from cache, not download") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_mmproj(hf_repo = requested_repo) + + assert out == str(snap / "mmproj-F16.gguf") + def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path): backend = LlamaCppBackend() downloaded: list[str] = [] @@ -264,7 +474,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), @@ -279,6 +489,48 @@ class TestGgufVariantFileResolution: assert downloaded == files assert out == "/fake/org/repo/model-Q4_K_M-00001-of-00002.GGUF" + def test_download_refetches_split_gguf_when_shards_span_snapshots(self, monkeypatch, hf_cache): + # The cached main shard lives in an older snapshot; its sibling shard is only + # in a newer, separate snapshot. Reusing the main shard alone would leave + # llama.cpp unable to resolve the sibling, so the whole set must be re-fetched + # together (co-located) rather than served split across snapshot dirs. + backend = LlamaCppBackend() + repo = "org/split" + files = [ + "model-Q4_K_M-00001-of-00002.gguf", + "model-Q4_K_M-00002-of-00002.gguf", + ] + _build_cache(hf_cache, repo, {files[0]: 4}, snapshot_sha = "a" * 40) + _build_cache(hf_cache, repo, {files[1]: 4}, snapshot_sha = "b" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "Q4_K_M") + + assert downloaded == files + assert out == f"/fake/{repo}/{files[0]}" + def _siblings(items: dict[str, int]): """Mock ``hf_model_info(...).siblings`` payload.""" @@ -315,6 +567,21 @@ class TestIterHfCacheSnapshots: out = list(_iter_hf_cache_snapshots("unsloth/multi")) assert [p.name for p in out] == ["b" * 40, "a" * 40] + def test_skips_snapshot_when_mtime_is_unavailable(self, hf_cache, monkeypatch): + stale = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) + good = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) + original_stat = Path.stat + + def flaky_stat(self, *args, **kwargs): + if self == stale: + raise FileNotFoundError(str(self)) + return original_stat(self, *args, **kwargs) + + monkeypatch.setattr(Path, "stat", flaky_stat) + + out = list(_iter_hf_cache_snapshots("unsloth/multi")) + assert out == [good] + def test_repo_id_match_is_case_insensitive(self, hf_cache): _build_cache(hf_cache, "unsloth/Foo-GGUF", {"Foo-Q4_K_M.gguf": 1}) # Lookup with different org/name casing still resolves @@ -347,6 +614,87 @@ class TestListGgufVariantsFromCache: assert _list_gguf_variants_from_hf_cache("unsloth/absent") is None +class TestCachedColocatedSplitMain: + def test_prefers_older_complete_snapshot_over_newer_partial(self, hf_cache): + # Newer snapshot has only shard 1; older snapshot has the complete set. The + # complete older snapshot must win so the split GGUF can load co-located. + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + old = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard1: 100, shard2: 100}, snapshot_sha = "a" * 40 + ) + new = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "b" * 40) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + main = _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) + assert main is not None + assert main.startswith(str(old)) + + def test_returns_none_when_shards_span_snapshots(self, hf_cache): + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + a = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "a" * 40) + b = _build_cache(hf_cache, "unsloth/split-GGUF", {shard2: 100}, snapshot_sha = "b" * 40) + os.utime(a, (1000, 1000)) + os.utime(b, (2000, 2000)) + + assert _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) is None + + +class TestResolveRepoIdCasing: + def test_maps_to_canonical_casing(self, monkeypatch): + monkeypatch.setattr( + "utils.paths.resolve_cached_repo_id_case", + lambda repo: "unsloth/Gemma-4-GGUF" if repo.lower() == "unsloth/gemma-4-gguf" else repo, + ) + # A companion download passed the resolved id reads the same cache entry + # as the main GGUF instead of missing it under the requested casing. + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/Gemma-4-GGUF" + + def test_passthrough_on_resolver_error(self, monkeypatch): + def boom(_repo): + raise RuntimeError("resolver unavailable") + + monkeypatch.setattr("utils.paths.resolve_cached_repo_id_case", boom) + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/gemma-4-gguf" + + def test_companion_only_newer_snapshot_does_not_shadow_real_variants(self, hf_cache): + # A newer snapshot holds only a vision projector fetched on demand, + # while the quant files live in an older snapshot. The newer snapshot + # must not shadow the real variants; the vision flag carries over. + old = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"vision-Q4_K_M.gguf": 100}, + snapshot_sha = "a" * 40, + ) + new = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"mmproj-vision-F16.gguf": 10}, + snapshot_sha = "b" * 40, + ) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert [v.quant for v in variants] == ["Q4_K_M"] + assert has_vision is True + + def test_companion_only_cache_returns_empty_variants_with_vision(self, hf_cache): + # Only a vision projector is cached anywhere: report the vision flag + # with an empty variant list rather than None. + _build_cache(hf_cache, "unsloth/vision-GGUF", {"mmproj-vision-F16.gguf": 10}) + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert variants == [] + assert has_vision is True + + class TestListGgufVariantsOffline: def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1}) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5d8458e5f0..281ca24281 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1617,36 +1617,60 @@ def _iter_hf_cache_snapshots(repo_id: str): cache_dir = Path(hf_constants.HF_HUB_CACHE) target = f"models--{repo_id.replace('/', '--')}".lower() - repo_dir: Optional[Path] = None + repo_dirs: list[Path] = [] try: if not cache_dir.is_dir(): return for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: - repo_dir = entry - break + repo_dirs.append(entry) except OSError: return - if repo_dir is None: + if not repo_dirs: return - snapshots = repo_dir / "snapshots" - try: - if not snapshots.is_dir(): - return - snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()] - except OSError: + snap_dirs: list[Path] = [] + for repo_dir in repo_dirs: + snapshots = repo_dir / "snapshots" + try: + if snapshots.is_dir(): + for snap_dir in snapshots.iterdir(): + try: + if snap_dir.is_dir(): + snap_dirs.append(snap_dir) + except OSError: + continue + except OSError: + continue + if not snap_dirs: return - snap_dirs.sort(key = lambda s: s.stat().st_mtime, reverse = True) - yield from snap_dirs + snap_dirs_with_mtime = [] + for snap_dir in snap_dirs: + try: + snap_dirs_with_mtime.append((snap_dir.stat().st_mtime, snap_dir)) + except OSError: + continue + snap_dirs_with_mtime.sort(key = lambda item: item[0], reverse = True) + yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime) def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: - """Variants from the local HF cache snapshot, or None if not cached.""" + """Variants from the local HF cache snapshot, or None if not cached. + + A newer snapshot can hold only a companion file (for example a vision + projector fetched on demand) while the quant files live in an older + snapshot. Returning the first snapshot that merely reports a vision flag + would shadow those real variants, so keep scanning older snapshots for + actual variants and carry the vision flag across snapshots. + """ + any_vision = False for snap in _iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snap)) - if variants or has_vision: - return variants, has_vision + any_vision = any_vision or has_vision + if variants: + return variants, any_vision + if any_vision: + return [], True return None diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 01014de4d3..477a47cc3d 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -597,6 +597,59 @@ def _loaded_models(base: str, key: str) -> list: return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) +_HF_REPO_ID_SEGMENT_RE = re.compile(r"^[A-Za-z0-9._-]+$") + + +def _is_hub_model_id(value: object) -> bool: + if not isinstance(value, str): + return False + text = value.strip() + if "\\" in text: + return False + if text.startswith(("/", "./", "../", "~")): + return False + if len(text) >= 2 and text[1] == ":" and text[0].isalpha(): + return False + # A hub id is exactly "namespace/name" over a restricted charset. Anything with + # extra path segments (e.g. a server-side relative path such as + # models/Llama/Foo.gguf on a remote Studio) is not a hub id and must not be + # casefold-matched against a differently cased path on a case-sensitive + # filesystem. This is host independent, unlike the existence probe below which + # cannot see a path that only exists on the server. + parts = text.split("/") + if len(parts) != 2: + return False + if any(part in ("", ".", "..") or not _HF_REPO_ID_SEGMENT_RE.match(part) for part in parts): + return False + try: + if Path(os.path.expanduser(text)).exists(): + return False + except OSError: + return False + return True + + +def _model_id_matches( + actual: object, + requested: object, + *, + allow_casefold: bool = True, +) -> bool: + if actual == requested: + return True + # Case-insensitive matching is only safe when the local existence probe in + # _is_hub_model_id is authoritative, i.e. against a loopback Studio on this host. + # Against a remote Studio a two-segment string is indistinguishable from a + # server-side relative path (e.g. Models/Foo vs models/foo), so casefolding it + # could attach to the wrong model on a case-sensitive server; defer to an exact + # match there and let the load endpoint resolve the requested path. + if not allow_casefold: + return False + if not (_is_hub_model_id(actual) and _is_hub_model_id(requested)): + return False + return str(actual).casefold() == str(requested).casefold() + + def _resolve_model( base: str, key: str, @@ -604,6 +657,9 @@ def _resolve_model( load: LoadOptions = LoadOptions(), ) -> dict: models = _loaded_models(base, key) + # Only casefold-match ids against a loopback Studio, where _is_hub_model_id's + # local existence probe can actually reject a server-side path; see the note there. + allow_casefold = is_loopback_url(base) # /v1/models reports the model id but not the active GGUF variant or runtime load # settings, so an id match alone can hide the wrong quant (Q8_0 serving while the # user asked for UD-Q4_K_XL). When the user passed any explicit load knob, defer to @@ -613,10 +669,21 @@ def _resolve_model( load_has_overrides = bool( load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel ) + # /v1/models also lists cached-but-unloaded catalog entries (loaded == False); + # matching one would skip /api/inference/load and leave the agent pointed at a + # model that is not resident, so only attach to an entry that is actually loaded. match = ( None if requested and load_has_overrides - else next((m for m in models if m["id"] == requested), None) + else next( + ( + m + for m in models + if _model_id_matches(m.get("id"), requested, allow_casefold = allow_casefold) + and m.get("loaded") is not False + ), + None, + ) ) if requested and match is None: typer.echo( @@ -651,7 +718,16 @@ def _resolve_model( if isinstance(loaded, dict): wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} models = _loaded_models(base, key) - match = next((m for m in models if m["id"] in wanted), None) + match = next( + ( + m + for m in models + if any( + _model_id_matches(m.get("id"), w, allow_casefold = allow_casefold) for w in wanted + ) + ), + None, + ) if match is not None: return match if requested: diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index bd964d5e54..ee5b442c27 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -457,6 +457,189 @@ def test_connect_codex_no_launch(fake_studio, tmp_path): assert (home / "unsloth_api.config.toml").exists() +def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, tmp_path): + result = CliRunner().invoke( + start.start_app, + [ + "codex", + "--no-launch", + "--model", + "unsloth/gemma-4-26b-a4b-it-gguf", + ], + ) + assert result.exit_code == 0, result.output + home = tmp_path / "agents" / "codex" + profile = _parse_toml((home / "unsloth_api.config.toml").read_text()) + assert profile["model"] == MODEL["id"] + + +def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch): + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/v1/models"): + return { + "data": [ + { + "id": "unsloth/gemma-4-E2B-it-GGUF" if state["loaded"] else "other/model", + "context_length": 131072, + } + ] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/gemma-4-E2B-it-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model( + BASE, + "sk-test", + "unsloth/gemma-4-e2b-it-gguf", + start.LoadOptions(gguf_variant = "UD-Q4_K_XL"), + ) + + assert entry["id"] == "unsloth/gemma-4-E2B-it-GGUF" + assert any(c[1].endswith("/api/inference/load") for c in calls) + + +def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch): + # A cached-but-unloaded catalog entry (loaded == False) that only case-differs must + # not be treated as ready; the load endpoint must still be called so the requested + # model becomes resident instead of the agent preflighting a different backend. + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [ + { + "id": "unsloth/Gemma-4-GGUF", + "loaded": state["loaded"], + "context_length": 131072, + } + ] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/Gemma-4-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf") + + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch): + # The mirror case: a loaded entry (loaded == True) that case-matches attaches with + # no /api/inference/load call. + calls = [] + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [{"id": "unsloth/Gemma-4-GGUF", "loaded": True, "context_length": 131072}] + } + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf") + + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert not any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_resolve_model_remote_studio_does_not_casefold_attach(monkeypatch): + # Against a remote Studio the local existence probe cannot see server-side paths, + # so a case-variant loaded id must NOT attach without a load: it could be a distinct + # server-side path on a case-sensitive host. The load endpoint resolves the request. + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [{"id": "unsloth/Gemma-4-GGUF", "loaded": True, "context_length": 131072}] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/Gemma-4-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model("http://10.0.0.5:8888", "sk-test", "unsloth/gemma-4-gguf") + + # The load endpoint was consulted (no casefold shortcut), and we still attach to the + # server's canonical id it reports back. + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_model_id_matching_does_not_casefold_local_paths(tmp_path): + existing_local = tmp_path / "Org" / "Foo" + existing_local.mkdir(parents = True) + + assert start._model_id_matches("Org/Foo", "org/foo") + assert not start._model_id_matches(str(existing_local), str(existing_local).lower()) + assert not start._model_id_matches("./Models/Foo", "./models/foo") + assert not start._model_id_matches(r".\Models\Foo", r".\models\foo") + # A server-side relative path (extra path segments) is not a hub id even when it + # does not exist on the CLI host, so it must not casefold-match a differently + # cased path on a case-sensitive server filesystem. + assert not start._is_hub_model_id("models/Llama/Foo.gguf") + assert not start._model_id_matches("models/Llama/Foo.gguf", "models/llama/foo.gguf") + # A genuine two-segment hub id still matches case-insensitively. + assert start._is_hub_model_id("unsloth/Gemma-3-4b-it-GGUF") + assert start._model_id_matches("unsloth/Gemma-3-4b-it-GGUF", "unsloth/gemma-3-4b-it-gguf") + # Casefolding is gated to loopback studios (allow_casefold). With it disabled (a + # remote studio, where a two-segment string could be a server-side path), even a + # genuine hub-id case variant must not match, so the load endpoint resolves it. + assert not start._model_id_matches( + "unsloth/Gemma-3-4b-it-GGUF", "unsloth/gemma-3-4b-it-gguf", allow_casefold = False + ) + assert start._model_id_matches("unsloth/Foo", "unsloth/Foo", allow_casefold = False) + + def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch): # Launch mode writes config to a throwaway temp CODEX_HOME and removes it after # the agent exits; the user's real ~/.codex is never the target. From f1a2621631d3adfd72a65dc9710290abb81c54db Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 03:09:52 -0700 Subject: [PATCH 059/113] Studio: show Hugging Face address on hover for Hub and online model rows (#6382) (#6928) * Studio: show Hugging Face address on hover for Hub and online model rows The model selector already shows an on-disk path tooltip on local rows, but Hub and online rows showed only the bare repo id, and nothing at all when there was no VRAM estimate. Add an optional hubUrl prop and a hubRepoUrl helper that mirrors localPathTooltip, and surface huggingface.co/ on hover for the Discover, search, and downloaded Hub rows. Local and VRAM tooltips are unchanged; the VRAM tooltip now also appends the address line. Closes #6382 * Studio: use a 700ms hover delay before the model-row tooltip Give the model-row hover tooltip (the Hugging Face address, plus the VRAM and local-path lines it shares) a 700ms open delay instead of showing it instantly, so it does not flash while sweeping the mouse down the list. * Fix/adjust GGUF tooltips for PR #6928 --------- Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- .../assistant-ui/model-selector/pickers.tsx | 62 ++++++++++++++----- 1 file changed, 45 insertions(+), 17 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 9890c5f574..ff0351883d 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -393,6 +393,7 @@ function ModelRow({ vramEst, gpuGb, tooltipText, + hubUrl, optionProps, onArrowDownIntoChildren, capabilities, @@ -409,6 +410,10 @@ function ModelRow({ vramEst?: number; gpuGb?: number; tooltipText?: ReactNode; + /** Hugging Face address (e.g. "huggingface.co/owner/name") for online/Hub + * rows; surfaced on hover so their repo id / URL is discoverable the same + * way local rows show an on-disk path. Omit to show no address line. */ + hubUrl?: string; optionProps?: ModelRowOptionProps; onArrowDownIntoChildren?: () => boolean; /** Capability override (HF rows have tags); falls back to name detection. */ @@ -546,30 +551,41 @@ function ModelRow({ ); - if (vramTooltipText) { - return ( - - {content} - - {label} - {vramTooltipText} - - - ); - } + // Optional Hugging Face address line for online/Hub rows, rendered under + // whichever tooltip shows so the repo id / URL is always visible on hover. + const hubUrlLine = hubUrl ? ( + + {hubUrl} + + ) : null; - if (tooltipText) { + const tooltipBody = vramTooltipText ? ( + <> + {label} + {vramTooltipText} + {hubUrlLine} + + ) : tooltipText ? ( + <> + {tooltipText} + {hubUrlLine} + + ) : hubUrl ? ( + <> + {label} + {hubUrlLine} + + ) : null; + + if (tooltipBody) { return ( - + {content} - {tooltipText} + {tooltipBody} ); @@ -1193,6 +1209,13 @@ function localPathTooltip(name: string, path: string): ReactNode { ); } +/** Hugging Face address for an online/Hub row, or undefined when the repo id is + * missing so the row shows no (empty) address line on hover. */ +function hubRepoUrl(id: string | null | undefined): string | undefined { + const trimmed = id?.trim(); + return trimmed ? `huggingface.co/${trimmed}` : undefined; +} + /** Whether a local model is an MLX build (name hint). MLX runs on Mac only, so * callers gate visibility on the host being a Mac. */ function localModelIsMlx(m: LocalModelInfo): boolean { @@ -2462,6 +2485,7 @@ export function HubModelPicker({
Date: Wed, 8 Jul 2026 03:13:32 -0700 Subject: [PATCH 060/113] Studio: fix currency and indentation edge cases in LaTeX rendering (#6957) * Studio: fix link, currency and indentation edge cases in LaTeX rendering Follow-up to #6914. Three fixes to studio/frontend/src/lib/latex.ts: - Skip reference-link definition URLs ([id]: url) during delimiter conversion, so escaped parens in such URLs are not rewritten as math. - Preserve the opener line's indentation when emitting a display $$ block, so a \[...\] inside a list item stays part of the list. - Stop a currency amount from pairing with a converted span's opening $, which swallowed the price into math (for example $5 + x \(y\)). * Exclude GFM footnote definitions from the reference-URL skip A footnote definition like [^1]: \(x\) had its body treated as a link destination, so leading math was left literal. Skip [^...] labels. * Merge overlapping link destination regions A reference-def token can nest inline-link spans (for example [1]: http://h/[a](b)/foo\(x\)), so the combined spans could overlap and isInRegion's binary search missed the outer one, rewriting the URL. Merge overlapping spans before the search. * Guard lineStart when the display opener is at index 0 Behavior is unchanged (lastIndexOf clamps a negative fromIndex to 0), but the explicit guard avoids relying on that implicit clamp. * Scope to indentation and currency fixes Drop the reference-link URL protection added earlier. It guards a case models effectively never emit (escaped parens in a reference-style URL), and approximating CommonMark reference definitions with a regex needs open-ended special-casing. Keep the two high-value fixes: preserve display math indentation (including multi-line bodies) inside a list item, and stop a currency amount from pairing with a converted span's opening dollar sign. --- studio/frontend/src/lib/latex.ts | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 86a9634048..edf9875602 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -173,7 +173,11 @@ function looksLikeMathBody(body: string): boolean { * (`**$X$**`, `__$X$__`) are always math: LLMs use that for "bold math" * and the heuristic would otherwise reject prose-shaped bodies like "90 - x". */ -function hasInlineMathCloser(content: string, offset: number): boolean { +function hasInlineMathCloser( + content: string, + offset: number, + mathRegions: Array<[number, number]>, +): boolean { const MAX_SPAN = 200; const limit = Math.min(content.length, offset + 1 + MAX_SPAN); for (let i = offset + 1; i < limit; i++) { @@ -181,6 +185,9 @@ function hasInlineMathCloser(content: string, offset: number): boolean { if (c === "\n") return false; if (c !== "$") continue; if (content[i - 1] === "\\") continue; + // A `$` opening a generated span (from `\(...\)`) is not a currency closer; + // pairing with it would swallow the price into math (`$5 + x \(y\)`). + if (isInRegion(i, mathRegions)) return false; if (content[i + 1] === "$") { i++; continue; @@ -294,7 +301,22 @@ function convertLatexDelimiters(content: string): { continue; } append(content.slice(last, match.index)); - const wrapped = isDisplay ? `\n$$\n${body}\n$$\n` : `$${body}$`; + let wrapped: string; + if (isDisplay) { + // Keep the opener's leading indentation so a `$$` block inside a list item + // stays in the container instead of breaking out at column 0. Only when the + // opener is whitespace-prefixed, so inline `text \[x\]` keeps column 0. + const lineStart = + match.index > 0 ? content.lastIndexOf("\n", match.index - 1) + 1 : 0; + const prefix = content.slice(lineStart, match.index); + const indent = /^\s*$/.test(prefix) ? prefix : ""; + // Indent every body line, not just the first, so multi-line display math + // (`\[a\nb\]`) stays wholly inside the container. + const inner = indent ? body.replace(/\n/g, `\n${indent}`) : body; + wrapped = `\n${indent}$$\n${indent}${inner}\n${indent}$$\n`; + } else { + wrapped = `$${body}$`; + } const start = append(wrapped); mathRegions.push([start, offset]); last = matchEnd; @@ -334,7 +356,7 @@ export function preprocessLaTeX(content: string): string { if (isInRegion(offset, mathRegions)) { return match; } - if (hasInlineMathCloser(text, offset)) { + if (hasInlineMathCloser(text, offset, mathRegions)) { return match; } return "\\" + match; From 38dacb8a1f49906d2fa23e4dace198a9c56b5220 Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 8 Jul 2026 18:25:26 +0800 Subject: [PATCH 061/113] Add MLX backend support for CLI unsloth train (#6709) * feat(studio): route CLI trainer to MLX backend * fix(studio): harden MLX trainer routing * fix(studio): harden MLX trainer adapter routing * test(studio): assert MLX CLI activation order * fix(studio): address MLX CLI review feedback * feat(cli): support MLX in legacy script * fix(cli): adapt MLX tokenizer for raw text * fix(cli): omit unsupported MLX eval batch arg * fix(cli): feed raw text to MLX trainer * Fix CLI MLX routing and Python 3.9 annotations Route the MLX backend through create_mlx_trainer_adapter so the torch-free Apple Silicon path never imports trainer.py (torch/unsloth/trl). Replace from __future__ import annotations with typing.Optional/Union so the CLI annotations stay Python 3.9 compatible without the unused-import lint hit. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Strip return_tensors from MLX raw-text tokenizer proxy On a torch-free MLX install, RawTextDataLoader calls the tokenizer with return_tensors='pt'; the callable proxy forwarded that to the HF tokenizer, which tried to build torch tensors and failed before training. Drop return_tensors so the MLX path returns plain token ids. * Tighten CLI MLX-backend comments --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/trainer.py | 31 +- studio/backend/core/training/training.py | 624 +++++++++++++++--- studio/backend/core/training/worker.py | 248 +++++-- .../backend/tests/test_training_preflight.py | 232 +++++++ unsloth-cli.py | 243 ++++--- unsloth_cli/commands/train.py | 46 +- 6 files changed, 1162 insertions(+), 262 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 20b2305a5a..958f8f4197 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -62,7 +62,6 @@ from loggers import get_logger import time from pathlib import Path from typing import Any, Dict, List, Optional, Callable -from dataclasses import dataclass import pandas as pd from datasets import Dataset from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset @@ -86,6 +85,11 @@ from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from .training import ( + TrainingProgress, + create_mlx_trainer_adapter, + should_use_mlx_training_backend, +) logger = get_logger(__name__) @@ -104,31 +108,16 @@ def _build_report_targets(training_args) -> list[str] | str: return report_to or "none" -@dataclass -class TrainingProgress: - """Training progress tracking""" - - epoch: float = 0 - step: int = 0 - total_steps: int = 0 - loss: Optional[float] = None - learning_rate: Optional[float] = None - is_training: bool = False - is_completed: bool = False - error: Optional[str] = None - status_message: str = "Ready to train" # Current stage - elapsed_seconds: Optional[float] = None - eta_seconds: Optional[float] = None - grad_norm: Optional[float] = None - num_tokens: Optional[int] = None - eval_loss: Optional[float] = None - - class UnslothTrainer: """ Unsloth Training Backend """ + def __new__(cls, *args, **kwargs): + if cls is UnslothTrainer and should_use_mlx_training_backend(): + return create_mlx_trainer_adapter(*args, **kwargs) + return super().__new__(cls) + def __init__(self): self.model = None self.tokenizer = None diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index f4233fcf04..2ddda19951 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -14,17 +14,19 @@ import json as _json import math import multiprocessing as mp import os +import platform import queue import re import shutil import threading import time +import traceback import structlog from datetime import datetime, timezone from loggers import get_logger -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Optional, Tuple, Any, TYPE_CHECKING +from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING if TYPE_CHECKING: import matplotlib.pyplot as plt @@ -98,6 +100,107 @@ def _coerce_optional_nonneg_float(name: str, value): return coerced +def is_apple_silicon_training_platform() -> bool: + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def is_mlx_training_device(device: Any) -> bool: + return ( + str(device).lower() == "mlx" + or str(device).lower().endswith(".mlx") + or getattr(device, "name", "").lower() == "mlx" + ) + + +def should_use_mlx_training_backend(*, device: Optional[Any] = None) -> bool: + if device is not None: + return is_mlx_training_device(device) + return is_apple_silicon_training_platform() + + +def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: + """Build the normalized worker config shared by Studio and the CLI adapter.""" + config = { + "model_name": values["model_name"], + "project_name": values.get("project_name"), + "training_type": values.get("training_type", "LoRA/QLoRA"), + "hf_token": values.get("hf_token", ""), + "load_in_4bit": values.get("load_in_4bit", True), + "max_seq_length": values.get("max_seq_length", 2048), + "vision_image_size": values.get("vision_image_size"), + "hf_dataset": values.get("hf_dataset", ""), + "local_datasets": values.get("local_datasets"), + "local_eval_datasets": values.get("local_eval_datasets"), + "format_type": values.get("format_type", ""), + "subset": values.get("subset"), + "train_split": values.get("train_split", "train"), + "eval_split": values.get("eval_split"), + "eval_steps": values.get("eval_steps", 0.00), + "dataset_streaming": values.get("dataset_streaming", False), + "dataset_slice_start": values.get("dataset_slice_start"), + "dataset_slice_end": values.get("dataset_slice_end"), + "custom_format_mapping": values.get("custom_format_mapping"), + "is_dataset_image": values.get("is_dataset_image", False), + "is_dataset_audio": values.get("is_dataset_audio", False), + "is_embedding": values.get("is_embedding", False), + "num_epochs": values.get("num_epochs", 3), + "learning_rate": values.get("learning_rate", "2e-4"), + "embedding_learning_rate": values.get("embedding_learning_rate"), + "batch_size": values.get("batch_size", 2), + "gradient_accumulation_steps": values.get("gradient_accumulation_steps", 4), + "warmup_steps": values.get("warmup_steps"), + "warmup_ratio": values.get("warmup_ratio"), + "max_steps": values.get("max_steps", 0), + "save_steps": values.get("save_steps", 0), + "weight_decay": values.get("weight_decay", 0.001), + "max_grad_norm": values.get("max_grad_norm", 0.0), + "max_grad_value": _coerce_optional_nonneg_float( + "max_grad_value", values.get("max_grad_value") + ), + "max_grad_leaf_norm": _coerce_optional_nonneg_float( + "max_grad_leaf_norm", values.get("max_grad_leaf_norm") + ), + "cast_norm_output_to_input_dtype": _coerce_optional_bool( + values.get("cast_norm_output_to_input_dtype"), True + ), + "random_seed": _coerce_seed(values.get("random_seed")), + "packing": values.get("packing", False), + "optim": values.get("optim", "adamw_8bit"), + "lr_scheduler_type": values.get("lr_scheduler_type", "linear"), + "use_lora": values.get("use_lora", True), + "lora_r": values.get("lora_r", 16), + "lora_alpha": values.get("lora_alpha", 16), + "lora_dropout": values.get("lora_dropout", 0.0), + "target_modules": values.get("target_modules"), + "gradient_checkpointing": values.get("gradient_checkpointing", "unsloth"), + "use_rslora": values.get("use_rslora", False), + "use_loftq": values.get("use_loftq", False), + "train_on_completions": values.get("train_on_completions", False), + "finetune_vision_layers": values.get("finetune_vision_layers", True), + "finetune_language_layers": values.get("finetune_language_layers", True), + "finetune_attention_modules": values.get("finetune_attention_modules", True), + "finetune_mlp_modules": values.get("finetune_mlp_modules", True), + "enable_wandb": values.get("enable_wandb", False), + "wandb_token": values.get("wandb_token"), + "wandb_project": values.get("wandb_project", "unsloth-training"), + "enable_tensorboard": values.get("enable_tensorboard", False), + "tensorboard_dir": values.get("tensorboard_dir", "runs"), + "resume_from_checkpoint": values.get("resume_from_checkpoint"), + "trust_remote_code": values.get("trust_remote_code", False), + "approved_remote_code_fingerprint": values.get("approved_remote_code_fingerprint"), + "subject": values.get("subject"), + "gpu_ids": values.get("gpu_ids"), + "s3_config": values.get("s3_config"), + "disable_xet": values.get("disable_xet", False), + } + for key in ("output_dir", "allow_external_output_dir"): + if key in values: + config[key] = values.get(key) + if config["training_type"] == "Full Finetuning": + config["load_in_4bit"] = False + return config + + _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") @@ -133,7 +236,7 @@ def _s3_dataset_name(s3_dataset: Any) -> Optional[str]: return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" -def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: +def _cleanup_cancelled_checkpoints(output_dir: Union[str, os.PathLike]) -> None: """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. Completed ``checkpoint-/`` dirs survive. Symlinked output_dir / children @@ -183,7 +286,7 @@ PLOT_HEIGHT = 3.5 @dataclass class TrainingProgress: - """Mirror of trainer.TrainingProgress so the parent never imports heavy ML modules.""" + """Shared training progress payload for Studio and backend-aware trainers.""" epoch: float = 0 step: int = 0 @@ -200,6 +303,423 @@ class TrainingProgress: num_tokens: Optional[int] = None eval_loss: Optional[float] = None peak_memory_gb: Optional[float] = None + output_dir: Optional[str] = None + + +class _MLXTrainerAdapter: + """Adapts the legacy UnslothTrainer API to the shared Studio MLX worker path.""" + + def __init__(self): + self.model = None + self.tokenizer = None + self.trainer = None + self.training_thread = None + self.training_progress = TrainingProgress() + self.progress_callbacks: list[Callable[[TrainingProgress], None]] = [] + self.is_training = False + self.should_stop = False + self.save_on_stop = True + self.load_in_4bit = True + self.output_dir = None + + self.is_cpt = False + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model_name = None + self.max_seq_length = None + + self._model_config: dict[str, Any] = {} + self._peft_config: dict[str, Any] = {} + self._dataset_config: dict[str, Any] = {} + self._event_queue: Optional[queue.Queue] = None + self._stop_queue: Optional[queue.Queue] = None + self._pump_thread: Optional[threading.Thread] = None + self._lock = threading.Lock() + + def _activate_transformers_for_model(self, model_name: str, hf_token: Optional[str]) -> None: + try: + from utils.transformers_version import activate_transformers_for_subprocess + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + logger.warning("MLX trainer adapter Transformers activation failed", error = str(exc)) + + def add_progress_callback(self, callback: Callable[[TrainingProgress], None]): + self.progress_callbacks.append(callback) + + def _update_progress(self, **kwargs): + with self._lock: + for key, value in kwargs.items(): + if hasattr(self.training_progress, key): + setattr(self.training_progress, key, value) + progress = self.training_progress + for callback in self.progress_callbacks: + try: + callback(progress) + except Exception: + pass + + def load_model( + self, + model_name: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + is_dataset_image: bool = False, + is_dataset_audio: bool = False, + trust_remote_code: bool = False, + full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, + ) -> bool: + self.model_name = model_name + self.max_seq_length = max_seq_length + self.load_in_4bit = load_in_4bit + self._audio_type = None + self._activate_transformers_for_model(model_name, hf_token) + try: + from utils.models import detect_audio_type, is_vision_model + + self._audio_type = detect_audio_type(model_name, hf_token) + if self._audio_type == "audio_vlm": + self.is_audio = False + self.is_audio_vlm = bool(is_dataset_audio) + self._audio_type = None + else: + self.is_audio = self._audio_type is not None + self.is_audio_vlm = False + vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False + self.is_vlm = not self.is_audio_vlm and vision and bool(is_dataset_image) + except Exception as exc: + logger.warning("MLX trainer adapter model type detection failed", error = str(exc)) + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model = object() + self.tokenizer = object() + self._model_config = { + "model_name": model_name, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "hf_token": hf_token or "", + "is_dataset_image": bool(is_dataset_image), + "is_dataset_audio": bool(is_dataset_audio), + "trust_remote_code": bool(trust_remote_code), + "gpu_ids": gpu_ids, + } + self._update_progress( + is_training = False, + is_completed = False, + error = None, + step = 0, + loss = 0.0, + epoch = 0, + status_message = f"Queued MLX model load: {model_name}", + ) + return True + + def prepare_model_for_training( + self, + use_lora: bool = True, + finetune_vision_layers: bool = True, + finetune_language_layers: bool = True, + finetune_attention_modules: bool = True, + finetune_mlp_modules: bool = True, + target_modules: Optional[Union[list, str]] = None, + lora_r: int = 16, + lora_alpha: int = 16, + lora_dropout: float = 0.0, + use_gradient_checkpointing: Union[str, bool] = "unsloth", + use_rslora: bool = False, + use_loftq: bool = False, + ) -> bool: + self._peft_config = { + "use_lora": bool(use_lora), + "lora_r": lora_r, + "lora_alpha": lora_alpha, + "lora_dropout": lora_dropout, + "target_modules": target_modules, + "gradient_checkpointing": use_gradient_checkpointing, + "use_rslora": bool(use_rslora), + "use_loftq": bool(use_loftq), + "finetune_vision_layers": bool(finetune_vision_layers), + "finetune_language_layers": bool(finetune_language_layers), + "finetune_attention_modules": bool(finetune_attention_modules), + "finetune_mlp_modules": bool(finetune_mlp_modules), + } + self._update_progress(status_message = "Queued MLX training setup") + return True + + def load_and_format_dataset( + self, + dataset_source: Optional[str], + format_type: str = "auto", + local_datasets: Optional[list[str]] = None, + local_eval_datasets: Optional[list[str]] = None, + custom_format_mapping: Optional[dict[str, Any]] = None, + subset: Optional[str] = None, + train_split: str = "train", + eval_split: Optional[str] = None, + dataset_streaming: bool = False, + eval_steps: float = 0.00, + dataset_slice_start: Optional[int] = None, + dataset_slice_end: Optional[int] = None, + is_cpt: bool = False, + s3_config: dict = None, + ) -> Optional[tuple]: + self._dataset_config = { + "hf_dataset": dataset_source or "", + "local_datasets": local_datasets, + "local_eval_datasets": local_eval_datasets, + "format_type": format_type or "", + "custom_format_mapping": custom_format_mapping, + "subset": subset, + "train_split": train_split or "train", + "eval_split": eval_split, + "dataset_streaming": bool(dataset_streaming), + "eval_steps": eval_steps or 0.0, + "dataset_slice_start": dataset_slice_start, + "dataset_slice_end": dataset_slice_end, + "s3_config": s3_config, + } + self.is_cpt = bool(is_cpt) + self._update_progress(status_message = "Queued MLX dataset load") + return ({"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, None) + + def start_training( + self, + dataset = None, + eval_dataset = None, + **training_args, + ) -> bool: + if self.is_training and self.training_thread and self.training_thread.is_alive(): + return False + if self._pump_thread and self._pump_thread.is_alive(): + self._pump_thread.join(timeout = 2.0) + if self._pump_thread.is_alive(): + self._update_progress(error = "Previous training event pump is still finalizing") + return False + if not self._model_config: + self._update_progress(error = "Model not loaded") + return False + if not self._dataset_config: + self._update_progress(error = "Dataset not loaded") + return False + if self.is_cpt: + self._update_progress( + error = "Continued Pretraining is not supported for MLX training yet.", + is_training = False, + is_completed = False, + ) + return False + + config = self._build_worker_config(training_args) + event_queue = queue.Queue() + stop_queue = queue.Queue() + self._event_queue = event_queue + self._stop_queue = stop_queue + self.should_stop = False + self.is_training = True + self.training_progress = TrainingProgress( + is_training = True, + status_message = "Initializing MLX training...", + ) + + self.training_thread = threading.Thread( + target = self._run_training_thread, + args = (config, event_queue, stop_queue), + daemon = True, + ) + self._pump_thread = threading.Thread( + target = self._pump_events, + args = (event_queue, self.training_thread), + daemon = True, + ) + self.training_thread.start() + self._pump_thread.start() + return True + + def _build_worker_config(self, training_args: dict[str, Any]) -> dict[str, Any]: + peft = { + "use_lora": True, + "lora_r": 16, + "lora_alpha": 16, + "lora_dropout": 0.0, + "target_modules": None, + "gradient_checkpointing": "unsloth", + "use_rslora": False, + "use_loftq": False, + "finetune_vision_layers": True, + "finetune_language_layers": True, + "finetune_attention_modules": True, + "finetune_mlp_modules": True, + **self._peft_config, + } + output_dir = training_args.get("output_dir") + if output_dir: + output_dir = os.path.abspath(os.path.expanduser(str(output_dir))) + values = { + **self._model_config, + **self._dataset_config, + **training_args, + "training_type": ( + "Continued Pretraining" + if self.is_cpt + else "LoRA/QLoRA" + if peft["use_lora"] + else "Full Finetuning" + ), + **peft, + "output_dir": output_dir, + "allow_external_output_dir": bool(output_dir), + } + config = _build_training_worker_config(values) + config["resolved_gpu_ids"] = None + config["gpu_selection"] = None + return config + + def _run_training_thread( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + try: + self._run_mlx_worker(config, event_queue, stop_queue) + except Exception as exc: + if event_queue is not None: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def _run_mlx_worker( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + from .worker import run_mlx_training_process + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + ) + + def _pump_events(self, event_queue: queue.Queue, training_thread: threading.Thread): + while True: + event = None + try: + event = event_queue.get(timeout = 0.25) + except queue.Empty: + pass + if event is not None: + self._handle_event(event) + continue + if not training_thread.is_alive(): + self._drain_events(event_queue) + with self._lock: + if self.training_progress.is_training: + self.training_progress.is_training = False + if self.should_stop: + self.training_progress.status_message = "Training stopped." + elif ( + not self.training_progress.error + and not self.training_progress.is_completed + ): + self.training_progress.error = "Training process exited unexpectedly" + self.is_training = False + self._event_queue = None + self._stop_queue = None + return + + def _drain_events(self, event_queue: Optional[queue.Queue] = None): + event_queue = event_queue or self._event_queue + if event_queue is None: + return + while True: + try: + self._handle_event(event_queue.get_nowait()) + except queue.Empty: + return + + def _handle_event(self, event: dict[str, Any]): + etype = event.get("type") + if etype == "status": + self._update_progress( + status_message = event.get("status_message") or event.get("message") or "" + ) + return + if etype == "progress": + self._update_progress( + step = event.get("step", self.training_progress.step), + epoch = event.get("epoch", self.training_progress.epoch), + loss = event.get("loss", self.training_progress.loss), + learning_rate = event.get("learning_rate", self.training_progress.learning_rate), + total_steps = event.get("total_steps", self.training_progress.total_steps), + elapsed_seconds = event.get( + "elapsed_seconds", + self.training_progress.elapsed_seconds, + ), + eta_seconds = event.get("eta_seconds", self.training_progress.eta_seconds), + grad_norm = event.get("grad_norm", self.training_progress.grad_norm), + num_tokens = event.get("num_tokens", self.training_progress.num_tokens), + eval_loss = event.get("eval_loss", self.training_progress.eval_loss), + peak_memory_gb = event.get("peak_memory_gb", self.training_progress.peak_memory_gb), + ) + return + if etype == "complete": + status_message = event.get("status_message") or "Training completed" + output_dir = event.get("output_dir") + was_cancelled = self.should_stop or status_message.strip().lower() in { + "training cancelled", + "training stopped", + } + self.output_dir = output_dir + self._update_progress( + is_training = False, + is_completed = not was_cancelled, + error = None, + status_message = status_message, + output_dir = output_dir, + ) + self.is_training = False + return + if etype == "error": + self._update_progress( + is_training = False, + is_completed = False, + error = event.get("error") or event.get("message") or "Training failed", + ) + self.is_training = False + return + + def stop_training(self, save: bool = True): + self.should_stop = True + self.save_on_stop = bool(save) + if self._stop_queue is not None: + self._stop_queue.put({"type": "stop", "save": save}) + status_message = ( + "Stopping training and saving checkpoint..." if save else "Cancelling training..." + ) + self._update_progress(status_message = status_message) + return True + + def get_training_progress(self) -> TrainingProgress: + pump_thread = self._pump_thread + training_thread = self.training_thread + if ( + pump_thread is not None + and pump_thread.is_alive() + and (training_thread is None or not training_thread.is_alive()) + and threading.current_thread() is not pump_thread + ): + pump_thread.join(timeout = 5.0) + if pump_thread is None or not pump_thread.is_alive(): + self._drain_events() + with self._lock: + return replace(self.training_progress) + + +def create_mlx_trainer_adapter(*args, **kwargs): + return _MLXTrainerAdapter(*args, **kwargs) class TrainingBackend: @@ -296,86 +816,7 @@ class TrainingBackend: # treat this fresh setup as a recoverable death. self._pump_running = False - # Build config dict for the subprocess - config = { - "model_name": kwargs["model_name"], - "project_name": kwargs.get("project_name"), - "training_type": kwargs.get("training_type", "LoRA/QLoRA"), - "hf_token": kwargs.get("hf_token", ""), - "load_in_4bit": kwargs.get("load_in_4bit", True), - "max_seq_length": kwargs.get("max_seq_length", 2048), - "vision_image_size": kwargs.get("vision_image_size"), - "hf_dataset": kwargs.get("hf_dataset", ""), - "local_datasets": kwargs.get("local_datasets"), - "local_eval_datasets": kwargs.get("local_eval_datasets"), - "format_type": kwargs.get("format_type", ""), - "subset": kwargs.get("subset"), - "train_split": kwargs.get("train_split", "train"), - "eval_split": kwargs.get("eval_split"), - "eval_steps": kwargs.get("eval_steps", 0.00), - "dataset_streaming": kwargs.get("dataset_streaming", False), - "dataset_slice_start": kwargs.get("dataset_slice_start"), - "dataset_slice_end": kwargs.get("dataset_slice_end"), - "custom_format_mapping": kwargs.get("custom_format_mapping"), - "is_dataset_image": kwargs.get("is_dataset_image", False), - "is_dataset_audio": kwargs.get("is_dataset_audio", False), - "is_embedding": kwargs.get("is_embedding", False), - "num_epochs": kwargs.get("num_epochs", 3), - "learning_rate": kwargs.get("learning_rate", "2e-4"), - "embedding_learning_rate": kwargs.get("embedding_learning_rate"), - "batch_size": kwargs.get("batch_size", 2), - "gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4), - "warmup_steps": kwargs.get("warmup_steps"), - "warmup_ratio": kwargs.get("warmup_ratio"), - "max_steps": kwargs.get("max_steps", 0), - "save_steps": kwargs.get("save_steps", 0), - "weight_decay": kwargs.get("weight_decay", 0.001), - "max_grad_norm": kwargs.get("max_grad_norm", 0.0), - "max_grad_value": _coerce_optional_nonneg_float( - "max_grad_value", kwargs.get("max_grad_value") - ), - "max_grad_leaf_norm": _coerce_optional_nonneg_float( - "max_grad_leaf_norm", kwargs.get("max_grad_leaf_norm") - ), - "cast_norm_output_to_input_dtype": _coerce_optional_bool( - kwargs.get("cast_norm_output_to_input_dtype"), True - ), - # MLX/CUDA/embedding workers need an int (transformers.set_seed(None) raises). - "random_seed": _coerce_seed(kwargs.get("random_seed")), - "packing": kwargs.get("packing", False), - "optim": kwargs.get("optim", "adamw_8bit"), - "lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"), - "use_lora": kwargs.get("use_lora", True), - "lora_r": kwargs.get("lora_r", 16), - "lora_alpha": kwargs.get("lora_alpha", 16), - "lora_dropout": kwargs.get("lora_dropout", 0.0), - "target_modules": kwargs.get("target_modules"), - "gradient_checkpointing": kwargs.get("gradient_checkpointing", "unsloth"), - "use_rslora": kwargs.get("use_rslora", False), - "use_loftq": kwargs.get("use_loftq", False), - "train_on_completions": kwargs.get("train_on_completions", False), - "finetune_vision_layers": kwargs.get("finetune_vision_layers", True), - "finetune_language_layers": kwargs.get("finetune_language_layers", True), - "finetune_attention_modules": kwargs.get("finetune_attention_modules", True), - "finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True), - "enable_wandb": kwargs.get("enable_wandb", False), - "wandb_token": kwargs.get("wandb_token"), - "wandb_project": kwargs.get("wandb_project", "unsloth-training"), - "enable_tensorboard": kwargs.get("enable_tensorboard", False), - "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), - "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), - "trust_remote_code": kwargs.get("trust_remote_code", False), - "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"), - "subject": kwargs.get("subject"), - "gpu_ids": kwargs.get("gpu_ids"), - "s3_config": kwargs.get("s3_config"), - # Flipped to True only by the HTTP-fallback respawn after a stall. - "disable_xet": kwargs.get("disable_xet", False), - } - - # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. - if config["training_type"] == "Full Finetuning": - config["load_in_4bit"] = False + config = _build_training_worker_config(kwargs) # Split GPU validation from placement around the VRAM hook: # * Explicit gpu_ids are validated here (raises -> the route returns 400 @@ -401,7 +842,7 @@ class TrainingBackend: ) defer_auto_selection = False - if _hw.DEVICE == _hw.DeviceType.MLX: + if should_use_mlx_training_backend(device = _hw.DEVICE): config["resolved_gpu_ids"] = None config["gpu_selection"] = None elif gpu_ids: @@ -1022,17 +1463,22 @@ class TrainingBackend: self._progress.is_training = True elif etype == "complete": - self._progress.is_training = False - self._progress.is_completed = True - self._output_dir = event.get("output_dir") msg = event.get("status_message", "Training completed") + stopped = self._should_stop or msg.strip().lower() in { + "training cancelled", + "training stopped", + } + self._progress.is_training = False + self._progress.is_completed = not stopped + self._output_dir = event.get("output_dir") + self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: db_action = "create_and_finalize" else: db_action = "finalize" db_action_kwargs = { - "status": "stopped" if self._should_stop else "completed", + "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, } diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 17dc1299ca..0ff4d517ed 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1309,14 +1309,18 @@ def _normalize_mlx_studio_scheduler(value): def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: - """Resolve Studio local dataset uploads without importing the GPU trainer.""" + """Resolve CLI paths and Studio local dataset uploads without importing the GPU trainer.""" from utils.paths import resolve_dataset_path all_files: list[str] = [] for dataset_file in file_paths or []: - file_path = ( - dataset_file if os.path.isabs(dataset_file) else str(resolve_dataset_path(dataset_file)) - ) + dataset_path = Path(os.path.expanduser(str(dataset_file))) + if dataset_path.is_absolute(): + file_path = str(dataset_path) + elif dataset_path.exists(): + file_path = str(dataset_path.resolve()) + else: + file_path = str(resolve_dataset_path(str(dataset_file))) file_path_obj = Path(file_path) if file_path_obj.is_dir(): @@ -1355,6 +1359,58 @@ def _mlx_local_dataset_loader_for_files(files: list[str]) -> str: raise ValueError(f"Unsupported dataset format: {files[0]}") +_MLX_WORKER_COMPLETE = "_mlx_worker_complete" + + +def _start_mlx_stop_poller(stop_queue): + import queue as _queue + import threading + + stop_save = [True] + stop_requested = [False] + trainer_ref = [None] + + def is_stop_requested(): + return stop_requested[0] + + def poll_stop(): + while True: + try: + msg = stop_queue.get(timeout = 0.25) + if msg and msg.get("type") == _MLX_WORKER_COMPLETE: + return + if msg and msg.get("type") == "stop": + stop_save[0] = msg.get("save", True) + stop_requested[0] = True + trainer = trainer_ref[0] + if trainer is not None: + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target = poll_stop, daemon = True) + stop_thread.start() + return stop_save, stop_requested, trainer_ref, is_stop_requested, stop_thread + + +def _resolve_mlx_output_dir(config, model_name): + from utils.paths import resolve_output_dir, default_run_dir_name + + output_dir = config.get("output_dir", "") + if not output_dir: + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + return str(resolve_output_dir(output_dir)) + if config.get("allow_external_output_dir"): + output_path = Path(output_dir).expanduser() + if not output_path.is_absolute(): + output_path = Path.cwd() / output_path + return str(output_path.resolve()) + return str(resolve_output_dir(output_dir)) + + def _run_mlx_training(event_queue, stop_queue, config): """Self-contained MLX training path for Apple Silicon. @@ -1363,8 +1419,6 @@ def _run_mlx_training(event_queue, stop_queue, config): """ import time import math - import threading - import queue as _queue from pathlib import Path def _send(event_type, **kwargs): @@ -1374,31 +1428,9 @@ def _run_mlx_training(event_queue, stop_queue, config): kwargs["message"] = sm event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) - _stop_save = [True] - _stop_requested = [False] - _trainer_ref = [None] - - def _is_stop_requested(): - return _stop_requested[0] - - def _poll_stop(): - while True: - try: - msg = stop_queue.get(timeout = 1.0) - if msg and msg.get("type") == "stop": - _stop_save[0] = msg.get("save", True) - _stop_requested[0] = True - trainer = _trainer_ref[0] - if trainer is not None: - trainer.stop_requested = True - return - except _queue.Empty: - continue - except (EOFError, OSError): - return - - stop_thread = threading.Thread(target = _poll_stop, daemon = True) - stop_thread.start() + _stop_save, _stop_requested, _trainer_ref, _is_stop_requested, _stop_thread = ( + _start_mlx_stop_poller(stop_queue) + ) _send("status", status_message = "Loading MLX libraries...") @@ -1804,21 +1836,14 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 5. Build output dir ── # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it - from utils.paths import resolve_output_dir, ensure_dir + from utils.paths import ensure_dir - output_dir = config.get("output_dir", "") - if not output_dir: - output_dir = build_default_output_dir_name( - model_name, - config.get("project_name"), - ) - output_dir = str(resolve_output_dir(output_dir)) + output_dir = _resolve_mlx_output_dir(config, model_name) ensure_dir(Path(output_dir)) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 if isinstance(eval_steps_val, float) and 0 < eval_steps_val < 1: - # Studio sometimes sends fraction-of-total-steps eval_steps_val = max(1, int(eval_steps_val * max_steps)) else: eval_steps_val = int(eval_steps_val) @@ -2043,12 +2068,27 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train(resume_from_checkpoint = resume_from_checkpoint) + _save_model = trainer.save_model + + def _skip_internal_final_save(*args, **kwargs): + raise ValueError("worker owns final save") + + trainer.save_model = _skip_internal_final_save + try: + trainer.train(resume_from_checkpoint = resume_from_checkpoint) + finally: + trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested and not _stop_save[0]: - # User clicked "Cancel" (save=False) — skip saving - _send("complete", output_dir = None, status_message = "Training cancelled") + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: _send("status", status_message = "Saving model...") mx.synchronize() @@ -2067,6 +2107,79 @@ def _run_mlx_training(event_queue, stop_queue, config): pass +def _is_current_process_apple_silicon() -> bool: + import platform + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def run_mlx_training_process( + *, + event_queue: Any, + stop_queue: Any, + config: dict, + transformers_activated: bool = False, +) -> None: + """MLX worker entrypoint shared by Studio subprocesses and the CLI adapter.""" + model_name = config["model_name"] + + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + + if not transformers_activated: + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE != _hw.DeviceType.MLX: + event_queue.put( + { + "type": "error", + "error": "MLX training requires Apple Silicon with the MLX backend available.", + "stack": "", + "ts": time.time(), + } + ) + return + + if config.get("is_dataset_audio"): + event_queue.put( + { + "type": "error", + "error": "Audio dataset training is not yet supported on Apple Silicon.", + "stack": "", + "ts": time.time(), + } + ) + return + + try: + try: + _run_mlx_training(event_queue, stop_queue, config) + finally: + try: + stop_queue.put({"type": _MLX_WORKER_COMPLETE}) + except (EOFError, OSError, ValueError): + pass + except Exception as exc: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None: """Subprocess entrypoint. Fresh Python — no stale module state. @@ -2141,36 +2254,26 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) + from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend + + mlx_backend_requested = is_apple_silicon_training_platform() + + mlx_transformers_activated = False + if mlx_backend_requested and _is_current_process_apple_silicon(): + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + mlx_transformers_activated = True + from utils.hardware import hardware as _hw _hw.detect_hardware() - if _hw.DEVICE == _hw.DeviceType.MLX: - if config.get("is_dataset_audio"): - event_queue.put( - { - "type": "error", - "error": "Audio dataset training is not yet supported on Apple Silicon.", - "stack": "", - "ts": time.time(), - } - ) - return - # Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.) - # Must happen before any transformers/mlx-lm imports in _run_mlx_training. - # Non-fatal: fall through with whatever version is installed, but log - # the failure instead of swallowing it (issue #6103). - _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) - try: - _run_mlx_training(event_queue, stop_queue, config) - except Exception as exc: - event_queue.put( - { - "type": "error", - "error": str(exc), - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - } - ) + if mlx_backend_requested or should_use_mlx_training_backend(device = _hw.DEVICE): + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + transformers_activated = mlx_transformers_activated, + ) return # ── 1. Activate correct transformers version BEFORE any ML imports ── @@ -2693,7 +2796,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) - from core.training.trainer import UnslothTrainer, TrainingProgress + from core.training.training import TrainingProgress + from core.training.trainer import UnslothTrainer from utils.paths import ( ensure_dir, resolve_output_dir, diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 54048a65dd..47c6669f8f 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -6,9 +6,15 @@ empty-chat-template crash) before train(). The real methods are bound onto a lig fake self so the production logic runs against controlled batches.""" import importlib +import json +import os +import queue +import subprocess import sys +import threading import types import unittest +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -184,5 +190,231 @@ class TestChatTemplateRendersEmpty(unittest.TestCase): self.assertFalse(s._chat_template_renders_empty()) +def _clear_trainer_module(package: str): + sys.modules.pop(f"{package}.trainer", None) + pkg = sys.modules.get(package) + if pkg is not None and hasattr(pkg, "trainer"): + delattr(pkg, "trainer") + + +def _set_training_platform(monkeypatch, package: str, backend: str): + training_mod = importlib.import_module(f"{package}.training") + from utils.hardware import hardware as hw + + monkeypatch.setattr(hw, "DEVICE", None) + monkeypatch.setattr( + training_mod.platform, + "system", + lambda: "Darwin" if backend == "mlx" else "Linux", + ) + monkeypatch.setattr( + training_mod.platform, + "machine", + lambda: "arm64" if backend == "mlx" else "x86_64", + ) + + +def _load_trainer_module( + monkeypatch, + backend: str, + package: str = "core.training", +): + _set_training_platform(monkeypatch, package, backend) + _clear_trainer_module(package) + if package in sys.modules: + importlib.reload(sys.modules[package]) + trainer_mod = importlib.import_module(f"{package}.trainer") + training_mod = importlib.import_module(f"{package}.training") + monkeypatch.setattr( + training_mod._MLXTrainerAdapter, + "_activate_transformers_for_model", + lambda self, model_name, hf_token: None, + ) + return trainer_mod + + +class _ExitedProc: + def join(self, timeout = None): + return None + + def is_alive(self): + return False + + +class _TerminableProc: + def __init__(self): + self.terminated = False + self._done = threading.Event() + + def join(self, timeout = None): + self._done.wait(timeout = timeout or 5) + + def is_alive(self): + return not self.terminated + + def terminate(self): + self.terminated = True + self._done.set() + + +def test_unsloth_trainer_dispatches_for_mlx_and_torch(monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + + mlx_trainer = trainer_mod.UnslothTrainer() + + assert type(mlx_trainer).__module__ == "core.training.training" + assert mlx_trainer.get_training_progress().status_message == "Ready to train" + + trainer_mod = _load_trainer_module(monkeypatch, "torch") + + assert trainer_mod.UnslothTrainer().__class__ is trainer_mod.UnslothTrainer + + +def test_cli_mlx_trainer_activates_before_importing_trainer(): + repo_root = Path(__file__).resolve().parents[3] + script = """ +import json +import sys +import unsloth_cli.commands.train as train_cmd +from studio.backend.core.training import training as training_mod +from utils.hardware import hardware as hw + +training_mod.platform.system = lambda: "Darwin" +training_mod.platform.machine = lambda: "arm64" +hw.DEVICE = None +events = [] + +def fake_activate(model_name, hf_token): + events.append({ + "model_name": model_name, + "trainer_loaded": "studio.backend.core.training.trainer" in sys.modules, + }) + +train_cmd._activate_mlx_transformers = fake_activate +trainer = train_cmd._create_cli_trainer("mlx-community/Qwen3-0.6B-4bit", None) +print(json.dumps({ + "trainer_module": type(trainer).__module__, + "events": events, +})) +""" + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + [str(repo_root), str(repo_root / "studio" / "backend"), env.get("PYTHONPATH", "")] + ) + result = subprocess.run( + [sys.executable, "-c", script], + cwd = repo_root, + env = env, + text = True, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + check = True, + ) + payload = json.loads(result.stdout) + + assert payload["trainer_module"] == "studio.backend.core.training.training" + assert payload["events"] == [ + {"model_name": "mlx-community/Qwen3-0.6B-4bit", "trainer_loaded": False} + ] + + +def test_mlx_adapter_builds_config_and_reports_completion(tmp_path, monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + captured = {} + + def fake_run_worker(config, event_queue, stop_queue): + captured["config"] = config + event_queue.put({"type": "progress", "step": 1, "total_steps": 1, "loss": 0.25}) + event_queue.put( + {"type": "complete", "status_message": "done", "output_dir": config["output_dir"]} + ) + + trainer = trainer_mod.UnslothTrainer() + monkeypatch.setattr(trainer, "_run_mlx_worker", fake_run_worker) + + assert trainer.load_model("mlx-community/Qwen3-0.6B-4bit", max_seq_length = 1024) + assert trainer.prepare_model_for_training(use_lora = False) + dataset, eval_dataset = trainer.load_and_format_dataset("org/dataset") + output_dir = tmp_path / "mlx-out" + + assert trainer.start_training( + dataset = dataset, + eval_dataset = eval_dataset, + output_dir = output_dir, + project_name = "Sales Assistant", + max_steps = 1, + learning_rate = 3e-4, + ) + trainer.training_thread.join(timeout = 5) + + progress = trainer.get_training_progress() + config = captured["config"] + assert progress.is_completed + assert progress.output_dir == str(output_dir.resolve()) + progress.status_message = "mutated" + assert trainer.get_training_progress().status_message == "done" + assert config["model_name"] == "mlx-community/Qwen3-0.6B-4bit" + assert config["project_name"] == "Sales Assistant" + assert config["hf_dataset"] == "org/dataset" + assert config["training_type"] == "Full Finetuning" + assert config["load_in_4bit"] is False + assert config["max_seq_length"] == 1024 + assert config["learning_rate"] == 3e-4 + assert config["output_dir"] == str(output_dir.resolve()) + assert config["allow_external_output_dir"] is True + + +def test_mlx_worker_helpers_cover_cli_paths(tmp_path, monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training.worker import ( + _resolve_mlx_local_dataset_files, + _resolve_mlx_output_dir, + ) + + dataset = tmp_path / "train.jsonl" + dataset.write_text('{"text":"hello"}\n', encoding = "utf-8") + monkeypatch.chdir(tmp_path) + + assert _resolve_mlx_local_dataset_files(["train.jsonl"]) == [str(dataset)] + assert _resolve_mlx_output_dir( + {"output_dir": "cli-out", "allow_external_output_dir": True}, + "mlx-community/Qwen3-0.6B-4bit", + ) == str((tmp_path / "cli-out").resolve()) + + +def test_run_mlx_training_process_applies_side_effects_before_hardware_detection(monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training import worker + from utils.hardware import hardware as hw + + order = [] + + def fake_activate(model_name, hf_token): + order.append(("activate", model_name, hf_token)) + + def fake_detect_hardware(): + order.append("detect") + hw.DEVICE = hw.DeviceType.CPU + return hw.DEVICE + + monkeypatch.delenv("HF_HUB_DISABLE_XET", raising = False) + monkeypatch.delenv("HF_HUB_ENABLE_HF_TRANSFER", raising = False) + monkeypatch.setattr(worker, "_activate_transformers_version_or_warn", fake_activate) + monkeypatch.setattr(hw, "detect_hardware", fake_detect_hardware) + + event_queue = queue.Queue() + worker.run_mlx_training_process( + event_queue = event_queue, + stop_queue = queue.Queue(), + config = {"model_name": "mlx-community/Gemma-4-12B", "disable_xet": True}, + ) + + event = event_queue.get_nowait() + assert order == [("activate", "mlx-community/Gemma-4-12B", None), "detect"] + assert os.environ["HF_HUB_DISABLE_XET"] == "1" + assert os.environ["HF_HUB_ENABLE_HF_TRANSFER"] == "0" + assert "MLX training requires Apple Silicon" in event["error"] + + if __name__ == "__main__": unittest.main() diff --git a/unsloth-cli.py b/unsloth-cli.py index 756efef0d0..be48893d8d 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -22,24 +22,179 @@ import argparse import os +def _is_mlx_backend(unsloth_module): + return bool(getattr(unsloth_module, "_IS_MLX", False)) + + +def _normalize_dtype(dtype, is_mlx): + if is_mlx and isinstance(dtype, str) and dtype.strip().lower() in {"", "none", "auto"}: + return None + return dtype + + +def _prepare_device_map(is_mlx): + if is_mlx: + return None, False + + from unsloth.models.loader_utils import prepare_device_map + return prepare_device_map() + + +class _CallableTokenizerProxy: + def __init__(self, tokenizer): + self._tokenizer = tokenizer + + def __getattr__(self, name): + return getattr(self._tokenizer, name) + + def __call__(self, text, *args, **kwargs): + # MLX/torch-free: never request torch tensors; keep plain python ids. + kwargs.pop("return_tensors", None) + wrapped = getattr(self._tokenizer, "_tokenizer", None) + if callable(wrapped): + return wrapped(text, *args, **kwargs) + + add_special_tokens = kwargs.get("add_special_tokens", False) + input_ids = self._tokenizer.encode(text, add_special_tokens = add_special_tokens) + return {"input_ids": input_ids} + + +def _tokenizer_for_raw_text_loader(tokenizer, is_mlx): + if not is_mlx or callable(tokenizer): + return tokenizer + return _CallableTokenizerProxy(tokenizer) + + +def _raw_text_loader_for_backend( + RawTextDataLoader, + tokenizer, + is_mlx, + chunk_size = 2048, + stride = 512, +): + return RawTextDataLoader( + _tokenizer_for_raw_text_loader(tokenizer, is_mlx), + chunk_size, + stride, + return_tokenized = not is_mlx, + ) + + +def _train_with_legacy_save_control(trainer, is_mlx): + if not is_mlx: + return trainer.train() + + original_save_model = getattr(trainer, "save_model", None) + if original_save_model is None: + return trainer.train() + + def skip_internal_final_save(*args, **kwargs): + raise ValueError("legacy unsloth-cli.py owns final save") + + trainer.save_model = skip_internal_final_save + try: + return trainer.train() + finally: + trainer.save_model = original_save_model + + +def _iter_quantization_methods(quantization): + if isinstance(quantization, list): + return quantization + return [quantization] + + +def _save_or_push_model(model, tokenizer, args, is_mlx): + if not args.save_model: + print("Warning: The model is not saved!") + return + + # Enter the GGUF branch when saving or pushing GGUF, so --push_gguf works + # without --save_gguf (the local save is guarded separately below). + if args.save_gguf or args.push_gguf: + if not args.save_gguf: + print("Warning: --save_gguf not set, pushing GGUF to hub without saving locally.") + for quantization_method in _iter_quantization_methods(args.quantization): + if args.save_gguf: + print(f"Saving model with quantization method: {quantization_method}") + model.save_pretrained_gguf( + args.save_path, + tokenizer, + quantization_method = quantization_method, + ) + if args.push_model or args.push_gguf: + model.push_to_hub_gguf( + args.hub_path, + tokenizer, + quantization_method = quantization_method, + token = args.hub_token, + ) + return + + if is_mlx: + model.save_pretrained_merged( + args.save_path, + tokenizer, + save_method = args.save_method, + push_to_hub = args.push_model, + repo_id = args.hub_path if args.push_model else None, + token = args.hub_token, + ) + return + + model.save_pretrained_merged(args.save_path, tokenizer, save_method = args.save_method) + if args.push_model: + model.push_to_hub_merged(args.hub_path, tokenizer, args.save_method, token = args.hub_token) + + +def _build_sft_config(SFTConfig, args, is_mlx, bf16_supported): + config_kwargs = dict( + per_device_train_batch_size = args.per_device_train_batch_size, + gradient_accumulation_steps = args.gradient_accumulation_steps, + warmup_steps = args.warmup_steps, + max_steps = args.max_steps, + learning_rate = args.learning_rate, + fp16 = not bf16_supported, + bf16 = bf16_supported, + logging_steps = args.logging_steps, + optim = args.optim, + weight_decay = args.weight_decay, + lr_scheduler_type = args.lr_scheduler_type, + seed = args.seed, + output_dir = args.output_dir, + report_to = args.report_to, + max_length = args.max_seq_length, + dataset_num_proc = 2, + packing = args.packing, + ) + if is_mlx: + if args.per_device_eval_batch_size != 4: + print("Warning: --per_device_eval_batch_size is ignored on MLX without eval data.") + else: + config_kwargs["per_device_eval_batch_size"] = args.per_device_eval_batch_size + return SFTConfig(**config_kwargs) + + def run(args): + import unsloth from unsloth import FastLanguageModel from datasets import load_dataset from transformers.utils import strtobool from trl import SFTTrainer, SFTConfig from unsloth import is_bfloat16_supported - from unsloth.models.loader_utils import prepare_device_map import logging from unsloth import RawTextDataLoader logging.getLogger("hf-to-gguf").setLevel(logging.WARNING) + is_mlx = _is_mlx_backend(unsloth) + # Load model and tokenizer - device_map, distributed = prepare_device_map() + device_map, distributed = _prepare_device_map(is_mlx) model, tokenizer = FastLanguageModel.from_pretrained( model_name = args.model_name, max_seq_length = args.max_seq_length, - dtype = args.dtype, + dtype = _normalize_dtype(args.dtype, is_mlx), load_in_4bit = args.load_in_4bit, device_map = device_map, ) @@ -92,11 +247,13 @@ def run(args): def load_dataset_smart(args): from transformers.utils import strtobool if args.raw_text_file: - loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride) + loader = _raw_text_loader_for_backend( + RawTextDataLoader, tokenizer, is_mlx, args.chunk_size, args.stride + ) dataset = loader.load_from_file(args.raw_text_file) elif args.dataset.endswith((".txt", ".md", ".json", ".jsonl")): # Auto-detect local raw text files - loader = RawTextDataLoader(tokenizer) + loader = _raw_text_loader_for_backend(RawTextDataLoader, tokenizer, is_mlx) dataset = loader.load_from_file(args.dataset) else: use_modelscope = strtobool(os.environ.get("UNSLOTH_USE_MODELSCOPE", "False")) @@ -115,27 +272,9 @@ def run(args): print("Data is formatted and ready!") # Configure training arguments - training_args = SFTConfig( - per_device_train_batch_size = args.per_device_train_batch_size, - per_device_eval_batch_size = args.per_device_eval_batch_size, - gradient_accumulation_steps = args.gradient_accumulation_steps, - warmup_steps = args.warmup_steps, - max_steps = args.max_steps, - learning_rate = args.learning_rate, - fp16 = not is_bfloat16_supported(), - bf16 = is_bfloat16_supported(), - logging_steps = args.logging_steps, - optim = args.optim, - weight_decay = args.weight_decay, - lr_scheduler_type = args.lr_scheduler_type, - seed = args.seed, - output_dir = args.output_dir, - report_to = args.report_to, - max_length = args.max_seq_length, - dataset_num_proc = 2, - ddp_find_unused_parameters = False if distributed else None, - packing = args.packing, - ) + training_args = _build_sft_config(SFTConfig, args, is_mlx, is_bfloat16_supported()) + if distributed: + training_args.ddp_find_unused_parameters = False # Initialize trainer trainer = SFTTrainer( @@ -145,57 +284,9 @@ def run(args): args = training_args, ) - trainer.train() + _train_with_legacy_save_control(trainer, is_mlx) - # Save model - if args.save_model: - # If args.quantization is a list, save once per quantization method - # Enter the GGUF branch when saving *or* pushing GGUF, so --push_gguf - # works even when --save_gguf is omitted (the local save is guarded - # separately below). - if args.save_gguf or args.push_gguf: - # Push-only GGUF (no --save_gguf) skips the local save; warn so it is not silent. - if not args.save_gguf: - print("Warning: --save_gguf not set, pushing GGUF to hub without saving locally.") - if isinstance(args.quantization, list): - for quantization_method in args.quantization: - if args.save_gguf: - print(f"Saving model with quantization method: {quantization_method}") - model.save_pretrained_gguf( - args.save_path, - tokenizer, - quantization_method = quantization_method, - ) - if args.push_model or args.push_gguf: - model.push_to_hub_gguf( - args.hub_path, - tokenizer, - quantization_method = quantization_method, - token = args.hub_token, - ) - else: - if args.save_gguf: - print(f"Saving model with quantization method: {args.quantization}") - model.save_pretrained_gguf( - args.save_path, - tokenizer, - quantization_method = args.quantization, - ) - if args.push_model or args.push_gguf: - model.push_to_hub_gguf( - args.hub_path, - tokenizer, - quantization_method = args.quantization, - token = args.hub_token, - ) - else: - model.save_pretrained_merged(args.save_path, tokenizer, args.save_method) - if args.push_model: - model.push_to_hub_merged( - args.hub_path, tokenizer, args.save_method, token = args.hub_token - ) - else: - print("Warning: The model is not saved!") + _save_or_push_model(model, tokenizer, args, is_mlx) if __name__ == "__main__": diff --git a/unsloth_cli/commands/train.py b/unsloth_cli/commands/train.py index 9d47a574d5..c52c2344b7 100644 --- a/unsloth_cli/commands/train.py +++ b/unsloth_cli/commands/train.py @@ -7,10 +7,42 @@ from typing import Optional import typer +from unsloth_cli._inference import ensure_studio_backend_path from unsloth_cli.config import Config, load_config from unsloth_cli.options import add_options_from_config +def _should_use_mlx_backend_for_cli() -> bool: + ensure_studio_backend_path() + from studio.backend.core.training.training import should_use_mlx_training_backend + return should_use_mlx_training_backend() + + +def _activate_mlx_transformers(model_name: str, hf_token: Optional[str]) -> None: + # Activate before any transformers import: adapter model-type detection imports utils.models. + ensure_studio_backend_path() + from utils.transformers_version import activate_transformers_for_subprocess + try: + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + typer.echo(f"Warning: failed to activate Transformers sidecar: {exc}", err = True) + + +def _create_cli_trainer(model_name: str, hf_token: Optional[str]): + if _should_use_mlx_backend_for_cli(): + _activate_mlx_transformers(model_name, hf_token) + # MLX is torch-free: use the lightweight adapter, not trainer.py (imports torch/unsloth/trl at load). + ensure_studio_backend_path() + from studio.backend.core.training.training import create_mlx_trainer_adapter + + return create_mlx_trainer_adapter() + + ensure_studio_backend_path() + from studio.backend.core.training.trainer import UnslothTrainer + + return UnslothTrainer() + + @add_options_from_config(Config) def train( config: Optional[Path] = typer.Option( @@ -39,6 +71,7 @@ def train( typer.echo(f"Error: {e}", err = True) raise typer.Exit(code = 2) + config_overrides = config_overrides or {} cfg.apply_overrides(**config_overrides) # CLI/env tokens take precedence; guard against unresolved typer.Option @@ -83,9 +116,7 @@ def train( ) raise typer.Exit(code = 2) - from studio.backend.core.training.trainer import UnslothTrainer - - trainer = UnslothTrainer() + trainer = _create_cli_trainer(cfg.model, hf_token) # Load model (trainer.is_vlm is set after this) if not trainer.load_model( @@ -124,13 +155,20 @@ def train( try: while trainer.training_thread and trainer.training_thread.is_alive(): + progress = trainer.get_training_progress() + if getattr(progress, "error", None): + break time.sleep(1) except KeyboardInterrupt: typer.echo("Stopping training (Ctrl+C detected)...") trainer.stop_training() finally: if trainer.training_thread: - trainer.training_thread.join() + progress = trainer.get_training_progress() + if getattr(progress, "error", None): + trainer.training_thread.join(timeout = 5) + else: + trainer.training_thread.join() final = trainer.get_training_progress() if getattr(final, "error", None): From 2a6abe2ff5c643ee853f2e0ef632b1e80d06918a Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 8 Jul 2026 18:25:39 +0800 Subject: [PATCH 062/113] feat(cli): support MLX distributed inference (#6845) * feat(cli): detect MLX distributed launch context * feat(mlx): wire distributed inference backend * feat(cli): broadcast MLX distributed chat turns * fix(cli): wait indefinitely for distributed chat turns * fix(cli): report MLX distributed load errors cleanly * fix(mlx): route distributed vlm through loader * fix(cli): detect inline MLX host JSON * fix(studio): harden distributed object sharing * fix(studio): select JACCL distributed backend * fix(cli): abort distributed error paths * Distinguish real stream errors from model text via GenStreamError in distributed CLI * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fail loud when MLX distributed init returns a singleton group The worker only reaches this block when distributed was explicitly requested. A singleton (size 1) group means the launch failed to form a real group (MLX built without distributed support, or an invalid launch env/hostfile); silently continuing leaves nonzero ranks looping forever on share_distributed_object. Raise instead so the surrounding handler returns a clear load error. * Tighten MLX distributed inference comments --------- Co-authored-by: Daniel Han Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../backend/core/inference/mlx_inference.py | 171 ++++++--- studio/backend/core/inference/orchestrator.py | 106 +++++- studio/backend/core/inference/worker.py | 123 ++++++- studio/backend/routes/inference.py | 46 +++ .../tests/test_mlx_inference_backend.py | 125 +++++++ unsloth_cli/_inference.py | 217 ++++++++++-- unsloth_cli/commands/chat.py | 178 +++++++--- unsloth_cli/commands/inference.py | 47 ++- unsloth_cli/tests/test_inference_chat.py | 330 +++++++++++++++++- 9 files changed, 1188 insertions(+), 155 deletions(-) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 62d268e15f..d84baa278d 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -5,6 +5,7 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm instead of torch/transformers for model loading and generation. """ +import os import threading from typing import Optional, Generator from core.inference.runtime_context import runtime_context_length @@ -41,6 +42,48 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): } +def _mlx_distributed_rank_size(group = None): + """Return ``(rank, world_size)`` for an optional MLX distributed group.""" + if group is None: + return 0, 1 + rank = int(group.rank()) + world_size = int(group.size()) + if world_size < 1: + raise ValueError(f"Invalid MLX distributed world_size={world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.") + return rank, world_size + + +def _mlx_distributed_backend_from_env(): + if os.environ.get("MLX_JACCL_COORDINATOR") and os.environ.get("MLX_IBV_DEVICES"): + return "jaccl" + return None + + +def _init_mlx_distributed(): + """Initialize MLX distributed state, falling back to singleton metadata.""" + import mlx.core as mx + + group = None + rank = 0 + world_size = 1 + distributed = getattr(mx, "distributed", None) + init = getattr(distributed, "init", None) if distributed is not None else None + if callable(init): + backend = _mlx_distributed_backend_from_env() + if backend is None: + group = init() + else: + try: + group = init(backend = backend) + except TypeError: + group = init() + if group is not None: + rank, world_size = _mlx_distributed_rank_size(group) + return group, rank, world_size + + def _make_mlx_presence_penalty_processor(penalty: float): """Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path. @@ -52,7 +95,7 @@ def _make_mlx_presence_penalty_processor(penalty: float): def _processor(tokens, logits): if state["prompt_len"] is None: - # First call = prompt only; latch its length. + # First call is prompt-only; latch its length. state["prompt_len"] = int(tokens.shape[0]) return logits generated = tokens[state["prompt_len"] :] @@ -61,22 +104,17 @@ def _make_mlx_presence_penalty_processor(penalty: float): import mlx.core as mx vocab = logits.shape[-1] - # Bound generated ids to the valid range [0, vocab) before they index - # logits. MLX does no bounds checking and out-of-bounds indexing is - # documented undefined behavior (crash / memory corruption), unlike the - # torch path's harmless negative wrap -- so this bound is load-bearing - # here and matches the torch filter seen[(seen >= 0) & (seen < vocab)]. - # MLX has no boolean-mask filtering (data-dependent output shape is - # unsupported), so instead of compacting the id list we route every - # out-of-range or negative id to a scratch slot at index ``vocab`` that - # is dropped before the subtract. That scratch slot can never collide - # with a real token, so real ids (including id 0) are penalized exactly - # once and stray ids are ignored. + # Bound ids to [0, vocab) before indexing logits: MLX does no bounds + # checking and out-of-bounds indexing is undefined behavior (crash / + # corruption), unlike torch's harmless negative wrap. MLX also lacks + # boolean-mask filtering, so out-of-range/negative ids route to a + # scratch slot at index vocab (dropped before the subtract) that never + # collides with a real token: real ids (including 0) are penalized + # once, strays ignored. valid = (generated >= 0) & (generated < vocab) safe = mx.where(valid, generated, vocab).astype(mx.int32) - # Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate - # ids are idempotent, so presence applies once per distinct token; the - # scratch column is discarded and the full-width subtract stays on-device. + # Scatter penalty into a (vocab + 1)-wide mask: duplicate ids are + # idempotent (presence applies once per token); scratch column dropped. mask = mx.zeros((vocab + 1,), dtype = logits.dtype) mask[safe] = penalty logits = logits - mask[:vocab] @@ -93,7 +131,7 @@ class MLXInferenceBackend: self.loaded_local_models = [] self.device = "mlx" self._generation_lock = threading.Lock() - # usage/timings of the latest generation; shipped on gen_done. + # usage/timings of the latest generation, shipped on gen_done. self.last_generation_stats = None self._model = None @@ -101,6 +139,9 @@ class MLXInferenceBackend: self._processor = None self._is_vlm = False self._config = {} + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 # Recorded for unload to release pinned memory back to the OS. self._memory_limits_applied = {} @@ -145,19 +186,26 @@ class MLXInferenceBackend: trust_remote_code = False, gpu_ids = None, dtype = None, + parallel_mode = None, + distributed_group = None, ) -> bool: import mlx.core as mx - # Keep the token so the native-template fallback can fetch a - # gated model's repo template later during generation. + # Keep the token so the native-template fallback can fetch a gated + # model's repo template during generation. self._hf_token = hf_token model_name = config.identifier if hasattr(config, "identifier") else str(config) is_vision = getattr(config, "is_vision", False) + distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group) + is_distributed = distributed_group is not None and distributed_size > 1 + self._distributed_group = distributed_group + self._distributed_rank = distributed_rank + self._distributed_world_size = distributed_size - # GGUF guard. GGUF models are served by llama-server in the parent - # process, not mlx-lm here. Reaching this with is_gguf=True means the - # route's first detection flaked (transient HF Hub) but the subprocess - # re-detected GGUF; raise loudly instead of a cryptic mlx_lm error. + # GGUF guard: GGUF is served by llama-server in the parent process, + # not mlx-lm. Reaching here with is_gguf=True means the route's + # detection flaked but the subprocess re-detected GGUF; raise loudly + # instead of a cryptic mlx_lm error. if getattr(config, "is_gguf", False): raise RuntimeError( f"MLXInferenceBackend cannot load GGUF model '{model_name}': " @@ -176,11 +224,26 @@ class MLXInferenceBackend: is_lora = getattr(config, "is_lora", False) logger.info( - "Loading %s via %s (is_lora=%s)", + "Loading %s via %s (is_lora=%s, distributed=%s, rank=%s/%s, mode=%s)", model_name, "mlx-vlm" if is_vision else "mlx-lm", is_lora, + is_distributed, + distributed_rank, + distributed_size, + parallel_mode, ) + if is_distributed and parallel_mode not in ("pipeline", "tensor"): + raise ValueError( + "Unsloth: distributed MLX inference requires parallel_mode='pipeline' " + "or parallel_mode='tensor'." + ) + if is_distributed and is_lora: + raise ValueError( + "Unsloth: distributed MLX inference for LoRA adapter repos " + "is not supported yet. Merge/export the adapter into an MLX model " + "before distributed inference." + ) try: from unsloth_zoo.mlx.loader import FastMLXModel @@ -190,14 +253,23 @@ class MLXInferenceBackend: "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon." ) from e + load_kwargs = { + "max_seq_length": max_seq_length, + "dtype": dtype, + "load_in_4bit": load_in_4bit, + "token": hf_token, + "trust_remote_code": trust_remote_code, + "text_only": False if is_vision else True, + } + if is_distributed: + if parallel_mode == "pipeline": + load_kwargs["pipeline_group"] = distributed_group + else: + load_kwargs["tensor_group"] = distributed_group + model, tokenizer_or_processor = FastMLXModel.from_pretrained( model_name, - max_seq_length = max_seq_length, - dtype = dtype, - load_in_4bit = load_in_4bit, - token = hf_token, - trust_remote_code = trust_remote_code, - text_only = False if is_vision else True, + **load_kwargs, ) if is_vision: @@ -217,8 +289,7 @@ class MLXInferenceBackend: self.models[model_name] = { # Per-model token for the native-template fallback (matches transformers). "hf_token": hf_token, - # Per-model consent for the native-template reload: re-use the exact - # trust_remote_code this model was loaded with (matches transformers). + # Per-model trust_remote_code reused by the native-template reload (matches transformers). "trust_remote_code": trust_remote_code, "model": self._model, "tokenizer": self._tokenizer, @@ -234,8 +305,7 @@ class MLXInferenceBackend: "has_audio_input": False, "context_length": runtime_context_length(self._model, max_seq_length), } - # Capture chat_template_info so the worker IPC reply ships it back and - # the route layer classifies capabilities like the other paths. + # Capture chat_template_info for the worker IPC reply and route capability classification. self._populate_chat_template_info(model_name) logger.info("Model %s loaded successfully", model_name) @@ -293,6 +363,9 @@ class MLXInferenceBackend: self._model = None self._tokenizer = None self._processor = None + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 if self.active_model_name == model_name: self.active_model_name = None gc.collect() @@ -320,8 +393,7 @@ class MLXInferenceBackend: max_new_tokens = 256, repetition_penalty = 1.0, cancel_event = None, - # Reasoning / tool kwargs forwarded by the route + worker; rendered via - # apply_chat_template_for_generation like the transformers path. + # Reasoning / tool kwargs, rendered via apply_chat_template_for_generation (transformers parity). tools = None, enable_thinking = None, reasoning_effort = None, @@ -334,7 +406,6 @@ class MLXInferenceBackend: # Reset so a failed run cannot surface stale stats. self.last_generation_stats = None - # Build messages with system prompt full_messages = [] if system_prompt: full_messages.append({"role": "system", "content": system_prompt}) @@ -351,7 +422,6 @@ class MLXInferenceBackend: {"type": "text", "text": content}, ] elif isinstance(content, list): - # Prepend image if not already present has_image = any( p.get("type") == "image" for p in content if isinstance(p, dict) ) @@ -429,11 +499,11 @@ class MLXInferenceBackend: if prompt is None: raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible") - # Same parity fix as the transformers backend: if the template dropped the - # requested tools, fall back to the native template so MLX text models keep - # advertising them. ``self._tokenizer`` is this entry's model_info tokenizer, - # so probe and native render share a renderer. (The VLM path renders via the - # processor for image tokens and is intentionally not wired here.) + # Parity with the transformers backend: if the template dropped the + # requested tools, fall back to the native template so MLX text models + # keep advertising them. self._tokenizer is this entry's tokenizer, so + # probe and native render share a renderer. (VLM renders via the + # processor for image tokens and is not wired here.) model_info = self.models.get(self.active_model_name, {}) prompt = render_with_native_template_fallback( formatted_prompt = prompt, @@ -455,7 +525,7 @@ class MLXInferenceBackend: min_p = float(min_p or 0.0), min_tokens_to_keep = 1, ) - # Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths). + # Repetition and/or presence penalty processors (GGUF/safetensors parity). logits_processors = [] if repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, @@ -496,7 +566,6 @@ class MLXInferenceBackend: ): final_response = response token_ids.append(response.token) - # Decode full sequence with skip_special_tokens cumulative = self._tokenizer.decode( token_ids, skip_special_tokens = True, @@ -544,8 +613,7 @@ class MLXInferenceBackend: ) # Pick the chat-template-aware caller: processors with their own - # apply_chat_template + chat_template (e.g. Qwen2.5-VL) use it - # directly; else fall back to the nested tokenizer. + # apply_chat_template + chat_template (e.g. Qwen2.5-VL), else the nested tokenizer. chat_target = self._processor if ( getattr(self._processor, "apply_chat_template", None) is None @@ -572,10 +640,9 @@ class MLXInferenceBackend: len(prompt), image is not None, ) - # mlx_vlm.stream_generate forwards **kwargs into generate_step, which - # builds the sampler + logits_processors internally. - # GOTCHA: generate_step expects ``temperature=`` (long form); ``temp=`` - # silently falls into **kwargs and is ignored, stuck at greedy 0.0. + # stream_generate forwards **kwargs into generate_step (builds the + # sampler + logits_processors internally). GOTCHA: generate_step expects + # temperature= (long form); temp= is silently ignored, stuck at greedy 0.0. vlm_kwargs = dict( max_tokens = max_new_tokens, temperature = temperature, @@ -589,7 +656,7 @@ class MLXInferenceBackend: ) if presence_penalty: # Presence needs a custom processor: pass the full list (repetition + - # presence) instead of the repetition_penalty shortcut so both apply once. + # presence) instead of the repetition_penalty shortcut so both apply. from mlx_lm.sample_utils import make_logits_processors _vlm_processors = [] @@ -634,7 +701,7 @@ class MLXInferenceBackend: cancel_event = None, **gen_kwargs, ) -> Generator[str, None, None]: - # MLX LoRA adapter toggling not yet supported — generate normally + # MLX LoRA adapter toggling not yet supported; generate normally yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) def reset_generation_state(self): diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index cf5d24c367..4fa0d3ed26 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -50,6 +50,18 @@ _DISPATCH_DRAIN_TIMEOUT = 5.0 _UNLOAD_GEN_LOCK_TIMEOUT = 15.0 +class GenStreamError(str): + """A stream chunk carrying a real backend/generation error, not model text. + + Subclasses str so existing display/logging consumers are unaffected, while + callers that must abort a distributed run on error (raise_on_streamed_error) + can distinguish a real error from model output whose visible text starts with + "Error:" by checking isinstance(chunk, GenStreamError). + """ + + __slots__ = () + + class InferenceOrchestrator: """ Inference backend orchestrator — subprocess-based. @@ -482,13 +494,13 @@ class InferenceOrchestrator: initial_resp_queue = self._resp_queue while True: if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: - yield f"Error: {self._subprocess_crash_message(crash_context)}" + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") return resp = read_one(read_timeout) if resp is None: # Check subprocess health if not self._ensure_subprocess_alive(): - yield f"Error: {self._subprocess_crash_message(crash_context)}" + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") return continue @@ -498,7 +510,7 @@ class InferenceOrchestrator: # Subprocess-level error (no request_id); request-scoped failures # arrive as gen_error below. if rtype == "error" and not resp.get("request_id"): - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return if rtype == "token": @@ -513,7 +525,7 @@ class InferenceOrchestrator: stats_holder["stats"] = resp.get("stats") return elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return # ------------------------------------------------------------------ @@ -640,11 +652,11 @@ class InferenceOrchestrator: GPU work stays serialized; this only avoids orchestrator lock contention. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return # Latch the target model so the recheck below can detect a switch that completed # between _start_dispatcher and mailbox registration (mirrors the locked path's @@ -655,7 +667,7 @@ class InferenceOrchestrator: # so without this early-out a compare request would enqueue a generate on the # outgoing model and delay the switch. if self._unload_pending: - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under @@ -727,7 +739,7 @@ class InferenceOrchestrator: # _stop_dispatcher joins the dispatcher, which itself takes that lock. if orphaned_dispatcher: self._stop_dispatcher() - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return try: @@ -735,7 +747,7 @@ class InferenceOrchestrator: except RuntimeError as exc: with self._mailbox_lock: self._mailboxes.pop(request_id, None) - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return def read_mailbox(timeout): @@ -813,6 +825,59 @@ class InferenceOrchestrator: self._stop_dispatcher() return True + def share_distributed_object( + self, + obj, + timeout: Optional[float] = 300.0, + ): + """Share a small object through the worker's MLX distributed group.""" + if not self._ensure_subprocess_alive(): + raise RuntimeError("Inference subprocess is not running") + + self._wait_dispatcher_idle() + with self._mailbox_lock: + if self._mailboxes: + raise RuntimeError( + "Cannot share distributed objects while compare requests are active" + ) + request_id = str(uuid.uuid4()) + cmd = { + "type": "share_object", + "request_id": request_id, + "object": obj, + } + + with self._gen_lock: + self._send_cmd(cmd) + deadline = None if timeout is None else time.monotonic() + timeout + while deadline is None or time.monotonic() < deadline: + remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("sharing chat turn")) + continue + + rtype = resp.get("type", "") + rid = resp.get("request_id") + if rid and rid != request_id: + logger.debug( + "Skipping response for request_id=%s while sharing request_id=%s", + rid, + request_id, + ) + continue + if rtype == "shared": + return resp.get("object") + if rtype == "share_error": + raise RuntimeError(resp.get("error", "Failed to share object")) + if rtype == "error": + raise RuntimeError(resp.get("error", "Subprocess error")) + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for distributed object share") + # ------------------------------------------------------------------ # Public API — same interface as InferenceBackend # ------------------------------------------------------------------ @@ -828,6 +893,8 @@ class InferenceOrchestrator: approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, subject: Optional[str] = None, + tensor_parallel: bool = False, + mlx_distributed: bool = False, ) -> bool: """Load a model for inference. @@ -853,6 +920,11 @@ class InferenceOrchestrator: "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "subject": subject, "gpu_ids": gpu_ids, + "tensor_parallel": bool(tensor_parallel), + "mlx_distributed": bool(mlx_distributed), + "mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline") + if mlx_distributed + else None, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( gpu_ids, @@ -1338,11 +1410,11 @@ class InferenceOrchestrator: readers don't consume each other's tokens off the shared resp_queue. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return expected_model = self.active_model_name @@ -1359,7 +1431,7 @@ class InferenceOrchestrator: # so we never generate on the wrong one. if self._unload_pending or self.active_model_name != expected_model: # Won the lock handoff during a switch; don't start on the outgoing model. - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None @@ -1385,7 +1457,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( @@ -1544,10 +1616,10 @@ class InferenceOrchestrator: ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return expected_model = self.active_model_name @@ -1556,7 +1628,7 @@ class InferenceOrchestrator: # cleared or swapped the model while we waited. if self._unload_pending or self.active_model_name != expected_model: # Won the lock handoff during a switch; don't start on the outgoing model. - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return request_id = str(uuid.uuid4()) @@ -1583,7 +1655,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index d4b102e422..05dee39283 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker. from __future__ import annotations import base64 +import json from loggers import get_logger import os import queue as _queue @@ -26,6 +27,9 @@ from typing import Any logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +_SHARE_OBJECT_MAX_BYTES = 1 << 20 +_SHARE_OBJECT_ERROR_SIZE = -1 + # studio/backend root, prepended to sys.path so the spawned subprocess can # import the utils/core packages. _BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent) @@ -75,6 +79,17 @@ def _send_response(resp_queue: Any, response: dict) -> None: logger.error("Failed to send response: %s", exc) +def _encode_share_object(obj: Any) -> bytes: + data = json.dumps(obj, separators = (",", ":"), ensure_ascii = False).encode("utf-8") + if len(data) > _SHARE_OBJECT_MAX_BYTES: + raise ValueError("Distributed object share payload is too large") + return data + + +def _decode_share_object(data: Any) -> Any: + return json.loads(bytes(data.tolist()).decode("utf-8")) + + def _clean_token(value: str | None) -> str | None: """Normalize an HF token: blank or whitespace-only becomes None.""" return value if value and value.strip() else None @@ -329,14 +344,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) try: - success = backend.load_model( - config = mc, - max_seq_length = config.get("max_seq_length", 2048), - load_in_4bit = load_in_4bit, - hf_token = hf_token, - trust_remote_code = trust_remote_code, - gpu_ids = config.get("resolved_gpu_ids"), - ) + load_kwargs = { + "config": mc, + "max_seq_length": config.get("max_seq_length", 2048), + "load_in_4bit": load_in_4bit, + "hf_token": hf_token, + "trust_remote_code": trust_remote_code, + "gpu_ids": config.get("resolved_gpu_ids"), + } + if getattr(backend, "device", None) == "mlx": + load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode") + load_kwargs["distributed_group"] = config.get("_mlx_distributed_group") + success = backend.load_model(**load_kwargs) finally: heartbeat_stop.set() @@ -521,6 +540,67 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: ) +def _handle_share_object(backend, cmd: dict, resp_queue: Any) -> None: + """Share a small Python object across MLX distributed ranks.""" + request_id = cmd.get("request_id", "") + group = getattr(backend, "_distributed_group", None) + rank = int(getattr(backend, "_distributed_rank", 0) or 0) + world_size = int(getattr(backend, "_distributed_world_size", 1) or 1) + obj = cmd.get("object") + + try: + if group is None or world_size <= 1: + shared = obj + else: + import mlx.core as mx + if rank == 0: + if obj is None: + mx.eval(mx.distributed.all_sum(mx.array(0), group = group)) + shared = None + else: + try: + data = mx.array(_encode_share_object(obj), dtype = mx.uint8) + except Exception: + mx.eval( + mx.distributed.all_sum( + mx.array(_SHARE_OBJECT_ERROR_SIZE), + group = group, + ) + ) + raise + mx.eval(mx.distributed.all_sum(mx.array(data.size), group = group)) + mx.eval(mx.distributed.all_sum(data, group = group)) + shared = obj + else: + size = int(mx.distributed.all_sum(mx.array(0), group = group).item()) + if size == _SHARE_OBJECT_ERROR_SIZE: + raise RuntimeError("Failed to share distributed object") + if size == 0: + shared = None + else: + data = mx.zeros(size, dtype = mx.uint8) + data = mx.distributed.all_sum(data, group = group) + shared = _decode_share_object(data) + _send_response( + resp_queue, + { + "type": "shared", + "request_id": request_id, + "object": shared, + }, + ) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "share_error", + "request_id": request_id, + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + }, + ) + + def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None: """Handle TTS audio generation — returns WAV bytes + sample_rate.""" request_id = cmd.get("request_id", "") @@ -720,9 +800,29 @@ def run_inference_process( exc, ) try: - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed backend = MLXInferenceBackend() + if config.get("mlx_distributed"): + group, rank, size = _init_mlx_distributed() + config["_mlx_distributed_group"] = group + if size <= 1: + # A singleton group (MLX built without distributed support, + # or an invalid launch env/hostfile) would leave nonzero ranks + # looping forever on share_distributed_object. Fail the load + # instead of silently continuing without sharding. + raise RuntimeError( + "MLX distributed launch requested but initialized a singleton " + "group (size 1). Ensure the installed MLX has distributed " + "support and the launch environment/hostfile is valid, or run " + "without distributed." + ) + logger.info( + "MLX distributed initialized in worker: rank=%s size=%s mode=%s", + rank, + size, + config.get("mlx_parallel_mode"), + ) _send_response( resp_queue, {"type": "status", "message": "Loading model..."}, @@ -764,6 +864,8 @@ def run_inference_process( if _drain_skip_generate(cmd, resp_queue, drain_event): continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) @@ -977,6 +1079,9 @@ def run_inference_process( continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) + elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ce755ae1fb..d9901a5b2e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -40,6 +40,43 @@ def _positive_int_or_none(value: Any) -> Optional[int]: return value_int if value_int > 0 else None +def _nonnegative_int_or_none(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + value_int = int(value) + except (TypeError, ValueError): + return None + return value_int if value_int >= 0 else None + + +_MLX_MPI_DISTRIBUTED_ENV_PAIRS = ( + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +) + + +def _mlx_distributed_launch_detected() -> bool: + if _nonnegative_int_or_none(os.environ.get("MLX_RANK")) is not None: + world_size = _positive_int_or_none(os.environ.get("MLX_WORLD_SIZE")) + if world_size is not None and world_size > 1: + return True + return bool( + os.environ.get("MLX_HOSTFILE") + or os.environ.get("MLX_IBV_DEVICES") + or os.environ.get("MLX_JACCL_COORDINATOR") + or (os.environ.get("NCCL_HOST_IP") and os.environ.get("NCCL_PORT")) + ) + return any( + _nonnegative_int_or_none(os.environ.get(rank_env)) is not None + and (_positive_int_or_none(os.environ.get(size_env)) or 0) > 1 + for rank_env, size_env in _MLX_MPI_DISTRIBUTED_ENV_PAIRS + ) + + def _install_httpcore_asyncgen_silencer() -> None: """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. @@ -3426,6 +3463,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre status_code = 400, detail = "gpu_ids is not supported for GGUF models yet.", ) + if not config.is_gguf and _mlx_distributed_launch_detected(): + raise HTTPException( + status_code = 400, + detail = ( + "Studio does not support distributed MLX inference under " + "mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio " + "without the distributed launcher." + ), + ) # Effective quantization (LoRA can flip 4-bit -> 16-bit); guard + load reuse it. effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index ac4088fb25..55a3198a6b 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -4,6 +4,8 @@ import sys import types from types import SimpleNamespace +import pytest + class _DummyMetal: @staticmethod @@ -185,6 +187,129 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri assert isinstance(backend._tokenizer, _DummyTokenizer) +def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch): + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + from core.inference.mlx_inference import MLXInferenceBackend + + group = SimpleNamespace(size = lambda: 2, rank = lambda: 0) + config = SimpleNamespace(identifier = "fake/vlm", is_vision = True, is_lora = False) + for mode, group_key in (("tensor", "tensor_group"), ("pipeline", "pipeline_group")): + calls.clear() + assert MLXInferenceBackend().load_model(config, parallel_mode = mode, distributed_group = group) + _, kwargs = calls.pop() + assert kwargs["text_only"] is False and kwargs[group_key] is group + + calls.clear() + singleton = SimpleNamespace(size = lambda: 1, rank = lambda: 0) + assert MLXInferenceBackend().load_model( + config, parallel_mode = "tensor", distributed_group = singleton + ) + assert not {"tensor_group", "pipeline_group"} & set(calls.pop()[1]) + + config = SimpleNamespace(identifier = "fake/adapter", is_vision = False, is_lora = True) + with pytest.raises(ValueError, match = "LoRA adapter repos"): + MLXInferenceBackend().load_model(config, parallel_mode = "tensor", distributed_group = group) + + +@pytest.mark.parametrize("accepts_backend", (True, False)) +def test_mlx_distributed_init_selects_jaccl_backend(monkeypatch, accepts_backend): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import _init_mlx_distributed + + group = SimpleNamespace(rank = lambda: 1, size = lambda: 2) + calls = [] + + def _init(**kwargs): + calls.append(kwargs) + if kwargs and not accepts_backend: + raise TypeError("backend keyword unsupported") + return group + + sys.modules["mlx.core"].distributed = SimpleNamespace(init = _init) + monkeypatch.setenv("MLX_JACCL_COORDINATOR", "127.0.0.1:12345") + monkeypatch.setenv("MLX_IBV_DEVICES", "/tmp/devices.json") + + assert _init_mlx_distributed() == (group, 1, 2) + assert calls == ([{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}]) + + +def test_worker_share_object_receives_distributed_payload(monkeypatch): + from core.inference import worker + + shared_obj = {"type": "turn", "text": "hi"} + payload = worker._encode_share_object(shared_obj) + + def _array(value): + val = value.item() if hasattr(value, "item") else value + return SimpleNamespace( + item = lambda: val, + tolist = lambda: list(val) if hasattr(val, "__iter__") else [val], + ) + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.uint8 = "uint8" + mlx_core.array = _array + mlx_core.zeros = lambda *_a, **_k: _array([]) + + def _all_sum(value, group = None): + value = value.item() if hasattr(value, "item") else value + return _array(len(payload)) if value == 0 else _array(payload) + + mlx_core.distributed = SimpleNamespace(all_sum = _all_sum) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 1, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": None}, + SimpleNamespace(put = responses.append), + ) + + response = responses[0] + assert response["object"] == shared_obj + + +def test_worker_share_object_oversize_notifies_peers(monkeypatch): + from core.inference import worker + + calls = [] + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.array = lambda value, **_kwargs: SimpleNamespace(item = lambda: value) + mlx_core.eval = lambda value: value + mlx_core.distributed = SimpleNamespace( + all_sum = lambda value, group = None: calls.append(value.item()) or value + ) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + monkeypatch.setattr(worker, "_SHARE_OBJECT_MAX_BYTES", 8) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 0, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": {"text": "too long"}}, + SimpleNamespace(put = responses.append), + ) + + assert calls == [worker._SHARE_OBJECT_ERROR_SIZE] + assert responses[0]["type"] == "share_error" + + # Regression: generate_chat_response must accept the four template kwargs # (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route # layer can forward UI toggles. The old signature raised diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index e8a3414d73..c3b188710e 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -4,9 +4,11 @@ """Model loading and streaming shared by `inference` and `chat`.""" import asyncio +import json import os import re import sys +from contextlib import contextmanager, redirect_stderr, redirect_stdout from pathlib import Path from typing import List, Optional @@ -14,10 +16,18 @@ import typer _THINK_OPEN = "" _THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", re.DOTALL) +_STREAMED_ERROR_PREFIX = "Error: " # Cloudflare (in front of remote Studio proxies like RunPod) 403s the default # "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request. _USER_AGENT = "unsloth-cli" +_MPI_ENV_PAIRS = ( + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +) # Built lazily; urllib stays function-local to match this module. _no_redirect_opener = None @@ -61,6 +71,108 @@ def configure_quiet_logging() -> None: os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") +def _parse_nonnegative_int(value: Optional[str]) -> Optional[int]: + if value is None: + return None + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed if parsed >= 0 else None + + +def _first_mpi_env_pair() -> tuple[Optional[int], Optional[int]]: + for rank_name, size_name in _MPI_ENV_PAIRS: + rank = _parse_nonnegative_int(os.environ.get(rank_name)) + world_size = _parse_nonnegative_int(os.environ.get(size_name)) + if rank is not None and world_size is not None and world_size > 1 and rank < world_size: + return rank, world_size + return None, None + + +def _json_rank_count_from_env(name: str) -> Optional[int]: + value = os.environ.get(name) + if not value: + return None + try: + if value.lstrip().startswith(("[", "{")): + data = json.loads(value) + else: + with open(value, "r") as f: + data = json.load(f) + except (OSError, json.JSONDecodeError): + return None + if isinstance(data, list): + return len(data) + if isinstance(data, dict) and isinstance(data.get("hosts"), list): + return len(data["hosts"]) + return None + + +def mlx_distributed_info() -> tuple[bool, int, Optional[int]]: + """Return launch-context metadata without initializing MLX distributed.""" + rank = _parse_nonnegative_int(os.environ.get("MLX_RANK")) + world_size = _parse_nonnegative_int(os.environ.get("MLX_WORLD_SIZE")) + if rank is not None: + if ( + world_size is not None + and world_size > 1 + and rank < world_size + and os.environ.get("NCCL_HOST_IP") + and os.environ.get("NCCL_PORT") + ): + return True, rank, world_size + inferred_size = _json_rank_count_from_env("MLX_HOSTFILE") + if inferred_size is not None and inferred_size > 1 and rank < inferred_size: + return True, rank, inferred_size + inferred_size = _json_rank_count_from_env("MLX_IBV_DEVICES") + if ( + inferred_size is not None + and inferred_size > 1 + and rank < inferred_size + and os.environ.get("MLX_JACCL_COORDINATOR") + ): + return True, rank, inferred_size + return False, 0, None + + mpi_rank, mpi_world_size = _first_mpi_env_pair() + return mpi_rank is not None, mpi_rank or 0, mpi_world_size + + +def mlx_distributed_uses_mpi() -> bool: + """Whether the current distributed context was launched through MPI.""" + return ( + _parse_nonnegative_int(os.environ.get("MLX_RANK")) is None + and _first_mpi_env_pair()[0] is not None + ) + + +@contextmanager +def quiet_if_nonzero_mlx_rank(): + """Silence parent and child-process stdout/stderr on nonzero ranks.""" + if mlx_distributed_info()[1] == 0: + yield + return + + sys.stdout.flush() + sys.stderr.flush() + saved_stdout_fd = os.dup(1) + saved_stderr_fd = os.dup(2) + with open(os.devnull, "w") as devnull: + try: + os.dup2(devnull.fileno(), 1) + os.dup2(devnull.fileno(), 2) + with redirect_stdout(devnull), redirect_stderr(devnull): + yield + finally: + sys.stdout.flush() + sys.stderr.flush() + os.dup2(saved_stdout_fd, 1) + os.dup2(saved_stderr_fd, 2) + os.close(saved_stdout_fd) + os.close(saved_stderr_fd) + + def visible_text(text: str, show_thinking: bool) -> str: if show_thinking: return text @@ -120,6 +232,21 @@ def collect_stream(stream, show_thinking: bool) -> str: return visible_text(raw, show_thinking) +def raise_on_streamed_error(stream): + # Match real backend errors by type (GenStreamError), not the "Error:" text + # prefix, so a completion whose text opens with "Error:" is not misread as a + # failure that aborts a distributed run. + try: + ensure_studio_backend_path() + from core.inference.orchestrator import GenStreamError + except Exception: + GenStreamError = None + for chunk in stream: + if GenStreamError is not None and isinstance(chunk, GenStreamError): + raise RuntimeError(str(chunk)[len(_STREAMED_ERROR_PREFIX) :].strip() or "Unknown error") + yield chunk + + def render_columns( left_label: str, left_text: str, @@ -200,6 +327,19 @@ class ChatBackend: except Exception: pass + def share_distributed_object( + self, + obj, + *, + timeout = 300.0, + ): + if self._kind != "unsloth" or not hasattr(self._backend, "share_distributed_object"): + raise RuntimeError( + "Distributed MLX chat requires the Unsloth MLX backend; " + f"backend '{self._kind}' cannot broadcast chat turns." + ) + return self._backend.share_distributed_object(obj, timeout = timeout) + def resolve_model_config(model: str, *, hf_token: Optional[str]): ensure_studio_backend_path() @@ -293,36 +433,59 @@ def load_chat_backend( fresh_backend uses a private orchestrator so a second model (compare's base column) can run alongside the main one. """ - if model_config is None: - model_config = resolve_model_config(model, hf_token = hf_token) + with quiet_if_nonzero_mlx_rank(): + is_mlx_distributed, rank, _world_size = mlx_distributed_info() + if model_config is None: + model_config = resolve_model_config(model, hf_token = hf_token) - typer.echo(f"Loading {model}", err = True) + if is_mlx_distributed and model_config.is_gguf: + if rank == 0: + typer.echo( + "Distributed MLX inference does not support GGUF/llama.cpp models. " + "Use a non-GGUF MLX model under mlx.launch, or run GGUF without " + "mlx.launch.", + err = True, + ) + raise typer.Exit(code = 1) - if model_config.is_gguf: - return _load_gguf_backend( - model_config, - hf_token = hf_token, - max_seq_length = max_seq_length, - tensor_parallel = tensor_parallel, - llama_extra_args = llama_extra_args, - ) + if rank == 0: + typer.echo(f"Loading {model}", err = True) - if fresh_backend: - ensure_studio_backend_path() - from core.inference import InferenceOrchestrator - backend = InferenceOrchestrator() - else: - ensure_studio_backend_path() - from core.inference import get_inference_backend - backend = get_inference_backend() - if not backend.load_model( - config = model_config, - max_seq_length = max_seq_length, - load_in_4bit = load_in_4bit, - hf_token = hf_token, - ): - typer.echo("Model load failed", err = True) - raise typer.Exit(code = 1) + if model_config.is_gguf: + return _load_gguf_backend( + model_config, + hf_token = hf_token, + max_seq_length = max_seq_length, + tensor_parallel = tensor_parallel, + llama_extra_args = llama_extra_args, + ) + + if fresh_backend: + ensure_studio_backend_path() + from core.inference import InferenceOrchestrator + backend = InferenceOrchestrator() + else: + ensure_studio_backend_path() + from core.inference import get_inference_backend + backend = get_inference_backend() + try: + loaded = backend.load_model( + config = model_config, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + hf_token = hf_token, + tensor_parallel = tensor_parallel, + mlx_distributed = is_mlx_distributed, + ) + except Exception as exc: + if not is_mlx_distributed: + raise + if rank == 0: + typer.echo(str(exc) or "Model load failed", err = True) + raise typer.Exit(code = 1) + if not loaded: + typer.echo("Model load failed", err = True) + raise typer.Exit(code = 1) return ChatBackend("unsloth", backend) diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py index d3bfbbf96b..bc4a72f36c 100644 --- a/unsloth_cli/commands/chat.py +++ b/unsloth_cli/commands/chat.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import sys from typing import List, Optional import typer @@ -12,6 +13,10 @@ from unsloth_cli._inference import ( connect_studio_server, ensure_studio_backend_path, load_chat_backend, + mlx_distributed_info, + mlx_distributed_uses_mpi, + quiet_if_nonzero_mlx_rank, + raise_on_streamed_error, render_columns, resolve_model_config, stream_markdown, @@ -107,6 +112,20 @@ def _compare_needs_second_model() -> bool: return False +def _drain_available_stdin() -> None: + """Drain already-buffered launcher stdin on nonzero distributed ranks.""" + try: + import os + from select import select + + fd = sys.stdin.fileno() + while select([fd], [], [], 0)[0]: + if not os.read(fd, 8192): + break + except Exception: + return + + def _pick_trained_model(console) -> str: ensure_studio_backend_path() from utils.models import scan_trained_models @@ -158,7 +177,8 @@ def chat( "--tensor-parallel/--no-tensor-parallel", help = ( "Split a GGUF across GPUs by tensor (--split-mode tensor) instead " - "of by layer. Ignored for non-GGUF models." + "of by layer. Under non-MPI mlx.launch, select MLX tensor " + "parallel mode instead of pipeline mode." ), ), llama_extra_args: Optional[List[str]] = typer.Option( @@ -195,15 +215,43 @@ def chat( console = Console() err = Console(stderr = True) + is_mlx_distributed, rank, _world_size = mlx_distributed_info() + should_print = rank == 0 + + if is_mlx_distributed and mlx_distributed_uses_mpi(): + if should_print: + err.print( + "Distributed `unsloth chat` with MPI needs rank-0 prompt broadcast, " + "which is not enabled yet. Use a non-MPI MLX launcher backend " + "such as ring/JACCL for now.", + style = "red", + markup = False, + ) + raise typer.Exit(code = 1) if model is None: + if is_mlx_distributed: + if should_print: + err.print( + "Distributed `unsloth chat` requires an explicit model id or path.", + style = "red", + markup = False, + ) + raise typer.Exit(code = 1) model = _pick_trained_model(console) # Resolve first so --compare can be rejected before the slow load. - model_config = resolve_model_config(model, hf_token = hf_token) + with quiet_if_nonzero_mlx_rank(): + model_config = resolve_model_config(model, hf_token = hf_token) compare_blocked = _compare_blocked_reason(model_config) + if is_mlx_distributed: + compare_blocked = ( + "distributed MLX chat does not support compare mode yet because it " + "would need a second distributed worker group on the same ranks" + ) if compare and compare_blocked: - err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False) + if should_print: + err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False) raise typer.Exit(code = 1) load_opts = dict( @@ -215,9 +263,11 @@ def chat( ) # Prefer a running Studio server: instant starts, model shared with the UI. - chat_backend = None if no_server else connect_studio_server(model, **load_opts) + chat_backend = ( + None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts) + ) server_mode = chat_backend is not None - if server_mode: + if server_mode and should_print: console.print( "(Studio server connected — model stays warm after /exit)", style = "bright_black", @@ -242,23 +292,26 @@ def chat( return True base_id = model_config.base_model if not base_id: - console.print( - "(compare unavailable: this adapter doesn't record its base model)", - style = "yellow", - ) + if should_print: + console.print( + "(compare unavailable: this adapter doesn't record its base model)", + style = "yellow", + ) return False - console.print( - f"(loading base model {base_id} for compare — keeps two models in memory)", - style = "bright_black", - markup = False, - ) + if should_print: + console.print( + f"(loading base model {base_id} for compare — keeps two models in memory)", + style = "bright_black", + markup = False, + ) try: # Use the same precision as the tuned model for fair comparison base_load_opts = dict(load_opts) # Copy original options base_load_opts["load_in_4bit"] = _get_base_load_in_4bit(model_config) base_backend = load_chat_backend(base_id, fresh_backend = True, **base_load_opts) except Exception as exc: - err.print(f"(base model load failed: {exc})", style = "red", markup = False) + if should_print: + err.print(f"(base model load failed: {exc})", style = "red", markup = False) return False return True @@ -267,7 +320,7 @@ def chat( def generate(backend = None, use_adapter = None): # Reads messages and show_thinking live, so /reset and /think apply. - return (backend or chat_backend).stream( + stream = (backend or chat_backend).stream( messages, system_prompt = system_prompt, temperature = temperature, @@ -278,22 +331,48 @@ def chat( enable_thinking = show_thinking, use_adapter = use_adapter, ) + return raise_on_streamed_error(stream) if is_mlx_distributed else stream - console.print() - console.print(f"Chatting with {name}", style = "bold green", markup = False) - console.print(_HELP, style = "bright_black") + if should_print: + console.print() + console.print(f"Chatting with {name}", style = "bold green", markup = False) + console.print(_HELP, style = "bright_black") # legacy_windows: pre-VT consoles print raw ANSI as ←[1;36m garbage. - you_prompt = _you_prompt(console.is_terminal and not console.legacy_windows) + you_prompt = ( + _you_prompt(console.is_terminal and not console.legacy_windows) if should_print else "" + ) assistant_label = "[bold magenta]Assistant:[/bold magenta]" try: while True: - try: - user = input(you_prompt).strip() - except (EOFError, KeyboardInterrupt): - console.print() - break + if should_print: + try: + user = input(you_prompt).strip() + except (EOFError, KeyboardInterrupt): + if should_print: + console.print() + user = "/exit" + turn = {"type": "turn", "text": user} + else: + turn = None + + if is_mlx_distributed: + try: + turn = chat_backend.share_distributed_object(turn, timeout = None) + if not should_print: + _drain_available_stdin() + except Exception as exc: + if should_print: + err.print( + f"\n(error sharing chat turn: {exc})", + style = "red", + markup = False, + ) + raise typer.Exit(code = 1) + if not turn: + continue + user = str(turn.get("text", "")).strip() if not user: continue @@ -301,55 +380,69 @@ def chat( break if user == "/reset": messages = [] - console.print("(history cleared)", style = "bright_black") + if should_print: + console.print("(history cleared)", style = "bright_black") continue if user == "/think": show_thinking = not show_thinking - state = "on" if show_thinking else "off" - console.print(f"(thinking {state})", style = "bright_black") + if should_print: + state = "on" if show_thinking else "off" + console.print(f"(thinking {state})", style = "bright_black") continue if user == "/compare": if compare_blocked: - console.print(f"(compare unavailable: {compare_blocked})", style = "yellow") + if should_print: + console.print(f"(compare unavailable: {compare_blocked})", style = "yellow") continue if not compare_mode and dual_compare and not load_base_for_compare(): continue compare_mode = not compare_mode - state = "on" if compare_mode else "off" - console.print(f"(compare {state})", style = "bright_black") + if should_print: + state = "on" if compare_mode else "off" + console.print(f"(compare {state})", style = "bright_black") continue if user in ("/help", "/?"): - console.print(_HELP, style = "bright_black") + if should_print: + console.print(_HELP, style = "bright_black") continue messages.append({"role": "user", "content": user}) try: if compare_mode: - console.print("(comparing base vs tuned…)", style = "bright_black") + if should_print: + console.print("(comparing base vs tuned…)", style = "bright_black") if dual_compare: base_text = collect_stream(generate(backend = base_backend), show_thinking) tuned_text = collect_stream(generate(), show_thinking) else: base_text = collect_stream(generate(use_adapter = False), show_thinking) tuned_text = collect_stream(generate(use_adapter = True), show_thinking) - console.print() - render_columns( - "base", base_text, f"{name} (tuned)", tuned_text, console = console - ) + if should_print: + console.print() + render_columns( + "base", base_text, f"{name} (tuned)", tuned_text, console = console + ) # History continues as the tuned model; base is just the reference. answer = tuned_text else: - console.print(assistant_label) - answer = stream_markdown(generate(), show_thinking, console = console) + if should_print: + console.print(assistant_label) + answer = stream_markdown(generate(), show_thinking, console = console) + else: + answer = collect_stream(generate(), show_thinking) except KeyboardInterrupt: # Ctrl-C aborts this answer only; drop the unanswered turn. - console.print("\n(interrupted)", style = "bright_black") + if should_print: + console.print("\n(interrupted)", style = "bright_black") messages.pop() continue except Exception as exc: - err.print(f"\n(error: {exc})", style = "red", markup = False) + if should_print: + err.print(f"\n(error: {exc})", style = "red", markup = False) messages.pop() + if is_mlx_distributed: + raise typer.Exit(code = 1) continue messages.append( @@ -359,4 +452,5 @@ def chat( chat_backend.close() if base_backend is not None: base_backend.close() - err.print("\nBye.", style = "bright_black") + if should_print: + err.print("\nBye.", style = "bright_black") diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py index 1401fa9e9e..84a126163e 100644 --- a/unsloth_cli/commands/inference.py +++ b/unsloth_cli/commands/inference.py @@ -6,9 +6,13 @@ from typing import List, Optional import typer from unsloth_cli._inference import ( + collect_stream, configure_quiet_logging, connect_studio_server, load_chat_backend, + mlx_distributed_info, + mlx_distributed_uses_mpi, + raise_on_streamed_error, stream_to_stdout, ) @@ -36,7 +40,8 @@ def inference( "--tensor-parallel/--no-tensor-parallel", help = ( "Split a GGUF across GPUs by tensor (--split-mode tensor) instead " - "of by layer. Ignored for non-GGUF models." + "of by layer. Under non-MPI mlx.launch, select MLX tensor " + "parallel mode instead of pipeline mode." ), ), llama_extra_args: Optional[List[str]] = typer.Option( @@ -69,8 +74,20 @@ def inference( if not verbose: configure_quiet_logging() - # A running Studio server keeps the model warm between runs, which is - # exactly what a one-shot command wants. + is_mlx_distributed, rank, _world_size = mlx_distributed_info() + if is_mlx_distributed and mlx_distributed_uses_mpi(): + if rank == 0: + typer.echo( + "Distributed `unsloth inference` with MPI is not supported by " + "the current subprocess backend. Use a non-MPI MLX launcher " + "backend such as ring/JACCL for now.", + err = True, + ) + raise typer.Exit(code = 1) + + # A running Studio server keeps the model warm between runs. Under + # mlx.launch, every rank must enter the local MLX path instead of rank 0 + # alone talking to a server. load_opts = dict( hf_token = hf_token, max_seq_length = max_seq_length, @@ -78,7 +95,9 @@ def inference( tensor_parallel = tensor_parallel, llama_extra_args = llama_extra_args, ) - chat_backend = None if no_server else connect_studio_server(model, **load_opts) + chat_backend = ( + None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts) + ) if chat_backend is None: chat_backend = load_chat_backend(model, **load_opts) try: @@ -92,7 +111,23 @@ def inference( repetition_penalty = repetition_penalty, enable_thinking = think, ) - typer.echo("Assistant:") - stream_to_stdout(stream, show_thinking = think) + if is_mlx_distributed: + stream = raise_on_streamed_error(stream) + if rank == 0: + typer.echo("Assistant:") + try: + stream_to_stdout(stream, show_thinking = think) + except RuntimeError as exc: + if not is_mlx_distributed: + raise + typer.echo(f"Error: {exc}", err = True) + raise typer.Exit(code = 1) + else: + try: + collect_stream(stream, show_thinking = think) + except RuntimeError: + if not is_mlx_distributed: + raise + raise typer.Exit(code = 1) finally: chat_backend.close() diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py index 013b7c5a81..56633408fb 100644 --- a/unsloth_cli/tests/test_inference_chat.py +++ b/unsloth_cli/tests/test_inference_chat.py @@ -26,6 +26,8 @@ from unsloth_cli._inference import ( ChatBackend, HttpChatBackend, collect_stream, + mlx_distributed_info, + mlx_distributed_uses_mpi, render_columns, visible_text, ) @@ -39,6 +41,16 @@ class _FakeConfig: path = None +_EXPECTED_MPI_ENV_PAIRS = [ + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +] +_IGNORED_DISTRIBUTED_ENV_PAIRS = [("SLURM_PROCID", "SLURM_NTASKS")] + + def _chat_app(): cli = typer.Typer() cli.command()(chatmod.chat) @@ -53,6 +65,39 @@ def _inference_app(): return cli +def _clear_mlx_distributed_env(monkeypatch): + for name in ( + "MLX_RANK", + "MLX_HOSTFILE", + "MLX_WORLD_SIZE", + "MLX_IBV_DEVICES", + "MLX_JACCL_COORDINATOR", + "NCCL_HOST_IP", + "NCCL_PORT", + *(rank for rank, _size in _EXPECTED_MPI_ENV_PAIRS + _IGNORED_DISTRIBUTED_ENV_PAIRS), + *(size for _rank, size in _EXPECTED_MPI_ENV_PAIRS + _IGNORED_DISTRIBUTED_ENV_PAIRS), + ): + monkeypatch.delenv(name, raising = False) + + +def _set_mlx_nccl_env( + monkeypatch, + *, + rank: str = "0", + size: str = "2", +): + monkeypatch.setenv("MLX_RANK", rank) + monkeypatch.setenv("MLX_WORLD_SIZE", size) + monkeypatch.setenv("NCCL_HOST_IP", "127.0.0.1") + monkeypatch.setenv("NCCL_PORT", "12345") + + +@pytest.fixture(autouse = True) +def _isolate_mlx_distributed_env(monkeypatch): + _clear_mlx_distributed_env(monkeypatch) + monkeypatch.delenv("HF_TOKEN", raising = False) + + def test_visible_text_passthrough_when_shown(): text = "reasoninganswer" assert visible_text(text, show_thinking = True) == text @@ -101,6 +146,46 @@ def test_inference_exposes_gguf_runtime_options(): assert "--llama-extra-arg" in (getattr(extra, "param_decls", None) or []) +def test_mlx_distributed_info_reads_launch_env(monkeypatch, tmp_path): + _clear_mlx_distributed_env(monkeypatch) + assert mlx_distributed_info() == (False, 0, None) + assert mlx_distributed_uses_mpi() is False + + monkeypatch.setenv("MLX_RANK", "1") + monkeypatch.setenv("MLX_WORLD_SIZE", "2") + assert mlx_distributed_info() == (False, 0, None) + monkeypatch.setenv("NCCL_HOST_IP", "127.0.0.1") + monkeypatch.setenv("NCCL_PORT", "12345") + assert mlx_distributed_info() == (True, 1, 2) + assert mlx_distributed_uses_mpi() is False + + _clear_mlx_distributed_env(monkeypatch) + ring_hostfile = tmp_path / "ring.json" + ring_hostfile.write_text('[["127.0.0.1:5000"], ["127.0.0.1:5001"]]\n') + monkeypatch.setenv("MLX_RANK", "0") + monkeypatch.setenv("MLX_HOSTFILE", str(ring_hostfile)) + assert mlx_distributed_info() == (True, 0, 2) + assert mlx_distributed_uses_mpi() is False + + _clear_mlx_distributed_env(monkeypatch) + monkeypatch.setenv("MLX_RANK", "1") + monkeypatch.setenv("MLX_IBV_DEVICES", '[["node-a"], ["node-b"]]') + monkeypatch.setenv("MLX_JACCL_COORDINATOR", "node-a:12345") + assert mlx_distributed_info() == (True, 1, 2) + assert mlx_distributed_uses_mpi() is False + + _clear_mlx_distributed_env(monkeypatch) + monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "1") + monkeypatch.setenv("OMPI_COMM_WORLD_SIZE", "2") + assert mlx_distributed_info() == (True, 1, 2) + assert mlx_distributed_uses_mpi() is True + + _clear_mlx_distributed_env(monkeypatch) + monkeypatch.setenv("MLX_RANK", "bad") + monkeypatch.setenv("MLX_WORLD_SIZE", "-3") + assert mlx_distributed_info() == (False, 0, None) + + def test_chat_command_is_registered_with_options(): params = inspect.signature(chatmod.chat).parameters assert "model" in params @@ -751,7 +836,6 @@ def test_chat_server_mode_compare_loads_base_locally(monkeypatch): assert result.exit_code == 0, result.output assert "(compare on)" in result.output - # Only the base model loaded locally, on its own private backend. assert base_loads == [("fake/base", True)] assert streamed == ["base", "tuned"] assert set(closed) == {"http", "base"} @@ -785,6 +869,248 @@ def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch): assert result.exit_code == 0, result.output assert loads == [("tuned-run", False), ("fake/base", True)] - # Both models answered the turn, via plain generation (no adapter toggle). assert ("base", None) in streamed and ("tuned", None) in streamed assert set(closed) == {"tuned", "base"} + + +@pytest.mark.parametrize( + ("chunk_kind", "expected_exit"), + [ + ("answer", 0), + ("model_text_error", 0), + ("real_error", 1), + ], +) +def test_inference_under_mlx_launch_handles_stream(monkeypatch, chunk_kind, expected_exit): + from unsloth_cli.commands import inference as infermod + from unsloth_cli._inference import ensure_studio_backend_path + + ensure_studio_backend_path() + from core.inference.orchestrator import GenStreamError + + if chunk_kind == "answer": + chunks = ["answer"] + elif chunk_kind == "model_text_error": + # Model output whose visible text starts with "Error:" must not abort. + chunks = ["Error: printed by the model, not a backend failure"] + else: + chunks = [GenStreamError("Error: generation failed")] + + loads, closed = [], [] + + class _FakeBackend: + def stream(self, messages, **kwargs): + return iter(chunks) + + def close(self): + closed.append(True) + + _set_mlx_nccl_env(monkeypatch, rank = "0") + monkeypatch.setattr( + infermod, + "connect_studio_server", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")), + ) + monkeypatch.setattr( + infermod, + "load_chat_backend", + lambda model, **kwargs: (loads.append((model, kwargs)), _FakeBackend())[1], + ) + + result = CliRunner().invoke( + _inference_app(), + ["fake-model", "hello", "--tensor-parallel"], + ) + + assert result.exit_code == expected_exit, result.output + assert loads[0][1]["tensor_parallel"] is True + if chunk_kind == "real_error": + assert "generation failed" in result.output + + +def test_chat_under_mlx_launch_nonzero_rank_drains_stdin(monkeypatch): + drains, closed = [], [] + turns = iter( + [ + {"type": "turn", "text": "hi"}, + {"type": "turn", "text": "/exit"}, + ] + ) + + class _FakeChatBackend: + def share_distributed_object( + self, + obj, + *, + timeout = 300.0, + ): + assert obj is None + return next(turns) + + def stream(self, messages, **kwargs): + return iter(["hidden"]) + + def close(self): + closed.append(True) + + _set_mlx_nccl_env(monkeypatch, rank = "1") + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr( + chatmod, + "connect_studio_server", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")), + ) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend()) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + monkeypatch.setattr(chatmod, "_drain_available_stdin", lambda: drains.append(True)) + + result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n") + + assert result.exit_code == 0, result.output + assert "Chatting with" not in result.output + assert drains == [True, True] + assert closed == [True] + + +def test_chat_under_mlx_launch_rank0_bypasses_studio_and_prints(monkeypatch): + loads, shares, closed = [], [], [] + + class _FakeChatBackend: + def share_distributed_object( + self, + obj, + *, + timeout = 300.0, + ): + shares.append((obj, timeout)) + return obj + + def stream(self, messages, **kwargs): + return iter(["hello"]) + + def close(self): + closed.append(True) + + _set_mlx_nccl_env(monkeypatch, rank = "0") + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr( + chatmod, + "connect_studio_server", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")), + ) + monkeypatch.setattr( + chatmod, + "load_chat_backend", + lambda model, **kwargs: (loads.append((model, kwargs)), _FakeChatBackend())[1], + ) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + + result = CliRunner().invoke( + _chat_app(), + ["fake-model", "--tensor-parallel"], + input = "hi\n/exit\n", + ) + + assert result.exit_code == 0, result.output + assert "Chatting with fake-model" in result.output + assert "hello" in result.output + assert loads and loads[0][0] == "fake-model" + assert loads[0][1]["tensor_parallel"] is True + assert shares == [ + ({"type": "turn", "text": "hi"}, None), + ({"type": "turn", "text": "/exit"}, None), + ] + + +@pytest.mark.parametrize( + ("stream_error", "expected_exit"), + [("exception", 1), ("chunk", 1), ("model_text", 0)], +) +def test_chat_under_mlx_launch_exits_on_generation_error(monkeypatch, stream_error, expected_exit): + from unsloth_cli._inference import ensure_studio_backend_path + + ensure_studio_backend_path() + from core.inference.orchestrator import GenStreamError + + closed = [] + + class _FakeChatBackend: + def share_distributed_object( + self, + obj, + *, + timeout = 300.0, + ): + return obj + + def stream(self, messages, **kwargs): + if stream_error == "exception": + raise RuntimeError("generation failed") + if stream_error == "model_text": + # Plain model text starting with "Error:" must not abort the run. + return iter(["Error: printed by the model"]) + return iter([GenStreamError("Error: generation failed")]) + + def close(self): + closed.append(True) + + _set_mlx_nccl_env(monkeypatch, rank = "0") + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr( + chatmod, + "connect_studio_server", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")), + ) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend()) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + + result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n") + + assert result.exit_code == expected_exit + if expected_exit: + assert "generation failed" in result.output + assert closed == [True] + + +def test_load_chat_backend_forwards_mlx_distributed_options(monkeypatch): + import unsloth_cli._inference as inference + + calls = [] + + class _FakeBackend: + def load_model(self, **kwargs): + calls.append(kwargs) + return True + + class _FakeModelConfig: + is_gguf = False + + @classmethod + def from_identifier(cls, **_kwargs): + return cls() + + fake_backend = _FakeBackend() + fake_inference = types.ModuleType("core.inference") + fake_inference.get_inference_backend = lambda: fake_backend + fake_utils = types.ModuleType("utils") + fake_utils.__path__ = [] + fake_models = types.ModuleType("utils.models") + fake_models.ModelConfig = _FakeModelConfig + + _set_mlx_nccl_env(monkeypatch, rank = "0") + monkeypatch.setitem(sys.modules, "core", types.ModuleType("core")) + monkeypatch.setitem(sys.modules, "core.inference", fake_inference) + monkeypatch.setitem(sys.modules, "utils", fake_utils) + monkeypatch.setitem(sys.modules, "utils.models", fake_models) + monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None) + + inference.load_chat_backend( + "fake-model", + hf_token = None, + max_seq_length = 2048, + load_in_4bit = True, + tensor_parallel = True, + ) + + assert calls[0]["tensor_parallel"] is True + assert calls[0]["mlx_distributed"] is True From 934f879043b289bc6fd699ccfc13f0f363cd655d Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 8 Jul 2026 18:25:50 +0800 Subject: [PATCH 063/113] feat(mlx): route trainer callbacks (#6929) --- tests/python/test_mlx_public_trainer_api.py | 58 +++++++++++++++++---- unsloth/__init__.py | 50 ++++++++++++++++++ 2 files changed, 99 insertions(+), 9 deletions(-) diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py index 2c33f86af1..89f304c76d 100644 --- a/tests/python/test_mlx_public_trainer_api.py +++ b/tests/python/test_mlx_public_trainer_api.py @@ -115,6 +115,9 @@ def test_mlx_training_arguments_accept_trl_style_kwargs(): def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): """Implemented and falsey inert compatibility kwargs should stay quiet.""" unsloth = _import_mlx_unsloth() + supported_eval_kwargs = {} + if "eval_strategy" in unsloth._MLX_TRAINING_CONFIG_FIELDS: + supported_eval_kwargs = {"eval_strategy": "no", "eval_delay": 1} with warnings.catch_warnings(record = True) as caught: warnings.simplefilter("always") @@ -125,12 +128,16 @@ def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): remove_unused_columns = False, assistant_only_loss = False, completion_only_loss = False, + **supported_eval_kwargs, ) assert args.warmup_steps == 2 assert args.padding_free is False assert args.remove_unused_columns is False assert args.completion_only_loss is False + if supported_eval_kwargs: + assert args.eval_strategy == "no" + assert args.eval_delay == 1 assert caught == [] @@ -705,17 +712,10 @@ def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs(): ) -def test_mlx_trainer_rejects_metrics_and_callbacks(): - """Trainer hooks should fail because MLXTrainer cannot honor them yet.""" +def test_mlx_trainer_rejects_compute_metrics(): + """compute_metrics is still unsupported by MLXTrainer.""" unsloth = _import_mlx_unsloth() - with pytest.raises(NotImplementedError, match = "callbacks"): - unsloth.UnslothTrainer( - model = _DummyModel(), - tokenizer = None, - train_dataset = [], - callbacks = [object()], - ) with pytest.raises(NotImplementedError, match = "compute_metrics"): unsloth.UnslothTrainer( model = _DummyModel(), @@ -725,6 +725,46 @@ def test_mlx_trainer_rejects_metrics_and_callbacks(): ) +def test_mlx_trainer_accepts_callbacks(): + """Callbacks are routed to MLXTrainer when the zoo backend supports them.""" + unsloth = _import_mlx_unsloth() + from transformers import TrainerCallback + + if not unsloth._mlx_trainer_supports_kwarg("callbacks"): + pytest.skip("requires unsloth-zoo MLXTrainer callback support") + + class Callback(TrainerCallback): + pass + + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [Callback()], + ) + assert any(isinstance(cb, Callback) for cb in trainer.callback_handler.callbacks) + + +def test_mlx_trainer_rejects_callbacks_with_old_zoo(monkeypatch): + """Older unsloth-zoo builds should fail clearly instead of TypeError.""" + unsloth = _import_mlx_unsloth() + from transformers import TrainerCallback + + monkeypatch.setattr( + unsloth, + "_mlx_trainer_supports_kwarg", + lambda name: name != "callbacks", + ) + + with pytest.raises(NotImplementedError, match = "callbacks require"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [TrainerCallback()], + ) + + def test_mlx_trainer_rejects_custom_data_collator(): """MLXTrainer owns batching; custom SFT data collators must not be ignored.""" unsloth = _import_mlx_unsloth() diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 04cc600725..de5cd0f61b 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -102,6 +102,7 @@ if _IS_MLX: ) from _e import dataclasses as _dataclasses + import inspect as _inspect import importlib.machinery as _machinery import sys as _sys import types as _types @@ -109,6 +110,30 @@ if _IS_MLX: __version__ = unsloth_zoo.__version__ DEVICE_TYPE = "mlx" + _MLX_TRAINER_ACCEPTS_VAR_KWARGS = False + _MLX_TRAINER_SUPPORTED_KWARGS = frozenset() + try: + _MLX_TRAINER_INIT_PARAMETERS = _inspect.signature(MLXTrainer.__init__).parameters + _MLX_TRAINER_ACCEPTS_VAR_KWARGS = any( + param.kind is _inspect.Parameter.VAR_KEYWORD + for param in _MLX_TRAINER_INIT_PARAMETERS.values() + ) + _MLX_TRAINER_SUPPORTED_KWARGS = frozenset( + name + for name, param in _MLX_TRAINER_INIT_PARAMETERS.items() + if name != "self" + and param.kind + in ( + _inspect.Parameter.POSITIONAL_OR_KEYWORD, + _inspect.Parameter.KEYWORD_ONLY, + ) + ) + except (TypeError, ValueError): + pass + + def _mlx_trainer_supports_kwarg(name): + """Return whether the installed zoo MLXTrainer accepts a kwarg.""" + return _MLX_TRAINER_ACCEPTS_VAR_KWARGS or name in _MLX_TRAINER_SUPPORTED_KWARGS def _is_mlx_cuda_device_target(device): """Return True when a torch .to/.cuda target asks for CUDA on MLX.""" @@ -966,6 +991,7 @@ if _IS_MLX: "args", "formatting_func", "processor", + "callbacks", ) _TRL_SFT_TRAINER_POSITIONAL_KWARGS = ( "model", @@ -985,6 +1011,29 @@ if _IS_MLX: ) _MLX_TRAINER_KWARGS = frozenset(_MLX_TRAINER_POSITIONAL_KWARGS) + def _filter_supported_mlx_trainer_kwargs(trainer_kwargs): + """Drop inert/empty kwargs unsupported by this zoo MLXTrainer.""" + unsupported = { + key: value + for key, value in trainer_kwargs.items() + if not _mlx_trainer_supports_kwarg(key) + } + names = sorted( + key for key, value in unsupported.items() if _is_meaningful_mlx_extra_value(value) + ) + if names: + subject = ", ".join(names) + verb = "requires" if len(names) == 1 else "require" + raise NotImplementedError( + "Unsloth MLX: " + f"{subject} {verb} an unsloth-zoo build with " + "matching MLXTrainer support. Upgrade unsloth-zoo together " + "with unsloth." + ) + for key in unsupported: + trainer_kwargs.pop(key, None) + return trainer_kwargs + def _is_mlx_native_text_collator(collator): """HF pad/copy collators are redundant on MLX; match by class name.""" for klass in type(collator).__mro__: @@ -1162,6 +1211,7 @@ if _IS_MLX: trainer_kwargs, config_kwargs, ignored_kwargs = _split_mlx_trainer_kwargs(kwargs) _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs) + trainer_kwargs = _filter_supported_mlx_trainer_kwargs(trainer_kwargs) trainer_kwargs["args"] = _coerce_mlx_training_args( trainer_kwargs.get("args"), config_kwargs, From 07c8bbbf5a48ee059f2f0e7767f664e24c8b0bf6 Mon Sep 17 00:00:00 2001 From: marcandrelarochelle Date: Wed, 8 Jul 2026 07:05:03 -0400 Subject: [PATCH 064/113] (GRPO) Fix PEFT replacement for TRL >= 1.7.0, add missing compute_aux_loss for TRL >= 1.7.0 (#6904) * Fix PEFT replacement for TRL >= 1.7.0, add missing compute_aux_loss for TRL 1.7.0 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix GRPO for TRL >= 1.7.0: PEFT ref-adapter removal and return arity rl.py: for trl >= 1.7.0, scope the PEFT removal regex to the ref-adapter block only by anchoring the end on ref_param.data.copy_(param.data), so it no longer also deletes the following gradient-checkpointing enable_input_require_grads() block. Neutralize TRL 1.7.0's `if _is_quantized_model:` bf16 cast the same way the existing is_loaded_in_4bit cast is handled. rl_replacements.py: initialize _extra_moe_kwargs before use (it was referenced before assignment whenever compute_aux_loss was passed) and only request output_router_logits when the aux loss is actually wanted. rl_replacements.py: _get_per_token_logps_and_entropies now returns a 3-tuple (logps, entropies, aux_loss) for trl >= 1.7.0 and a 2-tuple for older TRL, matching how every TRL call site unpacks the result. Without this, TRL 1.7.x _generate_and_score_completions unpacks 3 values from a 2-tuple and raises "not enough values to unpack (expected 3, got 2)". * Return zero aux_loss placeholder and drop inference-mode aux collection * GRPO TRL >= 1.7.0: reject router aux-loss opt-in at init; drop zero aux placeholder Unsloth's optimized GRPO forward cannot compute the MoE router auxiliary loss. Previously an explicit opt-in (router_aux_loss_coef > 0) returned a fabricated zero, silently training without the requested load-balancing penalty. Now reject it at trainer init with a clear NotImplementedError, and return None (not zero) for the aux slot of TRL's 3-tuple. Default stays off (coef 0), so the common path is unaffected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO hidden-states fallback: free ModelOutput before chunked log-softmax The old/ref logprob fallback binds the full ModelOutput (which holds every layer's hidden_states when output_hidden_states=True) and kept it alive across chunked_hidden_states_selective_log_softmax, an avoidable OOM on large models. Extract logits then del outputs in both the text and VLM branches. * Version-compat CI: proactively catch TRL GRPO breakage The existing TRL canary is a static symbol/source grep: it verifies symbols exist but is blind to structural changes (TRL 1.7.0's 2->3-tuple per-token-logps return arity and restructured PEFT ref-adapter block, which the fix in this PR addresses, both slipped past it because the methods still existed). Two additions: - test_trl_grpo_pinned_symbols.py: extend TRL_TAGS to 1.5/1.6/1.7 and pin the exact source-string contracts the rl.py / rl_replacements.py transforms depend on for TRL >= 1.7.0 (PEFT elif ref-adapter block + enable_input_require_grads survival, if _is_quantized_model, aux_loss_enabled anchor, compute_aux_loss arity). A future TRL change fails on main a few days before the PyPI release. - test_trl_grpo_fake_run.py + a version-compat-ci job: fake-CUDA run that drives the real GRPO/SFT/DPO source-transform patchers against latest + main TRL on a CPU-only runner (no training) and asserts the generated Unsloth trainer still satisfies the transform contracts. Catches behavioral regressions the grep cannot see. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fake-run test: use a normal Version import for the aux gate * version-compat CI: fix fake-run job gate + torch-absent collection - Drop the invalid job-level matrix if (matrix is not available in jobs..if -> 'Unrecognized named-value: matrix' fails the whole workflow). Use a single job that runs vs TRL latest always and re-runs vs TRL main only on schedule/dispatch via a step-level github.event_name guard. Validated with actionlint. - Module-level skip the fake-run test when torch is absent so daily-fresh-fetch (pytest-only, collects tests/version_compat/) does not crash on the top-level spoof import. * fake-run test: do not skip on import failure unsloth/trl are installed in the grpo-fake-run job, so a failing import is the import-time drift this canary must catch. Keep only the not-installed find_spec skips; let a real import error fail the test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO arity gate: regex downgrade + fail loud + CI coverage The TRL < 1.7.0 per-token-logps return downgrade was an exact-string replace anchored on the full return line incl. its comment, so a reformat (e.g. pre-commit) could silently no-op it and ship a 3-tuple to older TRL. Switch to a regex tolerant of comment/whitespace drift, and raise if the anchor stops matching (re.subn count != 1) instead of failing silently. Add a monkeypatched trl_version unit test asserting both arities, since CI only installs TRL >= 1.7.0 and never exercised the downgrade otherwise. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fake-run: give SFT/DPO a real contract, not just ast-parse The SFT/DPO fake patch runs only checked the generated trainer parses. Also assert the shared QLoRA _is_quantized_model bf16 cast is neutralized (TRL 1.7's spelling, present in both sft_trainer and dpo_trainer), so a structural TRL change to that block is caught for SFT/DPO too, not just GRPO. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO PEFT ref-adapter removal: lower gate to the TRL 1.4.0 floor The elif is_peft_model(model) and args.beta != 0.0: ref-adapter block was introduced in TRL 1.4.0 and is unchanged through 1.7.x, but the removal was gated at >= 1.7.0, so for 1.4 <= TRL < 1.7 the transform fell through to the 0.27 branch (which matches the older if is_peft_available()... form) and silently no-oped: a PEFT + beta != 0 GRPO run then computed the KL reference from the copied ref adapter instead of the base model. Lower the gate to 1.4.0 and keep the 1.7.0-only router aux-loss fail-fast nested. Widen the pinned-symbol contract test to run from 1.4.0 so the covered versions are actually exercised. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .github/workflows/version-compat-ci.yml | 77 ++++++ .../version_compat/test_trl_grpo_fake_run.py | 246 ++++++++++++++++++ .../test_trl_grpo_pinned_symbols.py | 110 ++++++++ unsloth/models/rl.py | 36 ++- unsloth/models/rl_replacements.py | 42 ++- 5 files changed, 504 insertions(+), 7 deletions(-) create mode 100644 tests/version_compat/test_trl_grpo_fake_run.py diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index e492d21e99..b15d5bfa25 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -285,6 +285,83 @@ jobs: tests/vllm_compat/test_extended_module_imports.py \ -v --tb=short + # Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike + # the static symbol/source greps above, this drives unsloth's actual + # source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only + # runner under the tests/conftest.py spoof harness -- no GPU, no training. + # Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple + # per-token-logps return, restructured PEFT ref-adapter block) by asserting + # the generated Unsloth trainer still satisfies the transform contracts. + grpo-fake-run: + name: GRPO fake-run (latest + main TRL, CPU spoof) + runs-on: ubuntu-latest + timeout-minutes: 18 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - name: Clone unsloth-zoo @ main + run: | + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install CPU torch + ecosystem + TRL latest + run: | + python -m pip install --upgrade pip + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10' + # Ecosystem floors unsloth needs; TRL itself is installed last so it + # can pull the transformers/peft it requires. + pip install \ + 'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \ + 'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \ + 'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow + pip install --upgrade trl + pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo" + pip install --no-deps -e ./unsloth + - name: Fake-run vs TRL latest + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + -v --tb=short + # `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge + # TRL break does not red every PR. github.event_name is valid in a step if. + - name: Fake-run vs TRL main (scheduled / dispatch only) + if: ${{ github.event_name != 'pull_request' }} + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + pip install --upgrade "git+https://github.com/huggingface/trl" + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + -v --tb=short + # Daily-only: same suites but with --strict on importable upstream # tags. Schedule-only so PR jobs stay fast; cron tolerates a flake. daily-fresh-fetch: diff --git a/tests/version_compat/test_trl_grpo_fake_run.py b/tests/version_compat/test_trl_grpo_fake_run.py new file mode 100644 index 0000000000..c79c9955da --- /dev/null +++ b/tests/version_compat/test_trl_grpo_fake_run.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Fake-CUDA GRPO patch run against the *installed* TRL (CPU-only, no training). + +The static symbol/source-string canaries (test_trl_grpo_pinned_symbols.py) +grep raw TRL source; they never execute unsloth's transforms. This test drives +the real pipeline: under the aggressive CUDA spoof it imports unsloth and calls +`_patch_trl_rl_trainers_impl`, which reads the installed GRPOTrainer via +inspect.getsource, applies every rl.py/rl_replacements.py rewrite, and compiles +the result into an UnslothGRPOTrainer. A structural TRL change that slips past +the greps (e.g. TRL 1.7.0's 2->3-tuple return arity, or a restructured PEFT +ref-adapter block) surfaces here as a transform error, a broken generated +source, or a violated contract -- with no GPU and no training run. + +Meant to run in CI against `trl==latest` and `trl @ main` (see +version-compat-ci.yml). The tests/conftest.py harness pre-loads device_type +with DEVICE_COUNT=0 so unsloth's kernel init takes the CPU-safe path. +""" + +from __future__ import annotations + +import ast +import importlib +import importlib.machinery +import importlib.util +import inspect +import sys +import types +from pathlib import Path + +import pytest + + +# daily-fresh-fetch collects tests/version_compat/ with only pytest installed; +# the spoof and the rest of this module need the real torch runtime. Skip the +# whole module cleanly when torch is absent rather than crashing collection. +if importlib.util.find_spec("torch") is None: + pytest.skip("torch not installed; fake-run needs the real runtime", allow_module_level = True) + +# Apply the spoof BEFORE any unsloth-touching import (mirrors +# tests/vllm_compat/test_extended_module_imports.py). +_SPOOF_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_SPOOF_DIR)) +import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402 + +_spoof.apply() + + +def _stub_module(name: str, attrs: dict | None = None) -> None: + if name in sys.modules: + return + m = types.ModuleType(name) + m.__spec__ = importlib.machinery.ModuleSpec(name = name, loader = None, origin = "") + for k, v in (attrs or {}).items(): + setattr(m, k, v) + sys.modules[name] = m + + +_stub_module("torchcodec") + + +def _trl_version(): + import trl + from packaging.version import Version + return Version(trl.__version__.split("+")[0]) + + +def _patch_grpo_and_get_source() -> str: + """Run the GRPO patcher against the installed TRL and return the generated + UnslothGRPOTrainer source. Calls the impl (not the try/except wrapper) so a + transform/compile regression surfaces as a hard error instead of a silent + no-op.""" + import trl.trainer.grpo_trainer as _g + + from unsloth.models import rl as _rl + + _rl._patch_trl_rl_trainers_impl("grpo_trainer") + patched = _g.GRPOTrainer + assert patched.__name__ == "UnslothGRPOTrainer", ( + f"GRPO patch silently no-oped: trl.trainer.grpo_trainer.GRPOTrainer is " + f"{patched.__name__!r}, expected 'UnslothGRPOTrainer' (transform failed " + f"or dispatch key drifted on this TRL)" + ) + # The transformed body (__init__ rewrites, injected per-token-logps) lives in + # the generated module's `_UnslothGRPOTrainer` base + module-level funcs, not + # the thin UnslothGRPOTrainer subclass -- read the whole generated module. + mod = inspect.getmodule(patched) + return inspect.getsource(mod) if mod is not None else inspect.getsource(patched) + + +@pytest.fixture(scope = "module") +def generated_grpo_source(): + if importlib.util.find_spec("unsloth") is None: + pytest.skip("unsloth not installed") + if importlib.util.find_spec("trl") is None: + pytest.skip("trl not installed") + # Do NOT swallow import errors: unsloth is installed here, so a failing + # `import unsloth` is exactly the import-time TRL/transformers drift this + # canary must surface as a failure, not a skip. + import unsloth # noqa: F401 -- _gpu_init bootstrap under spoof + + return _patch_grpo_and_get_source() + + +def test_grpo_patch_generates_valid_source(generated_grpo_source): + """The generated UnslothGRPOTrainer must be syntactically valid Python.""" + ast.parse(generated_grpo_source) + + +def test_grpo_patch_aux_fail_fast_injected(generated_grpo_source): + """TRL >= 1.7.0: rl.py injects a fail-fast for the unsupported MoE router + aux-loss opt-in right after `self.aux_loss_enabled = ...`.""" + from packaging.version import Version + + if _trl_version() < Version("1.7.0"): + pytest.skip("aux_loss_enabled / router_aux_loss_coef are TRL >= 1.7.0") + assert "does not compute the MoE router auxiliary loss" in generated_grpo_source, ( + "aux fail-fast raise missing from generated trainer; rl.py's " + "aux_loss_enabled .replace() anchor did not match this TRL" + ) + + +def test_grpo_patch_three_tuple_return(generated_grpo_source): + """TRL >= 1.7.0 call sites unpack a 3-tuple from + _get_per_token_logps_and_entropies; the injected replacement must return + (logps, entropies, aux_loss).""" + from packaging.version import Version + if _trl_version() >= Version("1.7.0"): + assert "return logprobs.detach(), entropies, aux_loss" in generated_grpo_source, ( + "3-tuple per-token-logps return missing; the arity version-gate in " + "rl_replacements.py did not emit the >=1.7.0 form" + ) + else: + assert ( + "return logprobs.detach(), entropies, aux_loss" not in generated_grpo_source + ), "2-tuple TRL got the 3-tuple return; arity gate mis-fired" + + +def test_grpo_patch_preserves_grad_checkpointing_block(generated_grpo_source): + """The tightened PR #6904 PEFT regex must remove only the ref-adapter init, + not the following enable_input_require_grads gradient-checkpointing block.""" + from packaging.version import Version + + if _trl_version() < Version("1.7.0"): + pytest.skip("ref-adapter elif block is the TRL >= 1.7.0 shape") + assert "enable_input_require_grads" in generated_grpo_source, ( + "gradient-checkpointing enable_input_require_grads() block was swallowed " + "by the PEFT-removal regex (over-reach regression)" + ) + + +def test_grpo_patch_neutralizes_ref_adapter_and_qlora_cast(generated_grpo_source): + """TRL >= 1.7.0: the ref-adapter copy and the hardcoded QLoRA bf16 cast must + both be gone from the generated trainer.""" + from packaging.version import Version + + if _trl_version() < Version("1.7.0"): + pytest.skip("targets the TRL >= 1.7.0 PEFT / _is_quantized_model shapes") + assert ( + "ref_param.data.copy_(param.data)" not in generated_grpo_source + ), "TRL's PEFT ref-adapter init survived; rl.py peft_pattern re.sub no-oped" + assert ( + "if _is_quantized_model:" not in generated_grpo_source + ), "TRL's hardcoded QLoRA bf16 cast survived; rl.py neutralization no-oped" + + +# SFT / DPO: the same source-transform patcher runs on them (a fake patch run, +# no training), so a structural TRL change can break generation. Assert the patch +# produces a valid, importable Unsloth trainer AND that the shared QLoRA +# `_is_quantized_model` bf16 cast is neutralized (TRL 1.7's spelling), which the +# patcher applies to every trainer. Catches "and or others" beyond GRPO. + + +def _patch_and_get_source(trainer_file: str, trainer_cls: str) -> str: + if importlib.util.find_spec("unsloth") is None or importlib.util.find_spec("trl") is None: + pytest.skip("unsloth or trl not installed") + # Let a real import failure fail the test (import-time drift is the target). + import unsloth # noqa: F401 + import trl.trainer # noqa: F401 + + from unsloth.models import rl as _rl + + _rl._patch_trl_rl_trainers_impl(trainer_file) + mod = importlib.import_module(f"trl.trainer.{trainer_file}") + patched = getattr(mod, trainer_cls) + assert patched.__name__ == f"Unsloth{trainer_cls}", ( + f"{trainer_cls} patch silently no-oped on this TRL " + f"(got {patched.__name__!r}); source-transform dispatch drifted" + ) + gen = inspect.getmodule(patched) + src = inspect.getsource(gen) if gen is not None else inspect.getsource(patched) + ast.parse(src) + return src + + +def _assert_quantized_cast_neutralized(src: str, trainer_cls: str) -> None: + from packaging.version import Version + if _trl_version() < Version("1.7.0"): + pytest.skip("pre-1.7.0 spells the QLoRA cast differently (is_loaded_in_4bit)") + assert "if _is_quantized_model:" not in src, ( + f"{trainer_cls}: TRL's hardcoded QLoRA bf16 cast survived; the shared " + f"rl.py `if _is_quantized_model:` -> `if False:` neutralization no-oped" + ) + + +def test_sft_patch_generates_valid_source(): + src = _patch_and_get_source("sft_trainer", "SFTTrainer") + _assert_quantized_cast_neutralized(src, "SFTTrainer") + + +def test_dpo_patch_generates_valid_source(): + src = _patch_and_get_source("dpo_trainer", "DPOTrainer") + _assert_quantized_cast_neutralized(src, "DPOTrainer") + + +# The installed TRL in CI is always >= 1.7.0, so the < 1.7.0 return-arity +# downgrade is never exercised by the fake-run above. Lock both arities by +# monkeypatching rl_replacements.trl_version and re-generating the injected +# _get_per_token_logps_and_entropies source directly (no TRL install needed). +def test_per_token_logps_arity_gate_both_directions(monkeypatch): + if importlib.util.find_spec("unsloth") is None: + pytest.skip("unsloth not installed") + import unsloth # noqa: F401 + from packaging.version import Version + + from unsloth.models import rl_replacements as _rlr + + gate = _rlr.grpo_trainer__get_per_token_logps_and_entropies + + # >= 1.7.0: 3-tuple return kept. + monkeypatch.setattr(_rlr, "trl_version", Version("1.7.0"), raising = False) + src_new = gate("_get_per_token_logps_and_entropies", None) + assert ( + "return logprobs.detach(), entropies, aux_loss" in src_new + ), "3-tuple return missing for TRL >= 1.7.0" + + # < 1.7.0: aux_loss element dropped -> 2-tuple. A no-op downgrade must raise + # (fail loud), never silently ship a 3-tuple to older TRL. + monkeypatch.setattr(_rlr, "trl_version", Version("1.6.0"), raising = False) + src_old = gate("_get_per_token_logps_and_entropies", None) + assert ( + "return logprobs.detach(), entropies # logps, entropies" in src_old + ), "2-tuple return missing for TRL < 1.7.0" + assert ( + "entropies, aux_loss" not in src_old + ), "aux_loss element still present in the TRL < 1.7.0 downgrade" diff --git a/tests/version_compat/test_trl_grpo_pinned_symbols.py b/tests/version_compat/test_trl_grpo_pinned_symbols.py index 834692e2dd..2f935dde49 100644 --- a/tests/version_compat/test_trl_grpo_pinned_symbols.py +++ b/tests/version_compat/test_trl_grpo_pinned_symbols.py @@ -51,10 +51,27 @@ TRL_TAGS = [ "v1.2.0", "v1.3.0", "v1.4.0", + "v1.5.0", + "v1.5.1", + "v1.6.0", + "v1.7.0", # anchor: first release unsloth's TRL>=1.7.0 GRPO patch targets + "v1.7.1", # current PyPI latest "main", ] +def _tag_ge(tag: str, floor: str) -> bool: + """True if `tag` is `main` or a version >= `floor` (e.g. "1.7.0").""" + if tag == "main": + return True + from packaging.version import Version + + try: + return Version(tag.lstrip("v")) >= Version(floor) + except Exception: + return False + + # unsloth/trainer.py + unsloth/models/rl.py rebind these top-level names. @@ -537,3 +554,96 @@ def test_trl_truncate_with_protected_tokens_optional(tag: str): assert src is not None has_it = "truncate_with_protected_tokens" in src _ = has_it # informational; pass either way. + + +# 24-27. TRL >= 1.7.0 GRPO source contracts. Unlike the has_def existence +# checks above, these pin the exact source strings unsloth/models/rl.py and +# rl_replacements.py transform for TRL >= 1.7.0 (the window PR #6904 fixes). +# The 1.7.0 break was invisible to the existence checks because the methods +# still existed -- only their internal structure / return arity changed. If +# TRL restructures one of these, the transform silently no-ops (or the +# generated trainer breaks), so failing here on `main` gives a few-day lead. + + +@pytest.mark.parametrize("tag", TRL_TAGS) +def test_trl_grpo_peft_ref_adapter_block_contract(tag: str): + """rl.py (trl>=1.4.0) strips TRL's PEFT ref-adapter init with a re.DOTALL + regex anchored on `elif is_peft_model(model) and args.beta != 0.0:` ... + `ref_param.data.copy_(param.data)`. Both anchors must exist (else the + regex no-ops and the ref adapter is created under Unsloth), and the + following `enable_input_require_grads` gradient-checkpointing block must + remain present -- the tightened regex must NOT swallow it (PR #6904). The + `elif` block shape appeared in TRL 1.4.0, so this contract runs from there.""" + if not _tag_ge(tag, "1.4.0"): + pytest.skip( + f"{tag}: pre-1.4.0 uses the `if is_peft_available()...` form (rl.py 0.27 branch)" + ) + src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") + assert src is not None + assert "elif is_peft_model(model) and args.beta != 0.0:" in src, ( + f"{tag}: PEFT ref-adapter `elif` anchor gone; unsloth/models/rl.py " + f"peft_pattern re.sub no-ops and TRL's ref adapter init runs under Unsloth" + ) + assert "ref_param.data.copy_(param.data)" in src, ( + f"{tag}: `ref_param.data.copy_(param.data)` end-anchor gone; " + f"unsloth/models/rl.py peft_pattern loses its DOTALL end match" + ) + assert "enable_input_require_grads" in src, ( + f"{tag}: `enable_input_require_grads` block gone from grpo_trainer.py; " + f"the tightened PR #6904 regex assumed it follows the ref-adapter block" + ) + + +@pytest.mark.parametrize("tag", TRL_TAGS) +def test_trl_grpo_quantized_model_cast_contract(tag: str): + """rl.py (trl>=1.7.0) neutralizes TRL's hardcoded QLoRA bf16 cast + `if _is_quantized_model:` -> `if False:`. A rename leaves the cast active, + which ignores the user's dtype and breaks GradScaler with fp16=True.""" + if not _tag_ge(tag, "1.7.0"): + pytest.skip(f"{tag}: pre-1.7.0 spells the cast differently (is_loaded_in_4bit)") + src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") + assert src is not None + assert "if _is_quantized_model:" in src, ( + f"{tag}: `if _is_quantized_model:` gone; unsloth/models/rl.py cannot " + f"neutralize TRL's hardcoded QLoRA bf16 cast and it runs under Unsloth" + ) + + +@pytest.mark.parametrize("tag", TRL_TAGS) +def test_trl_grpo_aux_loss_enabled_contract(tag: str): + """rl.py (trl>=1.7.0) appends a fail-fast after + `self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0` so an + explicit MoE router-aux opt-in errors instead of silently training without + the penalty (the optimized forward cannot compute it). A change to this + line drops the guard silently (PR #6904).""" + if not _tag_ge(tag, "1.7.0"): + pytest.skip(f"{tag}: aux_loss_enabled / router_aux_loss_coef added in TRL 1.7.0") + src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") + assert src is not None + assert "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0" in src, ( + f"{tag}: `aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0` " + f"changed; unsloth/models/rl.py's fail-fast .replace() anchor no-ops" + ) + + +@pytest.mark.parametrize("tag", TRL_TAGS) +def test_trl_grpo_per_token_logps_aux_arity_contract(tag: str): + """TRL 1.7.0 added `compute_aux_loss` to + _get_per_token_logps_and_entropies and made every call site unpack a + 3-tuple. rl_replacements.py version-gates its injected replacement to emit + a 3-tuple for trl>=1.7.0 (2-tuple below). This is the exact change the + has_def existence checks miss: the method still exists, only its arity + changed. If TRL drops/renames the aux return, the gate needs revisiting.""" + if not _tag_ge(tag, "0.20.0"): + pytest.skip(f"{tag}: pre-0.20 uses legacy _get_per_token_logps (2-tuple, no aux)") + src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") + assert src is not None + assert has_def(src, "_get_per_token_logps_and_entropies", "func"), ( + f"{tag}: _get_per_token_logps_and_entropies missing on TRL >=0.20; " + f"unsloth's per-token-logps injection dispatch key no longer matches" + ) + if _tag_ge(tag, "1.7.0"): + assert "compute_aux_loss" in src, ( + f"{tag}: TRL >=1.7.0 dropped `compute_aux_loss`; the 3-tuple " + f"injection gate in unsloth/models/rl_replacements.py must be revisited" + ) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index b5cadf2dea..eeab8fbaca 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1651,7 +1651,36 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): RLTrainer_source = re.sub(pattern, new_options, RLTrainer_source, flags = re.DOTALL) - if trl_version >= Version("0.27.0"): + if trl_version >= Version("1.4.0"): + # The `elif is_peft_model(model) and args.beta != 0.0:` ref-adapter block + # was introduced in TRL 1.4.0 and is used through 1.7.x. Remove only that + # block, anchored on the final ref_param copy so we do NOT also swallow the + # following gradient-checkpointing enable_input_require_grads() block. + peft_pattern = ( + r"\s*elif is_peft_model\(model\) and args\.beta != 0\.0:" + r".*?" + r"ref_param\.data\.copy_\(param\.data\)" + ) + + replacement_comment = ( + "\n # PEFT initialization logic removed via script for trl >= 1.4.0\n" + ) + + RLTrainer_source = re.sub( + peft_pattern, replacement_comment, RLTrainer_source, flags = re.DOTALL + ) + + if trl_version >= Version("1.7.0"): + # router_aux_loss_coef / aux_loss_enabled were added in TRL 1.7.0. Unsloth's + # optimized GRPO forward cannot compute the MoE router aux loss, so reject + # explicit opt-in (router_aux_loss_coef > 0) at init rather than silently ignoring it. + RLTrainer_source = RLTrainer_source.replace( + "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0", + "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0\n" + ' if self.aux_loss_enabled: raise NotImplementedError("Unsloth GRPO does not compute the MoE router auxiliary loss; set router_aux_loss_coef = 0 (the Unsloth default).")', + ) + + elif trl_version >= Version("0.27.0"): peft_pattern = ( r"\s*if is_peft_available\(\) and is_peft_model\(model\) and args\.beta != 0\.0:" r".*?" @@ -1689,6 +1718,11 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): 'if getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False):', "if False:", ) + # TRL >= 1.7.0 spells the same QLoRA bf16 cast as `if _is_quantized_model:`. + RLTrainer_source = RLTrainer_source.replace( + "if _is_quantized_model:", + "if False:", + ) if RLTrainer_name == "SFTTrainer": original_text = ( diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 098950de08..ffb845b04f 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -1203,6 +1203,8 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": self._autocast_dtype = torch.float16 + compute_aux_loss = kwargs.get("compute_aux_loss", None) + pixel_values, image_grid_thw = ( kwargs.get("pixel_values", None), kwargs.get("image_grid_thw", None), @@ -1846,7 +1848,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): _extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): if pixel_values is None: - logits_chunk = unwrapped_model( + outputs = unwrapped_model( input_ids = input_ids_chunk, attention_mask = attention_mask_chunk, pixel_values = pixel_values_chunk, @@ -1854,7 +1856,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): pixel_attention_mask = pixel_attention_mask_chunk, image_sizes = image_sizes_chunk, **_extra_vision_kwargs, - ).logits + ) + + logits_chunk = outputs.logits + del outputs # free hidden_states before chunked log-softmax completion_input_ids_chunk = input_ids_chunk[ :, -(logits_to_keep + max_left_pad) : @@ -1876,7 +1881,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): else: # Essentially, for VLMs we do not go via the optimized path in models/, # so we don't encounter the Flash Attn left-padding issue. - logits_chunk = unwrapped_model( + outputs = unwrapped_model( input_ids = input_ids_chunk, attention_mask = attention_mask_chunk, pixel_values = pixel_values_chunk, @@ -1885,7 +1890,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_sizes = image_sizes_chunk, logits_to_keep = logits_to_keep + 1, **_extra_vision_kwargs, - ).logits + ) + + logits_chunk = outputs.logits + del outputs # free hidden_states before chunked log-softmax logits_chunk = logits_chunk[:, :-1, :] completion_input_ids_chunk = input_ids_chunk[:, -logits_to_keep:] @@ -1914,11 +1922,15 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): all_logprobs_list.append(logprobs_chunk) if logprobs is None: # padded fallback when packing was not used logprobs = torch.cat(all_logprobs_list, dim = 0) + entropies = None os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" - - return logprobs.detach(), entropies # logps, entropies + # aux loss is unused: it is off by default (router_aux_loss_coef set to 0 in models/rl.py) + # and explicit opt-in is rejected at trainer init, so this is always None (kept in the + # return for TRL >= 1.7.0's 3-tuple contract). + aux_loss = None + return logprobs.detach(), entropies, aux_loss # logps, entropies, aux_loss # input_ids = input_ids[:, -logits_to_keep:] # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. # See https://github.com/huggingface/trl/issues/2770 @@ -1937,6 +1949,24 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): # return logps # compute logprobs for the input tokens function = inspect.getsource(_get_per_token_logps_and_entropies) + if trl_version < Version("1.7.0"): + # TRL < 1.7.0 unpacks (logps, entropies) at every call site; TRL >= 1.7.0 + # always unpacks (logps, entropies, aux_loss). Drop the aux_loss element so + # the return arity matches the installed TRL. Regex tolerates comment / + # whitespace drift on the return line; fail loud if the anchor ever stops + # matching rather than silently shipping a 3-tuple to older TRL. + new_function, n = re.subn( + r"return (logprobs\.detach\(\), entropies), aux_loss[^\n]*", + r"return \1 # logps, entropies", + function, + ) + if n != 1: + raise RuntimeError( + "Unsloth GRPO: could not downgrade the per-token-logps return to a " + f"2-tuple for TRL {trl_version} (matched {n} times, expected 1). The " + "return line changed; update the arity gate in rl_replacements.py." + ) + function = new_function return function From 0e1ed88bb8161d0cb048d46d2b71e50d925eeb46 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 04:06:28 -0700 Subject: [PATCH 065/113] version-compat CI: fake CPU training runs for SFT/GRPO/DPO (#6965) * version-compat CI: fake CPU training runs for SFT/GRPO/DPO Adds a runtime layer on top of the patch-run canary: actually runs trainer.train() for a couple of steps on a CPU-only runner under the CUDA spoof, wrapping a plain tiny HF model in the Unsloth-patched trainer. Exercises the real train() loop (collation, generation, the injected _get_per_token_logps_and_entropies, loss, backward, optimizer) so a TRL or transformers change that breaks the loop at runtime -- not just the source structure -- surfaces here. No GPU, no meaningful numerics. Needs a chain of small CPU shims (eager torch.compile, dynamo suppress, cuda tensor-alloc redirect to CPU, model.for_training/for_inference equivalents) documented inline. Does not exercise Unsloth's Triton/GPU kernels (CPU can't). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cpu fake-train: force adamw_torch + disable dynamo for CPU runner On a real CPU-build torch runner (GitHub CI) two things bit that a CUDA-build torch with GPUs hidden masked locally: - The default optimizer is adamw_8bit (bitsandbytes), whose is_on_gpu() check dies on CPU tensors. Force optim=adamw_torch in all three configs. - import unsloth reinstalls the real torch.compile over the eager passthrough, so the GRPO hot path (chunked_selective_log_softmax) actually compiles and inductor picks the spoofed CUDA device, crashing on device props (gcnArchName). Re-apply the eager passthrough after import and flip torch._dynamo.config.disable so every @torch.compile runs eager at call time. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cpu fake-train: write checkpoints under pytest tmp_path Use pytest's tmp_path for each trainer's output_dir instead of a hardcoded relative temp/ci_* path, so a local pytest run does not leave untracked dirs in the repo tree and the tests are CWD-independent. * version-compat CI: disable dynamo at process level for the fake-run job Set TORCHDYNAMO_DISABLE / TORCH_COMPILE_DISABLE in the fake-run step env so dynamo/inductor is off before conftest.py's early import unsloth, not only via the per-test runtime shim. Defense in depth on the GPU-less runner: the GRPO hot path never compiles regardless of when its functions were decorated. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/version-compat-ci.yml | 9 + .../version_compat/test_trl_fake_train_cpu.py | 285 ++++++++++++++++++ 2 files changed, 294 insertions(+) create mode 100644 tests/version_compat/test_trl_fake_train_cpu.py diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index b15d5bfa25..6becccc90a 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -339,12 +339,18 @@ jobs: env: UNSLOTH_IS_PRESENT: '1' UNSLOTH_COMPILE_DISABLE: '1' + # Disable dynamo/inductor at the process level, before conftest.py's early + # `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner + # (defense in depth; the CPU fake-train also flips this at runtime). + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python run: | cd unsloth python -c "import trl; print('Resolved TRL', trl.__version__)" PYTHONPATH=. python -m pytest \ tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ -v --tb=short # `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge # TRL break does not red every PR. github.event_name is valid in a step if. @@ -353,6 +359,8 @@ jobs: env: UNSLOTH_IS_PRESENT: '1' UNSLOTH_COMPILE_DISABLE: '1' + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python run: | pip install --upgrade "git+https://github.com/huggingface/trl" @@ -360,6 +368,7 @@ jobs: python -c "import trl; print('Resolved TRL', trl.__version__)" PYTHONPATH=. python -m pytest \ tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ -v --tb=short # Daily-only: same suites but with --strict on importable upstream diff --git a/tests/version_compat/test_trl_fake_train_cpu.py b/tests/version_compat/test_trl_fake_train_cpu.py new file mode 100644 index 0000000000..4dae696282 --- /dev/null +++ b/tests/version_compat/test_trl_fake_train_cpu.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Fake CPU training runs for the Unsloth-patched SFT / GRPO / DPO trainers. + +The patch-run canary (test_trl_grpo_fake_run.py) only compiles + inspects the +generated trainer source. This goes one layer deeper: it actually runs +`trainer.train()` for a couple of steps on a CPU-only runner, under the CUDA +spoof, wrapping a plain (tiny, random-weight) HF model in the Unsloth-patched +trainer. That exercises the real train() loop at runtime -- data collation, +generation (GRPO), the injected `_get_per_token_logps_and_entropies`, loss, +backward, optimizer -- so a TRL or transformers change that breaks the loop +(not just the source structure) surfaces here. No GPU, no meaningful numerics. + +What it does NOT cover: Unsloth's Triton/GPU-optimized model kernels (the +FastLanguageModel fast path) cannot run on CPU, so this validates the +trainer-transform + orchestration layer with a standard forward, not the +optimized kernels. +""" + +from __future__ import annotations + +import os + +# CPU-only: no torch.compile / dynamo (it reaches into the CUDA accelerator), no +# Unsloth kernel compile, no mixed precision. Must be set before torch/unsloth. +os.environ.setdefault("UNSLOTH_COMPILE_DISABLE", "1") +os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") +os.environ.setdefault("ACCELERATE_MIXED_PRECISION", "no") + +import importlib +import importlib.util +import sys +from pathlib import Path + +import pytest + + +# torch is needed for everything below (daily-fresh-fetch collects this dir with +# only pytest installed); skip the whole module cleanly when it is absent. +if importlib.util.find_spec("torch") is None: + pytest.skip( + "torch not installed; fake CPU train needs the real runtime", allow_module_level = True + ) + +# Apply the CUDA spoof before any unsloth-touching import. +_SPOOF_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_SPOOF_DIR)) +import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402 + +_spoof.apply() + +import torch # noqa: E402 + + +# The generated GRPO trainer hard-decorates hot functions with @torch.compile, +# which dynamo processes even under the disable env vars, reaching into +# torch.accelerator (real CUDA) on a GPU-less box. Make torch.compile an eager +# passthrough before unsloth generates/imports the trainer -- same logic, no +# dynamo. (An eager CPU run is exactly what we want here.) +def _eager_compile( + model = None, + *args, + **kwargs, +): + if callable(model): + return model + return lambda fn: fn + + +torch.compile = _eager_compile + +# Belt-and-suspenders: if any @torch.compile still routes through dynamo, let it +# fall back to eager instead of crashing, and stop its stream-capture probe from +# reaching torch.accelerator -> real CUDA on a GPU-less box. +try: + import torch._dynamo # noqa: E402 + torch._dynamo.config.suppress_errors = True +except Exception: + pass +if hasattr(torch, "accelerator"): + torch.accelerator.is_available = lambda *a, **k: False + + +# Redirect any `device="cuda"` tensor allocation / `.to("cuda")` / `.cuda()` to +# CPU. The aggressive spoof deliberately keeps real allocators, but a fake CPU +# train needs cuda-targeted ops (e.g. inductor's init_gpu_context does +# `torch.empty(1, device="cuda")`) to land on CPU instead of erroring. +def _is_cuda_dev(d): + try: + return d is not None and torch.device(d).type == "cuda" + except Exception: + return False + + +for _name in ( + "empty", + "zeros", + "ones", + "full", + "tensor", + "arange", + "randn", + "rand", + "randint", + "empty_like", + "zeros_like", + "ones_like", +): + _orig = getattr(torch, _name, None) + if _orig is None: + continue + + def _redir( + *args, + _orig = _orig, + **kwargs, + ): + if _is_cuda_dev(kwargs.get("device")): + kwargs["device"] = "cpu" + return _orig(*args, **kwargs) + + setattr(torch, _name, _redir) + +_orig_to = torch.Tensor.to + + +def _to_cpu(self, *args, **kwargs): + args = tuple("cpu" if _is_cuda_dev(a) else a for a in args) + if _is_cuda_dev(kwargs.get("device")): + kwargs["device"] = "cpu" + return _orig_to(self, *args, **kwargs) + + +torch.Tensor.to = _to_cpu +torch.Tensor.cuda = lambda self, *a, **k: self + +# Extra CUDA stubs the aggressive spoof lacks, needed to walk a real train(): +# Adam's _cuda_graph_capture_health_check() probes stream capture. +torch.cuda.is_current_stream_capturing = lambda *a, **k: False +try: + import torch.cuda.graphs as _cg # noqa: E402 + _cg._cuda_isCurrentStreamCapturing = lambda *a, **k: False +except Exception: + pass + +# A broken libmlx.so in the shared site-packages crashes transformers' Mac-only +# is_mlx_array probe on Linux; disable it. +try: + import transformers.utils.generic as _g # noqa: E402 + _g._is_mlx_available = False +except Exception: + pass + + +# Dense (non-MoE) tiny model on purpose: MoE models route through Unsloth's +# grouped_gemm Triton kernel, which is CUDA-only and cannot run on a CPU runner. +_MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM" + + +def _load_plain(): + """Tiny plain HF model + tokenizer on CPU. Skips (not fails) if the model + cannot be fetched -- that is a network/hub issue, not an unsloth regression.""" + from transformers import AutoModelForCausalLM, AutoTokenizer + + try: + tok = AutoTokenizer.from_pretrained(_MODEL) + model = AutoModelForCausalLM.from_pretrained(_MODEL, dtype = torch.float32) + except OSError as e: # hub unreachable / model missing + pytest.skip(f"could not fetch {_MODEL} (network/hub): {str(e)[:150]}") + if tok.pad_token is None: + tok.pad_token = tok.eos_token + # Unsloth's GRPO path calls model.for_training()/for_inference() (added by + # FastLanguageModel). A plain HF model lacks them; supply minimal train/eval + # equivalents so the loop proceeds without the optimized wrapper. + if not hasattr(model, "for_training"): + model.for_training = lambda *a, **k: model.train() + if not hasattr(model, "for_inference"): + model.for_inference = lambda *a, **k: model.eval() + return model.to("cpu"), tok + + +@pytest.fixture(autouse = True) +def _require_stack(): + global torch # the `import torch._dynamo` below would otherwise shadow it as local + if importlib.util.find_spec("unsloth") is None or importlib.util.find_spec("trl") is None: + pytest.skip("unsloth or trl not installed") + # A real import failure is a regression we want to surface, so do not guard it. + import unsloth # noqa: F401 -- patches TRL trainers to the Unsloth variants + + # `import unsloth` reinstalls the real torch.compile (overwriting the eager + # passthrough set at module load), so the GRPO hot path (chunked_selective_ + # log_softmax) would really compile -- and inductor picks the spoofed CUDA + # device, crashing on device props (`gcnArchName`). Re-apply the eager + # passthrough and flip dynamo's call-time kill switch so every @torch.compile + # runs eager regardless of when it was decorated. CPU eager is what we want. + torch.compile = _eager_compile + try: + import torch._dynamo # noqa: E402 + torch._dynamo.config.disable = True + except Exception: + pass + + +def test_sft_trains_on_cpu(tmp_path): + from datasets import Dataset + from trl import SFTConfig, SFTTrainer + + assert SFTTrainer.__name__ == "UnslothSFTTrainer", "SFT patch did not apply" + model, tok = _load_plain() + ds = Dataset.from_list([{"text": "The quick brown fox jumps over the lazy dog."}] * 8) + cfg = SFTConfig( + output_dir = str(tmp_path / "ci_sft"), + per_device_train_batch_size = 2, + max_steps = 2, + logging_steps = 1, + report_to = "none", + save_strategy = "no", + use_cpu = True, + max_length = None, + padding_free = False, + dataset_text_field = "text", + fp16 = False, + bf16 = False, + optim = "adamw_torch", + ) + SFTTrainer(model = model, processing_class = tok, args = cfg, train_dataset = ds).train() + + +def test_grpo_trains_on_cpu(tmp_path): + from datasets import Dataset + from trl import GRPOConfig, GRPOTrainer + + assert GRPOTrainer.__name__ == "UnslothGRPOTrainer", "GRPO patch did not apply" + model, tok = _load_plain() + ds = Dataset.from_list([{"prompt": "hi there"}] * 4) + cfg = GRPOConfig( + output_dir = str(tmp_path / "ci_grpo"), + per_device_train_batch_size = 2, + num_generations = 2, + max_steps = 2, + max_completion_length = 8, + logging_steps = 1, + report_to = "none", + temperature = 1.0, + beta = 0.0, + save_strategy = "no", + use_cpu = True, + use_vllm = False, + fp16 = False, + bf16 = False, + optim = "adamw_torch", + ) + GRPOTrainer( + model = model, + processing_class = tok, + reward_funcs = [lambda completions, **k: [float(len(c)) for c in completions]], + args = cfg, + train_dataset = ds, + ).train() + + +def test_dpo_trains_on_cpu(tmp_path): + from datasets import Dataset + from trl import DPOConfig, DPOTrainer + + assert DPOTrainer.__name__ == "UnslothDPOTrainer", "DPO patch did not apply" + model, tok = _load_plain() + ds = Dataset.from_list( + [{"prompt": "Hi", "chosen": " hello friend", "rejected": " go away"}] * 8 + ) + cfg = DPOConfig( + output_dir = str(tmp_path / "ci_dpo"), + per_device_train_batch_size = 2, + max_steps = 2, + logging_steps = 1, + report_to = "none", + save_strategy = "no", + use_cpu = True, + beta = 0.1, + fp16 = False, + bf16 = False, + optim = "adamw_torch", + ) + DPOTrainer(model = model, processing_class = tok, args = cfg, train_dataset = ds).train() From 6ef09361800a6268ac6ae2f89e37c771e5e516f1 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:25:42 +0100 Subject: [PATCH 066/113] Fix OpenClaw start default to local TUI (#6937) * fix: launch OpenClaw local TUI by default * Fix/adjust OpenClaw launch paths for PR #6937 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Default OpenClaw to the local TUI only on a bare invocation The first-arg startswith('-') branch rewrote passthrough globals into a broken command: OpenClaw's grammar is openclaw [--dev] [--profile ] , so 'unsloth start openclaw --profile test' became 'openclaw tui --local --profile test', but tui does not accept --profile (or --dev), so the invocation failed. A leading '--flag value' is ambiguous between a global (--profile test) and a tui option (--message hi), so it cannot be reinterpreted safely. Default to the local TUI only when no passthrough args are given, and forward everything else verbatim so OpenClaw parses it under its own grammar. The bare-launch default (the point of this change) is preserved; explicit subcommands and global flags pass through. --------- Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: Wasim Yousef Said --- .github/scripts/agent-guides-drive.sh | 4 ++-- unsloth_cli/commands/start.py | 12 +++++++++++- unsloth_cli/tests/test_start.py | 24 ++++++++++++++++++++++++ 3 files changed, 37 insertions(+), 3 deletions(-) diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index d430d2c172..defdb498c7 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -376,7 +376,7 @@ case "$MODE" in hermes) patch_hermes_tools none invoke_via_connect "$OUT" -z "$PROMPT" ;; openclaw) patch_openclaw_agent notools - invoke_via_connect "$OUT" agent --local --agent ci \ + CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$OUT" agent --local --agent ci \ --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; *) invoke_via_connect "$OUT" "$PROMPT" ;; esac @@ -449,7 +449,7 @@ case "$MODE" in fi ;; opencode) invoke_via_connect "$out" run "$prompt" ;; hermes) invoke_via_connect "$out" -z "$prompt" ;; - openclaw) invoke_via_connect "$out" agent --local --agent ci \ + openclaw) CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$out" agent --local --agent ci \ --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;; *) invoke_via_connect "$out" "$prompt" ;; esac diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 477a47cc3d..764f5c7963 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -1568,7 +1568,17 @@ def openclaw( serve = serve, launch = launch, ) - command = ["openclaw", *ctx.args] + openclaw_args = list(ctx.args) + # Default a bare `unsloth start openclaw` to the local TUI. Anything the caller + # passes through is forwarded verbatim so OpenClaw parses it under its own grammar + # (openclaw [global-flags] [options]): an explicit subcommand, a global + # flag that must precede the command such as --profile/--dev, or a tui option. We + # cannot reinterpret those safely because a leading "--flag value" is ambiguous + # between a global (`--profile test`) and a tui option (`--message hi`); prepending + # `tui --local` would break the global form, so only the empty case is defaulted. + if not openclaw_args: + openclaw_args = ["tui", "--local"] + command = ["openclaw", *openclaw_args] install_hint = ( "iwr -useb https://openclaw.ai/install.ps1 | iex" if os.name == "nt" diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index ee5b442c27..18cb40f18d 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -1634,10 +1634,34 @@ def test_connect_openclaw_no_launch(fake_studio, tmp_path): config = json.loads(config_path.read_text()) assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert _launch_command(result.output) == ["openclaw", "tui", "--local"] # OpenAI /v1/chat/completions works on either backend — no GGUF gate. assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) +def test_connect_openclaw_no_launch_keeps_explicit_subcommand(fake_studio): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch", "crestodian"]) + assert result.exit_code == 0, result.output + assert _launch_command(result.output) == ["openclaw", "crestodian"] + + +def test_connect_openclaw_no_launch_passes_global_flags_through(fake_studio): + # OpenClaw globals (openclaw [--dev] [--profile ] ) precede the + # command, and tui does not accept them, so any passthrough args must be forwarded + # verbatim rather than rewritten into `openclaw tui --local `. + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch", "--profile", "test"]) + assert result.exit_code == 0, result.output + assert _launch_command(result.output) == ["openclaw", "--profile", "test"] + + +def test_connect_openclaw_no_launch_keeps_explicit_tui(fake_studio): + result = CliRunner().invoke( + start.start_app, ["openclaw", "--no-launch", "tui", "--message", "hi"] + ) + assert result.exit_code == 0, result.output + assert _launch_command(result.output) == ["openclaw", "tui", "--message", "hi"] + + # ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── From e86b7874d433ea1c5a2a49063c07932c36aa63ce Mon Sep 17 00:00:00 2001 From: ErenAta16 Date: Wed, 8 Jul 2026 15:26:50 +0300 Subject: [PATCH 067/113] feat: detect installed coding agent CLIs in Studio settings (#6909) * feat: detect installed coding agent CLIs in Studio settings The API-keys panel only ever showed the "claude" flavor of the `unsloth start` command, so anyone using Codex, OpenCode, OpenClaw, Hermes, or Pi had to manually rewrite the copied command by hand. Add a backend check that looks for each agent's CLI binary on PATH (shutil.which, mirroring the pattern already used elsewhere in studio/backend/utils) and expose it as GET /api/settings/coding-agents. The API-keys panel now renders a picker for all six supported agents, marks the ones it finds installed, and defaults to one of those instead of always falling back to claude. Includes unit tests for the detection helper. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review feedback on coding-agent detection Three fixes from PR review: - detect_installed_coding_agents now treats a PATH lookup failure as "not installed" instead of letting it bubble up and break the settings endpoint; added a regression test for it. - CodingAgentsResponse.agents is now typed as an immutable tuple instead of a list built from one, matching CODING_AGENTS itself. - Fixed a race in the API-keys panel: picking an agent while the installed-CLI check is still in flight could get silently overwritten once that check resolved. A ref now tracks whether the user has made a manual choice, so the auto-detected default only applies before that happens. * Address Codex feedback: GGUF gating and remote-detection scope - codex refuses to launch against a non-GGUF (transformers-backed) model (unsloth_cli's _require_gguf_for_codex), so auto-defaulting to it produced a copy-pasteable command that fails immediately whenever the loaded model isn't GGUF. Add useActiveModelIsGguf() (looks up the active checkpoint in the chat runtime store) and a correction effect that steers the auto-pick away from codex unless the loaded model qualifies, without ever touching a choice the user made by hand. - Detection runs via shutil.which on the Studio backend host, which isn't the same machine as the browser in a tunnel/remote session. Reword the 'installed'/'detected' copy to say so explicitly when the tunnel URL is in use, instead of implying the check ran on the viewer's own device. * Rework auto-default per review: loopback gating + inline GGUF check Replaces the previous approach with the exact shape discussed on the PR: - Export isLoopbackHost/normalizeHost from agent-command.ts. The detection endpoint runs shutil.which on the Studio backend, which only describes the browser's own machine when the base this panel targets resolves to loopback. For a LAN or tunnel/remote base, gate the whole thing off -- don't mark anything as "detected" and don't let it drive the default -- instead of just relabeling the copy. - Drop the separate GGUF-correction effect and useActiveModelIsGguf hook. Read useChatRuntimeStore.getState().activeGgufVariant inline inside the existing detection effect's .then() (so it doesn't need to sit in the effect's deps), and pick the first detected agent that isn't codex unless the loaded model is GGUF, leaving the existing default untouched when no compatible agent is detected. Verified both branches (loopback vs LAN/tunnel base, gguf vs non-gguf, manual pick preserved, no-compatible-agent fallback) with a standalone port of the .then() logic. * Address latest Codex findings: stale detection, model swap, cache - Clear detectedAgents (and skip the network call entirely) when the panel leaves a loopback base, instead of leaving a previous loopback detection result marked 'installed' for a command that now targets a LAN/tunnel/ remote host. - Add a separate, network-free correction effect keyed on the live activeGgufVariant: if codex was auto-picked while a GGUF model was loaded and the user then switches to a transformers-backed model while this panel stays mounted, steer away from codex instead of leaving a command that unsloth_cli's _require_gguf_for_codex will now reject. Never touches a manual pick. - Drop coding-agents.ts's module-lifetime cache. Installed-CLI detection is environment state, not a persisted setting, so a stale positive/negative from before the user installed something (or reopened the tab) is worse than one extra cheap local API call per mount; keep only the in-flight de-dupe for concurrent callers. Verified the correction-effect logic (gguf->non-gguf swap with/without a fallback, still-gguf no-op, manual pick never overridden) with a standalone port of the effect. * Make the codex/GGUF auto-pick symmetric in both directions The correction effect only steered away from codex when the model stopped being GGUF; it never steered back toward codex if the model became GGUF *after* a non-GGUF-gated fallback had already picked something else (e.g. codex is the only detected CLI, a transformers model is loaded so the selection correctly falls back to the claude default, then the user loads a GGUF model while the panel stays mounted -- codex never gets reconsidered). Consolidate into one effect that re-derives the preferred detected agent from scratch whenever detectedAgents or activeGgufVariant changes, in either direction, instead of only reacting to the codex-specific downgrade case. The fetch effect now only populates detectedAgents/availableAgents; this effect is the single source of truth for what gets auto-picked from that list. Never overrides a manual choice. Verified both transition directions plus the manual-pick-survives and initial-detection cases with a standalone port of the derivation logic. * Reset the auto-pick to the default when it stops being trustworthy Two more real gaps from the latest Codex pass on d988f52: - The unified derivation effect only handled the case where a *different* detected agent could take over. If codex was the only detected agent and auto-picked while a GGUF model was loaded, then the model stopped being GGUF, 'preferred' came back undefined and the effect silently left the selection on codex -- exactly the command unsloth_cli's _require_gguf_for_codex now rejects. Fall back to DEFAULT_AGENT in that case instead of leaving it untouched. - Leaving a loopback base cleared detectedAgents (so the 'installed' badges correctly disappear) but left whatever agent had been auto-picked from that now-stale, server-side-only detection still selected. Reset to DEFAULT_AGENT there too, unless the user picked by hand. Introduces a shared DEFAULT_AGENT constant instead of repeating the "claude" literal at each reset site. Verified all five cases (both new resets, both manual-pick-survives variants, and the existing multi-detected-agent fallback still preferring another compatible agent over resetting) with a standalone port of the effects. * Derive GGUF-ness from the actual loaded state, not just the variant string activeGgufVariant only covers an HF-repo GGUF pick (a specific quant variant string). A direct local .gguf file -- custom folder, LM Studio, or drag-drop -- is just as much a GGUF the codex preflight (unsloth_cli's _require_gguf_for_codex) would accept, but it never has a "variant" to report, so it read as non-GGUF here even though /api/inference/status correctly reports is_gguf: true for it. That mismatch could leave a Codex-only install not auto-selected, or reset an auto-picked Codex, for a model that actually supports it. Combined activeGgufVariant with activeNativePathToken (covers the drag-drop/picked-file case) and ggufContextLength (only ever populated when the backend last reported is_gguf: true for the active model, see applyActiveModelStatusToStore) so all three paths a model can be GGUF through are covered, matching the same is_gguf-or-equivalent check hasGgufSource already applies to a staged pick elsewhere in this codebase. * Clear stale native-path token on a non-GGUF status refresh When a native (drag-dropped or picked) GGUF was loaded and the backend later switches to a transformers model outside the UI load path, refresh() adopts the new /api/inference/status via setCheckpoint and applyActiveModelStatusToStore. Those reset activeGgufVariant and ggufContextLength but never clear activeNativePathToken, so the isGguf OR stays true after the switch and a Codex-only detection auto-selects unsloth start codex for a non-GGUF model its preflight rejects. Drop activeNativePathToken in applyActiveModelStatusToStore whenever the status is non-GGUF. A real GGUF load reports is_gguf: true, so its token is preserved (the load path owns it); only a non-GGUF status clears it. * Add the AGPL-3.0 header to the new studio contract test * Fix/adjust agent detection for PR #6909 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> --- studio/backend/routes/settings.py | 14 ++ studio/backend/tests/test_coding_agents.py | 50 +++++ studio/backend/utils/coding_agents.py | 39 ++++ .../lib/apply-inference-status-to-store.ts | 25 ++- .../features/settings/api/coding-agents.ts | 45 +++++ .../settings/components/agent-command.ts | 4 +- .../settings/components/usage-examples.tsx | 175 +++++++++++++++++- studio/frontend/src/i18n/locales/en.ts | 2 + .../test_chat_response_details_ui_contract.py | 3 + ...usage_examples_agent_detection_contract.py | 21 +++ 10 files changed, 361 insertions(+), 17 deletions(-) create mode 100644 studio/backend/tests/test_coding_agents.py create mode 100644 studio/backend/utils/coding_agents.py create mode 100644 studio/frontend/src/features/settings/api/coding-agents.ts create mode 100644 tests/studio/test_usage_examples_agent_detection_contract.py diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 914699f540..bbee374334 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -32,6 +32,7 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents from utils.openai_auto_switch_settings import ( DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, @@ -174,6 +175,19 @@ def update_helper_precache( return _helper_precache_response(enabled) +class CodingAgentsResponse(BaseModel): + # All agents `unsloth start` supports, in the CLI's declared order. + agents: tuple[str, ...] = CODING_AGENTS + # Subset of `agents` whose CLI binary was found on PATH; the frontend uses + # this to default the API-keys panel to a command the user can run as-is. + detected: list[str] + + +@router.get("/coding-agents", response_model = CodingAgentsResponse) +def get_coding_agents(current_subject: str = Depends(get_current_subject)) -> CodingAgentsResponse: + return CodingAgentsResponse(detected = detect_installed_coding_agents()) + + @router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) def get_openai_auto_switch( current_subject: str = Depends(get_current_subject), diff --git a/studio/backend/tests/test_coding_agents.py b/studio/backend/tests/test_coding_agents.py new file mode 100644 index 0000000000..b19da1dded --- /dev/null +++ b/studio/backend/tests/test_coding_agents.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for coding-agent CLI detection used by the API-keys settings panel.""" + +from unittest.mock import patch + +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents + + +def test_matches_unsloth_start_subcommands(): + # Each entry must be an actual `unsloth start ` subcommand name + # (unsloth_cli/commands/start.py). Spelled out here rather than imported + # from that module, which pulls in the CLI's heavier dependencies. + assert CODING_AGENTS == ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def test_detects_only_agents_present_on_path(): + installed = {"claude", "opencode"} + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: f"/usr/bin/{name}" if name in installed else None, + ): + assert detect_installed_coding_agents() == ["claude", "opencode"] + + +def test_returns_empty_list_when_nothing_is_installed(): + with patch("utils.coding_agents.shutil.which", return_value = None): + assert detect_installed_coding_agents() == [] + + +def test_preserves_declared_order_regardless_of_path_lookup_order(): + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: name if name in ("pi", "claude", "hermes") else None, + ): + assert detect_installed_coding_agents() == ["claude", "hermes", "pi"] + + +def test_treats_a_path_lookup_error_as_not_installed(): + # An advisory check: shutil.which raising for one entry (e.g. a permission + # error walking a PATH directory) should not take down the whole endpoint, + # and should not stop the remaining agents from being checked. + def flaky_which(name: str): + if name == "codex": + raise OSError("permission denied") + return name if name == "claude" else None + + with patch("utils.coding_agents.shutil.which", side_effect = flaky_which): + assert detect_installed_coding_agents() == ["claude"] diff --git a/studio/backend/utils/coding_agents.py b/studio/backend/utils/coding_agents.py new file mode 100644 index 0000000000..f7dd2f8357 --- /dev/null +++ b/studio/backend/utils/coding_agents.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Detect which `unsloth start ` coding-agent CLIs are on PATH. + +The web UI only ever shows the user the "claude" flavor of the `unsloth start` +command (see agent-command.ts), leaving anyone using Codex, OpenCode, and the +other supported agents to manually edit the copied command. This module gives +the frontend a way to ask which of those CLIs are actually installed so it can +default to one the user can run immediately. +""" + +import shutil + +# Keep in sync with the `unsloth start ` subcommands defined in +# unsloth_cli/commands/start.py. Each entry is the exact executable name that +# subcommand launches, so a hit here means `unsloth start ` can find the +# binary on PATH without the user installing anything first. +CODING_AGENTS: tuple[str, ...] = ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def _is_on_path(agent: str) -> bool: + # shutil.which is documented to return None on a miss, but PATH lookups can + # still raise (e.g. a permission error while probing a directory entry); + # this is an advisory check, so a lookup failure should read as "not + # installed" instead of breaking the settings endpoint. + try: + return shutil.which(agent) is not None + except OSError: + return False + + +def detect_installed_coding_agents() -> list[str]: + """Return the subset of CODING_AGENTS whose CLI binary is on PATH. + + Order follows CODING_AGENTS, not discovery order, so callers can treat the + first entry as the preferred default among the installed agents. + """ + return [agent for agent in CODING_AGENTS if _is_on_path(agent)] diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 9386650fee..60788c23bd 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -3,16 +3,19 @@ import { getInferenceStatus } from "../api/chat-api"; import { mergeBackendRecommendedInference } from "../presets/preset-policy"; +import { clampReasoningEffortToLevels } from "../provider-capabilities"; import { CHAT_REASONING_ENABLED_KEY, - loadOptionalBool, type ReasoningEffort, type ReasoningStyle, + loadOptionalBool, resolveToolsEnabledOnLoad, useChatRuntimeStore, } from "../stores/chat-runtime-store"; -import { isMultimodalResponse, type InferenceStatusResponse } from "../types/api"; -import { clampReasoningEffortToLevels } from "../provider-capabilities"; +import { + type InferenceStatusResponse, + isMultimodalResponse, +} from "../types/api"; import type { ChatModelSummary } from "../types/runtime"; type LocalReasoningEffort = Extract; @@ -31,7 +34,10 @@ export function normalizeSpeculativeType( return "ngram"; } if (s === "mtp+ngram") return "mtp+ngram"; - const parts = s.split(",").map((p) => p.trim()).filter(Boolean); + const parts = s + .split(",") + .map((p) => p.trim()) + .filter(Boolean); const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); const hasNgram = parts.some( (p) => p === "ngram" || p === "ngram-mod" || p === "ngram-simple", @@ -197,6 +203,12 @@ export function applyActiveModelStatusToStore( ggufContextLength: currentGgufContextLength, ggufMaxContextLength, ggufNativeContextLength, + // A non-GGUF status must also drop a stale native-path token: without this the + // isGguf OR (activeGgufVariant || activeNativePathToken || ggufContextLength) + // stays true after switching from a native GGUF to a transformers model, so a + // Codex-only detection would auto-select for a model its preflight rejects. A real + // GGUF load reports is_gguf: true, so its token is preserved (the load path owns it). + ...(status.is_gguf ? {} : { activeNativePathToken: null }), modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false, defaultChatTemplate: nextDefaultChatTemplate, loadedIsMultimodal: isMultimodalResponse(status), @@ -245,7 +257,7 @@ export function applyActiveModelStatusToStore( const mid = checkpointId.toLowerCase(); if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) { const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); - if (sizeMatch && parseFloat(sizeMatch[1]) < 9) { + if (sizeMatch && Number.parseFloat(sizeMatch[1]) < 9) { reasoningDefault = false; } } @@ -281,8 +293,7 @@ export async function tryAdoptServerActiveModel(): Promise { } // Re-check after the await: keep a checkpoint the user picked meanwhile. - const previousCheckpoint = - useChatRuntimeStore.getState().params.checkpoint; + const previousCheckpoint = useChatRuntimeStore.getState().params.checkpoint; if (previousCheckpoint) { return true; } diff --git a/studio/frontend/src/features/settings/api/coding-agents.ts b/studio/frontend/src/features/settings/api/coding-agents.ts new file mode 100644 index 0000000000..ae371b2d3a --- /dev/null +++ b/studio/frontend/src/features/settings/api/coding-agents.ts @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type CodingAgentsInfo = { + // Every agent `unsloth start` supports, in the CLI's declared order. + agents: string[]; + // Subset of `agents` whose CLI binary was found on PATH by the backend. + detected: string[]; +}; + +type ApiCodingAgentsInfo = { + agents: string[]; + detected: string[]; +}; + +// Which CLIs are on PATH is environment state, not a persisted setting -- it +// can change any time the user installs something new, so this only +// de-duplicates concurrent in-flight calls (e.g. React strict-mode's double +// mount) rather than caching the result across the module's lifetime. Every +// fresh call (each time a settings panel mounts) re-checks PATH for real. +let inFlightInfo: Promise | null = null; + +function fromApi(info: ApiCodingAgentsInfo): CodingAgentsInfo { + return { agents: info.agents, detected: info.detected }; +} + +async function fetchCodingAgents(): Promise { + const res = await authFetch("/api/settings/coding-agents"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load installed coding agents"), + ); + } + return fromApi(await res.json()); +} + +export async function loadCodingAgents(): Promise { + inFlightInfo ??= fetchCodingAgents().finally(() => { + inFlightInfo = null; + }); + return inFlightInfo; +} diff --git a/studio/frontend/src/features/settings/components/agent-command.ts b/studio/frontend/src/features/settings/components/agent-command.ts index 9e87922970..38b2d73c3b 100644 --- a/studio/frontend/src/features/settings/components/agent-command.ts +++ b/studio/frontend/src/features/settings/components/agent-command.ts @@ -12,7 +12,7 @@ const DEFAULT_AGENT = "claude"; // URL.hostname brackets IPv6 literals (`new URL("http://[::1]:8888").hostname` is // "[::1]"), so strip the brackets before matching the bare "::1" loopback rules below. -function normalizeHost(host: string): string { +export function normalizeHost(host: string): string { const lower = host.toLowerCase(); return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower; } @@ -26,7 +26,7 @@ function isDefaultLocalHost(host: string): boolean { } // Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8. -function isLoopbackHost(host: string): boolean { +export function isLoopbackHost(host: string): boolean { if (host === "localhost" || host === "::1") return true; const octets = host.split("."); return ( diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index c43dd0f219..b3396c8d9b 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -16,6 +16,7 @@ import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { useChatRuntimeStore } from "@/features/chat"; import { useT } from "@/i18n"; import type { TranslationKey } from "@/i18n"; +import { isTauri } from "@/lib/api-base"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { Tick02Icon } from "@/lib/tick-icon"; import { cn } from "@/lib/utils"; @@ -25,14 +26,15 @@ import { InformationCircleIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useEffect, useMemo, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { Streamdown } from "streamdown"; +import { loadCodingAgents } from "../api/coding-agents"; import { type OpenAIAutoSwitchSettings, loadOpenAIAutoSwitchSettings, updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; -import { buildAgentCommand } from "./agent-command"; +import { buildAgentCommand, isLoopbackHost, normalizeHost } from "./agent-command"; type ExampleType = | "curl" @@ -114,6 +116,30 @@ const DOC_LINKS = [ { label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" }, ]; +// Falls back to this list until the backend's installed-CLI check resolves; +// kept in sync with the `unsloth start ` subcommands and with +// CODING_AGENTS in studio/backend/utils/coding_agents.py. +const DEFAULT_AGENTS = [ + "claude", + "codex", + "openclaw", + "opencode", + "hermes", + "pi", +]; +// The agent selection resets to this whenever an auto-pick is no longer +// trustworthy (leaving loopback, or the only compatible detected agent +// stops being compatible) rather than lingering on a stale choice. +const DEFAULT_AGENT = "claude"; +const AGENT_LABELS: Record = { + claude: "Claude Code", + codex: "Codex", + openclaw: "OpenClaw", + opencode: "OpenCode", + hermes: "Hermes", + pi: "Pi", +}; + const j = (s: string): string => JSON.stringify(s); const shSingle = (s: string): string => s.replace(/'/g, "'\\''"); const psSingle = (s: string): string => s.replace(/'/g, "''"); @@ -399,6 +425,17 @@ function useLoadedModelName(): string { }, [checkpoint, ggufVariant]); } +// Backend PATH detection is only safe in the desktop app, where the UI owns +// the local backend. A browser loopback URL may be an SSH/local port forward. +function canUseLocalAgentDetection(base: string): boolean { + if (!isTauri) return false; + try { + return isLoopbackHost(normalizeHost(new URL(base).hostname)); + } catch { + return false; + } +} + const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [ typeof unslothLightTheme, typeof unslothDarkTheme, @@ -443,7 +480,18 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); const [copiedAgent, setCopiedAgent] = useState(false); + const [agent, setAgent] = useState(DEFAULT_AGENT); + const [availableAgents, setAvailableAgents] = + useState(DEFAULT_AGENTS); + const [detectedAgents, setDetectedAgents] = useState([]); + // True once the user has picked an agent themselves; guards the detection + // effect below from clobbering that choice if it resolves afterward. + const agentPickedByUserRef = useRef(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); + const origin = typeof window !== "undefined" ? window.location.origin : ""; + const base = + useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); + const localAgentDetection = canUseLocalAgentDetection(base); // null while loading; the same setting the General tab exposes (shared cache). const [autoSwitch, setAutoSwitch] = useState( null, @@ -454,6 +502,78 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { void fetchDeviceType({ force: true }); }, []); + // Fetching is the only job of this effect: populate availableAgents/ + // detectedAgents (or clear them). Which agent gets auto-picked from that + // list is derived separately below, so it can react to the loaded model + // changing too, not just a fresh fetch. + useEffect(() => { + // Browser loopback URLs can be SSH/local forwards, so only the desktop app + // may use backend PATH checks to mark or auto-pick local agents. + if (!localAgentDetection) { + setDetectedAgents([]); + // A previously auto-picked agent was only ever verified against the + // Studio backend's PATH, which is meaningless now that this panel no + // longer targets a loopback base -- don't leave it selected, but + // never touch a choice the user made by hand. + if (!agentPickedByUserRef.current) { + setAgent(DEFAULT_AGENT); + } + return; + } + + let cancelled = false; + void loadCodingAgents() + .then((info) => { + if (cancelled) return; + setAvailableAgents(info.agents); + setDetectedAgents(info.detected); + }) + .catch(() => { + // Best-effort: keep the default agent list and let the user pick manually. + }); + return () => { + cancelled = true; + }; + }, [localAgentDetection]); + + // Single source of truth for the auto-picked agent, re-derived whenever + // the detected list or the loaded model's GGUF-ness changes -- in either + // direction. `codex` needs a GGUF model (unsloth_cli's + // _require_gguf_for_codex exits otherwise), so it's only preferred once + // the loaded model actually qualifies; loading a GGUF model *after* a + // non-GGUF-gated fallback picked something else re-steers back to codex + // just as loading a non-GGUF model steers away from it. Never overrides a + // choice the user made by hand. + // activeGgufVariant alone only covers an HF-repo GGUF pick (a specific + // quant variant string) -- a direct local .gguf file (custom folder / + // LM Studio / drag-drop) is just as much a GGUF the codex preflight would + // accept, but never has a "variant" to report, and would otherwise read as + // non-GGUF here. activeNativePathToken covers the drag-drop/picked-file + // case; ggufContextLength is only ever populated when the backend's + // /api/inference/status last reported is_gguf: true for the active model + // (see applyActiveModelStatusToStore), so together these three cover every + // path a model can be GGUF through, matching the same is_gguf-or-equivalent + // check hasGgufSource applies to a staged pick. + const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + const activeNativePathToken = useChatRuntimeStore((s) => s.activeNativePathToken); + const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + useEffect(() => { + if (agentPickedByUserRef.current) return; + if (detectedAgents.length === 0) return; + const isGguf = + activeGgufVariant != null || activeNativePathToken != null || ggufContextLength != null; + const preferred = detectedAgents.find((a) => a !== "codex" || isGguf); + if (preferred) { + setAgent(preferred); + } else if (agent === "codex" && !isGguf) { + // codex was auto-picked while a GGUF model was active and it's the + // only detected agent; now that the model isn't GGUF anymore, nothing + // detected is actually runnable, so fall back to the default instead + // of leaving a codex command unsloth_cli will reject. + setAgent(DEFAULT_AGENT); + } + }, [agent, detectedAgents, activeGgufVariant, activeNativePathToken, ggufContextLength]); + useEffect(() => { let cancelled = false; void loadOpenAIAutoSwitchSettings() @@ -470,9 +590,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const model = useLoadedModelName(); const key = apiKey || KEY_PLACEHOLDER; - const origin = typeof window !== "undefined" ? window.location.origin : ""; - const base = - useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); const autoSwitchOn = autoSwitch?.enabled ?? false; const snippets = useMemo( @@ -481,8 +598,8 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); // Agent command must target the server the panel shows, not the :8888 default. const agentCommand = useMemo( - () => buildAgentCommand(base, key, os), - [base, key, os], + () => buildAgentCommand(base, key, os, agent), + [base, key, os, agent], ); const osAware = OS_AWARE[lang]; @@ -710,6 +827,42 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {t("settings.apiKeys.codingAgentsHint")} +
+ {availableAgents.map((id) => { + const installed = detectedAgents.includes(id); + const active = agent === id; + return ( + + ); + })} +
{agentCommand} @@ -727,7 +880,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
- {t("settings.apiKeys.codingAgentsSwap")} + {detectedAgents.length > 0 + ? t("settings.apiKeys.codingAgentsDetectedHint", { + agents: detectedAgents + .map((id) => AGENT_LABELS[id] ?? id) + .join(", "), + }) + : t("settings.apiKeys.codingAgentsSwap")}
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index b67fd5ca1d..a9b5d839b3 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -445,6 +445,8 @@ export const en = { codingAgentsHint: "Launch a coding agent against this server. It uses the loaded model; a local server mints an API key automatically, a remote one includes it in the command.", codingAgentsSwap: "Swap claude for codex, openclaw, opencode, hermes, or pi.", + codingAgentDetected: "Installed on this machine", + codingAgentsDetectedHint: "Detected on this machine: {agents}.", relativeNever: "never", relativeJustNow: "just now", relativeHoursAgo: "{count}h ago", diff --git a/tests/studio/test_chat_response_details_ui_contract.py b/tests/studio/test_chat_response_details_ui_contract.py index 04301a0de6..89d3000629 100644 --- a/tests/studio/test_chat_response_details_ui_contract.py +++ b/tests/studio/test_chat_response_details_ui_contract.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + """Static contract for the chat response-details action and metadata.""" from __future__ import annotations diff --git a/tests/studio/test_usage_examples_agent_detection_contract.py b/tests/studio/test_usage_examples_agent_detection_contract.py new file mode 100644 index 0000000000..0dc16f0d30 --- /dev/null +++ b/tests/studio/test_usage_examples_agent_detection_contract.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Static contract for API usage-example agent detection scope.""" + +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +USAGE_EXAMPLES_TSX = REPO / "studio/frontend/src/features/settings/components/usage-examples.tsx" + + +def test_agent_detection_requires_desktop_scope(): + src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8") + assert 'import { isTauri } from "@/lib/api-base"' in src + assert "function canUseLocalAgentDetection(base: string): boolean" in src + helper = src[src.find("function canUseLocalAgentDetection") : src.find("const SHIKI_THEMES")] + assert "if (!isTauri) return false" in helper + assert "isLoopbackHost(normalizeHost(new URL(base).hostname))" in helper + assert "const localAgentDetection = canUseLocalAgentDetection(base)" in src + assert "if (!localAgentDetection)" in src + assert "}, [localAgentDetection]);" in src From 41dd95ea0a588343fc010ee24af4837cc8c08b98 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 05:33:16 -0700 Subject: [PATCH 068/113] Studio: don't pin transformers before the training worker activates the 5.x sidecar (#6968) * Studio: keep transformers off sys.modules until the training worker activates the sidecar The training worker (core/training/worker.py:run_training_process) decides the per-worker Xet env flip during preflight by importing utils/hf_xet_fallback.py, which eagerly imported unsloth_zoo at module load. unsloth_zoo's __init__ imports transformers, so the default transformers 4.57.x was cached in sys.modules before activate_transformers_for_subprocess prepended the 5.x sidecar to sys.path. Since activation only edits sys.path, the already cached module won, and 5.x models failed to load their tokenizer or config: - Qwen3.5 / GLM-4.7 (tokenizer_class TokenizersBackend): "Tokenizer class TokenizersBackend does not exist or is not currently imported." - gemma-4: "... is not supported yet in transformers==4.57.6." Fix: load the shared unsloth_zoo backend lazily (only when a heavy download helper is first used, which is after activation). child_should_disable_xet and the DEFAULT_* constants are defined locally so importing the shim stays light. The download wrappers, the DownloadStallError class, start_watchdog and get_hf_download_state resolve the shared backend on first use, and the degraded no-unsloth_zoo fallback is preserved. Tests: - test_hf_xet_fallback.py: existing suite kept green via the restored _shared_* seam; the GPU-init retry test now triggers the lazy load explicitly; new guard asserts importing child_should_disable_xet does not import transformers/unsloth_zoo. - test_training_worker_import_discipline.py: new invariant test that the worker preflight imports leave transformers unimported, so this class of regression cannot return silently. Runs in studio-backend-ci (CPU only, no network/GPU/weights). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: CPU-only guard that activation switches transformers to the model's sidecar version Adds test_worker_activates_correct_transformers.py: runs the real worker preflight (from utils.hf_xet_fallback import child_should_disable_xet) plus the real tier detection and activate_transformers_for_subprocess for a transformers-5.x model (Qwen3.5, tier 530), then asserts the in-process transformers actually switched to the 5.x sidecar. A stale pre-activation import leaves 4.57.x pinned and fails the assertion, which is exactly the TokenizersBackend regression (#6951). Self-contained CUDA spoof (mirrors tests/_zoo_aggressive_cuda_spoof.py) forces unsloth_zoo down its full, transformers-importing init path on a GPU-less runner; without it unsloth_zoo degrades and never preloads transformers, masking the bug. A one-line stub sidecar stands in for the 5.x venv, so no GPU, network, weights, or real sidecar are needed. Passes on this fix, fails on buggy main. * Studio: load the repo's canonical CUDA spoof in the correct-version guard Load tests/_zoo_aggressive_cuda_spoof.py (the committed spoof the consolidated CI already relies on) as the single source of truth so the guard matches CI and stays robust on a CPU-only torch wheel, where a partial hand-rolled spoof could miss a torch.cuda call and let the unsloth_zoo import raise (masking the bug). Falls back to a minimal inline spoof for a standalone studio checkout. Verified: passes on this fix, fails on buggy main, and the fallback path passes when the spoof file is absent. * Studio: declare the lazily-resolved xet names so ruff F822 stays green DownloadStallError, start_watchdog and get_hf_download_state are provided via the module __getattr__ (PEP 562), so ruff F822 flagged them as undefined names in __all__ and the Source-lint / pre-commit checks went red. Add annotation-only declarations (no value bound, so __getattr__ still resolves them lazily to the shared unsloth_zoo backend) to mark them defined for the linter while keeping F822 active for the rest of __all__. * Studio: tighten comments on the sidecar-activation fix and its tests * Studio: mirror the new MLX-dispatch preflight import in the import-discipline guard The worker preflight now also runs 'from core.training.training import is_apple_silicon_training_platform, should_use_mlx_training_backend' before it activates the transformers sidecar. Add that import (guarded) to the guard's preflight snippet so the invariant test stays a faithful mirror: a future change that makes core.training.training pull transformers/unsloth_zoo eagerly would then be caught too. Verified clean on the current tree (no leak). --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/tests/test_hf_xet_fallback.py | 42 ++- .../test_training_worker_import_discipline.py | 81 ++++ ...t_worker_activates_correct_transformers.py | 155 ++++++++ studio/backend/utils/hf_xet_fallback.py | 353 +++++++++++------- 4 files changed, 490 insertions(+), 141 deletions(-) create mode 100644 studio/backend/tests/test_training_worker_import_discipline.py create mode 100644 studio/backend/tests/test_worker_activates_correct_transformers.py diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 4d73213d15..2fff744b64 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -287,7 +287,9 @@ def test_degrades_when_shared_helper_import_raises_importerror(): def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim - retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails.""" + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails. + The backend loads lazily (first use of a heavy helper), so this triggers the load explicitly + before asserting the retry/degrade behavior.""" import importlib import os @@ -321,11 +323,15 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): sys.meta_path.insert(0, finder) try: degraded = importlib.import_module("utils.hf_xet_fallback") - # First attempt without the light env, then a retry with it set. + # Import is light (lazy backend); unsloth_zoo not loaded yet. + assert seen_env == [], seen_env + # First use of a heavy helper triggers the load (attempt without the light env, then a retry + # with it set); accessing DownloadStallError drives it via __getattr__. + stall_error = degraded.DownloadStallError assert seen_env == [None, "1"], seen_env # Both attempts raised -> Studio still boots in degraded mode. - assert issubclass(degraded.DownloadStallError, RuntimeError) - # The env override must not leak past the import. + assert issubclass(stall_error, RuntimeError) + # The env override must not leak past the load. assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None finally: sys.meta_path.remove(finder) @@ -333,3 +339,31 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): sys.modules.update(saved) if saved_shim is not None: sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_importing_child_should_disable_xet_stays_light(monkeypatch): + """Regression guard for the stale-transformers-sidecar bug: importing the shim (and + ``child_should_disable_xet``) must NOT pull in ``transformers``/``unsloth_zoo``. The worker calls + this at startup to decide the Xet env flip BEFORE activating the sidecar; an eager import here + would cache the default transformers 4.57.x in sys.modules, defeating the sidecar sys.path prepend + and breaking 5.x models (Qwen3.5/GLM/gemma-4).""" + import importlib + + for name in [ + m + for m in list(sys.modules) + if m == "transformers" + or m.startswith("transformers.") + or m == "unsloth_zoo" + or m.startswith("unsloth_zoo.") + or m == "utils.hf_xet_fallback" + ]: + monkeypatch.delitem(sys.modules, name, raising = False) + + mod = importlib.import_module("utils.hf_xet_fallback") + # The lightweight decision works without the heavy backend. + assert mod.child_should_disable_xet({"disable_xet": True}) is True + assert mod.child_should_disable_xet({}) is False + # And nothing heavy was imported as a side effect. + assert "transformers" not in sys.modules, "importing the shim must not import transformers" + assert "unsloth_zoo" not in sys.modules, "importing the shim must not import unsloth_zoo" diff --git a/studio/backend/tests/test_training_worker_import_discipline.py b/studio/backend/tests/test_training_worker_import_discipline.py new file mode 100644 index 0000000000..a047c91704 --- /dev/null +++ b/studio/backend/tests/test_training_worker_import_discipline.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: the training worker must not import ``transformers`` before it activates the +transformers sidecar. + +``core/training/worker.py:run_training_process`` runs a preflight (Xet decision, logging, hardware +detection) and only THEN calls ``_activate_transformers_version`` -> ``activate_transformers_for_subprocess``, +which prepends the correct ``.venv_t5_*`` (5.x) sidecar to ``sys.path``. Because activation only edits +``sys.path``, it is a no-op for any module already cached in ``sys.modules``. So if the preflight imports +``transformers`` (directly or transitively via ``unsloth_zoo``), the default 4.57.x gets pinned before +the sidecar is on the path -- and 5.x models (Qwen3.5, GLM-4.7, gemma-4) then fail to load their +tokenizer/config ("Tokenizer class TokenizersBackend does not exist"). + +This regression shipped once when ``utils/hf_xet_fallback.py`` eagerly imported ``unsloth_zoo`` (which +imports ``transformers``) at module load; the worker imports that shim during preflight to decide the +Xet env flip (see issue #6951). This test locks the invariant in a fresh interpreter. It is CPU-only, +needs no network/GPU/weights/sidecars, so it runs in the standard ``studio-backend-ci`` matrix. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend + +# Mirrors run_training_process's imports that run BEFORE _activate_transformers_version (worker.py); +# keep in sync. torch-dependent imports are optional (a no-torch CI shard skips them) but must still +# not drag in transformers. +_PREFLIGHT_SNIPPET = r""" +import sys + +# worker.py: from utils.hf_xet_fallback import child_should_disable_xet (+ call it) +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) + +# worker.py: from loggers.config import LogConfig +from loggers.config import LogConfig # noqa: F401 + +# worker.py: from utils.hardware import hardware (imports torch, not transformers) +try: + from utils.hardware import hardware as _hw # noqa: F401 +except Exception: + pass # torch may be absent in a no-torch shard; the invariant below still applies + +# worker.py: from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend +# (the MLX-dispatch preflight; must also stay clear of transformers). Guarded because it may pull +# unsloth/trl, absent in a minimal shard -- but a partial import that leaked transformers would still +# be caught by the assertion below. +try: + from core.training.training import ( # noqa: F401 + is_apple_silicon_training_platform as _is_apple, + should_use_mlx_training_backend as _use_mlx, + ) +except Exception: + pass + +leaked_tf = sorted(m for m in sys.modules if m == "transformers" or m.startswith("transformers.")) +leaked_zoo = sorted(m for m in sys.modules if m == "unsloth_zoo" or m.startswith("unsloth_zoo.")) +assert not leaked_tf, f"transformers imported during worker preflight (before sidecar activation): {leaked_tf}" +assert not leaked_zoo, f"unsloth_zoo imported during worker preflight (before sidecar activation): {leaked_zoo}" +print("PREFLIGHT_CLEAN") +""" + + +def test_worker_preflight_does_not_import_transformers(): + """A fresh interpreter running the worker's pre-activation imports must leave ``transformers`` + (and ``unsloth_zoo``) unimported, so the 5.x sidecar prepend is not defeated by a stale module.""" + result = subprocess.run( + [sys.executable, "-c", _PREFLIGHT_SNIPPET], + cwd = str(_BACKEND_DIR), + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight imported transformers before sidecar activation.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + assert "PREFLIGHT_CLEAN" in result.stdout, result.stdout diff --git a/studio/backend/tests/test_worker_activates_correct_transformers.py b/studio/backend/tests/test_worker_activates_correct_transformers.py new file mode 100644 index 0000000000..fe7b8dd25a --- /dev/null +++ b/studio/backend/tests/test_worker_activates_correct_transformers.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: after the training worker runs its preflight and then activates the transformers +sidecar, the in-process ``transformers`` must be the sidecar version the model requires -- not the +default 4.57.x that the base environment ships. + +The CPU-only "does it choose the correct transformers version" guard, stronger than the pure +import-order check in ``test_training_worker_import_discipline.py``: it runs the REAL tier detection +(``get_transformers_tier``) and REAL activation (``activate_transformers_for_subprocess``) for a +transformers-5.x model (Qwen3.5, tier 530) and asserts the version actually switched. It catches the +whole failure family at once: + + * a stale pre-activation ``transformers`` import (the #6951 / ``TokenizersBackend`` regression: an + already-cached 4.57.x defeats the sidecar's ``sys.path`` prepend), + * a wrong tier selected for a 5.x model, and + * activation not actually swapping the resident module. + +Why the CUDA spoof matters (verified): ``unsloth_zoo``'s eager ``import transformers`` only happens on +its full, GPU-present init path. On a GPU-less runner it silently degrades and never preloads +transformers -- which would MASK the stale-import bug (the check would falsely pass). Spoofing +``torch.cuda`` so ``unsloth_zoo`` believes a GPU is present forces the real init path, exposing the +regression on CPU CI. The spoof mirrors ``tests/_zoo_aggressive_cuda_spoof.py`` but is inlined so the +test is self-contained in the ``studio-backend-ci`` matrix (whose conftest does not apply the shared +spoof). No GPU/network/weights/real sidecar needed: a one-line stub sidecar stands in for the 5.x venv, +so we only assert activation lands on it. + +Proven: passes on the fixed tree (active == 5.3.0) and fails on the buggy tree (active == 4.57.x) on +a simulated GPU-less runner. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend +# Canonical CUDA spoof at the repo root (studio/backend -> studio -> repo root). Loaded by the +# subprocess when present (matches the consolidated CI); absent in a standalone studio checkout, where +# the subprocess falls back to a minimal inline spoof. +_SPOOF_PATH = _BACKEND_DIR.parent.parent / "tests" / "_zoo_aggressive_cuda_spoof.py" + +# Runs in a fresh interpreter with cwd == studio/backend so ``utils.*`` resolves like the worker. +# STUB_HOME (a pytest tmp dir) holds a throwaway ``.venv_t5_530`` sidecar exporting transformers 5.3.0. +_SNIPPET = r""" +import os, sys +sys.path.insert(0, os.getcwd()) + +# CUDA spoof so unsloth_zoo takes its full, transformers-importing init path on a GPU-less runner. +# Without it unsloth_zoo degrades and never preloads transformers, which would MASK the stale-import +# regression under test (verified). Prefer the repo's canonical spoof (single source of truth, and the +# one the consolidated CI already relies on); fall back to a minimal inline spoof so this also works in +# a standalone studio checkout. If torch is absent the fixed tree still passes below; the bug just +# would not be exposable in that shard. +try: + import torch # noqa: F401 + _sp = os.environ.get("SPOOF_PATH") + if _sp and os.path.exists(_sp): + import importlib.util + _spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", _sp) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + _mod.apply() + else: + torch.cuda.is_available = lambda: True + torch.cuda.device_count = lambda: 1 + torch.cuda.current_device = lambda: 0 + torch.cuda.get_device_capability = lambda *a, **k: (8, 0) + torch.cuda.get_device_name = lambda *a, **k: "NVIDIA A100-SPOOFED" + torch.cuda.is_bf16_supported = lambda *a, **k: True + class _Props: + name = "NVIDIA A100-SPOOFED" + major = 8 + minor = 0 + total_memory = 80 * 1024**3 + multi_processor_count = 108 + torch.cuda.get_device_properties = lambda *a, **k: _Props() + torch.cuda.mem_get_info = lambda *a, **k: (0, 80 * 1024**3) +except Exception: + pass +os.environ["UNSLOTH_IS_PRESENT"] = "1" + +# Stub 5.x sidecar: activation only edits sys.path, so a package that merely exports __version__ is +# enough to prove the resident transformers switched to it. +home = os.environ["STUB_HOME"] +pkg = os.path.join(home, ".venv_t5_530", "transformers") +os.makedirs(pkg, exist_ok = True) +with open(os.path.join(pkg, "__init__.py"), "w") as f: + f.write('__version__ = "5.3.0"\n') +os.environ["UNSLOTH_STUDIO_HOME"] = home + +# Faithful worker preflight (worker.py: from utils.hf_xet_fallback import child_should_disable_xet). +# This is the exact stale-import trigger: on the buggy tree it pulls unsloth_zoo -> transformers 4.57.x +# into sys.modules BEFORE activation. +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) +_tf = sys.modules.get("transformers") +preload = _tf.__version__ if _tf is not None else None + +# Real tier detection + real activation, with the 530 sidecar pointed at the stub above. +import utils.transformers_version as tv +tv._VENV_T5_530_DIR = os.path.join(home, ".venv_t5_530") +tv._ensure_venv_t5_530_exists = lambda: True +tier = tv.get_transformers_tier("Qwen/Qwen3.5-9B", None) +tv.activate_transformers_for_subprocess("Qwen/Qwen3.5-9B", None) + +import transformers +print(f"RESULT tier={tier} preload={preload} active={transformers.__version__}") +""" + + +def _parse(stdout: str) -> dict[str, str]: + for line in stdout.splitlines(): + if line.startswith("RESULT "): + return dict(kv.split("=", 1) for kv in line.split()[1:]) + return {} + + +def test_worker_activates_correct_transformers_version(tmp_path): + """The worker's real preflight + activation for a transformers-5.x model (Qwen3.5, tier 530) must + leave the in-process ``transformers`` on the 5.x sidecar. A stale pre-activation import leaves the + default 4.57.x pinned and fails this assertion -- exactly the #6951 ``TokenizersBackend`` regression.""" + result = subprocess.run( + [sys.executable, "-c", _SNIPPET], + cwd = str(_BACKEND_DIR), + env = { + **__import__("os").environ, + "STUB_HOME": str(tmp_path), + **({"SPOOF_PATH": str(_SPOOF_PATH)} if _SPOOF_PATH.exists() else {}), + }, + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight + activation harness crashed.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + parsed = _parse(result.stdout) + assert parsed, f"No RESULT line.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + + # Correct tier chosen for a transformers-5.x model (pure, deterministic; no network/GPU). + assert parsed["tier"] == "530", ( + f"Wrong transformers tier for Qwen3.5 (expected 530, got {parsed['tier']}). " + "Tier detection regressed." + ) + + # Activation must actually swap the resident transformers to the sidecar version. If a preflight + # import cached 4.57.x first, the sidecar prepend is a no-op and this stays 4.57.x -- the bug. + assert parsed["active"] == "5.3.0", ( + "Sidecar activation did NOT switch the in-process transformers to the model's 5.x version " + f"(active={parsed['active']}, preloaded-before-activation={parsed['preload']}). A pre-activation " + "transformers import (directly or via unsloth_zoo) defeated the sidecar; 5.x models (Qwen3.5, " + "GLM-4.7, gemma-4) then fail with 'Tokenizer class TokenizersBackend does not exist'. See #6951." + ) diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 2dd2247396..9bc4a60fad 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -6,6 +6,16 @@ Re-exports the shared API and injects Studio's marker-aware cache purge (``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` marker semantics on the HTTP retry. + +Import discipline: ``unsloth_zoo``'s ``__init__`` eagerly imports ``transformers``. The workers +import this shim at startup (to decide the per-worker Xet env flip) *before* activating the model's +``transformers`` sidecar. Activation only prepends the sidecar to ``sys.path``, so a ``transformers`` +already cached in ``sys.modules`` (via an eager ``unsloth_zoo`` import here) wins -- pinning the +default 4.57.x and regressing Qwen3.5 / GLM-4.7 / gemma-4 training with +``Tokenizer class TokenizersBackend does not exist``. So the shared backend is loaded **lazily** +(``_load_shared``), only on first use of a heavy download helper, i.e. after the sidecar is active. +``child_should_disable_xet`` and the ``DEFAULT_*`` constants are defined locally so importing them +never triggers the heavy load. """ from __future__ import annotations @@ -13,161 +23,230 @@ from __future__ import annotations import threading from typing import Any, Callable, Optional -_shared_import_error = None -try: - import unsloth_zoo.hf_xet_fallback as _shared - _shared_available = True -except Exception as _exc: # noqa: BLE001 - any import failure must degrade, not crash - # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less Studio - # host. The download helper needs none of it, so retry via the light UNSLOTH_ZOO_DISABLE_GPU_INIT - # path before giving up. - _shared_import_error = _exc - import os as _os +# Defaults mirror unsloth_zoo.hf_xet_fallback; plain literals so they resolve (including as +# default args below) without importing unsloth_zoo/transformers. +DEFAULT_GRACE_PERIOD = 10.0 +DEFAULT_HEARTBEAT_INTERVAL = 30.0 +DEFAULT_STALL_TIMEOUT = 180.0 - _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") - _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" - try: - import unsloth_zoo.hf_xet_fallback as _shared - _shared_available = True - _shared_import_error = None - except Exception as _exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF downloads - _shared_import_error = _exc2 - _shared_available = False - finally: - if _prev_gpu_init is None: - _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) - else: - _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init +# --- lazy shared-backend loader ---------------------------------------------------------------- +_shared: Any = None +_shared_available: Optional[bool] = None # None = not yet attempted +_shared_import_error: Optional[BaseException] = None +_load_lock = threading.Lock() -if _shared_available: - # Bind by assignment so each public name shares one module-level binding with the degraded branch. - DEFAULT_GRACE_PERIOD = _shared.DEFAULT_GRACE_PERIOD - DEFAULT_HEARTBEAT_INTERVAL = _shared.DEFAULT_HEARTBEAT_INTERVAL - DEFAULT_STALL_TIMEOUT = _shared.DEFAULT_STALL_TIMEOUT - DownloadStallError = _shared.DownloadStallError - child_should_disable_xet = _shared.child_should_disable_xet - get_hf_download_state = _shared.get_hf_download_state - start_watchdog = _shared.start_watchdog - _shared_hf_hub_download_with_xet_fallback = _shared.hf_hub_download_with_xet_fallback - _shared_snapshot_download_with_xet_fallback = _shared.snapshot_download_with_xet_fallback -else: - # Degrade instead of crashing Studio: plain HF downloads, stall watchdog disabled. Thin stubs, - # not a second copy of the orchestration; recovery returns once unsloth_zoo is upgraded. - import logging as _logging - _logging.getLogger(__name__).warning( - "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " - "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " - "re-enable automatic Xet -> HTTP download recovery.", - _shared_import_error, - ) +def _load_shared() -> bool: + """Import ``unsloth_zoo.hf_xet_fallback`` on demand; return True if available. Deferred so + importing this module at worker startup does not pull transformers in before the sidecar is + activated. Degrades (returns False) rather than crashing when unsloth_zoo is unavailable.""" + global _shared, _shared_available, _shared_import_error + if _shared_available is not None: + return _shared_available + with _load_lock: + if _shared_available is not None: + return _shared_available + try: + import unsloth_zoo.hf_xet_fallback as shared - DEFAULT_HEARTBEAT_INTERVAL = 30.0 - DEFAULT_STALL_TIMEOUT = 180.0 - DEFAULT_GRACE_PERIOD = 10.0 + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc: # noqa: BLE001 - any import failure must degrade, not crash + # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less + # host. The download helper needs none of it, so retry via UNSLOTH_ZOO_DISABLE_GPU_INIT. + _shared_import_error = exc + import os as _os - class DownloadStallError(RuntimeError): - """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" + try: + import unsloth_zoo.hf_xet_fallback as shared - def child_should_disable_xet(config: dict) -> bool: - return bool(config.get("disable_xet")) + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF + _shared_import_error = exc2 + _shared_available = False + import logging as _logging - def get_hf_download_state(*args: Any, **kwargs: Any) -> None: - return None # unmeasurable -> the (absent) watchdog never fires + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " + "re-enable automatic Xet -> HTTP download recovery.", + _shared_import_error, + ) + return False + finally: + if _prev_gpu_init is None: + _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) + else: + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init - def start_watchdog( - *, - on_heartbeat: "Optional[Callable[[str], None]]" = None, - interval: float = DEFAULT_HEARTBEAT_INTERVAL, - xet_disabled: bool = False, - **kwargs: Any, - ) -> "threading.Event": - # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline - # is not tripped during a long download. - stop = threading.Event() - if on_heartbeat is None: - return stop - transport = "https" if xet_disabled else "xet" - def _beat() -> None: - while not stop.wait(interval): - try: - on_heartbeat(f"Downloading ({transport} transport)...") - except Exception: - pass +def child_should_disable_xet(config: dict) -> bool: + """Single source of truth for the per-worker Xet env flip (mirrors + ``unsloth_zoo.hf_xet_fallback.child_should_disable_xet``). Deliberately lightweight: importing or + calling it must NOT pull in unsloth_zoo/transformers, so the worker can decide before activating + the transformers sidecar (see the module docstring).""" + return bool(config.get("disable_xet")) - threading.Thread( - target = _beat, - daemon = True, - name = "hf-xet-degraded-heartbeat", - ).start() + +# --- degraded stubs (used only when unsloth_zoo is unavailable) ------------------------------- +class _DegradedDownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + + +def _degraded_get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires + + +def _degraded_start_watchdog( + *, + on_heartbeat: "Optional[Callable[[str], None]]" = None, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + xet_disabled: bool = False, + **kwargs: Any, +) -> "threading.Event": + # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline + # is not tripped during a long download. + stop = threading.Event() + if on_heartbeat is None: return stop + transport = "https" if xet_disabled else "xet" - def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: - return cancel_event is not None and cancel_event.is_set() + def _beat() -> None: + while not stop.wait(interval): + try: + on_heartbeat(f"Downloading ({transport} transport)...") + except Exception: + pass - def _shared_hf_hub_download_with_xet_fallback( - repo_id: str, - filename: str, - token: Optional[str], - *, - repo_type: str = "model", - revision: Optional[str] = None, - cache_dir: Optional[str] = None, - force_download: bool = False, - cancel_event: "Optional[threading.Event]" = None, - **_ignored: Any, - ) -> str: - # Keep the cancellation contract: do not start or return a download once cancelled. - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() + return stop - from huggingface_hub import hf_hub_download - path = hf_hub_download( - repo_id = repo_id, - filename = filename, - token = token, - repo_type = repo_type, - revision = revision, - cache_dir = cache_dir, - force_download = force_download, - ) - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - return path +def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() - def _shared_snapshot_download_with_xet_fallback( - repo_id: str, - *, - revision: Optional[str] = None, - token: Optional[str] = None, - repo_type: str = "model", - cache_dir: Optional[str] = None, - allow_patterns: Optional[Any] = None, - ignore_patterns: Optional[Any] = None, - force_download: bool = False, - cancel_event: "Optional[threading.Event]" = None, - **_ignored: Any, - ) -> str: - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - from huggingface_hub import snapshot_download +def _degraded_hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + repo_type: str = "model", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + # Keep the cancellation contract: do not start or return a download once cancelled. + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") - path = snapshot_download( - repo_id = repo_id, - repo_type = repo_type, - revision = revision, - token = token, - cache_dir = cache_dir, - allow_patterns = allow_patterns, - ignore_patterns = ignore_patterns, - force_download = force_download, - ) - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - return path + from huggingface_hub import hf_hub_download + + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + token = token, + repo_type = repo_type, + revision = revision, + cache_dir = cache_dir, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +def _degraded_snapshot_download_with_xet_fallback( + repo_id: str, + *, + revision: Optional[str] = None, + token: Optional[str] = None, + repo_type: str = "model", + cache_dir: Optional[str] = None, + allow_patterns: Optional[Any] = None, + ignore_patterns: Optional[Any] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + + from huggingface_hub import snapshot_download + + path = snapshot_download( + repo_id = repo_id, + repo_type = repo_type, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +# --- lazy attribute access for the heavy shared API ------------------------------------------- +# ``DownloadStallError`` (class identity matters for ``except``), ``start_watchdog`` and +# ``get_hf_download_state`` come from the shared backend when available, else the degraded stubs. +# Resolved via PEP 562 ``__getattr__`` so ``from utils.hf_xet_fallback import X`` triggers the load +# only for these heavy names, not for ``child_should_disable_xet`` / ``DEFAULT_*``. +_DEGRADED_ATTRS = { + "DownloadStallError": _DegradedDownloadStallError, + "start_watchdog": _degraded_start_watchdog, + "get_hf_download_state": _degraded_get_hf_download_state, +} + +# Annotation-only declarations for the three names above: they bind NO value, so lookup still misses +# and PEP 562 ``__getattr__`` resolves them lazily -- but ruff/pyflakes see them as defined, so listing +# them in ``__all__`` does not trip F822 (while F822 still catches a real typo elsewhere in the list). +DownloadStallError: type +start_watchdog: Any +get_hf_download_state: Any + + +def __getattr__(name: str) -> Any: + if name in _DEGRADED_ATTRS: + if _load_shared(): + return getattr(_shared, name) + return _DEGRADED_ATTRS[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +# Indirection seam the public wrappers call (and tests monkeypatch): lazy-load the shared backend, +# then dispatch to it or the degraded stub. The ``_shared_*`` names preserve the pre-refactor contract. +def _shared_hf_hub_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.hf_hub_download_with_xet_fallback + if _load_shared() + else _degraded_hf_hub_download_with_xet_fallback + ) + return impl(*args, **kwargs) + + +def _shared_snapshot_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.snapshot_download_with_xet_fallback + if _load_shared() + else _degraded_snapshot_download_with_xet_fallback + ) + return impl(*args, **kwargs) __all__ = [ From fcb1152c76417c9ae6d6c649a5036a36110c69c1 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Wed, 8 Jul 2026 09:34:59 -0300 Subject: [PATCH 069/113] Studio: source CPU llama.cpp prebuilts from unslothai/llama.cpp (#6311) * Studio: source CPU llama.cpp prebuilts from the unslothai fork * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: reject unknown Linux CPU arches and keep ROCm-tooling hosts off the CPU prebuilt * Studio: extend the resolve-prebuilt ROCm-tooling guard to Windows * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: let ROCm-SDK-only CPU hosts take the fork CPU prebuilt * Studio: accept windows-arm64 prebuilt kind and refresh stale fork-routing comments * Studio: correct stale fork-routing comments and --resolve-prebuilt help * Refresh stale ggml-org routing comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .../workflows/studio-windows-update-smoke.yml | 6 +- .../tests/test_install_resolve_prebuilt.py | 105 +++-------- .../backend/tests/test_llama_cpp_freshness.py | 6 +- studio/backend/tests/test_llama_cpp_update.py | 7 +- studio/install_llama_prebuilt.py | 63 +++---- studio/setup.ps1 | 23 +-- studio/setup.sh | 74 ++------ .../install/test_llama_pr_force_and_source.py | 14 +- tests/studio/install/test_pr4562_bugfixes.py | 18 +- tests/studio/install/test_rocm_support.py | 100 +++++----- tests/studio/install/test_selection_logic.py | 171 +++++++++++++++++- 11 files changed, 322 insertions(+), 265 deletions(-) diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 888b3d70a3..5b92f1a3e0 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -6,9 +6,9 @@ # windows-latest runner: # # 1. install.ps1 --local --no-torch installs Studio AND auto-fetches -# the prebuilt llama.cpp Windows binary (llama-bNNNN-bin-win-cpu- -# x64 from ggml-org/llama.cpp). Hitting the source-build fallback -# is treated as an Unsloth bug -- Studio must always pick the +# the prebuilt llama.cpp Windows binary (app--windows-x64-cpu +# from unslothai/llama.cpp). Hitting the source-build fallback is +# treated as an Unsloth bug -- Studio must always pick the # prebuilt on Windows. # 2. unsloth studio update --local is idempotent. Two consecutive # runs both report "prebuilt up to date and validated", no diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index e9941d9e62..b825172a63 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode. +"""install_llama_prebuilt.py: the --resolve-prebuilt probe (plans against the fork +by default; --published-repo overrides). These back the in-app update for source-build (markerless) installs: the backend asks the installer whether an official prebuilt exists for this host without @@ -24,9 +25,7 @@ if str(_studio) not in sys.path: ilp = importlib.import_module("install_llama_prebuilt") -if not hasattr(ilp, "published_repo_for_host") or not hasattr( - ilp, "resolve_simple_install_release_plans" -): +if not hasattr(ilp, "resolve_simple_install_release_plans"): pytest.skip("PR symbols not present - check branch", allow_module_level = True) FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp @@ -56,73 +55,6 @@ def _host(**kw): return ilp.HostInfo(**base) -def test_published_repo_for_host(): - # CPU-only Linux (x64 and arm64) -> ggml-org upstream. - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) - == UPSTREAM - ) - # GPU Linux -> fork. - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) - == FORK - ) - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK - # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). - assert ( - ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) - == UPSTREAM - ) - # GPU Windows -> fork. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) - ) - == FORK - ) - # macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS). - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - # Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling). - assert ( - ilp.published_repo_for_host( - _host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True - ) - == FORK - ) - # The tooling hint is Linux-only: Windows CPU stays on ggml-org. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True), - linux_amd_tooling_present = True, - ) - == UPSTREAM - ) - - -def test_macos_intel_and_arm_both_route_to_fork(): - # macOS uses the unslothai fork's own Mac prebuilts for BOTH arm64 and Intel; - # there is no longer any upstream-on-macOS default path, so the obsolete - # pre-macOS-26 pin (b9415) is gone. - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_x86_64 = True, machine = "x86_64") - ) - == FORK - ) - - def test_macos_upstream_pin_only_for_explicit_pre26_upstream(): pre26 = _host( system = "Darwin", @@ -188,15 +120,13 @@ def test_resolve_prebuilt_unavailable(monkeypatch, capsys): assert out["repo"] == FORK -def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): - # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the - # fork so a HIP source build is not offered an upstream CPU prebuilt. - monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) - monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") +def _run_resolve_capture_host(monkeypatch, capsys): + """Drive --resolve-prebuilt and return the host the resolver was handed.""" seen = {} def _resolver(tag, host, repo, published_release_tag): seen["repo"] = repo + seen["host"] = host raise ilp.PrebuiltFallback("no asset") monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) @@ -207,10 +137,33 @@ def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): ) assert ilp.main() == ilp.EXIT_SUCCESS out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + return seen, out + + +def test_resolve_prebuilt_cpu_linux_routes_to_fork(monkeypatch, capsys): + # CPU-only Linux host (no GPU): the dispatch routes to the fork, which now + # ships the CPU prebuilt -- it no longer falls back to ggml-org upstream. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == FORK assert out["repo"] == FORK +def test_resolve_prebuilt_rocm_sdk_only_host_still_offered_cpu(monkeypatch, capsys): + # A CPU-only host that merely has ROCm/HIP SDK tools on PATH (no AMD GPU, so + # detect_host leaves has_rocm False) is a valid CPU-prebuilt target. The probe + # must NOT reclassify it as ROCm from tool presence alone and suppress the CPU + # bundle -- that would deny the fork CPU prebuilt to a legitimate CPU source + # build. The host is left CPU-only and resolves against the fork. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr( + ilp.shutil, "which", lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == FORK + assert seen["host"].has_rocm is False + + # Blackwell floor is sm_100 (data-center B100/B200, B300/GB300), below consumer # sm_120 -- 120 wrongly excluded data-center hosts from the prebuilt selection. diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index 2a2e113585..08e1334ac9 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -137,9 +137,9 @@ def test_read_install_marker_finds_windows_cmake_layout(tmp_path): @pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"]) def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo): - # The freshness check queries whichever release repo the marker records, - # so CUDA (unslothai), CPU/macOS (ggml-org), and ROCm all get the right - # "latest" tag. + # The freshness check queries whichever release repo the marker records: + # new installs record the fork, legacy CPU/macOS markers still say ggml-org, + # and both must get the right "latest" tag. install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b9000", published_repo = repo) bin_path = _fake_binary(install_dir, layout = "cmake") diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 4ceffbf75b..5138e90471 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -594,9 +594,10 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): - # CPU installs come from ggml-org. Re-running into the same install-dir/repo - # reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU - # detection) is reserved for setup.sh's arm64 rescue and must not appear here. + # Legacy CPU installs recorded a ggml-org marker (new installs use the fork). + # Re-running into the same install-dir/repo reproduces the same CPU bundle; + # --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's + # arm64 rescue and must not appear here. cmd = _capture_install_cmd( monkeypatch, tmp_path, diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index e40cb3083e..6c75e6c394 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -165,9 +165,9 @@ def env_int( # errors. Only use "master" temporarily when the latest release is missing # support for a new model architecture. DEFAULT_LLAMA_TAG = os.environ.get("UNSLOTH_LLAMA_TAG", "latest") -# Default published repo for prebuilt release resolution. Linux uses -# Unsloth prebuilts; setup.sh/setup.ps1 pass --published-repo explicitly -# for macOS/Windows to override with ggml-org/llama.cpp when needed. +# Default published repo for prebuilt release resolution. Every host plans +# its prebuilt against the Unsloth fork; setup.sh/setup.ps1 pass it via +# --published-repo. ggml-org is reachable only via an explicit override. DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" DEFAULT_PUBLISHED_TAG = os.environ.get("UNSLOTH_LLAMA_RELEASE_TAG") DEFAULT_PUBLISHED_MANIFEST_ASSET = os.environ.get( @@ -3135,21 +3135,6 @@ def _apply_host_overrides( return host -def published_repo_for_host(host: HostInfo, *, linux_amd_tooling_present: bool = False) -> str: - """The release repo setup.sh / setup.ps1 pick for this host: macOS always the - fork (ggml-org macOS bundles need too-new macOS); else CPU-only Linux/Windows - -> ggml-org upstream (the fork ships no CPU bundle) and any usable GPU (NVIDIA - or ROCm) -> the fork. linux_amd_tooling_present mirrors setup.sh routing Linux - hosts that expose AMD tooling (rocminfo/amd-smi/hipconfig/hipinfo) to the fork - even when the probe cannot confirm an active GPU. Mirrors the shell routing.""" - if host.is_macos: - return DEFAULT_PUBLISHED_REPO - has_gpu = ( - host.has_usable_nvidia or host.has_rocm or (host.is_linux and linux_amd_tooling_present) - ) - return DEFAULT_PUBLISHED_REPO if has_gpu else UPSTREAM_REPO - - def pick_windows_cuda_runtime(host: HostInfo) -> str | None: if not host.driver_cuda_version: return None @@ -4015,6 +4000,9 @@ def resolve_release_asset_choice( published_choice = published_rocm_choice_for_host(release, host, "windows-rocm") else: published_choice = published_asset_choice_for_kind(release, "windows-cpu") + elif host.is_windows and host.is_arm64: + # Windows arm64 has no GPU prebuilt, so it always takes the CPU bundle. + published_choice = published_asset_choice_for_kind(release, "windows-arm64") elif host.is_macos and host.is_arm64: published_choice = published_asset_choice_for_kind(release, "macos-arm64") elif host.is_macos and host.is_x86_64: @@ -6127,8 +6115,13 @@ def _linux_published_attempts(host: HostInfo, bundle: PublishedReleaseBundle) -> # CPU-only host. A usable-NVIDIA host never reaches here -- if its CUDA # selection produced nothing we want an empty attempt list so the caller # source-builds with CUDA, not a CPU-only binary silently installed on a - # GPU host (mirrors the ROCm branch, and Windows NVIDIA). - cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu") + # GPU host (mirrors the ROCm branch, and Windows NVIDIA). Only x86_64 and + # arm64 have a CPU bundle; any other Linux arch (ppc64le, riscv64, s390x) + # has none, so leave attempts empty and source-build rather than hand it + # the x86_64 linux-cpu binary (the Linux preflight checks libraries, not + # ELF arch, so a wrong-arch binary would not be caught). + kind = "linux-cpu" if host.is_x86_64 else "linux-arm64" if host.is_arm64 else None + cpu_choice = published_asset_choice_for_kind(bundle, kind) if kind else None if cpu_choice is not None: attempts.append(cpu_choice) return attempts @@ -6143,9 +6136,9 @@ def _fork_manifest_release_plans( max_release_fallbacks: int = DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS, ) -> tuple[str, list[InstallReleasePlan]]: """Manifest-reading branch of resolve_simple_install_release_plans, used for - the fork's bundles whose GPU/arch coverage lives in - llama-prebuilt-manifest.json rather than in the filename: arm64 CUDA, Windows - CUDA, per-gfx ROCm, and macOS. Linux x64 takes the faster filename path.""" + every fork host: all of the fork's bundles describe their GPU/arch coverage + in llama-prebuilt-manifest.json rather than in the asset filename (CPU, + x64/arm64 CUDA, Windows CUDA, per-gfx ROCm, and macOS).""" requested_tag = normalized_requested_llama_tag(llama_tag) allow_older_release_fallback = requested_tag == "latest" and not published_release_tag release_limit = max(1, max_release_fallbacks) @@ -6714,8 +6707,8 @@ def install_prebuilt( log( f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install" ) - # Single resolver: linux-x64 takes the fast filename path internally, - # every other fork host reads the manifest. + # Single resolver: every fork host selects from the release manifest; + # an explicit ggml-org override selects by asset filename instead. requested_tag, release_plans = resolve_simple_install_release_plans( llama_tag, host, @@ -6903,8 +6896,8 @@ def parse_args() -> argparse.Namespace: const = "latest", help = ( "Report whether an official prebuilt exists for this host without " - "downloading. Picks the host's published repo when --published-repo " - "is left at the default. Use --output-format json." + "downloading. Plans against --published-repo (defaults to the " + "fork). Use --output-format json." ), ) parser.add_argument( @@ -6992,24 +6985,16 @@ def main() -> int: return EXIT_SUCCESS if args.resolve_prebuilt is not None: - # Host-aware "is a prebuilt available" probe, no download. A default repo - # means "pick the repo for this host"; PrebuiltFallback == source build. + # Host-aware "is a prebuilt available" probe, no download. Every host now + # plans against the fork (args.published_repo defaults to it); an explicit + # --published-repo overrides. PrebuiltFallback == source build. host = _apply_host_overrides( detect_host(), override_has_rocm = args.has_rocm, override_rocm_gfx = args.rocm_gfx, force_cpu = args.cpu_fallback, ) - # setup.sh routes Linux hosts with AMD tooling to the fork even when no GPU - # is probed; mirror that so a HIP source build is not offered a CPU prebuilt. - amd_tooling = host.is_linux and any( - shutil.which(t) for t in ("rocminfo", "amd-smi", "hipconfig", "hipinfo") - ) - repo = ( - published_repo_for_host(host, linux_amd_tooling_present = amd_tooling) - if args.published_repo == DEFAULT_PUBLISHED_REPO - else args.published_repo - ) + repo = args.published_repo try: _requested, plans = resolve_simple_install_release_plans( args.resolve_prebuilt, host, repo, args.published_release_tag or "" diff --git a/studio/setup.ps1 b/studio/setup.ps1 index bb1e88cc4e..07dcb17335 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -3088,12 +3088,11 @@ $LlamaCppDir = Join-Path $UnslothHome "llama.cpp" $NeedLlamaSourceBuild = $false $SkipPrebuiltInstall = $false $RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { $DefaultLlamaTag } -# GPU Windows (CUDA / ROCm) installs the fork's app-* prebuilts; CPU-only stays -# on ggml-org (the fork ships no windows-cpu bundle). Mirrors setup.sh's routing. -# A resolved gfx arch counts as a GPU host even when $HasROCm is false (Adrenalin -# driver only, no HIP runtime): the fork's per-gfx bundle ships its own runtime, -# so route there instead of ggml-org / a CPU build. -$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) { "unslothai/llama.cpp" } else { "ggml-org/llama.cpp" } +# Every host installs the fork's app-* prebuilts now: GPU Windows (CUDA / ROCm) +# already did, and the fork now also ships the CPU bundles for Windows x64 and +# arm64 (windows-cpu / windows-arm64). ggml-org artifacts are no longer used by +# default. Mirrors setup.sh's routing. +$HelperReleaseRepo = "unslothai/llama.cpp" $LlamaPr = if ($env:UNSLOTH_LLAMA_PR) { $env:UNSLOTH_LLAMA_PR.Trim() } else { "" } $LlamaPrForce = if ($env:UNSLOTH_LLAMA_PR_FORCE) { $env:UNSLOTH_LLAMA_PR_FORCE.Trim() } else { $DefaultLlamaPrForce } @@ -3283,11 +3282,13 @@ if ($LocalLlamaCppLinked) { # treat a valid ROCm install as mismatched. A name-inferred gfx # arch (Adrenalin-only, no confirmed runtime) still counts as # ROCm-capable -- the ROCm prebuilt bundles its own runtime, - # mirroring the --rocm-gfx forward below. NOTE: this block is - # currently inert -- write_prebuilt_metadata does not persist an - # install_kind key, so $existingKind is always null. If that changes, - # add the remaining host kinds (e.g. windows-arm64) before relying on it. - $expectedKinds = if ($HasROCm -or $script:ROCmGfxArch) { @("windows-rocm", "windows-hip") } elseif ($HasNvidiaSmi) { @("windows-cuda") } else { @("windows-cpu") } + # mirroring the --rocm-gfx forward below. The CPU branch covers both + # the x64 windows-cpu and arm64 windows-arm64 bundles (Windows arm64 + # has no GPU prebuilt). NOTE: this block is currently inert -- + # write_prebuilt_metadata does not persist an install_kind key, so + # $existingKind is always null; keep $expectedKinds in sync with the + # kinds install_llama_prebuilt.py installs before relying on it. + $expectedKinds = if ($HasROCm -or $script:ROCmGfxArch) { @("windows-rocm", "windows-hip") } elseif ($HasNvidiaSmi) { @("windows-cuda") } else { @("windows-cpu", "windows-arm64") } if ($existingKind -and ($existingKind -notin $expectedKinds)) { substep "Removing mismatched llama.cpp install (found '$existingKind', need one of: $($expectedKinds -join ', '))..." Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue diff --git a/studio/setup.sh b/studio/setup.sh index 6a74cd2296..d244e3cdcf 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1176,63 +1176,19 @@ _REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}" _HOST_SYSTEM="$(uname -s 2>/dev/null || true)" _HOST_MACHINE="$(uname -m 2>/dev/null || true)" -# Pick the release repo install_llama_prebuilt.py plans against. -# The fork ships CUDA (Linux x64/arm64, Windows), ROCm (Linux/Windows) and -# macOS bundles. Only the plain CPU/Vulkan bundles still come from ggml-org, so -# CPU-only Linux (x86_64 and arm64) routes there; GPU Linux, Windows and macOS -# use unslothai. -_LINUX_HAS_GPU=false -# Route to the fork only for a usable GPU. NVIDIA counts only when a device is -# actually enumerated and not hidden via CUDA_VISIBLE_DEVICES=""/-1 -# (_setup_nvidia_usable, from _setup_has_usable_nvidia_gpu above) -- mirroring -# install_llama_prebuilt.py's has_usable_nvidia. Mere nvidia-smi presence -# (CPU-only CUDA-toolkit containers, broken drivers) or a hidden GPU therefore -# takes the ggml-org CPU prebuilt instead of a slow source build. AMD is -# deliberately left on tooling presence, not usability: an unusable NVIDIA host -# has a good CPU prebuilt to fall back to, whereas tightening AMD would regress -# ROCm hosts exposing only hipconfig/hipinfo into an unnecessary CPU build. -if [ "$_setup_nvidia_usable" = true ]; then - _LINUX_HAS_GPU=true -else - for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo; do - if command -v "$_GPU_TOOL" >/dev/null 2>&1; then - _LINUX_HAS_GPU=true - break - fi - done -fi +# Pick the release repo install_llama_prebuilt.py plans against. Every host this +# installer supports now pulls its llama.cpp prebuilt from the unslothai fork: it +# ships the CUDA (Linux x64/arm64, Windows), ROCm (Linux/Windows) and macOS +# bundles, plus the CPU bundles for Linux/Windows on both x86_64 and arm64. +# ggml-org artifacts are no longer used by default. +_HELPER_RELEASE_REPO="unslothai/llama.cpp" # UNSLOTH_ROCM_GFX_ARCH may be set on a host where no probe fired, so the override # nested in the AMD-detected branch above never ran and _setup_gfx is still empty. -# Honour it here so the routing guard below and the --rocm-gfx forwarding both see -# it (install_llama_prebuilt.py reads the same env var as the --rocm-gfx default). -if [ "$_setup_nvidia_usable" != true ] && [ -z "${_setup_gfx:-}" ] && [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then +# Honour it here so the --rocm-gfx forwarding below still sees it +# (install_llama_prebuilt.py reads the same env var as the --rocm-gfx default). +if [ "${_setup_nvidia_usable:-}" != true ] && [ -z "${_setup_gfx:-}" ] && [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then _setup_gfx="${UNSLOTH_ROCM_GFX_ARCH}" fi -# A resolved/forwarded gfx arch (UNSLOTH_ROCM_GFX_ARCH) means an AMD GPU even when -# no ROCm tooling is on PATH; route it to the fork so the per-gfx prebuilt is -# picked instead of ggml-org / a source build. -if [ "$_LINUX_HAS_GPU" = false ] && [ -n "${_setup_gfx:-}" ]; then - _LINUX_HAS_GPU=true -fi - -if [ "$_HOST_SYSTEM" = "Linux" ] \ - && [ "$_HOST_MACHINE" = "x86_64" ] \ - && [ "$_LINUX_HAS_GPU" = false ]; then - _HELPER_RELEASE_REPO="ggml-org/llama.cpp" -elif [ "$_HOST_SYSTEM" = "Linux" ] \ - && { [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; } \ - && [ "$_LINUX_HAS_GPU" = false ]; then - # CPU-only Linux ARM64 (Ampere Altra, Raspberry Pi 5, GitHub - # `ubuntu-24.04-arm`, CPU-only Jetson rescue mode, ...). The fork ships no - # arm64 CPU bundle, so without this branch the prebuilt resolver returns 0 - # attempts and the installer falls back to a source build. ggml-org ships - # llama-bNNNN-bin-ubuntu-arm64.tar.gz from at least b9072 onward. - _HELPER_RELEASE_REPO="ggml-org/llama.cpp" -else - # GPU Linux (x64 CUDA/ROCm, arm64 CUDA), Windows (CUDA/ROCm), and macOS. - _HELPER_RELEASE_REPO="unslothai/llama.cpp" -fi -unset _GPU_TOOL _LLAMA_PR="${UNSLOTH_LLAMA_PR:-}" _SKIP_PREBUILT_INSTALL=false _LLAMA_PR_FORCE="${UNSLOTH_LLAMA_PR_FORCE:-${_DEFAULT_LLAMA_PR_FORCE}}" @@ -1939,19 +1895,19 @@ else fi # end _SKIP_GGUF_BUILD check # ── arm64 Linux GPU: CPU prebuilt as a last resort ── -# arm64 Linux with a GPU has no CUDA prebuilt anywhere (the unslothai fork is -# x64 only; ggml-org ships no Linux CUDA build), so it source-builds for the -# GPU above. If that produced no binary, install ggml-org's arm64 CPU prebuilt -# instead of leaving the host without llama.cpp. +# An arm64 Linux GPU host source-builds for the GPU above. If that produced no +# binary, install the fork's arm64 CPU prebuilt (app--linux-arm64-cpu.tar.gz) +# instead of leaving the host without llama.cpp. --cpu-fallback drops the GPU +# attributes so the CPU bundle is selected rather than re-attempting CUDA. if [ "$_LLAMA_CPP_DEGRADED" = true ] \ && [ "$_HOST_SYSTEM" = "Linux" ] \ && { [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; }; then - substep "GPU source build unavailable; trying ggml-org arm64 CPU prebuilt..." + substep "GPU source build unavailable; trying arm64 CPU prebuilt..." _ARM64_CPU_CMD=( python "$SCRIPT_DIR/install_llama_prebuilt.py" --install-dir "$LLAMA_CPP_DIR" --llama-tag "$_REQUESTED_LLAMA_TAG" - --published-repo "ggml-org/llama.cpp" + --published-repo "unslothai/llama.cpp" --cpu-fallback ) # Trust the installer's exit code: it validates the server before exiting 0, diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py index 8f40c9d720..4ff8c349c3 100644 --- a/tests/studio/install/test_llama_pr_force_and_source.py +++ b/tests/studio/install/test_llama_pr_force_and_source.py @@ -362,8 +362,11 @@ class TestSourcePatternsSh: assert '_LLAMA_SOURCE="${_DEFAULT_LLAMA_SOURCE}"' in self.content def test_release_repo_override_removed(self): + # No env-based release-repo override, and CPU-only hosts no longer fall + # back to ggml-org -- every host now routes to the fork. assert "UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp" not in self.content - assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' in self.content + assert '_HELPER_RELEASE_REPO="unslothai/llama.cpp"' in self.content + assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' not in self.content def test_force_compile_skips_prebuilt_resolution_early(self): assert 'if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then' in self.content @@ -425,12 +428,11 @@ class TestSourcePatternsPs1: assert "$LlamaSource = $DefaultLlamaSource" in self.content def test_release_repo_override_removed(self): - # Repo chosen by GPU detection (GPU -> fork, CPU -> ggml-org), no env override. + # No env-based release-repo override; every host now routes to the fork + # (the CPU-only ggml-org fallback was removed), mirroring setup.sh. assert "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" not in self.content - assert ( - "$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) " - '{ "unslothai/llama.cpp" } else { "ggml-org/llama.cpp" }' in self.content - ) + assert '$HelperReleaseRepo = "unslothai/llama.cpp"' in self.content + assert "$HelperReleaseRepo = if (" not in self.content def test_force_compile_skips_prebuilt_resolution_early(self): assert 'if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {' in self.content diff --git a/tests/studio/install/test_pr4562_bugfixes.py b/tests/studio/install/test_pr4562_bugfixes.py index a8cec0cb85..0d2b092924 100644 --- a/tests/studio/install/test_pr4562_bugfixes.py +++ b/tests/studio/install/test_pr4562_bugfixes.py @@ -658,15 +658,21 @@ class TestSourceCodePatterns: assert "_HELPER_RELEASE_REPO}/releases/latest" not in content assert "ggml-org/llama.cpp/releases/latest" not in content - def test_setup_sh_routes_to_fork_only_on_usable_gpu(self): - """Linux routing gates NVIDIA on GPU usability, not nvidia-smi presence, so - CPU-only/hidden-GPU hosts get the ggml CPU prebuilt. Guards the old presence-only loop.""" + def test_setup_sh_routes_every_host_to_fork(self): + """CPU-only Linux (the last ggml-org artifact consumer) now routes to the + fork like every other host, so the release-repo decision is unconditional. + Guards against a silent reintroduction of a ggml-org CPU routing branch. + GPU usability detection (used for PyTorch / source decisions) must stay.""" content = SETUP_SH.read_text() + assert '_HELPER_RELEASE_REPO="unslothai/llama.cpp"' in content + assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' not in content + # Usability gating (not routing) still distinguishes a hidden GPU. assert '[ "$_setup_nvidia_usable" = true ]' in content assert "CUDA_VISIBLE_DEVICES" in content - # nvidia-smi must NOT be back in the bare presence loop. - assert "for _GPU_TOOL in nvidia-smi" not in content - assert "for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo" in content + # The GPU-tooling probe (PR #4562) stays: ROCm detection goes through + # command -v, not a bare presence loop that mishandled a hidden nvidia-smi. + assert "command -v rocminfo" in content + assert "command -v amd-smi" in content def test_setup_sh_reports_installed_prebuilt_release(self): """Shell wrapper should report the installed prebuilt release from metadata.""" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index bfc8132683..5cabf41f57 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -3164,21 +3164,25 @@ class TestRocmGfxForwarding: assert "--rocm-gfx" in source assert "$script:ROCmGfxArch" in source - def test_setup_sh_routes_inferred_gfx_to_fork(self): - # An inferred gfx arch must route to the fork even without ROCm tooling. - # Pin the specific guard (a bare "${_setup_gfx:-}" also appears elsewhere). + def test_setup_sh_routes_unconditionally_to_fork(self): + # CPU-only hosts no longer fall back to ggml-org -- the release-repo + # decision is an unconditional fork assignment now. Pin the line text. source = _SETUP_SH_PATH.read_text(encoding = "utf-8") - assert '[ "$_LINUX_HAS_GPU" = false ] && [ -n "${_setup_gfx:-}" ]' in source + assert '_HELPER_RELEASE_REPO="unslothai/llama.cpp"' in source + assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' not in source - def test_setup_ps1_routes_inferred_gfx_to_fork(self): - # On Windows, a resolved $script:ROCmGfxArch counts as a fork install - # even when $HasROCm is false (Adrenalin-only, no HIP runtime). + def test_setup_ps1_routes_unconditionally_to_fork(self): + # Same on Windows: the fork now ships the windows-cpu / windows-arm64 + # bundles, so $HelperReleaseRepo is an unconditional fork assignment. source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") - assert "$HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch" in source + assert '$HelperReleaseRepo = "unslothai/llama.cpp"' in source + assert "$HelperReleaseRepo = if (" not in source - # The assertions above pin the guard *text*; the tests below *execute* the - # real routing block and assert the resolved repo, so a refactor that keeps - # the literal but breaks the inferred-gfx -> fork decision is still caught. + # The text pins above guard the literal. The tests below *execute* the real + # routing line from setup.sh / setup.ps1 and assert the resolved release repo, + # so a refactor that reintroduces a conditional (or a ggml-org branch) is still + # caught. Inputs are varied -- CPU-only, inferred/forwarded gfx, usable NVIDIA -- + # to prove no host slips back onto ggml-org. No GPU, no tooling, no network. @staticmethod def _resolve_setup_sh_repo( @@ -3187,15 +3191,18 @@ class TestRocmGfxForwarding: setup_gfx, rocm_gfx_arch_env = "", ): - """Run setup.sh's routing block under bash with PATH emptied (no ROCm tooling) and return _HELPER_RELEASE_REPO.""" + """Run setup.sh's release-repo routing block under bash and return the + resolved _HELPER_RELEASE_REPO. PATH is emptied so any stray tooling probe + misses; routing is unconditional, so the GPU inputs only prove no branch + reroutes a host to ggml-org.""" import shutil bash = shutil.which("bash") if bash is None: pytest.skip("bash not available") source = _SETUP_SH_PATH.read_text(encoding = "utf-8") - start = source.index("\n_LINUX_HAS_GPU=false\n") + 1 - end = source.index("\nunset _GPU_TOOL", start) + len("\nunset _GPU_TOOL") + start = source.index('\n_HELPER_RELEASE_REPO="unslothai/llama.cpp"\n') + 1 + end = source.index("\n_LLAMA_PR=", start) block = source[start:end] assert "_HELPER_RELEASE_REPO" in block, "setup.sh routing anchors not found" env = { @@ -3217,25 +3224,31 @@ class TestRocmGfxForwarding: assert result.returncode == 0, result.stderr return result.stdout.strip() - def test_setup_sh_inferred_gfx_resolves_to_fork(self): - # Only a name-inferred gfx arch -> route to the fork's per-gfx prebuilt - # (not ggml-org). x64 and arm64 share the fork branch. - assert self._resolve_setup_sh_repo("x86_64", False, "gfx1100") == "unslothai/llama.cpp" - assert self._resolve_setup_sh_repo("aarch64", False, "gfx1100") == "unslothai/llama.cpp" - - def test_setup_sh_env_forwarded_gfx_resolves_to_fork(self): - # No probe fired but UNSLOTH_ROCM_GFX_ARCH is set: adopt the env arch - # and route to the fork, same as name-inference. - repo = self._resolve_setup_sh_repo("x86_64", False, "", rocm_gfx_arch_env = "gfx1100") - assert repo == "unslothai/llama.cpp" - - def test_setup_sh_cpu_host_still_resolves_to_ggml(self): - # A real CPU host (no GPU, no inferred gfx, no env override) must keep routing to ggml-org. - assert self._resolve_setup_sh_repo("x86_64", False, "") == "ggml-org/llama.cpp" + @pytest.mark.parametrize( + "machine, nvidia_usable, setup_gfx, env_gfx", + [ + ("x86_64", False, "", ""), # plain CPU host (used to take ggml-org) + ("aarch64", False, "", ""), # plain CPU arm64 host (used to take ggml-org) + ("x86_64", False, "gfx1100", ""), # name-inferred gfx + ("x86_64", False, "", "gfx1100"), # env-forwarded gfx + ("x86_64", True, "", ""), # usable NVIDIA + ], + ) + def test_setup_sh_routing_block_always_resolves_to_fork( + self, machine, nvidia_usable, setup_gfx, env_gfx + ): + assert ( + self._resolve_setup_sh_repo( + machine, nvidia_usable, setup_gfx, rocm_gfx_arch_env = env_gfx + ) + == "unslothai/llama.cpp" + ) @staticmethod - def _resolve_setup_ps1_repo(has_nvidia, has_rocm, gfx_arch): - """Run setup.ps1's $HelperReleaseRepo selection under pwsh and return the resolved repo.""" + def _resolve_setup_ps1_repo(): + """Run setup.ps1's $HelperReleaseRepo assignment under pwsh and return the + resolved repo. The assignment is unconditional now, so there are no host + inputs to vary.""" import shutil pwsh = shutil.which("pwsh") @@ -3243,21 +3256,11 @@ class TestRocmGfxForwarding: pytest.skip("pwsh not available") source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") line = next( - ( - ln - for ln in source.splitlines() - if ln.strip().startswith("$HelperReleaseRepo = if (") - ), + (ln for ln in source.splitlines() if ln.strip().startswith("$HelperReleaseRepo =")), None, ) assert line is not None, "$HelperReleaseRepo selection not found in setup.ps1" - harness = ( - f"$HasNvidiaSmi = ${'true' if has_nvidia else 'false'}\n" - f"$HasROCm = ${'true' if has_rocm else 'false'}\n" - f"$script:ROCmGfxArch = '{gfx_arch}'\n" - f"{line}\n" - "Write-Output $HelperReleaseRepo" - ) + harness = f"{line}\nWrite-Output $HelperReleaseRepo" result = subprocess.run( [pwsh, "-NoProfile", "-Command", harness], capture_output = True, @@ -3267,13 +3270,10 @@ class TestRocmGfxForwarding: assert result.returncode == 0, result.stderr return result.stdout.strip() - def test_setup_ps1_inferred_gfx_resolves_to_fork(self): - # Adrenalin-only host: $HasROCm false but gfx inferred -> fork's windows-rocm bundle. - assert self._resolve_setup_ps1_repo(False, False, "gfx1100") == "unslothai/llama.cpp" - - def test_setup_ps1_cpu_host_still_resolves_to_ggml(self): - # No NVIDIA, no ROCm, no inferred gfx -> CPU host stays on ggml-org. - assert self._resolve_setup_ps1_repo(False, False, "") == "ggml-org/llama.cpp" + def test_setup_ps1_routing_resolves_to_fork(self): + # Windows routing is unconditional now: CPU-only Windows (x64 and arm64) + # uses the fork's windows-cpu / windows-arm64 bundles, not ggml-org. + assert self._resolve_setup_ps1_repo() == "unslothai/llama.cpp" # TEST: _pick_rocm_gfx_target -- visible-device selection from rocminfo output. diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index 464aa9a578..92f08c646b 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -1720,8 +1720,10 @@ class TestResolveInstallAttempts: assert approved.release_tag == "llama-prebuilt-latest" def test_linux_cpu_fork_without_bundle_raises_no_upstream_fallback(self, monkeypatch): - # CPU-only Linux on the fork must not fall back to the ggml-org CPU asset; with - # no fork CPU bundle the resolver raises rather than reaching upstream. + # A CPU-only Linux host on the fork never falls back to the ggml-org CPU + # asset. CPU-only Linux now routes to the fork, but if a release manifest + # happens to ship no CPU bundle the resolver raises rather than quietly + # reaching for an upstream asset. host = make_host( has_usable_nvidia = False, has_physical_nvidia = False, @@ -1846,6 +1848,151 @@ class TestResolveInstallAttempts: assert attempts[0].name == asset_name assert attempts[0].source_label == "published" + @pytest.mark.parametrize( + "system, machine, asset_name, install_kind, bundle_profile", + [ + # CPU-only Linux x64 -> fork linux-cpu (was ggml-org ubuntu-x64). + ("Linux", "x86_64", "app-b9625-linux-x64-cpu.tar.gz", "linux-cpu", "linux-cpu-x64"), + # CPU-only Linux arm64 -> fork linux-arm64 (was ggml-org ubuntu-arm64). + ( + "Linux", + "aarch64", + "app-b9625-linux-arm64-cpu.tar.gz", + "linux-arm64", + "linux-cpu-arm64", + ), + # CPU-only Windows arm64 -> fork windows-arm64 (was ggml-org win-cpu-arm64). + ( + "Windows", + "arm64", + "app-b9625-windows-arm64-cpu.zip", + "windows-arm64", + "windows-cpu-arm64", + ), + ], + ) + def test_cpu_host_prefers_published_fork_asset( + self, monkeypatch, system, machine, asset_name, install_kind, bundle_profile + ): + # CPU-only hosts now select the fork's CPU bundle from the manifest and + # must never query ggml-org upstream assets. Windows x64 CPU is covered + # separately by test_windows_cpu_prefers_published_asset. + host = make_host( + system = system, + machine = machine, + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + release = make_release( + [ + make_artifact( + asset_name, + install_kind = install_kind, + runtime_line = None, + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = bundle_profile, + rank = 1000, + ) + ], + release_tag = "llama-prebuilt-latest", + upstream_tag = "b9625", + assets = {asset_name: f"https://published.example/{asset_name}"}, + ) + checksums = make_checksums_with_source( + [asset_name], + release_tag = release.release_tag, + upstream_tag = "b9625", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: (_ for _ in ()).throw( + AssertionError("fork CPU host must not query upstream assets") + ), + ) + + _requested_tag, resolved_tag, attempts, _approved = resolve_install_attempts( + "latest", + host, + "unslothai/llama.cpp", + "", + ) + + assert resolved_tag == "b9625" + assert attempts[0].name == asset_name + assert attempts[0].install_kind == install_kind + assert attempts[0].source_label == "published" + + def test_cpu_only_unsupported_arch_source_builds(self, monkeypatch): + # A CPU-only Linux host that is neither x86_64 nor arm64 (ppc64le, + # riscv64, s390x) has no compatible CPU bundle. It must source-build, not + # receive the x86_64 linux-cpu binary (the Linux preflight checks libs, + # not ELF arch, so a wrong-arch binary would slip through). + host = make_host( + machine = "ppc64le", + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + assert not host.is_x86_64 and not host.is_arm64 + x64_asset = "app-b9625-linux-x64-cpu.tar.gz" + release = make_release( + [ + make_artifact( + x64_asset, + install_kind = "linux-cpu", + runtime_line = None, + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = "linux-cpu-x64", + rank = 1000, + ) + ], + release_tag = "llama-prebuilt-latest", + upstream_tag = "b9625", + assets = {x64_asset: f"https://published.example/{x64_asset}"}, + ) + checksums = make_checksums_with_source( + [x64_asset], + release_tag = release.release_tag, + upstream_tag = "b9625", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + + with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt asset was found"): + resolve_install_attempts("latest", host, "unslothai/llama.cpp", "") + def test_macos_prefers_published_asset(self, monkeypatch): host = make_host( system = "Darwin", @@ -3447,8 +3594,9 @@ class TestLinuxArm64ForkFallsBackToSource: assert plans == ["plan"] def test_arm64_cpu_on_ggml_org_is_not_blocked(self, monkeypatch): - # CPU-only arm64 routes to ggml-org, so the guard must not fire; it reaches the - # iterator (empty here -> generic message). + # ggml-org is reachable only via an explicit --published-repo override now, + # but the guard must still not fire on arm64 there; it reaches the iterator + # (empty here -> generic message). monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", @@ -3474,7 +3622,7 @@ class TestLinuxArm64ForkFallsBackToSource: class TestCpuFallback: - """--cpu-fallback drops GPU attributes so the host's OS/arch CPU prebuilt is selected, letting an arm64 GPU host install ggml-org's arm64 CPU build when its source build produced no binary.""" + """--cpu-fallback drops GPU attributes so the host's OS/arch CPU prebuilt is selected, letting an arm64 GPU host install the fork's arm64 CPU bundle when its source build produced no binary.""" _SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh" @@ -3553,10 +3701,15 @@ class TestCpuFallback: def test_setup_sh_has_arm64_cpu_prebuilt_fallback(self): source = self._SETUP_SH.read_text(encoding = "utf-8") - assert "--cpu-fallback" in source - # Fallback targets ggml-org (only repo with an arm64 Linux build), gated on a - # degraded arm64 source build. - assert "ggml-org/llama.cpp" in source + # The arm64 GPU last-resort CPU fallback now pulls the fork's arm64 CPU + # bundle (app--linux-arm64-cpu.tar.gz), not ggml-org's, and is gated + # on a degraded source build for arm64. + start = source.index("_ARM64_CPU_CMD=(") + end = source.index(")", start) + block = source[start:end] + assert "--cpu-fallback" in block + assert '--published-repo "unslothai/llama.cpp"' in block + assert '--published-repo "ggml-org/llama.cpp"' not in block assert "_LLAMA_CPP_DEGRADED" in source From d0c8d550a6db537583dbb78dd186e56ff103d2fa Mon Sep 17 00:00:00 2001 From: Tai An Date: Wed, 8 Jul 2026 05:38:06 -0700 Subject: [PATCH 070/113] fix(studio/hub): apply repo_id length limit per segment, not whole string (#6946) (#6953) * fix(studio/hub): apply repo_id length limit per segment, not whole string is_valid_repo_id() applied the 96-char limit to the full "namespace/repo_name" string, so a repo with a valid (<=96 char) name but a long combined id was falsely rejected. Match huggingface_hub.validate_repo_id by checking the length per segment instead. Fixes #6946. * Fix long repo id state filenames * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../backend/hub/tests/test_model_services.py | 46 +++++++++++++++++++ studio/backend/hub/utils/paths.py | 9 +++- studio/backend/hub/utils/state_dir.py | 33 +++++++++++-- 3 files changed, 82 insertions(+), 6 deletions(-) diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index f05d8359ec..44701c0b64 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -105,6 +105,10 @@ def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id): assert paths.is_valid_repo_id(repo_id) +def test_repo_id_validation_accepts_max_length_namespaced_repo(): + assert paths.is_valid_repo_id(f"{'a' * 96}/{'b' * 96}") + + @pytest.mark.parametrize( "repo_id", [ @@ -121,6 +125,48 @@ def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id): assert not paths.is_valid_repo_id(repo_id) +def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + + path = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M") + + assert path is not None + assert path.name == "models--owner--repo--variant--q4_k_m.json" + + +@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64]) +def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + repo_id = f"{'a' * 96}/{'b' * 96}" + + assert paths.is_valid_repo_id(repo_id) + assert download_manifest.write_cancel_marker("model", repo_id, variant, "http") + assert download_manifest.write_manifest( + "model", + repo_id, + variant, + [download_manifest.ExpectedFile(path = "model.gguf", size = 1)], + "http", + ) + + marker_path = state_dir.marker_path("model", repo_id, variant) + manifest_path = state_dir.manifest_path("model", repo_id, variant) + + assert marker_path is not None + assert manifest_path is not None + assert "--sha256-" in marker_path.name + assert len(marker_path.name.encode("utf-8")) <= 255 + assert len(f".{marker_path.name}.tmp-00000000".encode("utf-8")) <= 255 + assert download_manifest.has_cancel_marker("model", repo_id, variant) + assert download_manifest.read_manifest("model", repo_id, variant) is not None + assert list(download_manifest.iter_variant_markers("model", repo_id)) == [ + (variant, marker_path) + ] + assert list(download_manifest.iter_variant_manifests("model", repo_id)) == [ + (variant, manifest_path) + ] + + class _RecordingLogger: def __init__(self): self.warnings = [] diff --git a/studio/backend/hub/utils/paths.py b/studio/backend/hub/utils/paths.py index afcb0b41dc..5435202565 100644 --- a/studio/backend/hub/utils/paths.py +++ b/studio/backend/hub/utils/paths.py @@ -181,15 +181,20 @@ def is_valid_repo_id(repo_id: str) -> bool: """Validate Hugging Face ``repo_name`` or ``namespace/repo_name`` IDs.""" if not repo_id or repo_id != repo_id.strip(): return False - if len(repo_id) > _MAX_REPO_ID_LENGTH or repo_id.endswith(".git"): + if repo_id.endswith(".git"): return False if "--" in repo_id or ".." in repo_id: return False segments = repo_id.split("/") if len(segments) not in (1, 2): return False + # Match huggingface_hub.validate_repo_id: the 96-char limit applies per + # segment (repo name / namespace), not to the whole "namespace/repo_name" + # string, so long-but-valid repo names are not falsely rejected. return all( - segment not in ("", ".", "..") and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None + segment not in ("", ".", "..") + and len(segment) <= _MAX_REPO_ID_LENGTH + and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None for segment in segments ) diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py index a304477a3d..183e934724 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -11,8 +11,9 @@ cache lifecycle. Two subdirectories: manifests/ .json per-download expected-files manifest cancelled/ .json per-download cancel marker -The ```` mirrors HF's cache dir naming so a state file can be -eyeballed next to the on-disk repo it describes: +The ```` mirrors HF's cache dir naming while the resulting manifest, +cancel-marker, and atomic-write temp filenames fit common filesystem basename +limits. Very long repo IDs use a stable hash in the state key: models---- full snapshot models------variant-- GGUF variant @@ -49,6 +50,11 @@ _MANIFESTS_SUBDIR = "manifests" _CANCELLED_SUBDIR = "cancelled" _WORKERS_SUBDIR = "workers" _SAFE_VARIANT_FRAGMENT = re.compile(r"^[a-z0-9._-]{1,64}$") +_MAX_STATE_BASENAME_BYTES = 255 +_STATE_EXTENSION = ".json" +# _atomic_write_json writes "..tmp-<8hex>" beside the final file. +_ATOMIC_WRITE_TMP_OVERHEAD = len(".") + len(".tmp-") + 8 +_MAX_VARIANT_FRAGMENT_LENGTH = 64 def state_root() -> Optional[Path]: @@ -84,16 +90,35 @@ def repo_cache_basename(repo_type: RepoType, repo_id: str) -> str: return f"{repo_type}s--{repo_id.replace('/', '--')}".lower() +def _filename_bytes(name: str) -> int: + return len(name.encode("utf-8")) + + +def _state_filename_fits(entry_key: str) -> bool: + filename = f"{entry_key}{_STATE_EXTENSION}" + return _filename_bytes(filename) + _ATOMIC_WRITE_TMP_OVERHEAD <= _MAX_STATE_BASENAME_BYTES + + +def _state_repo_key(repo_type: RepoType, repo_id: str) -> str: + base = repo_cache_basename(repo_type, repo_id) + variant_prefix = f"{base}--variant--" + longest_variant_key = f"{variant_prefix}{'x' * _MAX_VARIANT_FRAGMENT_LENGTH}" + if _state_filename_fits(longest_variant_key): + return base + digest = hashlib.sha256(base.encode("utf-8")).hexdigest()[:32] + return f"{repo_type}s--sha256-{digest}" + + def variant_filename_prefix(repo_type: RepoType, repo_id: str) -> str: """Lowercased prefix every variant-keyed state file for this repo shares. The single source the download_manifest enumerators match against, so the scheme in :func:`_entry_key` cannot drift from them silently.""" - return f"{repo_cache_basename(repo_type, repo_id)}--variant--" + return f"{_state_repo_key(repo_type, repo_id)}--variant--" def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str: - base = repo_cache_basename(repo_type, repo_id) + base = _state_repo_key(repo_type, repo_id) if not variant: return base normalized_variant = variant.strip().lower() From 62a6eb2a3df395e2e0e94218cfc963318f5c7392 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 05:57:44 -0700 Subject: [PATCH 071/113] MoE LoRA: auto-target per-expert Linear experts (gpt-oss 4bit) instead of leaving them frozen (#6936) * models: auto-target per-expert Linear MoE experts for LoRA (gpt-oss 4bit) MoE checkpoints whose experts are stored as per-expert nn.Linear ModuleLists could not receive expert LoRA. gpt-oss bnb-4bit is the canonical case: its experts live at mlp.experts.gate_up_projs. and mlp.experts.down_projs. as per-expert Linear4bit modules, not a fused nn.Parameter. The target_parameters path only handles the fused nn.Parameter layout, and the plain gate_proj/up_proj/down_proj leaf names do not match the per-expert indices, so get_peft_model attached LoRA to attention only and left every expert frozen (0 of 1536 on gpt-oss-20b) even though the grouped bnb-4bit training forward exists. Add get_moe_target_modules, the module-LoRA counterpart of get_moe_target_parameters: it detects per-expert Linear ModuleLists under an experts container and returns their suffix target_modules names (gate_up_projs. / down_projs.). get_peft_model in both llama.py and vision.py extends target_modules with these, handling the explicit leaf-list form and the regex form (auto / all-linear / scoped). It is gated on the same MLP-in-scope condition as the parameter path, so an attention-only request still skips the experts. Also gate get_moe_target_parameters on the fused parameter actually existing, so a per-expert-Linear layout no longer produces a dead target_parameters path or a misleading "Enabling LoRA on MoE parameters" line; those experts are handled through target_modules instead. Validated on gpt-oss-20b-unsloth-bnb-4bit (transformers 5.5.0): experts attach (1536 modules, trainable 0.036 percent to 1.65 percent) across the default, None and all-linear paths; training memorizes and the LoRA adapter reproduces exactly after a cold reload in a fresh process. No regression: fused-parameter MoEs (Qwen3-30B-A3B-4bit), non-MoE models, and attention-only requests are unaffected (get_moe_target_modules returns an empty list). Merging these per-expert adapters into a merged_16bit checkpoint is handled by a companion unsloth-zoo change (saving_utils folds each per-expert delta into the fused gate_up_proj / down_proj tensor). With both, the LoRA adapter and the merged_16bit checkpoint reload the trained behavior identically. * models: scope per-expert MoE targets, keep repeat get_peft_model idempotent, warn on old zoo Address review of the per-expert Linear MoE targeting: - Scope get_moe_target_modules to the requested projection leaves (gate/up map to the gate_up ModuleList, down maps to the down ModuleList), so a narrowed request such as target_modules=["down_proj"] no longer also trains gate_up_projs, matching get_moe_target_parameters. - Detect experts through a PEFT-wrapped base_layer as well, and recompute the auto-added expert targets in the llama.py existing-adapter check, so a repeat get_peft_model call with the same arguments stays idempotent instead of raising on the saved expert targets. - Warn when the installed unsloth_zoo cannot fold these per-expert experts into a merged_16bit checkpoint (older releases keep the fused gate_up_proj / down_proj tensors and drop the per-expert deltas), so the expert LoRA is not silently lost on save_pretrained_merged; the fold lands in unsloth-zoo #885. The LoRA adapter itself is unaffected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/models/_utils.py | 119 +++++++++++++++++++++++++++++++++++++-- unsloth/models/llama.py | 18 ++++++ unsloth/models/vision.py | 29 ++++++++++ 3 files changed, 162 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1c75f8ce66..b68cb702b1 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -86,6 +86,8 @@ __all__ = [ "maybe_prefetch_hf_snapshot", "is_moe_model", "get_moe_target_parameters", + "get_moe_target_modules", + "warn_if_zoo_cannot_merge_moe_experts", "_select_moe_detection_targets", "make_fast_generate_wrapper", "_mark_unsloth_disable_data_parallel", @@ -4060,13 +4062,17 @@ def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str alternate_name = "experts.down_proj", ) - # gate_up_proj combines both gate_proj and up_proj in MoE - # Also match "gate_up_proj" directly since users may specify the fused name + # gate_up_proj combines gate_proj and up_proj; also match the fused name directly. + # Only target a fused expert Parameter that exists: per-expert Linear layouts + # (e.g. gpt-oss bnb-4bit) have no fused Parameter and are handled by + # get_moe_target_modules, so skip them rather than pass PEFT a dead path. if "gate_proj" in target_set or "up_proj" in target_set or "gate_up_proj" in target_set: - moe_params.append(gate_up_name) + if _moe_parameter_exists(model, gate_up_name): + moe_params.append(gate_up_name) if "down_proj" in target_set: - moe_params.append(down_name) + if _moe_parameter_exists(model, down_name): + moe_params.append(down_name) if moe_params: print( @@ -4077,6 +4083,111 @@ def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str return None +def _moe_parameter_exists(model, name: str) -> bool: + """True if ``name`` is an exact suffix of some parameter path on the model.""" + if not hasattr(model, "named_parameters"): + return False + try: + for parameter_name, _ in model.named_parameters(): + if parameter_name == name or parameter_name.endswith("." + name): + return True + except Exception: + return False + return False + + +def get_moe_target_modules(model, target_modules = None) -> List[str]: + """Per-expert ``target_modules`` suffixes for MoE models whose experts are stored + as per-expert ``nn.Linear`` ModuleLists rather than fused nn.Parameters. + + gpt-oss bnb-4bit is the canonical case (mlp.experts.gate_up_projs. / + down_projs. as Linear4bit): no fused Parameter, and the plain + gate/up/down_proj leaves do not match, so LoRA skips them. Returning the + per-expert suffixes makes PEFT attach via ordinary suffix matching (the + module-LoRA counterpart of get_moe_target_parameters). Returns [] for non-MoE, + fused-parameter MoEs, an absent per-expert layout, or a request that omits the + MLP experts (so an attention-only run does not train experts). + """ + if not is_moe_model(model): + return [] + if target_modules is None: + return [] + if isinstance(target_modules, str): + target_set = _moe_target_set_from_string(target_modules) + else: + target_set = { + target + for target in target_modules or () + if (isinstance(target, str) and "." not in target and target in _MOE_BROAD_MLP_TARGETS) + } + if not (target_set & _MOE_BROAD_MLP_TARGETS): + return [] + + if not hasattr(model, "named_modules"): + return [] + + # Scope the returned suffixes to the requested projection leaves, matching + # get_moe_target_parameters: gate_proj/up_proj/gate_up_proj map to the fused + # gate_up ModuleList (e.g. gate_up_projs); down_proj maps to the down ModuleList + # (e.g. down_projs). A down-only (or gate/up-only) request must not pull in the + # other projection. + want_gate_up = bool(target_set & {"gate_proj", "up_proj", "gate_up_proj"}) + want_down = "down_proj" in target_set + + targets = set() + for name, module in model.named_modules(): + if not isinstance(module, torch.nn.ModuleList) or len(module) == 0: + continue + parent, _, leaf = name.rpartition(".") + # ModuleList directly under an ``experts`` container, holding only Linear + # leaves (bnb Linear4bit / Linear8bitLt subclass nn.Linear). After PEFT has + # wrapped the experts the child is a LoRA layer whose ``base_layer`` is the + # Linear, so accept that too (keeps this idempotent across a re-wrapped model). + if not parent.endswith("experts"): + continue + if not all( + isinstance(child, torch.nn.Linear) + or isinstance(getattr(child, "base_layer", None), torch.nn.Linear) + for child in module + ): + continue + # Honor the requested subset: classify the ModuleList by projection role. + leaf_lower = leaf.lower() + is_down = "down" in leaf_lower + is_gate_up = (not is_down) and ("gate" in leaf_lower or "up" in leaf_lower) + if is_down and not want_down: + continue + if is_gate_up and not want_gate_up: + continue + # One entry per expert index; ``leaf.`` matches expert i in every layer. + for expert_index in range(len(module)): + targets.add(f"{leaf}.{expert_index}") + + return sorted(targets) + + +def warn_if_zoo_cannot_merge_moe_experts(): + """Warn once when the installed unsloth_zoo cannot fold per-expert Linear MoE LoRA + into a merged_16bit checkpoint. Older zoo releases keep the fused gate_up_proj / + down_proj tensors and drop the per-expert gate_up_projs. / down_projs. deltas, + so save_pretrained_merged("merged_16bit") would silently lose the expert training + (the LoRA adapter itself still saves and reloads correctly).""" + try: + from unsloth_zoo import saving_utils as _saving_utils + + # _fold_perexpert_lora_into_fused is the helper that folds these experts. + if hasattr(_saving_utils, "_fold_perexpert_lora_into_fused"): + return + except Exception: + return # cannot introspect zoo -> stay quiet rather than false-alarm + logger.warning_once( + "Unsloth: the installed unsloth_zoo will not fold these per-expert experts into " + "a merged_16bit checkpoint, so save_pretrained_merged('merged_16bit') would drop " + "the expert LoRA. Upgrade unsloth_zoo to merge them; saving the LoRA adapter is " + "unaffected." + ) + + def _select_moe_detection_targets( original_target_modules, scoped_target_modules, diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index a1da099758..1f43f61443 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3105,6 +3105,11 @@ class FastLlamaModel: new_target_modules = list(target_modules) + list( modules_to_save if modules_to_save is not None else [] ) + # Per-expert Linear MoE experts (e.g. gpt-oss bnb-4bit) were auto-added to the + # saved target_modules when the adapter was first created. Recompute them so a + # repeat get_peft_model call with the same args stays idempotent instead of + # tripping the mismatch below. No-op for non per-expert-Linear models. + new_target_modules += get_moe_target_modules(model, target_modules) # Now check! new_target_modules = set(new_target_modules) @@ -3331,6 +3336,19 @@ class FastLlamaModel: if target_parameters is None: target_parameters = get_moe_target_parameters(model, target_modules) + # Per-expert Linear expert layouts (e.g. gpt-oss bnb-4bit) are Linear modules, + # not fused Parameters, so target them via target_modules. No-op otherwise. + _moe_module_targets = get_moe_target_modules(model, target_modules) + if _moe_module_targets: + _added = [t for t in _moe_module_targets if t not in final_modules] + final_modules.extend(_added) + if _added: + print( + f"Unsloth: Detected MoE model with per-expert Linear experts. " + f"Enabling LoRA on {len(_added)} expert projection modules." + ) + warn_if_zoo_cannot_merge_moe_experts() + if finetune_last_n_layers is not None and layers_to_transform is None: from .vision import _get_total_transformer_layers _total_layers = _get_total_transformer_layers(model) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 0a68a49fee..0235d80e93 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1807,6 +1807,35 @@ class FastBaseModel: ) target_parameters = get_moe_target_parameters(model, _moe_targets) + # Per-expert Linear expert layouts (e.g. gpt-oss bnb-4bit) target experts via + # target_modules, not fused Parameters. Extend either form PEFT accepts: a leaf + # list (explicit) or a regex string (auto / all-linear / scoped). No-op otherwise. + _moe_module_detect = _select_moe_detection_targets( + _moe_detect_target, + target_modules, + finetune_mlp_modules = finetune_mlp_modules, + finetune_language_layers = finetune_language_layers, + ) + _moe_module_targets = get_moe_target_modules(model, _moe_module_detect) + if _moe_module_targets: + if isinstance(target_modules, (list, tuple)): + target_modules = list(target_modules) + [ + target for target in _moe_module_targets if target not in target_modules + ] + elif isinstance(target_modules, str): + _expert_leaves = sorted({t.rsplit(".", 1)[0] for t in _moe_module_targets}) + _expert_alt = ( + r".*\.experts\.(?:" + + "|".join(re.escape(leaf) for leaf in _expert_leaves) + + r")\.\d+" + ) + target_modules = f"(?:{target_modules})|(?:{_expert_alt})" + print( + f"Unsloth: Detected MoE model with per-expert Linear experts. " + f"Enabling LoRA on {len(_moe_module_targets)} expert projection modules." + ) + warn_if_zoo_cannot_merge_moe_experts() + if finetune_last_n_layers is not None and layers_to_transform is None: _total_layers = _get_total_transformer_layers(model) if _total_layers is not None and _total_layers > 0: From 03cbe211a38b78448e5844f25f49b65a449d8b10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Eric=20=20=F0=9F=87=A7=F0=9F=87=B7?= Date: Wed, 8 Jul 2026 10:38:10 -0300 Subject: [PATCH 072/113] Studio: fix flash-attn and torchao install on Blackwell (sm_100+) GPUs (Closes #6961) (#6970) * fix: Remove moot has_blackwell_gpu() function Fixes unslothai/unsloth#6961. This function skipped flash-attn on Blackwell GPUs because no prebuilt wheel existed; Dao-AILab now ships one and url_exists() already gates resolution. Co-Authored-By: Claude Opus 4.8 * fix: use torchao 0.17.0 for Blackwell Fixes #6961. Torchao 0.16.0's cpp extensions are built against CUDA 12, so on a CUDA-13 torch (cu130 / Blackwell) they fail to load with "libcudart.so.12: cannot open shared object file". Select 0.17.0 there instead: its cpp targets torch 2.11, so it is skipped cleanly rather than crashing. CUDA-12 / ROCm / CPU torch 2.10 keeps 0.16.0 and its working kernels. Co-Authored-By: Claude Opus 4.8 * Condense torchao version-selection comments (no behavior change) * Support torch 2.11 in the Studio installer via the torch2.10 prebuilt wheels Map torch 2.11 to the torch2.10 prebuilt wheels for flash-attn, causal-conv1d, and mamba through wheel_utils.prebuilt_wheel_torch_mm, applied in direct_wheel_url (filename) and flash_attn_wheel_url (version). Those torch2.10 CUDA wheels load and pass each project's own test suite on torch 2.11 (verified on B200), so a torch 2.11 environment gets the prebuilt accelerators instead of skipping or building from source. Raise _CUDA_TORCH_PKG_SPEC to <2.12.0 (torchvision <0.27.0, torchaudio <2.12.0) so the CUDA torch repair path can install torch 2.11, where torchao 0.17's cpp kernels load cleanly. Add tests for the mapping. * Keep has_blackwell_gpu as a False stub for future arch gating * Restore has_blackwell_gpu as a return-False probe kept for future arch gating Keep the nvidia-smi compute_cap detection and its two call sites, but short-circuit with return False at the top so flash-attn is no longer skipped on Blackwell (sm_100+ now has prebuilt wheels and url_exists gates resolution). Drop the early return to re-enable arch-based detection later. --------- Co-authored-by: Claude Opus 4.8 Co-authored-by: Daniel Han --- .../tests/test_mlx_training_worker_config.py | 1 - studio/backend/tests/test_torchao_select.py | 17 +- .../tests/test_training_worker_flash_attn.py | 23 -- studio/backend/utils/wheel_utils.py | 26 ++- studio/install_python_stack.py | 70 +++--- .../test_flash_attn_install_python_stack.py | 203 ++++-------------- 6 files changed, 123 insertions(+), 217 deletions(-) diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 14fc0933d0..503bae8da3 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -37,7 +37,6 @@ def _load_worker_module(): for name in ( "direct_wheel_url", "flash_attn_wheel_url", - "has_blackwell_gpu", "install_wheel", "probe_torch_wheel_env", "url_exists", diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 2d3dc5fbff..e4775a10a6 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -32,16 +32,23 @@ def _load_module(monkeypatch): @pytest.mark.parametrize( "torch_version, expected", [ - # torch 2.10 (the reported bug: cu130 resolves 2.10.0) -> 0.16.0, - # independent of the local +cuXXX/+rocm/+cpu suffix or patch level. - ("2.10.0+cu130", "torchao==0.16.0"), + # torch 2.10 on CUDA <= 12 -> 0.16.0 (its cpp is built for torch 2.10.0 and + # loads against the CUDA-12 PyPI wheel). Independent of patch level. + ("2.10.0+cu128", "torchao==0.16.0"), + ("2.10.0+cu126", "torchao==0.16.0"), ("2.10.0+rocm6.4", "torchao==0.16.0"), ("2.10.0+cpu", "torchao==0.16.0"), ("2.10.1", "torchao==0.16.0"), ("2.10.0", "torchao==0.16.0"), - # Pre-release / dev / rc builds: the minor is cleaned of non-digits. + # torch 2.10 on CUDA >= 13 (Blackwell / cu130): 0.16.0's CUDA-12 cpp can't + # load against a CUDA-13 torch (libcudart.so.12 error), so use 0.17.0. + ("2.10.0+cu130", "torchao==0.17.0"), + ("2.10.0+cu140", "torchao==0.17.0"), + # Pre-release / dev / rc builds: the minor is cleaned of non-digits; the + # CUDA tag still decides 0.16.0 vs 0.17.0. ("2.10.0rc1", "torchao==0.16.0"), - ("2.10.0.dev20250804+cu130", "torchao==0.16.0"), + ("2.10.0.dev20250804+cu130", "torchao==0.17.0"), + ("2.10.0.dev20250804+cu128", "torchao==0.16.0"), ("2.10rc1", "torchao==0.16.0"), # torch 2.11 (reachable via ROCm rocm7.2) and forward -> 0.17.0. ("2.11.0+cu130", "torchao==0.17.0"), diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 3c5d6cd094..7e7fc1af48 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -59,7 +59,6 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -88,7 +87,6 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -141,27 +139,6 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): worker._sp.run.assert_not_called() -def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): - statuses: list[str] = [] - install_mock = mock.Mock() - - monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True) - monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) - monkeypatch.setattr( - worker, - "_send_status", - lambda queue, message: statuses.append(message), - ) - - worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536) - - install_mock.assert_not_called() - assert len(statuses) == 1 - assert "Blackwell" in statuses[0] - - def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch): install_mock = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 98697df83c..1b5926fd49 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -26,11 +26,14 @@ FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/rele def has_blackwell_gpu() -> bool: """Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell). - Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels - fail to load, so callers use this to skip the flash-attn install path. Cached - for the process lifetime; tests mocking nvidia-smi must call + Cached for the process lifetime; tests mocking nvidia-smi must call ``has_blackwell_gpu.cache_clear()`` first. """ + # Detection disabled for now: Dao-AILab ships Blackwell (sm_100+) flash-attn + # wheels and url_exists() already gates resolution, so we no longer skip + # flash-attn on Blackwell. The nvidia-smi probe below is kept for possible + # future arch-based gating; drop this early return to re-enable it. + return False exe = shutil.which("nvidia-smi") if not exe: return False @@ -117,6 +120,19 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non return env +# torch 2.11 has no native prebuilt wheels for flash-attn / causal-conv1d / mamba +# yet, but their torch 2.10 CUDA wheels load and pass the projects' own test suites +# on torch 2.11 (verified on B200: FA2 fwd/bwd, causal-conv1d, and mamba selective +# scan all match reference). Reuse the torch 2.10 wheels on torch 2.11 so a 2.11 +# install still gets these prebuilt accelerators instead of building from source. +_PREBUILT_WHEEL_TORCH_MM = {"2.11": "2.10"} + + +def prebuilt_wheel_torch_mm(torch_mm: str) -> str: + """Map a torch major.minor to the one whose prebuilt accelerator wheels to use.""" + return _PREBUILT_WHEEL_TORCH_MM.get(torch_mm, torch_mm) + + def direct_wheel_url( *, filename_prefix: str, @@ -130,7 +146,7 @@ def direct_wheel_url( filename = ( f"{filename_prefix}-{package_version}" - f"+cu{env['cuda_major']}torch{env['torch_mm']}" + f"+cu{env['cuda_major']}torch{prebuilt_wheel_torch_mm(env['torch_mm'])}" f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}" f"-{env['platform_tag']}.whl" ) @@ -152,7 +168,7 @@ def flash_attn_package_version(torch_mm: str) -> str | None: def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None: if env is None: return None - package_version = flash_attn_package_version(env["torch_mm"]) + package_version = flash_attn_package_version(prebuilt_wheel_torch_mm(env["torch_mm"])) if package_version is None: return None return direct_wheel_url( diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e033a56a0a..19c492deaa 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -103,33 +103,47 @@ _PYTORCH_WHL_BASE = ( os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" ).rstrip("/") -# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are -# pinned to the torch<2.11 family rather than left bare: the install uses an -# exclusive --index-url (no PyPI fallback), so a bare name could resolve a -# torchvision built against a different torch major (e.g. 0.27 for torch 2.12) -# and fail at runtime with an ABI mismatch. Same bounds as the _default ROCm -# spec above, which targets the same torch family. +# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed: its +# torchao 0.17 cpp kernels load cleanly (0.16 crashes on cu130), and the flash-attn +# / causal-conv1d / mamba torch2.10 wheels load and pass their upstream suites on +# 2.11 (see wheel_utils._PREBUILT_WHEEL_TORCH_MM). torchvision/torchaudio are pinned +# (not bare) because the install uses an exclusive --index-url (no PyPI fallback), so +# a bare name could resolve one built against a different torch major (e.g. 0.27 for +# torch 2.12) and fail at runtime with an ABI mismatch. _CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = ( - "torch>=2.4,<2.11.0", - "torchvision>=0.19,<0.26.0", - "torchaudio>=2.4,<2.11.0", + "torch>=2.4,<2.12.0", + "torchvision>=0.19,<0.27.0", + "torchaudio>=2.4,<2.12.0", ) -# torchao's C++ extensions are built against ONE exact torch release; a newer -# torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions -# due to incompatible torch version ...") and fall back to slow Python. Because -# the torch pin above is a range (and every CUDA index now tops out at torch -# 2.10), the torch actually installed drifts ahead of a fixed torchao pin. So -# pick the torchao whose build matches the torch in the venv. Table: pytorch/ao#2919. -# torch 2.9.x -> torchao 0.14.0 (today's pin; built for torch 2.9.0) -# torch 2.10.x -> torchao 0.16.0 (built for torch 2.10.0) -# torch 2.11.x -> torchao 0.17.0 (built for torch 2.11.0; reachable via ROCm rocm7.2) -# Unknown/older torch keeps the conservative default (no regression vs today). +# torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch +# mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails +# to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a +# range, so match torchao to the installed torch (table: pytorch/ao#2919): +# 2.9.x -> 0.14.0 +# 2.10.x, CUDA<=12 -> 0.16.0 (cpp built for 2.10, loads via the CUDA-12 wheel) +# 2.10.x, CUDA>=13 -> 0.17.0 (cu130: 0.16.0's CUDA-12 cpp crashes on load; 0.17.0 +# targets torch 2.11 so its cpp is cleanly skipped, not crashed) +# 2.11.x -> 0.17.0 (reachable via CUDA or ROCm rocm7.2) +# Unknown/older torch keeps the conservative default. _TORCHAO_DEFAULT_SPEC = "torchao==0.14.0" -_TORCHAO_BY_TORCH_MINOR: dict[int, str] = { - 10: "torchao==0.16.0", - 11: "torchao==0.17.0", -} +_TORCHAO_TORCH_210_SPEC = "torchao==0.16.0" +_TORCHAO_TORCH_210_CUDA13_SPEC = "torchao==0.17.0" +_TORCHAO_TORCH_211_PLUS_SPEC = "torchao==0.17.0" +# torch 2.10 built against CUDA >= this major can't load 0.16.0's CUDA-12 cpp. +_TORCHAO_CUDA13_MIN_MAJOR = 13 + + +def _cuda_major_from_torch_version(torch_version: str) -> int | None: + """Extract the CUDA major from a torch local version tag, e.g. '2.10.0+cu130' + -> 13, '2.10.0+cu128' -> 12. Returns None for rocm/cpu/tagless builds.""" + local = str(torch_version).split("+", 1) + if len(local) < 2 or not local[1].startswith("cu"): + return None + digits = re.sub(r"[^0-9].*", "", local[1][2:]) # 'cu130' -> '130' + if not digits: + return None + return int(digits) // 10 # '130' -> 13, '128' -> 12, '118' -> 11 def _select_torchao_spec(torch_version: str | None) -> str: @@ -151,8 +165,14 @@ def _select_torchao_spec(torch_version: str | None) -> str: if major != 2: return _TORCHAO_DEFAULT_SPEC if minor >= 11: - return _TORCHAO_BY_TORCH_MINOR[11] # newest known build; covers 2.11+ - return _TORCHAO_BY_TORCH_MINOR.get(minor, _TORCHAO_DEFAULT_SPEC) + return _TORCHAO_TORCH_211_PLUS_SPEC # newest known build; covers 2.11+ + if minor == 10: + # cu130+ can't load 0.16.0's CUDA-12 cpp; use 0.17.0 (cpp skipped, not crashed). + cuda_major = _cuda_major_from_torch_version(str(torch_version)) + if cuda_major is not None and cuda_major >= _TORCHAO_CUDA13_MIN_MAJOR: + return _TORCHAO_TORCH_210_CUDA13_SPEC + return _TORCHAO_TORCH_210_SPEC + return _TORCHAO_DEFAULT_SPEC def _probe_installed_torch_version() -> str | None: diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py index 26ff03505a..bf3ed57788 100644 --- a/tests/python/test_flash_attn_install_python_stack.py +++ b/tests/python/test_flash_attn_install_python_stack.py @@ -13,102 +13,35 @@ sys.path.insert(0, str(STUDIO_DIR)) sys.path.insert(0, str(STUDIO_DIR / "backend")) import install_python_stack as ips -from backend.utils import wheel_utils +from utils import wheel_utils -def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess: - return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "") +class TestPrebuiltWheelTorchMapping: + def test_torch_211_maps_to_torch210(self): + assert wheel_utils.prebuilt_wheel_torch_mm("2.11") == "2.10" + def test_other_versions_pass_through(self): + for torch_mm in ("2.9", "2.10", "2.12"): + assert wheel_utils.prebuilt_wheel_torch_mm(torch_mm) == torch_mm -class TestHasBlackwellGpu: - def setup_method(self): - wheel_utils.has_blackwell_gpu.cache_clear() - - def teardown_method(self): - wheel_utils.has_blackwell_gpu.cache_clear() - - def test_returns_false_when_nvidia_smi_missing(self): - with mock.patch.object(wheel_utils.shutil, "which", return_value = None): - assert wheel_utils.has_blackwell_gpu() is False - - def test_returns_true_for_sm_100(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")), - ): - assert wheel_utils.has_blackwell_gpu() is True - - def test_returns_true_for_sm_120(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")), - ): - assert wheel_utils.has_blackwell_gpu() is True - - def test_returns_true_for_sm_121(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")), - ): - assert wheel_utils.has_blackwell_gpu() is True - - def test_returns_false_for_sm_90(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")), - ): - assert wheel_utils.has_blackwell_gpu() is False - - def test_returns_false_for_sm_89(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")), - ): - assert wheel_utils.has_blackwell_gpu() is False - - def test_mixed_gpus_with_one_blackwell_returns_true(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object( - wheel_utils.subprocess, - "run", - return_value = _smi_result("8.0\n10.0\n"), - ), - ): - assert wheel_utils.has_blackwell_gpu() is True - - def test_returns_false_when_nvidia_smi_fails(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object( - wheel_utils.subprocess, - "run", - return_value = _smi_result("", returncode = 1), - ), - ): - assert wheel_utils.has_blackwell_gpu() is False - - def test_returns_false_on_subprocess_timeout(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object( - wheel_utils.subprocess, - "run", - side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10), - ), - ): - assert wheel_utils.has_blackwell_gpu() is False - - def test_returns_false_on_malformed_output(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object( - wheel_utils.subprocess, - "run", - return_value = _smi_result("not-a-number\n\n"), - ), - ): - assert wheel_utils.has_blackwell_gpu() is False + def test_direct_wheel_url_reuses_torch210_on_211(self): + # causal-conv1d / mamba go through direct_wheel_url; torch 2.11 reuses the + # torch2.10 wheel filename just like flash-attn does. + url = wheel_utils.direct_wheel_url( + filename_prefix = "causal_conv1d", + package_version = "1.6.1", + release_tag = "v1.6.1.post4", + release_base_url = "https://example.test/download", + env = { + "python_tag": "cp313", + "torch_mm": "2.11", + "cuda_major": "13", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + }, + ) + assert url is not None + assert "causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url class TestFlashAttnWheelSelection: @@ -118,9 +51,24 @@ class TestFlashAttnWheelSelection: def test_torch_29_maps_to_v283(self): assert ips._select_flash_attn_version("2.9") == "2.8.3" - def test_unsupported_torch_has_no_wheel_mapping(self): + def test_torch_211_has_no_native_version_entry(self): + # The raw version table has no torch2.11-tagged wheel; the URL builder + # reuses the torch2.10 wheel instead (see test_torch_211_reuses_torch210_wheel). assert ips._select_flash_attn_version("2.11") is None + def test_torch_211_reuses_torch210_wheel(self): + url = ips._build_flash_attn_wheel_url( + { + "python_tag": "cp313", + "torch_mm": "2.11", + "cuda_major": "13", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + } + ) + assert url is not None + assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url + def test_exact_wheel_url_uses_full_env_tuple(self): url = ips._build_flash_attn_wheel_url( { @@ -333,83 +281,22 @@ class TestEnsureFlashAttn: mock_probe.assert_not_called() mock_install_wheel.assert_not_called() - def test_blackwell_gpu_skips_install_with_warning(self): - step_messages: list[tuple[str, str]] = [] - - def fake_step( - label: str, - value: str, - color_fn = None, - ): - step_messages.append((label, value)) - - with ( - mock.patch.object(ips, "NO_TORCH", False), - mock.patch.object(ips, "IS_WINDOWS", False), - mock.patch.object(ips, "IS_MACOS", False), - mock.patch.object(ips, "has_blackwell_gpu", return_value = True), - mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, - mock.patch.object(ips, "install_wheel") as mock_install_wheel, - mock.patch.object(ips, "_step", side_effect = fake_step), - mock.patch("subprocess.run", return_value = self._import_check()), - ): - ips._ensure_flash_attn() - - mock_probe.assert_not_called() - mock_install_wheel.assert_not_called() - assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages) - - def test_blackwell_gpu_on_windows_emits_blackwell_warning(self): - step_messages: list[tuple[str, str]] = [] - - def fake_step( - label: str, - value: str, - color_fn = None, - ): - step_messages.append((label, value)) - + def test_windows_skips_install_without_probing(self): + # flash-attn is Linux-only: on Windows the installer returns before + # probing the torch env or resolving a wheel (no Windows wheels are + # published upstream). with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", True), mock.patch.object(ips, "IS_MACOS", False), - mock.patch.object(ips, "has_blackwell_gpu", return_value = True), mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, mock.patch.object(ips, "install_wheel") as mock_install_wheel, - mock.patch.object(ips, "_step", side_effect = fake_step), mock.patch("subprocess.run", return_value = self._import_check()), ): ips._ensure_flash_attn() mock_probe.assert_not_called() mock_install_wheel.assert_not_called() - assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages) - - def test_non_blackwell_windows_does_not_emit_blackwell_warning(self): - step_messages: list[tuple[str, str]] = [] - - def fake_step( - label: str, - value: str, - color_fn = None, - ): - step_messages.append((label, value)) - - with ( - mock.patch.object(ips, "NO_TORCH", False), - mock.patch.object(ips, "IS_WINDOWS", True), - mock.patch.object(ips, "IS_MACOS", False), - mock.patch.object(ips, "has_blackwell_gpu", return_value = False), - mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, - mock.patch.object(ips, "install_wheel") as mock_install_wheel, - mock.patch.object(ips, "_step", side_effect = fake_step), - mock.patch("subprocess.run", return_value = self._import_check()), - ): - ips._ensure_flash_attn() - - mock_probe.assert_not_called() - mock_install_wheel.assert_not_called() - assert not any("Blackwell" in msg for _, msg in step_messages) class TestInstallPythonStackFlashAttnIntegration: From 38ea267124cc3c5a82b96fe9a7a140e4f4e566cc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 06:51:58 -0700 Subject: [PATCH 073/113] Versioning --- pyproject.toml | 6 +++--- unsloth/models/_utils.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 80b3d757e3..2b79121c82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "wheel>=0.42.0", "packaging", "numpy", @@ -94,7 +94,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "torchvision", "unsloth[triton]", ] @@ -579,7 +579,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b68cb702b1..fa0e0b1c49 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.7.1" +__version__ = "2026.7.2" __all__ = [ "SUPPORTS_BFLOAT16", From 3d41e5868d8aa46ece3b58de63f89be0a5d1d6b9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 07:22:54 -0700 Subject: [PATCH 074/113] Add has_blackwell_gpu to the mlx worker test's wheel_utils stub (#6980) worker.py imports has_blackwell_gpu from utils.wheel_utils, but _load_worker_module stubs utils.wheel_utils with a fixed name tuple that omitted it, so loading the worker raised ImportError (cannot import name 'has_blackwell_gpu') and Backend CI could not collect test_mlx_training_worker_config.py. Add the name to the stub so it matches worker.py's imports. --- studio/backend/tests/test_mlx_training_worker_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 503bae8da3..14fc0933d0 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -37,6 +37,7 @@ def _load_worker_module(): for name in ( "direct_wheel_url", "flash_attn_wheel_url", + "has_blackwell_gpu", "install_wheel", "probe_torch_wheel_env", "url_exists", From 116ce48c1a9d2b17596d66a247fe441551bdca91 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 07:26:10 -0700 Subject: [PATCH 075/113] Studio: allow CPU-only DiffusionGemma by granting the diffusion runner the CPU device (#6979) * Studio: allow CPU-only DiffusionGemma by granting the diffusion runner the CPU device * Studio: mark CPU-only DiffusionGemma as non-GPU-resident for training VRAM preflight * Studio: keep the CPU DiffusionGemma change minimal (revert VRAM-flag tweak; Metal hosts still hold unified memory) * Studio: keep CPU DiffusionGemma fallback fully CPU-masked so a masked GPU host does not re-expose GPU 0 --- studio/backend/core/inference/llama_cpp.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b07bd33076..f61402aa5c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -3971,7 +3971,11 @@ class LlamaCppBackend: # Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM # (capped at the training context). An explicit in-range n_ctx overrides it. maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0 - gpu = os.environ.get("DG_GPU", "0") + # No visible CUDA GPU: a genuine CPU host, or a GPU host masked with + # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child + # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. + cpu_only = self._effective_gpu_count() == 0 + gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") cmd = list(shim_cmd) + [ "--gguf", @@ -3991,6 +3995,12 @@ class LlamaCppBackend: # refuses to load unless UNSLOTH_IS_PRESENT is set (normally by `import # unsloth`). The shim never imports unsloth, so set it here as unsloth does. env["UNSLOTH_IS_PRESENT"] = "1" + # The shim's `import unsloth_zoo` aborts in get_device_type() ("needs a GPU") + # when no accelerator is visible, even though it only drives the CPU + # visual-server binary and does no torch GPU work. Allow the CPU device so the + # runner starts; the visual server still runs on the CPU llama.cpp build. + if cpu_only: + env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. From 1a274c488e9621281f86e2f6e85316a77e2c8070 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 07:51:53 -0700 Subject: [PATCH 076/113] Bump install.sh / install.ps1 pins to unsloth>=2026.7.2 and unsloth-zoo>=2026.7.2 (#6981) PyPI release unsloth 2026.7.2 is now live. Bumps the pinned floor in install.sh and install.ps1 from 2026.7.1 to 2026.7.2 for both unsloth and unsloth-zoo across all 5 install commands (no-torch / reinstall / upgrade / local / auto torch backend paths) so fresh installs resolve to the new wheel. Follows the same pattern as #5716. --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 9114f80af9..696f4e613a 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2155,7 +2155,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2169,7 +2169,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2235,7 +2235,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2247,7 +2247,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2275,7 +2275,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index 796d80e401..0acc9ec0be 100755 --- a/install.sh +++ b/install.sh @@ -2706,7 +2706,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2721,7 +2721,7 @@ if [ "$_MIGRATED" = true ]; then # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2925,7 +2925,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2943,7 +2943,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" + --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2975,7 +2975,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." From 5c2e53606e513cab3b698852b928b2b4484ada76 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:14:03 -0700 Subject: [PATCH 077/113] Studio: render thinking blocks for safetensors inference with prefilled templates (#6816) * Studio: render thinking blocks for safetensors inference with prefilled templates Reasoning templates like Qwen3.6 end the generation prompt with an open tag. skip_prompt streaming drops it, so the frontend never sees the opening tag and shows reasoning as plain text. Detect the prefill and re-emit it at the start of the stream on the transformers and MLX paths. Also stop stripping think tags in _clean_generated_text when a tokenizer marks them special. * Studio: guard think re-emit for special close tags, yield prefill early Address review feedback: - Guard: skip re-emitting the open when the tokenizer marks as a special token, since skip_special_tokens would strip the model's close tag and leave an unclosed block that swallows the answer. Falls back to plain text (pre-fix behaviour) for those tokenizers. - Yield the prefilled before the first token so the thinking block renders during prompt prefill instead of after the first generated token. - Drop the now-unnecessary _clean_generated_text think-tag exemption; the guard handles the special-token case at the source. No mainstream reasoning model (Qwen3.6, Qwen3, DeepSeek-R1, QwQ, GLM-4.6) marks think tags special, so behaviour is unchanged for them. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lyxot --- .../core/inference/chat_template_helpers.py | 37 ++++++++ studio/backend/core/inference/inference.py | 31 ++++++- .../backend/core/inference/mlx_inference.py | 20 ++++- .../tests/test_think_prefill_reemit.py | 89 +++++++++++++++++++ 4 files changed, 173 insertions(+), 4 deletions(-) create mode 100644 studio/backend/tests/test_think_prefill_reemit.py diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index dfd4c1c0bc..897db8262d 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -12,6 +12,43 @@ import json import logging from typing import Optional +_THINK_OPEN = "" +_THINK_CLOSE = "" + + +def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str: + """Return the trailing open ```` prefill of a rendered prompt. + + Reasoning templates (Qwen3.6, DeepSeek-R1-style) end the generation + prompt with ``\\n`` so the model starts reasoning immediately. + Because that opening tag is part of the *prompt*, skip_prompt streaming + never emits it, and the frontend's ````/```` parser shows + the reasoning as plain text instead of a thinking block. (The GGUF path + is unaffected: llama-server's reasoning parser returns + ``reasoning_content``, which gets re-wrapped in think tags.) + + Returns the exact prompt tail to re-emit at the start of the generated + stream (e.g. ``"\\n"``), or ``""`` when the prompt does not end + with an open think block, including the ``enable_thinking=False`` case + where templates prefill an already-closed ``\\n\\n``. + + ``special_tokens`` is the tokenizer's special-token list. If ```` + is one, the streamer's skip_special_tokens strips the model's closing tag, + so re-emitting the open would leave an unclosed block that swallows the + answer. In that case return ``""`` and fall back to plain text. + """ + if not prompt: + return "" + open_idx = prompt.rfind(_THINK_OPEN) + if open_idx == -1: + return "" + tail = prompt[open_idx:] + if _THINK_CLOSE in tail or tail.strip() != _THINK_OPEN: + return "" + if special_tokens and _THINK_CLOSE in set(special_tokens): + return "" + return tail + logger = logging.getLogger(__name__) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 167706f701..7e69e05124 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -1178,13 +1178,22 @@ class InferenceBackend: add_special_tokens = False, return_tensors = "pt", ).to(model.device) + prompt_text = input_text else: # Text-only path for a vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device) + prompt_text = formatted_prompt # Stream with TextIteratorStreamer + background thread try: + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill swallowed by skip_prompt (see + # generate_stream). + think_prefix = detect_think_prefill( + prompt_text, getattr(raw_tokenizer, "all_special_tokens", None) + ) from transformers import TextIteratorStreamer import threading @@ -1233,7 +1242,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False @@ -1467,6 +1480,16 @@ class InferenceBackend: from transformers import TextIteratorStreamer import threading + from core.inference.chat_template_helpers import detect_think_prefill + + # skip_prompt swallows an open prefilled by the template; + # re-emit it so the frontend can render the thinking block. + # gpt-oss emits its own tags via HarmonyTextStreamer. + think_prefix = ( + "" + if self._is_gpt_oss_model() + else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) + ) # gpt-oss models: HarmonyTextStreamer parses the multi-channel # harmony protocol into tags @@ -1550,7 +1573,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index d84baa278d..6287b184a6 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -485,6 +485,7 @@ class MLXInferenceBackend: from core.inference.chat_template_helpers import ( apply_chat_template_for_generation, + detect_think_prefill, render_with_native_template_fallback, ) @@ -518,6 +519,15 @@ class MLXInferenceBackend: hf_token = model_info.get("hf_token"), ) + # An open prefilled by the template lives in the prompt, not + # the generated tokens; re-emit it so the frontend renders the block. + think_prefix = detect_think_prefill( + prompt, getattr(self._tokenizer, "all_special_tokens", None) + ) + # Emit it before the first token so the block renders during prefill. + if think_prefix: + yield think_prefix + sampler = make_sampler( temp = temperature, top_p = top_p, @@ -570,7 +580,7 @@ class MLXInferenceBackend: token_ids, skip_special_tokens = True, ) - yield cumulative + yield think_prefix + cumulative if cancel_event and cancel_event.is_set(): break @@ -634,7 +644,13 @@ class MLXInferenceBackend: # mlx_vlm's stream_generate handles pixel_values (None for text-only) images = [image] if image is not None else None - cumulative = "" + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill from the prompt (see _generate_text). + cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) + # Emit it before the first token so the block renders during prefill. + if cumulative: + yield cumulative logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), diff --git a/studio/backend/tests/test_think_prefill_reemit.py b/studio/backend/tests/test_think_prefill_reemit.py new file mode 100644 index 0000000000..300ff92776 --- /dev/null +++ b/studio/backend/tests/test_think_prefill_reemit.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Unit tests for detect_think_prefill. + +Reasoning templates (Qwen3.6-style) end the generation prompt with an open +``\\n`` so the model starts reasoning immediately. skip_prompt +streaming drops that opening tag, so the safetensors/MLX paths must re-emit +it for the frontend's parser to render a thinking block. +""" + +import os +import sys + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.chat_template_helpers import detect_think_prefill + + +QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n" + + +def test_open_think_prefill_reemitted(): + """Qwen3.6-style enable_thinking=True prompt tail: \\n.""" + assert detect_think_prefill(QWEN_PROMPT + "\n") == "\n" + + +def test_bare_open_think_prefill_reemitted(): + """Prefill without trailing newline still detected.""" + assert detect_think_prefill(QWEN_PROMPT + "") == "" + + +def test_closed_think_prefill_not_reemitted(): + """enable_thinking=False prefills a closed, empty think block.""" + assert detect_think_prefill(QWEN_PROMPT + "\n\n\n\n") == "" + + +def test_prompt_without_think_untouched(): + """Non-reasoning templates produce no prefix.""" + assert detect_think_prefill(QWEN_PROMPT) == "" + + +def test_historical_think_blocks_ignored(): + """A closed think block in a prior assistant turn (preserve_thinking) + must not trigger re-emission when the generation tail is plain.""" + prompt = ( + "<|im_start|>user\nHi!<|im_end|>\n" + "<|im_start|>assistant\n\nprior reasoning\n\n\nHello!<|im_end|>\n" + "<|im_start|>user\nAgain?<|im_end|>\n<|im_start|>assistant\n" + ) + assert detect_think_prefill(prompt) == "" + + +def test_historical_blocks_plus_open_prefill(): + """Prior closed blocks plus a fresh open prefill: only the tail matters.""" + prompt = ( + "<|im_start|>assistant\n\nprior\n\n\nHello!<|im_end|>\n" + "<|im_start|>assistant\n\n" + ) + assert detect_think_prefill(prompt) == "\n" + + +def test_content_after_open_tag_not_reemitted(): + """If non-whitespace follows the tag it is not a plain prefill.""" + assert detect_think_prefill(QWEN_PROMPT + "\npartial reasoning") == "" + + +def test_empty_and_none_prompts(): + assert detect_think_prefill("") == "" + assert detect_think_prefill(None) == "" + + +def test_guard_suppresses_when_close_tag_is_special(): + """If is a special token, skip_special_tokens strips the model's + close tag, so re-emitting the open would leave an unclosed block. Guard off.""" + specials = ["<|im_end|>", "", ""] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "" + + +def test_guard_emits_when_think_not_special(): + specials = ["<|im_end|>", "<|endoftext|>"] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "\n" + + +def test_guard_default_and_empty_keep_emitting(): + assert detect_think_prefill(QWEN_PROMPT + "\n", None) == "\n" + assert detect_think_prefill(QWEN_PROMPT + "\n", []) == "\n" From 7a9fb4404e5c81ef5eb34de7d7944bd476c6192b Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:17:09 -0700 Subject: [PATCH 078/113] Remove API menu new badge (#6983) --- studio/frontend/src/components/app-sidebar.tsx | 3 --- 1 file changed, 3 deletions(-) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index f59a952b3a..823e420869 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -1572,9 +1572,6 @@ export function AppSidebar() { > {t("shell.navigation.api")} - - {t("common.new")} - } From 92c3e48529c8b7f96033f52f93845819b2ae53e3 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 8 Jul 2026 13:37:44 -0700 Subject: [PATCH 079/113] Fix BAD_MAPPINGS not redirecting the -unsloth-bnb-4bit dynamic quants (#6949) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .github/workflows/consolidated-tests-ci.yml | 1 + tests/test_bad_mappings_redirect.py | 47 +++++++++++++++++++++ unsloth/models/loader_utils.py | 5 +++ 3 files changed, 53 insertions(+) create mode 100644 tests/test_bad_mappings_redirect.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index ae4b386589..6ff3d19ba2 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -364,6 +364,7 @@ jobs: tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_bad_mappings_redirect.py \ tests/test_prefetch_snapshot_scope.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is diff --git a/tests/test_bad_mappings_redirect.py b/tests/test_bad_mappings_redirect.py new file mode 100644 index 0000000000..49dab2d98b --- /dev/null +++ b/tests/test_bad_mappings_redirect.py @@ -0,0 +1,47 @@ +"""Regression test for BAD_MAPPINGS redirecting oversized dynamic quants. + +get_model_name previously applied BAD_MAPPINGS only to the resolver's output, +but several listed names (the `-unsloth-bnb-4bit` dynamic quants, plus any name +the resolver doesn't map) come back as None, so their BAD_MAPPINGS entries were +dead and the oversized model loaded. Asserting over every entry catches all of +them. The mapper table and the resolver have no heavy imports of their own, +so we exec the import-free mapper module and ast-extract the resolver functions +rather than importing unsloth (which needs a GPU). +""" + +import ast +import os + +_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models") + + +def _load_get_model_name(): + mapper_ns = {} + with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f: + exec(compile(f.read(), "mapper.py", "exec"), mapper_ns) + + with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + namespace = dict(mapper_ns) + namespace["SUPPORTS_FOURBIT"] = True + namespace["_env_says_offline"] = lambda: True + namespace["_get_new_mapper"] = lambda: ({}, {}, {}) + + wanted = {"__get_model_name", "_resolve_with_mappers", "get_model_name"} + for node in tree.body: + if isinstance(node, ast.Assign) and any( + getattr(target, "id", None) == "BAD_MAPPINGS" for target in node.targets + ): + exec(compile(ast.Module([node], []), "", "exec"), namespace) + elif isinstance(node, ast.FunctionDef) and node.name in wanted: + exec(compile(ast.Module([node], []), node.name, "exec"), namespace) + + return namespace["get_model_name"], namespace["BAD_MAPPINGS"] + + +def test_bad_mappings_redirect_every_listed_name(): + get_model_name, bad_mappings = _load_get_model_name() + assert bad_mappings, "BAD_MAPPINGS should not be empty" + for name, expected in bad_mappings.items(): + assert get_model_name(name, load_in_4bit = True) == expected, name diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index d6d6ce877e..fa6282bcf6 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -239,6 +239,11 @@ def get_model_name( and new_model_name.lower() in BAD_MAPPINGS ): new_model_name = BAD_MAPPINGS[new_model_name.lower()] + elif new_model_name is None and model_name.lower() in BAD_MAPPINGS: + # Some bad names (e.g. the `-unsloth-bnb-4bit` dynamic quants) are keys + # of the mappers, not values, so the resolver returns None for them and + # the remap above is skipped; remap the input name directly instead. + new_model_name = BAD_MAPPINGS[model_name.lower()] if ( new_model_name is None From dc4618ce475554d99fab76daae9746baa62bf090 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 8 Jul 2026 13:40:39 -0700 Subject: [PATCH 080/113] Fix duplicate unsloth/gemma-2b-bnb-4bit mapper key routing the base 4bit repo to the instruct model (#6891) --- .github/workflows/consolidated-tests-ci.yml | 1 + tests/test_gemma_2b_mapper_key.py | 46 +++++++++++++++++++++ unsloth/models/mapper.py | 2 +- 3 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 tests/test_gemma_2b_mapper_key.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 6ff3d19ba2..1bb4c2bb58 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -366,6 +366,7 @@ jobs: tests/python/test_fast_language_model_text_only.py \ tests/test_bad_mappings_redirect.py \ tests/test_prefetch_snapshot_scope.py \ + tests/test_gemma_2b_mapper_key.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn diff --git a/tests/test_gemma_2b_mapper_key.py b/tests/test_gemma_2b_mapper_key.py new file mode 100644 index 0000000000..31edacfce4 --- /dev/null +++ b/tests/test_gemma_2b_mapper_key.py @@ -0,0 +1,46 @@ +"""Regression test for the duplicate ``unsloth/gemma-2b-bnb-4bit`` key in +``unsloth/models/mapper.py``. + +The 4bit instruction-tuned Gemma 2B entry was accidentally keyed with the base +model's repo name, so ``__INT_TO_FLOAT_MAPPER`` held two identical +``unsloth/gemma-2b-bnb-4bit`` keys. Python keeps only the last value for a +duplicate literal key, so the base 4bit repo resolved to the *instruct* model, +the base model lost its reverse (4x-faster) mapping, and +``unsloth/gemma-2b-it-bnb-4bit`` was never registered at all. + +``mapper.py`` has no imports, so we exec it directly and inspect the built +mappers without importing ``unsloth`` (which requires a GPU). +""" + +import os + +MAPPER_PATH = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models", "mapper.py") + + +def _load_mappers(): + with open(MAPPER_PATH) as f: + source = f.read() + namespace = {} + exec(compile(source, MAPPER_PATH, "exec"), namespace) + return namespace + + +def test_gemma_2b_base_and_instruct_4bit_are_distinct(): + namespace = _load_mappers() + int_to_float = namespace["INT_TO_FLOAT_MAPPER"] + float_to_int = namespace["FLOAT_TO_INT_MAPPER"] + + # The base 4bit repo must resolve to the base model, not the instruct one. + assert int_to_float["unsloth/gemma-2b-bnb-4bit"] == "unsloth/gemma-2b" + + # The instruct 4bit repo must be registered and resolve to the instruct model. + assert "unsloth/gemma-2b-it-bnb-4bit" in int_to_float + assert int_to_float["unsloth/gemma-2b-it-bnb-4bit"] == "unsloth/gemma-2b-it" + + # The base model must reverse-map back to the base 4bit repo. + assert float_to_int["unsloth/gemma-2b"] == "unsloth/gemma-2b-bnb-4bit" + assert float_to_int["google/gemma-2b"] == "unsloth/gemma-2b-bnb-4bit" + + # The instruct model must reverse-map to the instruct 4bit repo. + assert float_to_int["unsloth/gemma-2b-it"] == "unsloth/gemma-2b-it-bnb-4bit" + assert float_to_int["google/gemma-2b-it"] == "unsloth/gemma-2b-it-bnb-4bit" diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index 57c1e292c3..f3a0e1f9bb 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -134,7 +134,7 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/gemma-7b-it", "google/gemma-7b-it", ), - "unsloth/gemma-2b-bnb-4bit" : ( + "unsloth/gemma-2b-it-bnb-4bit" : ( "unsloth/gemma-2b-it", "google/gemma-2b-it", ), From 85a068cfe10f2f8bc214dbd2e0b82164419fc8de Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 8 Jul 2026 13:57:41 -0700 Subject: [PATCH 081/113] Fix to_sharegpt optional block rendering "None" for missing extra columns (#6827) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .../python/test_to_sharegpt_optional_none.py | 97 +++++++++++++++++++ unsloth/chat_templates.py | 13 ++- 2 files changed, 108 insertions(+), 2 deletions(-) create mode 100644 tests/python/test_to_sharegpt_optional_none.py diff --git a/tests/python/test_to_sharegpt_optional_none.py b/tests/python/test_to_sharegpt_optional_none.py new file mode 100644 index 0000000000..05fcb8a6a8 --- /dev/null +++ b/tests/python/test_to_sharegpt_optional_none.py @@ -0,0 +1,97 @@ +import ast +import re +from pathlib import Path + + +def _load_formatter_builders(): + # Extract _parse_combined_prompt and _create_formatter without importing + # unsloth (importing unsloth needs unsloth_zoo / a GPU). Both are pure + # Python and only use the `re` module. + source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + wanted = {"_parse_combined_prompt", "_create_formatter"} + funcs = [ + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted + ] + namespace = {"re": re} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace["_parse_combined_prompt"], namespace["_create_formatter"] + + +class _StubDataset: + def __init__(self, column_names): + self.column_names = column_names + + +def _render(merged_prompt, columns, batch): + parse, create = _load_formatter_builders() + possible_columns, final_optional_prompts = parse(merged_prompt, _StubDataset(columns)) + processor = create(possible_columns, final_optional_prompts, "text") + return processor(batch)["text"] + + +def test_optional_block_missing_second_column_does_not_render_none(): + # A [[...]] block may reference several columns; only the first gates the + # block. A later column that is None must not render as the literal "None". + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": [None]}, + ) + assert out[0] == "Location: Paris, end" + assert "None" not in out[0] + + +def test_optional_block_all_columns_present_unchanged(): + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": ["France"]}, + ) + assert out[0] == "Location: Paris, France end" + + +def test_optional_block_gating_column_empty_is_dropped(): + # When the gating (first) column is empty the whole block is omitted; this + # behaviour is unchanged by the None coercion. + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": [""], "country": ["France"]}, + ) + assert out[0] == "Location: end" + + +def test_single_column_optional_block_gated_out_on_none(): + # Single-column blocks were already gated correctly (the sole column is the + # gate); confirm they stay unaffected. + merged_prompt = "Name: [[{name}]]!" + out = _render(merged_prompt, ["name"], {"name": [None, "Bob"]}) + assert out == ["Name: !", "Name: Bob!"] + + +def test_required_column_none_does_not_render_none(): + # A required (non-[[...]]) column that is None must not render as the + # literal "None" either; coercion happens at the row source, so both the + # required and optional branches are covered. + merged_prompt = "Location: {city}, {country} end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": [None]}, + ) + assert out[0] == "Location: Paris, end" + assert "None" not in out[0] + + +def test_optional_block_falsy_but_present_gating_value_still_renders(): + # The gate keeps a block whenever the first column is not "". A falsy but + # real value (0) must not be treated as absent, so the block still renders. + merged_prompt = "Count: [[{n}]]!" + out = _render(merged_prompt, ["n"], {"n": [0]}) + assert out[0] == "Count: 0!" diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 169b2dbd0e..dd1e433471 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2185,7 +2185,16 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name texts = [] for row_idx in range(n_rows): - row_values = {column: examples[column][row_idx] for column in columns} + # Coerce missing (None) columns to "" so they do not render as the + # literal string "None" in the emitted text. In a [[...]] block only + # the first column gates the block, so a later column can still be + # None here; required columns can be None too. Coercing at the source + # covers both; since None is now "", the gate below only needs to + # test for "" (an empty first column still drops the block). + row_values = { + column: ("" if (value := examples[column][row_idx]) is None else value) + for column in columns + } formatter_values = {} for formatter_template in formatter_templates: @@ -2196,7 +2205,7 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name continue _, optional_name, prompt, needed_columns = formatter_template - if row_values[needed_columns[0]] not in (None, ""): + if row_values[needed_columns[0]] != "": prompt_values = {column: row_values[column] for column in needed_columns} formatter_values[optional_name] = prompt.format(**prompt_values) else: From 81f789ba85bb45c30fe7a8e60126112e07e3ddac Mon Sep 17 00:00:00 2001 From: ramisworld Date: Thu, 9 Jul 2026 09:32:51 +1200 Subject: [PATCH 082/113] Guard FP8 Triton launches with tensor device context (#6888) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- tests/test_fp8_device_context.py | 249 +++++++++++++++++++++++++++++++ unsloth/kernels/fp8.py | 87 ++++++----- 2 files changed, 300 insertions(+), 36 deletions(-) create mode 100644 tests/test_fp8_device_context.py diff --git a/tests/test_fp8_device_context.py b/tests/test_fp8_device_context.py new file mode 100644 index 0000000000..2eea35f4e6 --- /dev/null +++ b/tests/test_fp8_device_context.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import ast +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +FP8_SOURCE = REPO_ROOT / "unsloth" / "kernels" / "fp8.py" + + +class _FakeDeviceModule: + def __init__(self, device_count: int) -> None: + self._device_count = device_count + self.device_calls = [] + + def device_count(self) -> int: + return self._device_count + + def device(self, device): + self.device_calls.append(device) + return ("device-context", device) + + +class _FakeTorch: + Tensor = object + + def __init__( + self, + cuda_device_count: int, + xpu_device_count: int = 0, + ) -> None: + self.cuda = _FakeDeviceModule(cuda_device_count) + self.xpu = _FakeDeviceModule(xpu_device_count) + + +class _LaunchVisitor(ast.NodeVisitor): + def __init__(self) -> None: + self.guarded_launches: set[str] = set() + self.unguarded_launches: set[str] = set() + self._inside_fp8_device_context = 0 + + def visit_With(self, node: ast.With) -> None: + enters_context = any( + isinstance(item.context_expr, ast.Call) + and isinstance(item.context_expr.func, ast.Name) + and item.context_expr.func.id == "_fp8_triton_device_context" + for item in node.items + ) + if enters_context: + self._inside_fp8_device_context += 1 + for statement in node.body: + self.visit(statement) + if enters_context: + self._inside_fp8_device_context -= 1 + + def visit_Call(self, node: ast.Call) -> None: + launch_name = self._triton_launch_name(node) + if launch_name is not None: + if self._inside_fp8_device_context: + self.guarded_launches.add(launch_name) + else: + self.unguarded_launches.add(launch_name) + self.generic_visit(node) + + @staticmethod + def _triton_launch_name(node: ast.Call) -> str | None: + if isinstance(node.func, ast.Name) and node.func.id == "triton_quantize_fp8_block": + return node.func.id + if not isinstance(node.func, ast.Subscript): + return None + if not isinstance(node.func.value, ast.Name): + return None + return node.func.value.id + + +def _load_device_context_helper(fake_torch: _FakeTorch): + source = FP8_SOURCE.read_text() + tree = ast.parse(source) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == "_fp8_triton_device_context": + namespace = {"torch": fake_torch, "nullcontext": nullcontext} + exec(ast.get_source_segment(source, node), namespace) + return namespace["_fp8_triton_device_context"] + raise AssertionError("_fp8_triton_device_context was not found") + + +def test_fp8_device_context_selects_cuda_tensor_device_on_multi_gpu() -> None: + fake_torch = _FakeTorch(cuda_device_count = 2) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "cuda")) + + context = helper(tensor) + + assert context == ("device-context", tensor.device) + assert fake_torch.cuda.device_calls == [tensor.device] + + +def test_fp8_device_context_is_noop_for_single_cuda_device() -> None: + fake_torch = _FakeTorch(cuda_device_count = 1) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "cuda")) + + context = helper(tensor) + + assert isinstance(context, nullcontext) + assert fake_torch.cuda.device_calls == [] + + +def test_fp8_device_context_selects_xpu_tensor_device_on_multi_gpu() -> None: + fake_torch = _FakeTorch(cuda_device_count = 0, xpu_device_count = 2) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "xpu")) + + context = helper(tensor) + + assert context == ("device-context", tensor.device) + assert fake_torch.xpu.device_calls == [tensor.device] + + +def test_fp8_device_context_is_noop_for_single_xpu_device() -> None: + fake_torch = _FakeTorch(cuda_device_count = 0, xpu_device_count = 1) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "xpu")) + + context = helper(tensor) + + assert isinstance(context, nullcontext) + assert fake_torch.xpu.device_calls == [] + + +def test_fp8_device_context_is_noop_for_non_cuda_tensor() -> None: + fake_torch = _FakeTorch(cuda_device_count = 8) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "cpu")) + + context = helper(tensor) + + assert isinstance(context, nullcontext) + assert fake_torch.cuda.device_calls == [] + + +def test_fp8_triton_launches_enter_tensor_device_context() -> None: + tree = ast.parse(FP8_SOURCE.read_text()) + function_names = {node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)} + assert "_fp8_triton_device_context" in function_names + + visitor = _LaunchVisitor() + visitor.visit(tree) + + expected_launches = { + "weight_dequant_kernel", + "act_quant_kernel", + "_w8a8_block_fp8_matmul", + "triton_quantize_fp8_block", + } + assert expected_launches <= visitor.guarded_launches + assert not (expected_launches & visitor.unguarded_launches) + + +def _require_two_cuda_devices(): + torch = pytest.importorskip("torch") + pytest.importorskip("triton") + + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("requires at least two CUDA devices") + return torch + + +def test_weight_dequant_block_runs_on_tensor_device_when_current_device_differs() -> None: + torch = _require_two_cuda_devices() + from unsloth.kernels.fp8 import weight_dequant_block + + previous_device = torch.cuda.current_device() + try: + torch.cuda.set_device(0) + x = torch.arange(256 * 256, device = "cuda:1", dtype = torch.float32).reshape(256, 256) + scales = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device = "cuda:1", dtype = torch.float32) + + actual = weight_dequant_block(x, scales, block_size = 128, dtype = torch.float32) + + expanded_scales = scales.repeat_interleave(128, dim = 0).repeat_interleave(128, dim = 1) + expected = x * expanded_scales + + assert actual.device == x.device + assert torch.cuda.current_device() == 0 + torch.testing.assert_close(actual, expected) + finally: + torch.cuda.set_device(previous_device) + + +def test_act_quant_runs_on_tensor_device_when_current_device_differs() -> None: + torch = _require_two_cuda_devices() + if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("requires torch.float8_e4m3fn") + if torch.cuda.get_device_capability(1)[0] < 9: + pytest.skip("requires FP8-capable CUDA hardware") + + from unsloth.kernels.fp8 import act_quant + + previous_device = torch.cuda.current_device() + try: + torch.cuda.set_device(0) + x = torch.arange(256, device = "cuda:1", dtype = torch.float32).reshape(2, 128) + + y, scales = act_quant(x, block_size = 128) + + assert y.device == x.device + assert scales.device == x.device + assert torch.cuda.current_device() == 0 + finally: + torch.cuda.set_device(previous_device) + + +def test_w8a8_block_fp8_matmul_triton_runs_on_tensor_device_when_current_device_differs() -> None: + torch = _require_two_cuda_devices() + if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("requires torch.float8_e4m3fn") + if torch.cuda.get_device_capability(1)[0] < 9: + pytest.skip("requires FP8-capable CUDA hardware") + + from unsloth.kernels.fp8 import w8a8_block_fp8_matmul_triton + + previous_device = torch.cuda.current_device() + try: + torch.cuda.set_device(0) + A = torch.ones((128, 128), device = "cuda:1", dtype = torch.float32).to(torch.float8_e4m3fn) + B = torch.ones((128, 128), device = "cuda:1", dtype = torch.float32).to(torch.float8_e4m3fn) + As = torch.ones((128, 1), device = "cuda:1", dtype = torch.float32) + Bs = torch.ones((1, 1), device = "cuda:1", dtype = torch.float32) + + actual = w8a8_block_fp8_matmul_triton( + A, + B, + As, + Bs, + block_size = [128, 128], + output_dtype = torch.float32, + ) + + expected = torch.full((128, 128), 128.0, device = "cuda:1", dtype = torch.float32) + assert actual.device == A.device + assert torch.cuda.current_device() == 0 + torch.testing.assert_close(actual, expected) + finally: + torch.cuda.set_device(previous_device) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 935ffbb447..4efc4bd5d3 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import os +from contextlib import nullcontext import torch import torch.nn as nn import triton @@ -24,6 +25,15 @@ from unsloth_zoo.temporary_patches.common import torch_compile torch_matmul = torch.matmul + +def _fp8_triton_device_context(tensor: torch.Tensor): + if tensor.device.type == "cuda" and torch.cuda.device_count() > 1: + return torch.cuda.device(tensor.device) + if tensor.device.type == "xpu" and hasattr(torch, "xpu") and torch.xpu.device_count() > 1: + return torch.xpu.device(tensor.device) + return nullcontext() + + try: from transformers.integrations.finegrained_fp8 import FP8Linear except: @@ -95,7 +105,8 @@ def weight_dequant_block( triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"]), ) - weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size) + with _fp8_triton_device_context(x): + weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size) return y @@ -149,7 +160,8 @@ def act_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, tor def grid(meta): return (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),) - act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size) + with _fp8_triton_device_context(x): + act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size) return y, s @@ -274,32 +286,33 @@ def w8a8_block_fp8_matmul_triton( def grid(META): return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),) - _w8a8_block_fp8_matmul[grid]( - A, - B, - C, - As, - Bs, - M, - N, - K, - block_n, - block_k, - A.stride(-2), - A.stride(-1), - B.stride(1), - B.stride(0), - C.stride(-2), - C.stride(-1), - As.stride(-2), - As.stride(-1), - Bs.stride(1), - Bs.stride(0), - BLOCK_SIZE_M = BLOCK_SIZE_M, - BLOCK_SIZE_N = BLOCK_SIZE_N, - BLOCK_SIZE_K = BLOCK_SIZE_K, - GROUP_SIZE_M = 8, - ) + with _fp8_triton_device_context(A): + _w8a8_block_fp8_matmul[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + BLOCK_SIZE_M = BLOCK_SIZE_M, + BLOCK_SIZE_N = BLOCK_SIZE_N, + BLOCK_SIZE_K = BLOCK_SIZE_K, + GROUP_SIZE_M = 8, + ) return C @@ -311,13 +324,14 @@ def torchao_block_matmul( block_size: tuple[int, int], output_dtype: torch.dtype = torch.bfloat16, ): - out = torchao_blockwise_gemm( - act_q.contiguous(), - act_scale.contiguous(), - weight_q.contiguous(), - weight_scale.contiguous(), - block_size = block_size[1], - ) + with _fp8_triton_device_context(act_q): + out = torchao_blockwise_gemm( + act_q.contiguous(), + act_scale.contiguous(), + weight_q.contiguous(), + weight_scale.contiguous(), + block_size = block_size[1], + ) return out.to(output_dtype) @@ -540,7 +554,8 @@ class FP8_fbgemm_block_linear(torch.autograd.Function): f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {bs_n, bs_k}" ) - xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None) + with _fp8_triton_device_context(X): + xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None) # TODO: WARNING - diverges from baseline for high X values, producing # gibberish / high starting loss. Do not use until resolved; kept for a # future headstart. From 3b73cd88293728b3ba173e235d405cec40f7a42b Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Thu, 9 Jul 2026 08:33:03 +0530 Subject: [PATCH 083/113] Fix per-block ID collisions and add block cleanup for unstructured uploads (#6944) * unstructured block removal * Enhance unstructured block handling * Restrict block cleanup to upload UIDs * cleanup for seed block uploads * upload cleanup queue for unstructured blocks in recipe studio * Fix unstructured upload cleanup edge cases * Fix unstructured upload import ownership * Fix-unstructured-import-path-ownership * Guard failed-delete restore against stale block in unstructured drop zone * Drain queued upload cleanups when autosave is skipped --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 Co-authored-by: Daniel Han --- studio/backend/routes/data_recipe/seed.py | 37 ++++++ studio/backend/tests/test_data_recipe_seed.py | 97 +++++++++++++++ .../src/features/recipe-studio/api/index.ts | 10 ++ .../dialogs/seed/seed-dialog.tsx | 69 ++++++++++- .../dialogs/seed/unstructured-drop-zone.tsx | 37 ++++-- .../hooks/use-recipe-persistence.ts | 86 ++++++++++++-- .../recipe-studio/stores/recipe-studio.ts | 61 +++++++++- .../src/features/recipe-studio/types/index.ts | 2 + .../recipe-studio/utils/config-factories.ts | 43 +++++++ .../recipe-studio/utils/import/importer.ts | 112 +++++++++++------- .../import/parsers/seed-config-parser.ts | 14 ++- .../utils/payload/build-payload.ts | 3 + .../recipe-studio/utils/payload/types.ts | 2 + 13 files changed, 502 insertions(+), 71 deletions(-) diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 57a291291e..a5b75b7335 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -10,6 +10,7 @@ import binascii import json import os import re +import shutil from itertools import islice from pathlib import Path from typing import Any @@ -59,6 +60,9 @@ UNSTRUCTURED_ALLOWED_EXTS = {".pdf", ".docx", ".txt", ".md"} SEED_UPLOAD_DIR = seed_uploads_root() UNSTRUCTURED_UPLOAD_ROOT = unstructured_uploads_root() _SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$") +# Frontend-generated upload namespace (UUID4 hex). Legacy node ids (n1, ...) +# never match: those directories can be shared by several recipes. +_UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$") def _validate_safe_id(value: str, label: str) -> str: @@ -580,6 +584,39 @@ async def remove_unstructured_file(block_id: str, file_id: str): return {"status": "ok"} +@router.delete("/seed/unstructured-block/{block_id}") +async def remove_unstructured_block(block_id: str): + """Delete a block's upload directory; files on disk still count toward its quota. + + Only uid-namespaced directories may be bulk-deleted: they have exactly one + owning block. Legacy node-id directories (n1, ...) can be shared by other + recipes, so they are managed file-by-file instead. + """ + _validate_safe_id(block_id, "block_id") + if not _UPLOAD_UID_RE.match(block_id): + raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted") + + block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve() + if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()): + raise HTTPException(400, "Invalid block_id: outside upload root") + if not block_dir.exists(): + return {"status": "ok", "deleted": False} + + try: + shutil.rmtree(block_dir) + except OSError as exc: + raise log_and_http_error( + exc, + 500, + "failed to delete uploaded files", + event = "data_recipe.seed.unstructured_block_delete_failed", + log = logger, + ) from exc + if block_dir.exists(): + raise HTTPException(500, "failed to delete uploaded files") + return {"status": "ok", "deleted": True} + + @router.post("/seed/inspect-upload", response_model = SeedInspectResponse) def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse: if payload.file_ids is not None: diff --git a/studio/backend/tests/test_data_recipe_seed.py b/studio/backend/tests/test_data_recipe_seed.py index 09e22116ed..58bbd24061 100644 --- a/studio/backend/tests/test_data_recipe_seed.py +++ b/studio/backend/tests/test_data_recipe_seed.py @@ -124,3 +124,100 @@ def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, e assert result.status == "error" assert result.error == "Text extraction failed." assert _block_files(seed_route) == [] + + +_TEST_UPLOAD_UID = "0f" * 16 + + +def test_remove_unstructured_block_deletes_directory(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = _TEST_UPLOAD_UID) + assert _block_files(seed_route, _TEST_UPLOAD_UID) != [] + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": True} + assert not (seed_route.UNSTRUCTURED_UPLOAD_ROOT / _TEST_UPLOAD_UID).exists() + + +def test_remove_unstructured_block_missing_directory_is_ok(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": False} + + +def test_remove_unstructured_block_rejects_unsafe_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("../escape")) + + assert exc.value.status_code == 400 + + +def test_remove_unstructured_block_rejects_legacy_node_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = "n1") + assert _block_files(seed_route, "n1") != [] + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("n1")) + + assert exc.value.status_code == 400 + assert _block_files(seed_route, "n1") != [] + + +def test_remove_unstructured_block_rejects_symlink_escape(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + outside = tmp_path / "outside" + outside.mkdir() + (outside / "victim.txt").write_text("keep me") + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + root.mkdir(parents = True) + (root / _TEST_UPLOAD_UID).symlink_to(outside) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert exc.value.status_code == 400 + assert (outside / "victim.txt").exists() + + +def test_remove_unstructured_block_fails_if_directory_remains(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + block_dir = root / _TEST_UPLOAD_UID + block_dir.mkdir(parents = True) + (block_dir / "victim.txt").write_text("keep me") + + calls = [] + + def noop_rmtree(path, *args, **kwargs): + calls.append((path, args, kwargs)) + + monkeypatch.setattr(seed_route.shutil, "rmtree", noop_rmtree) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert calls + assert exc.value.status_code == 500 + assert block_dir.exists() + + +def test_total_upload_quota_is_scoped_per_block(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr(seed_route, "UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES", 10) + + first = _run_upload(seed_route, "a.txt", b"123456789") + assert first.status == "ok" + + with pytest.raises(seed_route.HTTPException) as exc: + _run_upload(seed_route, "b.txt", b"123") + assert exc.value.status_code == 413 + + # Another block starts with its own untouched budget. + other = _run_upload(seed_route, "c.txt", b"123", block_id = "other") + assert other.status == "ok" diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index 4fe24b1f6d..f4e8167cb3 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -494,3 +494,13 @@ export async function removeUnstructuredFile( throw new Error("Failed to remove file"); } } + +export async function removeUnstructuredBlock(blockId: string): Promise { + const res = await authFetch( + `${DATA_DESIGNER_API_BASE}/seed/unstructured-block/${encodeURIComponent(blockId)}`, + { method: "DELETE" }, + ); + if (!res.ok && res.status !== 404) { + throw new Error("Failed to remove uploaded files"); + } +} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index 0ed7eeed75..53f8566090 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -53,6 +53,11 @@ import { inspectSeedDataset, inspectSeedUpload, } from "../../api"; +import { useRecipeStudioStore } from "../../stores/recipe-studio"; +import { + makeUnstructuredUploadUid, + resolveUnstructuredUploadBlockId, +} from "../../utils/config-factories"; import { resolveImagePreview } from "../../utils/image-preview"; import type { GithubItemType, @@ -597,6 +602,41 @@ export function SeedDialog({ const mode = config.seed_source_type ?? "hf"; const previewEmpty = getPreviewEmptyStateCopy(mode); + const queueUploadCleanup = useRecipeStudioStore( + (state) => state.queueUploadCleanup, + ); + + // config.id collides across recipes (ids reset to n1 on import); use a + // stable per-block uid instead. Generate one synchronously so the first + // rendered drop zone cannot upload under a legacy node id. + const uploadUid = config.unstructured_upload_uid?.trim() ?? ""; + const unstructuredFileCount = config.unstructured_file_ids?.length ?? 0; + const generatedUploadUidRef = useRef(null); + if ( + mode === "unstructured" && + !uploadUid && + unstructuredFileCount === 0 && + generatedUploadUidRef.current === null + ) { + generatedUploadUidRef.current = makeUnstructuredUploadUid(); + } + const uploadBlockId = resolveUnstructuredUploadBlockId({ + configId: config.id, + uploadUid, + generatedUploadUid: generatedUploadUidRef.current, + unstructuredFileCount, + }); + + useEffect(() => { + if (mode !== "unstructured") return; + if (uploadUid) return; + if (unstructuredFileCount > 0) return; + const nextUid = + generatedUploadUidRef.current ?? makeUnstructuredUploadUid(); + generatedUploadUidRef.current = nextUid; + onUpdate({ unstructured_upload_uid: nextUid }); + }, [mode, uploadUid, unstructuredFileCount, onUpdate]); + const prevModeRef = useRef(mode); useEffect(() => { const prevMode = prevModeRef.current; @@ -720,6 +760,11 @@ export function SeedDialog({ subset: config.hf_subset?.trim() || undefined, preview_size: 10, }); + // Queue the block's upload directory for deletion after the next + // save; only uid-namespaced directories qualify (single owner). + if (uploadUid && unstructuredFileCount > 0) { + queueUploadCleanup(uploadUid); + } onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, @@ -730,6 +775,7 @@ export function SeedDialog({ hf_split: response.split ?? "", hf_subset: response.subset ?? "", local_file_name: "", + unstructured_upload_uid: "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -754,6 +800,11 @@ export function SeedDialog({ content_base64: payload, preview_size: 10, }); + // Queue the block's upload directory for deletion after the next + // save; only uid-namespaced directories qualify (single owner). + if (uploadUid && unstructuredFileCount > 0) { + queueUploadCleanup(uploadUid); + } onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, @@ -765,6 +816,7 @@ export function SeedDialog({ hf_subset: "", hf_split: "", local_file_name: localFile.name, + unstructured_upload_uid: "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -789,7 +841,7 @@ export function SeedDialog({ const { chunkSize, chunkOverlap } = resolveChunking(config); const response = await inspectSeedUpload({ - block_id: config.id, + block_id: uploadBlockId, file_ids: fileIds, file_names: fileNames, preview_size: 10, @@ -827,7 +879,18 @@ export function SeedDialog({ setIsInspecting(false); } }, - [config, getCurrentLoadKey, localFile, mode, onUpdate, unstructuredFiles], + [ + config, + getCurrentLoadKey, + localFile, + mode, + onUpdate, + queueUploadCleanup, + unstructuredFiles, + unstructuredFileCount, + uploadBlockId, + uploadUid, + ], ); useEffect(() => { @@ -997,7 +1060,7 @@ export function SeedDialog({ {mode === "unstructured" && ( (null); const filesRef = useRef(files); + const blockIdRef = useRef(blockId); + const mountedRef = useRef(true); const [isDragOver, setIsDragOver] = useState(false); useEffect(() => { filesRef.current = files; - }, [files]); + blockIdRef.current = blockId; + }, [files, blockId]); + useEffect(() => () => { + mountedRef.current = false; + }, []); const totalSize = files.reduce((sum, f) => sum + f.size, 0); @@ -134,15 +140,32 @@ export function UnstructuredDropZone({ if (entry.status === "uploading" && entry.abortController) { entry.abortController.abort(); } - if ( + const needsServerRemove = entry.id && entry.status === "ok" && - !deletedIdsRef.current.has(entry.id) - ) { - deletedIdsRef.current.add(entry.id); - void removeUnstructuredFile(blockId, entry.id).catch(() => {}); - } + !deletedIdsRef.current.has(entry.id); onFilesChange((prev) => prev.filter((_, i) => i !== index)); + if (!needsServerRemove) return; + deletedIdsRef.current.add(entry.id); + removeUnstructuredFile(blockId, entry.id).catch(() => { + // Skip if the drop zone unmounted or its block changed: the id no + // longer belongs here and restoring would leak it into another block. + if (!mountedRef.current || blockIdRef.current !== blockId) return; + // Still exists server-side (counts toward quota); restore it at its + // original position. + deletedIdsRef.current.delete(entry.id); + onFilesChange((prev) => { + const next = [...prev]; + next.splice(Math.min(index, next.length), 0, { + id: entry.id, + name: entry.name, + size: entry.size, + status: "ok", + error: "Remove failed — try again", + }); + return next; + }); + }); }, [blockId, onFilesChange], ); diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts index 0417a6dd22..2d91272469 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts @@ -4,11 +4,13 @@ import { useCallback, useEffect, useMemo, useState } from "react"; import { toastError, toastSuccess } from "@/shared/toast"; import { normalizeNonEmptyName } from "@/utils"; +import { removeUnstructuredBlock } from "../api"; import { buildSignature, copyTextToClipboard, formatSavedLabel, } from "../executions/execution-helpers"; +import { useRecipeStudioStore } from "../stores/recipe-studio"; import { importRecipePayload, type RecipeSnapshot } from "../utils/import"; import type { RecipePayloadResult } from "../utils/payload/types"; @@ -72,7 +74,10 @@ function stripApiKeys(value: unknown): unknown { !Array.isArray(output.env) ) { output.env = Object.fromEntries( - Object.keys(output.env as Record).map((envKey) => [envKey, ""]), + Object.keys(output.env as Record).map((envKey) => [ + envKey, + "", + ]), ); } return output; @@ -82,10 +87,7 @@ function inferHfRepoIdFromPath(pathValue: unknown): string { if (typeof pathValue !== "string") { return ""; } - const parts = pathValue - .trim() - .split("/") - .filter(Boolean); + const parts = pathValue.trim().split("/").filter(Boolean); if (parts.length >= 3 && parts[0] === "datasets") { return `${parts[1]}/${parts[2]}`; } @@ -126,8 +128,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { typeof ui?.seed_source_type === "string" ? ui.seed_source_type : null; const sourceType = typeof source?.seed_type === "string" ? source.seed_type : null; - const shouldResetHfState = - sourceType === "hf" || uiSourceType === "hf"; + const shouldResetHfState = sourceType === "hf" || uiSourceType === "hf"; const shouldResetLocalState = sourceType === "local" || sourceType === "unstructured" || @@ -144,6 +145,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { ui.seed_drop_columns = []; ui.seed_preview_rows = []; ui.local_file_name = ""; + ui.unstructured_upload_uid = ""; ui.unstructured_file_ids = []; ui.unstructured_file_names = []; ui.unstructured_file_sizes = []; @@ -165,6 +167,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { ui.seed_drop_columns = []; ui.seed_preview_rows = []; ui.local_file_name = ""; + ui.unstructured_upload_uid = ""; ui.unstructured_file_ids = []; ui.unstructured_file_names = []; ui.unstructured_file_sizes = []; @@ -174,6 +177,43 @@ function sanitizeSeedForShare(payload: unknown): unknown { return root; } +// Delete queued upload directories once a save stops referencing them, so a +// reload before autosave can never leave the saved recipe pointing at +// already-deleted files. Skips any uid the just-saved payload still uses. +function drainQueuedUploadCleanups( + savedPayload: RecipePayloadResult["payload"], +): void { + const pending = useRecipeStudioStore.getState().pendingUploadCleanups; + if (pending.length === 0) { + return; + } + const ui = + savedPayload && typeof savedPayload === "object" + ? (savedPayload as { ui?: Record }).ui + : undefined; + const savedUid = + ui && typeof ui.unstructured_upload_uid === "string" + ? ui.unstructured_upload_uid + : ""; + const ready = pending.filter((uid) => uid !== savedUid); + if (ready.length === 0) { + return; + } + for (const uid of ready) { + void removeUnstructuredBlock(uid) + .then(() => { + useRecipeStudioStore.setState((state) => ({ + pendingUploadCleanups: state.pendingUploadCleanups.filter( + (pendingUid) => pendingUid !== uid, + ), + })); + }) + .catch((error) => { + console.warn("Failed to clean up uploaded documents:", error); + }); + } +} + export function useRecipePersistence({ recipeId, initialRecipeName, @@ -202,8 +242,10 @@ export function useRecipePersistence({ () => buildSignature(normalizedWorkflowName, currentPayload), [currentPayload, normalizedWorkflowName], ); - const isDirty = savedSignature.length > 0 && currentSignature !== savedSignature; - const saveTone: SaveTone = !isDirty && Boolean(lastSavedAt) ? "success" : "error"; + const isDirty = + savedSignature.length > 0 && currentSignature !== savedSignature; + const saveTone: SaveTone = + !isDirty && Boolean(lastSavedAt) ? "success" : "error"; const savedAtLabel = formatSavedLabel(lastSavedAt); useEffect(() => { @@ -214,7 +256,9 @@ export function useRecipePersistence({ setLastSavedAt(initialSavedAt); setCopied(false); - const parsed = importRecipePayload(JSON.stringify(initialPayload)); + const parsed = importRecipePayload(JSON.stringify(initialPayload), { + preserveUnstructuredUploads: true, + }); if (parsed.snapshot) { loadRecipe(parsed.snapshot); } else { @@ -252,6 +296,7 @@ export function useRecipePersistence({ }); setLastSavedAt(result.updatedAt); setSavedSignature(buildSignature(nextName, currentPayload)); + drainQueuedUploadCleanups(currentPayload); } catch (error) { console.error("Save recipe failed:", error); toastError("Save failed", "Could not save recipe."); @@ -270,11 +315,28 @@ export function useRecipePersistence({ return () => window.clearTimeout(timeoutId); }, [isDirty, persistRecipe, saveLoading]); + // Drain queued cleanups even when autosave is skipped: a net-zero edit (add + // then remove an unstructured seed before the 800ms debounce) keeps isDirty + // false, so the autosave effect never drains and the queued uid leaks its + // upload dir. Not-dirty means currentPayload equals the saved recipe, and + // drain skips the uid it still references, so only dirs no saved recipe + // points at are deleted (keeps the save-first invariant). + useEffect(() => { + if (!initialRecipeReady || isDirty || saveLoading) { + return; + } + drainQueuedUploadCleanups(currentPayload); + }, [currentPayload, initialRecipeReady, isDirty, saveLoading]); + const copyRecipe = useCallback(async (): Promise => { setCopied(false); try { - const safePayload = sanitizeSeedForShare(stripApiKeys(payloadResult.payload)); - const ok = await copyTextToClipboard(JSON.stringify(safePayload, null, 2)); + const safePayload = sanitizeSeedForShare( + stripApiKeys(payloadResult.payload), + ); + const ok = await copyTextToClipboard( + JSON.stringify(safePayload, null, 2), + ); if (!ok) { throw new Error("Clipboard not available."); } diff --git a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts index 8659cad4fb..a1ff72ee9b 100644 --- a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts +++ b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts @@ -36,6 +36,7 @@ import { } from "../utils/handles"; import type { RecipeSnapshot } from "../utils/import"; import { getLayoutedElements } from "../utils/layout"; +import { makeUnstructuredUploadUid } from "../utils/config-factories"; import { centerModelInfraNodes, optimizeModelInfraEdgeHandles, @@ -76,6 +77,12 @@ type RecipeStudioState = { nextId: number; nextY: number; fitViewTick: number; + // Upload-uid directories whose owning block dropped them; server-side + // deletion is deferred until a save no longer references them, so a + // reload before autosave cannot leave a saved recipe pointing at + // deleted files. + pendingUploadCleanups: string[]; + queueUploadCleanup: (uid: string) => void; setSheetOpen: (open: boolean) => void; setSheetView: (view: SheetView) => void; setProcessors: (processors: RecipeProcessorConfig[]) => void; @@ -137,6 +144,7 @@ const INITIAL_STATE = { nextId: 3, nextY: 280, fitViewTick: 0, + pendingUploadCleanups: [], } satisfies Pick< RecipeStudioState, | "nodes" @@ -154,6 +162,7 @@ const INITIAL_STATE = { | "nextId" | "nextY" | "fitViewTick" + | "pendingUploadCleanups" >; function buildAddedNodeState( @@ -269,6 +278,20 @@ function isModelSemanticEdge( ); } +// Upload uid of a seed block whose server-side directory becomes orphaned +// when the block drops it. Only uid directories qualify (single owner); +// legacy node-id directories can be shared by other recipes. +function seedUploadCleanupUid(config: NodeConfig | undefined): string | null { + if (!config || config.kind !== "seed") { + return null; + } + const uid = config.unstructured_upload_uid?.trim(); + if (!uid || !config.unstructured_file_ids?.length) { + return null; + } + return uid; +} + export const useRecipeStudioStore = create((set, get) => ({ ...INITIAL_STATE, setSheetOpen: (open) => set({ sheetOpen: open }), @@ -278,6 +301,12 @@ export const useRecipeStudioStore = create((set, get) => ({ setDialogOpen: (open) => set({ dialogOpen: open }), setExecutionLocked: (locked) => set({ executionLocked: locked }), resetRecipe: () => set(INITIAL_STATE), + queueUploadCleanup: (uid) => + set((state) => + state.pendingUploadCleanups.includes(uid) + ? state + : { pendingUploadCleanups: [...state.pendingUploadCleanups, uid] }, + ), selectConfig: (id) => set({ activeConfigId: id, dialogOpen: false }), openConfig: (id) => set({ activeConfigId: id, dialogOpen: true }), setLayoutDirection: (direction) => @@ -383,7 +412,18 @@ export const useRecipeStudioStore = create((set, get) => ({ } return buildAddedNodeState(state, "sampler", type, position, openDialog); }), - addSeedNode: (type, position, openDialog = true) => + addSeedNode: (type, position, openDialog = true) => { + const current = get(); + if (!current.executionLocked) { + // The reset below clears the block's upload uid and file list; queue + // its server-side directory for deletion after the next save. + const uid = seedUploadCleanupUid( + Object.values(current.configs).find((config) => config.kind === "seed"), + ); + if (uid) { + current.queueUploadCleanup(uid); + } + } set((state) => { if (state.executionLocked) { return state; @@ -413,6 +453,8 @@ export const useRecipeStudioStore = create((set, get) => ({ hf_token: "", hf_endpoint: "https://huggingface.co", local_file_name: "", + unstructured_upload_uid: + nextSourceType === "unstructured" ? makeUnstructuredUploadUid() : "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -446,7 +488,8 @@ export const useRecipeStudioStore = create((set, get) => ({ activeConfigId: existing.id, dialogOpen: openDialog, }; - }), + }); + }, addLlmNode: (type, position, openDialog = true) => set((state) => { if (state.executionLocked) { @@ -699,6 +742,9 @@ export const useRecipeStudioStore = create((set, get) => ({ dialogOpen: false, sheetView: "root", fitViewTick: state.fitViewTick + 1, + // Queued cleanups belong to the previous recipe; draining them after + // a save of this one could delete files its saved payload still uses. + pendingUploadCleanups: [], })), setAuxNodePosition: (id, position) => set((state) => { @@ -786,6 +832,17 @@ export const useRecipeStudioStore = create((set, get) => ({ set(applyUpdate); }, onNodesChange: (changes) => { + const current = get(); + if (!current.executionLocked) { + for (const change of changes) { + if (change.type === "remove") { + const uid = seedUploadCleanupUid(current.configs[change.id]); + if (uid) { + current.queueUploadCleanup(uid); + } + } + } + } const applyNodesChange = (state: RecipeStudioState) => { if (state.executionLocked) { return state; diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index b8ed13f70b..9231c5f7a3 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -340,6 +340,8 @@ export type SeedConfig = { hf_token?: string; hf_endpoint?: string; local_file_name?: string; + // ui-only: stable per-block id for uploads, since node ids collide across imports + unstructured_upload_uid?: string; unstructured_file_ids?: string[]; unstructured_file_names?: string[]; unstructured_file_sizes?: number[]; diff --git a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts index 76fc2c38ee..d47bac8858 100644 --- a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts +++ b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts @@ -20,6 +20,46 @@ import type { } from "../types"; import { nextName } from "./naming"; +export function makeUnstructuredUploadUid(): string { + if (typeof globalThis.crypto?.randomUUID === "function") { + return globalThis.crypto.randomUUID().replace(/-/g, "").toLowerCase(); + } + if (typeof globalThis.crypto?.getRandomValues === "function") { + const bytes = new Uint8Array(16); + globalThis.crypto.getRandomValues(bytes); + return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join( + "", + ); + } + let uid = ""; + while (uid.length < 32) { + uid += Math.floor(Math.random() * 0x100000000) + .toString(16) + .padStart(8, "0"); + } + return uid.slice(0, 32); +} + +export function resolveUnstructuredUploadBlockId({ + configId, + uploadUid, + generatedUploadUid, + unstructuredFileCount, +}: { + configId: string; + uploadUid: string; + generatedUploadUid: string | null; + unstructuredFileCount: number; +}): string { + if (uploadUid) { + return uploadUid; + } + if (generatedUploadUid) { + return generatedUploadUid; + } + return unstructuredFileCount > 0 ? configId : ""; +} + export function makeSamplerConfig( id: string, samplerType: SamplerType, @@ -368,6 +408,9 @@ export function makeSeedConfig( hf_token: "", hf_endpoint: "https://huggingface.co", local_file_name: "", + ...(seedSourceType === "unstructured" + ? { unstructured_upload_uid: makeUnstructuredUploadUid() } + : {}), unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts index ac881d0373..54df2ffd50 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts @@ -16,11 +16,7 @@ import type { } from "../../types"; import { buildEdges } from "./edges"; import { isRecord, parseJson, readString } from "./helpers"; -import { - parseColumn, - parseModelConfig, - parseModelProvider, -} from "./parsers"; +import { parseColumn, parseModelConfig, parseModelProvider } from "./parsers"; import { parseSeedConfig } from "./parsers/seed-config-parser"; import { buildNodes, parseUi } from "./ui"; import type { ImportResult } from "./types"; @@ -43,6 +39,7 @@ type UiInput = { seed_drop_columns?: unknown; seed_preview_rows?: unknown; local_file_name?: unknown; + unstructured_upload_uid?: unknown; unstructured_file_ids?: unknown; unstructured_file_names?: unknown; unstructured_file_sizes?: unknown; @@ -51,6 +48,10 @@ type UiInput = { advanced_open_by_node?: unknown; }; +type ImportRecipePayloadOptions = { + preserveUnstructuredUploads?: boolean; +}; + type UiMarkdownNoteNode = { name: string; markdown: string; @@ -90,7 +91,7 @@ function parseProcessors(input: unknown): RecipeProcessorConfig[] { ? templateRaw : isRecord(templateRaw) ? JSON.stringify(templateRaw, null, 2) - : "{\n \"text\": \"{{ column_name }}\"\n}"; + : '{\n "text": "{{ column_name }}"\n}'; processors.push({ id: `p${index + 1}`, // biome-ignore lint/style/useNamingConvention: api schema @@ -135,9 +136,7 @@ function parseSeedDropColumns(input: unknown): string[] { return Array.from(values); } -function parseMcpProviders( - input: unknown, -): Map { +function parseMcpProviders(input: unknown): Map { const providers = new Map(); if (!Array.isArray(input)) { return providers; @@ -156,13 +155,12 @@ function parseMcpProviders( const args = Array.isArray(item.args) ? item.args.map((value) => String(value)) : []; - const envPairs = - isRecord(item.env) - ? Object.entries(item.env).map(([key, value]) => ({ - key: String(key), - value: String(value), - })) - : []; + const envPairs = isRecord(item.env) + ? Object.entries(item.env).map(([key, value]) => ({ + key: String(key), + value: String(value), + })) + : []; providers.set(name, { id: `mcp-${index + 1}`, name, @@ -209,7 +207,8 @@ function parseToolConfigs(input: unknown): Map { allow_tools: allowTools, // biome-ignore lint/style/useNamingConvention: api schema max_tool_call_turns: - item.max_tool_call_turns === null || item.max_tool_call_turns === undefined + item.max_tool_call_turns === null || + item.max_tool_call_turns === undefined ? "5" : String(item.max_tool_call_turns), // biome-ignore lint/style/useNamingConvention: api schema @@ -257,7 +256,9 @@ function parseUiMarkdownNoteNodes(input: unknown): UiMarkdownNoteNode[] { return noteNodes; } -function parseUiToolProfileNodes(input: unknown): Map> { +function parseUiToolProfileNodes( + input: unknown, +): Map> { const toolProfiles = new Map>(); if (!Array.isArray(input)) { return toolProfiles; @@ -312,9 +313,15 @@ function parseAdvancedOpenByNode(input: unknown): Record { return out; } -type AdvancedOpenConfig = LlmConfig | SamplerConfig | SeedConfig | ValidatorConfig; +type AdvancedOpenConfig = + | LlmConfig + | SamplerConfig + | SeedConfig + | ValidatorConfig; -function isAdvancedOpenConfig(config: NodeConfig): config is AdvancedOpenConfig { +function isAdvancedOpenConfig( + config: NodeConfig, +): config is AdvancedOpenConfig { return ( config.kind === "llm" || config.kind === "sampler" || @@ -350,7 +357,8 @@ function buildToolProfileConfig( .map((providerName) => mcpProvidersByName.get(providerName)) .flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])), // biome-ignore lint/style/useNamingConvention: ui schema - fetched_tools_by_provider: fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, + fetched_tools_by_provider: + fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, // biome-ignore lint/style/useNamingConvention: api schema allow_tools: [...(canonical.allow_tools ?? [])], // biome-ignore lint/style/useNamingConvention: api schema @@ -360,7 +368,10 @@ function buildToolProfileConfig( }; } -export function importRecipePayload(input: string): ImportResult { +export function importRecipePayload( + input: string, + options: ImportRecipePayloadOptions = {}, +): ImportResult { const parsed = parseJson(input); if (!parsed.data || !isRecord(parsed.data)) { return { @@ -369,9 +380,9 @@ export function importRecipePayload(input: string): ImportResult { }; } - const recipe = (isRecord(parsed.data.recipe) - ? parsed.data.recipe - : parsed.data) as RecipeInput; + const recipe = ( + isRecord(parsed.data.recipe) ? parsed.data.recipe : parsed.data + ) as RecipeInput; const ui = isRecord(parsed.data.ui) ? (parsed.data.ui as UiInput) : null; if (!Array.isArray(recipe.columns)) { @@ -410,21 +421,36 @@ export function importRecipePayload(input: string): ImportResult { .map((row) => ({ ...row })) : undefined; const uiLocalFileName = readString(ui?.local_file_name) ?? undefined; - // Preserve file IDs/names from saved recipes (cleared at share time by sanitizeSeedForShare) - const uiUnstructuredFileIds: string[] = Array.isArray(ui?.unstructured_file_ids) - ? (ui.unstructured_file_ids as string[]).filter((v): v is string => typeof v === "string") - : []; - const uiUnstructuredFileNames: string[] = Array.isArray(ui?.unstructured_file_names) - ? (ui.unstructured_file_names as string[]).filter((v): v is string => typeof v === "string") - : []; - const uiUnstructuredFileSizes: number[] = Array.isArray(ui?.unstructured_file_sizes) - ? (ui.unstructured_file_sizes as number[]).filter((v): v is number => typeof v === "number") - : []; + const preserveUnstructuredUploads = + options.preserveUnstructuredUploads === true; + const uiUnstructuredUploadUid = preserveUnstructuredUploads + ? (readString(ui?.unstructured_upload_uid) ?? undefined) + : undefined; + const uiUnstructuredFileIds: string[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_ids) + ? (ui.unstructured_file_ids as string[]).filter( + (v): v is string => typeof v === "string", + ) + : []; + const uiUnstructuredFileNames: string[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_names) + ? (ui.unstructured_file_names as string[]).filter( + (v): v is string => typeof v === "string", + ) + : []; + const uiUnstructuredFileSizes: number[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_sizes) + ? (ui.unstructured_file_sizes as number[]).filter( + (v): v is number => typeof v === "number", + ) + : []; const uiUnstructuredChunkSize = readStringNumber(ui?.unstructured_chunk_size); const uiUnstructuredChunkOverlap = readStringNumber( ui?.unstructured_chunk_overlap, ); - const uiAdvancedOpenByNode = parseAdvancedOpenByNode(ui?.advanced_open_by_node); + const uiAdvancedOpenByNode = parseAdvancedOpenByNode( + ui?.advanced_open_by_node, + ); const uiMarkdownNotes = parseUiMarkdownNoteNodes(ui?.nodes); const uiToolProfilesByName = parseUiToolProfileNodes(ui?.nodes); @@ -459,11 +485,13 @@ export function importRecipePayload(input: string): ImportResult { : payloadSeedDropColumns, seed_preview_rows: uiSeedPreviewRows, local_file_name: uiLocalFileName, + unstructuredUploadUid: uiUnstructuredUploadUid, unstructuredFileIds: uiUnstructuredFileIds, unstructuredFileNames: uiUnstructuredFileNames, unstructuredFileSizes: uiUnstructuredFileSizes, unstructured_chunk_size: uiUnstructuredChunkSize, unstructured_chunk_overlap: uiUnstructuredChunkOverlap, + preserveUnstructuredUploads, }); if (seedConfig) { applyAdvancedOpen(seedConfig, uiAdvancedOpenByNode); @@ -567,12 +595,7 @@ export function importRecipePayload(input: string): ImportResult { const { layouts, auxNodes, edges: uiEdges, layoutDirection } = parseUi(ui); const resolvedLayoutDirection = layoutDirection ?? "LR"; const nodes = buildNodes(configs, layouts); - const edges = buildEdges( - configs, - nameToId, - uiEdges, - resolvedLayoutDirection, - ); + const edges = buildEdges(configs, nameToId, uiEdges, resolvedLayoutDirection); const auxNodePositions = Object.fromEntries( auxNodes.flatMap((item) => { const llmId = nameToId.get(item.llm); @@ -583,10 +606,7 @@ export function importRecipePayload(input: string): ImportResult { }), ); - const maxY = nodes.reduce( - (acc, node) => Math.max(acc, node.position.y), - 0, - ); + const maxY = nodes.reduce((acc, node) => Math.max(acc, node.position.y), 0); return { errors: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts index 21eadb3195..939205fe6d 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts @@ -197,17 +197,26 @@ export function parseSeedConfig( seed_drop_columns?: string[]; seed_preview_rows?: Record[]; local_file_name?: string; + unstructuredUploadUid?: string; unstructuredFileIds?: string[]; unstructuredFileNames?: string[]; unstructuredFileSizes?: number[]; unstructured_chunk_size?: string; unstructured_chunk_overlap?: string; + preserveUnstructuredUploads?: boolean; }, ): SeedConfig | null { if (!seedConfigRaw) { return null; } - const parsed = parseSeedSettings(seedConfigRaw); + const parsed = { ...parseSeedSettings(seedConfigRaw) }; + if ( + parsed.seed_source_type === "unstructured" && + options?.preserveUnstructuredUploads !== true + ) { + parsed.hf_path = ""; + parsed.resolved_paths = []; + } let sourceType: SeedSourceType = "hf"; if (parsed.seed_source_type === "hf") { sourceType = "hf"; @@ -230,6 +239,9 @@ export function parseSeedConfig( ...(options?.local_file_name !== undefined ? { local_file_name: options.local_file_name } : {}), + ...(options?.unstructuredUploadUid + ? { unstructured_upload_uid: options.unstructuredUploadUid } + : {}), ...(options?.unstructuredFileIds !== undefined ? { unstructured_file_ids: options.unstructuredFileIds } : {}), diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts index 34c3c34274..9b2ad5b2b5 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts @@ -440,6 +440,9 @@ export function buildRecipePayload( unstructured_file_names: firstSeed.unstructured_file_names, unstructured_file_sizes: firstSeed.unstructured_file_sizes, }), + ...(firstSeed?.unstructured_upload_uid?.trim() && { + unstructured_upload_uid: firstSeed.unstructured_upload_uid, + }), ...(firstSeed && firstSeed.unstructured_chunk_size !== undefined && { unstructured_chunk_size: firstSeed.unstructured_chunk_size, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts index 902ea796b4..763e68c1fb 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts @@ -71,6 +71,8 @@ export type RecipePayload = { seed_preview_rows?: Record[]; local_file_name?: string; // biome-ignore lint/style/useNamingConvention: api schema + unstructured_upload_uid?: string; + // biome-ignore lint/style/useNamingConvention: api schema unstructured_file_ids?: string[]; // biome-ignore lint/style/useNamingConvention: api schema unstructured_file_names?: string[]; From 1b825213ea2ffe4774404f27795ad287634f6681 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Thu, 9 Jul 2026 00:16:05 -0700 Subject: [PATCH 084/113] Stabilize floating monitor drag (#6984) * Stabilize floating monitor drag * Restore floating monitor exit animation * Harden Windows Studio smoke checks * Keep API menu badge removed * Apply no-build-tools env overrides in-script The runner does not apply step-level env keys containing parentheses, so ProgramFiles(x86) kept its real value and Find-VsBuildTools still detected VS through vswhere. Set the overrides inside each pwsh step instead; child processes inherit them. The resolver step moves to pwsh because bash cannot export a variable named ProgramFiles(x86). * Reset chat UI session without a second browser context macOS runs Chromium with --single-process, where closing the last context tears down the whole browser, so the shutdown re-login died with TargetClosedError on new_page. Clear cookies and swap pages inside the same context instead, opening the replacement page before closing the old one. * Keep the no-build-tools Path filtered across session refreshes install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment rebuild the session Path from the Machine and User registry scopes, so the process-level filter could be undone mid-install and re-expose CMake. Filter those scopes in the Prepare step with normalized dir matching and restore them in cleanup. * Drop stale localStorage auth tokens before re-login Auth tokens live in localStorage, not cookies, and the login guest guard redirects on their mere presence. Remove them during the session reset so the /login navigation is deterministic instead of relying on the tolerated redirect bounce. --- .../studio-windows-inference-smoke.yml | 161 +++++++++---- .../frontend/src/components/app-sidebar.tsx | 17 +- .../src/components/floating-monitor.tsx | 223 ++++++++++-------- studio/frontend/src/features/chat/index.ts | 1 + .../frontend/src/features/settings/index.ts | 1 + studio/frontend/src/i18n/locales/en.ts | 1 - studio/frontend/src/i18n/locales/ja.ts | 1 - studio/frontend/src/i18n/locales/pt-br.ts | 1 - studio/frontend/src/i18n/locales/zh-CN.ts | 1 - tests/studio/playwright_chat_ui.py | 36 ++- 10 files changed, 277 insertions(+), 166 deletions(-) diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 0bc216d65a..dbb0f9ea6f 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -1334,42 +1334,75 @@ jobs: try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { } } - - name: Hide Visual Studio + CMake (simulate a host with no build tools) + - name: Prepare no-build-tools simulation shell: pwsh run: | $ErrorActionPreference = 'Stop' - # A Program Files dir can hold a transient handle (Defender / MSBuild node) - # so Rename-Item intermittently fails with "Access is denied"; retry to ride it out. - function Rename-WithRetry($Path, $NewName) { - for ($i = 1; $i -le 6; $i++) { - try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return } - catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { + [void] $blocked.Add( + [Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\')) + } + } } } - # Rename the Visual Studio install roots (incl. the Installer that holds - # vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss. - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - if (Test-Path -LiteralPath $d) { - Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff') - Write-Host "Hid VS: $d" - } + # Normalized comparison so registry spellings (trailing slash, + # unexpanded %VAR%) still match. + function Test-Blocked([string]$p) { + $n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\') + return $blocked.Contains($n) } - # Surgically rename each cmake executable on PATH (not its parent dir -- - # cmake can share a dir with other shims) so Get-Command cmake fails. - $hidden = @() - foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) { - if ($c.Source -and (Test-Path -LiteralPath $c.Source)) { - Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off') - $hidden += $c.Source - Write-Host "Hid cmake: $($c.Source)" - } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not (Test-Blocked $_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + # install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment + # rebuild the session Path from these scopes mid-install, so filter + # them too. Originals are saved for the cleanup step. + foreach ($scope in @('Machine', 'User')) { + $orig = [Environment]::GetEnvironmentVariable('Path', $scope) + if (-not $orig) { continue } + Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline + $kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';' + [Environment]::SetEnvironmentVariable('Path', $kept, $scope) + Write-Host "Filtered $scope Path scope." + } + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH<&1 | Tee-Object -FilePath logs/install.log @@ -1480,19 +1517,19 @@ jobs: [ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; } echo "Inference OK without Visual Studio: $CONTENT" - - name: Restore Visual Studio + CMake + - name: Clean no-build-tools simulation if: always() shell: pwsh run: | - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - $off = "$d.vsoff" - if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" } - } - if ($env:HIDDEN_CMAKE) { - foreach ($src in ($env:HIDDEN_CMAKE -split '\|')) { - if ($src -and (Test-Path -LiteralPath "$src.off")) { Rename-Item -LiteralPath "$src.off" -NewName (Split-Path $src -Leaf) } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + foreach ($scope in @('Machine', 'User')) { + $saved = Join-Path $root "orig-path-$scope.txt" + if (Test-Path -LiteralPath $saved) { + [Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope) + Write-Host "Restored $scope Path scope." } } + Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue - name: Stop Studio if: always() @@ -1540,21 +1577,34 @@ jobs: with: python-version: '3.12' - - name: Hide Visual Studio + - name: Prepare no-build-tools simulation shell: pwsh run: | $ErrorActionPreference = 'Stop' - # Retry the rename: a Program Files dir can hold a transient handle that - # makes Rename-Item intermittently fail with "Access is denied". - function Rename-WithRetry($Path, $NewName) { - for ($i = 1; $i -le 6; $i++) { - try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return } - catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { [void] $blocked.Add($dir) } + } } } - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" } - } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not $blocked.Contains($_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH< /tmp/resolve.json || { - echo "::error::resolver exited non-zero"; cat /tmp/resolve.json || true; exit 1; } - cat /tmp/resolve.json - echo "Prebuilt resolver ran with no Visual Studio present." + if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 } + python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json + if ($LASTEXITCODE -ne 0) { + Write-Host "::error::resolver exited non-zero" + if (Test-Path resolve.json) { Get-Content resolve.json } + exit 1 + } + Get-Content resolve.json + Write-Host "Prebuilt resolver ran with no Visual Studio present." - - name: Restore Visual Studio + - name: Clean no-build-tools simulation if: always() shell: pwsh run: | - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - $off = "$d.vsoff" - if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" } - } + Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue # ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ── pester: diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 823e420869..1a74b38524 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -81,7 +81,6 @@ import { TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; -import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, TooltipContent, @@ -97,6 +96,7 @@ import { createChatProject, deleteChatProject, deleteChatItem, + listStoredChatThreads, moveChatItemToProject, renameChatItem, renameChatProject, @@ -582,7 +582,14 @@ export function AppSidebar() { useEffect(() => { if (!pendingRename) return; const match = allChatItems.find((i) => i.id === pendingRename.id); - if (match && match.title === pendingRename.title) setPendingRename(null); + if (!match || match.title !== pendingRename.title) return; + queueMicrotask(() => { + setPendingRename((current) => + current?.id === pendingRename.id && current.title === pendingRename.title + ? null + : current, + ); + }); }, [allChatItems, pendingRename]); const [creatingProject, setCreatingProject] = useState(false); const [projectNameDraft, setProjectNameDraft] = useState(""); @@ -680,12 +687,6 @@ export function AppSidebar() { useState(null); const [deleteProjectFiles, setDeleteProjectFiles] = useState(false); - useEffect(() => { - if (confirmingDelete?.kind !== "project") { - setDeleteProjectFiles(false); - } - }, [confirmingDelete]); - async function commitDelete() { const target = confirmingDelete; if (!target) return; diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx index bce4bf2831..0a51875de9 100644 --- a/studio/frontend/src/components/floating-monitor.tsx +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -3,27 +3,35 @@ import { Button } from "@/components/ui/button"; import { Progress } from "@/components/ui/progress"; -import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store"; +import { useMonitorOverlayStore } from "@/features/settings"; import { useSystemInfo } from "@/hooks/use-system"; import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react"; -import { motion } from "motion/react"; -import { useRef } from "react"; +import { AnimatePresence, motion, useDragControls } from "motion/react"; +import { type PointerEvent, useMemo, useState } from "react"; function clampPercent(value: number): number { return Math.max(0, Math.min(100, value)); } function usageIndicatorClass(percent: number): string { - if (percent >= 90) return "bg-destructive"; - if (percent >= 70) return "bg-amber-500"; + if (percent >= 90) { + return "bg-destructive"; + } + if (percent >= 70) { + return "bg-amber-500"; + } return "bg-primary"; } function usageTextClass(percent: number): string { - if (percent >= 90) return "text-destructive"; - if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + if (percent >= 90) { + return "text-destructive"; + } + if (percent >= 70) { + return "text-amber-600 dark:text-amber-400"; + } return "text-primary"; } @@ -39,9 +47,18 @@ export function FloatingMonitor() { const { isOpen, setIsOpen } = useMonitorOverlayStore(); const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 }); - const constraintsRef = useRef(null); + const [constraintsElement, setConstraintsElement] = + useState(null); + const constraintsRef = useMemo( + () => ({ current: constraintsElement }), + [constraintsElement], + ); + const dragControls = useDragControls(); - if (!isOpen) return null; + function startDrag(event: PointerEvent) { + event.preventDefault(); + dragControls.start(event); + } const ramTotal = systemInfo.memory?.total_gb ?? 0; const ramAvailable = systemInfo.memory?.available_gb ?? 0; @@ -64,99 +81,109 @@ export function FloatingMonitor() { const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0; return ( -
- -
-
- - - {t("settings.resources.liveMonitor.title")} - -
-
-
- -
- - -
-
- - + {isOpen && ( +
-
-
- {t("settings.resources.liveMonitor.ram")} - - {Math.round(ramPercent)}% - -
-
- {formatGiB(ramUsed)} / {formatGiB(ramTotal)} -
- -
- - {hasGpu && ( -
-
- - {t("settings.resources.liveMonitor.vram")}{" "} - {devices.length > 1 - ? `(${devices.length} GPUs)` - : `(${devices[0].name ?? "GPU"})`} + +
+
+ + + {t("settings.resources.liveMonitor.title")} - +
+
- {Math.round(vramPercent)}% - + +
+ +
-
- {formatGiB(vramUsed)} / {formatGiB(vramTotal)} -
-
- )} - - -
+ + +
+
+ {t("settings.resources.liveMonitor.ram")} + + {Math.round(ramPercent)}% + +
+
+ {formatGiB(ramUsed)} / {formatGiB(ramTotal)} +
+ +
+ + {hasGpu && ( +
+
+ + {t("settings.resources.liveMonitor.vram")}{" "} + {devices.length > 1 + ? `(${devices.length} GPUs)` + : `(${devices[0].name ?? "GPU"})`} + + + {Math.round(vramPercent)}% + +
+
+ {formatGiB(vramUsed)} / {formatGiB(vramTotal)} +
+ +
+ )} + + +
+ )} + ); } diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 7cd9611c71..d070ed15de 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -36,6 +36,7 @@ export { ChatSearchDialog } from "./components/chat-search-dialog"; export { setTrainingCompareHandoff } from "./lib/training-compare-handoff"; export type { ProjectRecord } from "./types"; export { clearAllChats, countAllChats } from "./utils/clear-all-chats"; +export { listStoredChatThreads } from "./utils/chat-history-storage"; export { ArtifactCard } from "./artifacts/artifact-card"; export { useChatArtifactsStore, diff --git a/studio/frontend/src/features/settings/index.ts b/studio/frontend/src/features/settings/index.ts index 364ca1611f..3fefd8c63a 100644 --- a/studio/frontend/src/features/settings/index.ts +++ b/studio/frontend/src/features/settings/index.ts @@ -7,6 +7,7 @@ export { savePersonalization, } from "./api/personalization"; export { setTheme, useTheme } from "./stores/theme-store"; +export { useMonitorOverlayStore } from "./stores/monitor-overlay-store"; export type { Personalization, PersonalizationAppearance, diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index a9b5d839b3..e0cb8030ae 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -409,7 +409,6 @@ export const en = { description: "Access Unsloth via the OpenAI-compatible API.", readDocs: "Read the API docs", noAccess: "No API access yet.", - newBadge: "New", accessTokens: "Access tokens", loadError: "Couldn't load API access.", createError: "Couldn't create access token.", diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts index f2020f3cfb..08b8d4f4d3 100644 --- a/studio/frontend/src/i18n/locales/ja.ts +++ b/studio/frontend/src/i18n/locales/ja.ts @@ -296,7 +296,6 @@ export const ja = { description: "OpenAI互換 API を介して Unsloth にアクセスします。", readDocs: "API ドキュメントを読む", noAccess: "まだ API アクセス権がありません。", - newBadge: "新規", accessTokens: "アクセストークン", loadError: "API アクセス権を読み込めませんでした。", createError: "アクセストークンを作成できませんでした。", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts index c261e4ed0c..494a98ec32 100644 --- a/studio/frontend/src/i18n/locales/pt-br.ts +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -363,7 +363,6 @@ export const ptBR = { "Acesse o Unsloth por meio da API compatível com OpenAI.", readDocs: "Leia a documentação da API", noAccess: "Nenhum acesso à API ainda.", - newBadge: "Novo", accessTokens: "Tokens de acesso", loadError: "Não foi possível carregar o acesso à API.", createError: "Não foi possível criar o token de acesso.", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index f6dc265fc7..dda8e017ab 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -267,7 +267,6 @@ export const zhCN = { description: "通过兼容 OpenAI 的 API 以编程方式访问 Unsloth。", readDocs: "阅读 API 文档", noAccess: "还没有 API 访问权限。", - newBadge: "新", accessTokens: "访问 token", loadError: "无法加载 API 访问权限。", createError: "无法创建访问 token。", diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index 71297f9043..0698d8e0d2 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -1264,11 +1264,37 @@ with sync_playwright() as p: # placeholder, and /api/health goes unreachable shortly after. # ───────────────────────────────────────────────────── step("Shutdown via account menu") - # Re-login with NEW2 for a valid /api/shutdown token (CLI rotation - # invalidated the old one). The stale token can make the SPA auth guard - # abort this goto with ERR_ABORTED, or redirect to the same /login URL - # ("interrupted by another navigation"); resolve on domcontentloaded and - # tolerate either -- the pw-field wait below confirms we are on /login. + # Start fresh after the CLI rotation invalidates this browser session. + # Stay in the SAME context: macOS Chromium runs --single-process, where + # closing the last context kills the browser and a second context cannot + # be created. Open the new page before closing the old one; the context + # init script covers the new page. + try: + ctx.clear_cookies() + except Exception as exc: + info(f"WARN clearing stale session cookies failed: {exc!r}") + # Auth tokens live in localStorage, and /login's guest guard redirects on + # their mere presence, so drop them before navigating. + try: + page.evaluate( + "['unsloth_auth_token', 'unsloth_auth_refresh_token']" + ".forEach((key) => localStorage.removeItem(key))" + ) + except Exception as exc: + info(f"WARN clearing stale auth tokens failed: {exc!r}") + _fresh_page = ctx.new_page() + _fresh_page.set_default_timeout(60_000) + _fresh_page.on("pageerror", lambda e: page_errors.append(str(e))) + _fresh_page.on("console", _on_console) + try: + page.close() + except Exception: + pass + page = _fresh_page + + # Re-login with NEW2 for a valid /api/shutdown token. Route changes can + # still abort or interrupt this navigation, so the field wait below is the + # final confirmation that we reached /login. _tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation") try: page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000) From 8205d4c0819088a3c864fdd505ce1c1e6d72852d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 01:46:14 -0700 Subject: [PATCH 085/113] Retry the Studio UI shutdown re-login on transient goto timeout (#7027) * Retry the Studio UI shutdown re-login on transient goto timeout The Chat UI Playwright smoke intermittently failed at the pre-shutdown re-login: page.goto('/login') can hit a 60s TimeoutError on a slow runner even while the server is healthy, and the surrounding except only tolerated ERR_ABORTED / interrupted-navigation, so a plain timeout hard-failed the job. Wrap the re-login goto/wait/fill/submit in the same 3-attempt retry the change-password step already uses (recover_or_replace_page between tries, per-attempt fail screenshots, wait_for_health pre-gate). The composer wait stays outside the loop so a retry never re-navigates after login has set tokens (which would redirect to /chat via the guest guard); it remains the authoritative confirmation, so a genuinely broken login still fails. * Catch transient login-request failures and preserve error listeners on recovery Wait on the /api/auth/login POST inside the retry (via click_and_wait_for_response) so a transient 4xx/5xx is retried in-loop instead of surfacing only at the out-of-loop composer wait, matching the change-password step. When recover_or_replace_page swaps in a fresh page, re-attach the pageerror/console listeners so error tracking survives the replacement. --- tests/studio/playwright_chat_ui.py | 96 ++++++++++++++++++++++++++---- 1 file changed, 86 insertions(+), 10 deletions(-) diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index 0698d8e0d2..6a88b98c19 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -1296,16 +1296,92 @@ with sync_playwright() as p: # still abort or interrupt this navigation, so the field wait below is the # final confirmation that we reached /login. _tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation") - try: - page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000) - except Exception as exc: - if not any(t in str(exc) for t in _tolerated_nav): - raise - info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login") - pw_field = page.locator("#password") - pw_field.wait_for(state = "visible", timeout = 60_000) - pw_field.fill(NEW2) - page.locator('button[type="submit"]').click() + # A slow CI runner can make this re-login navigation time out even with the + # server healthy, so retry the whole goto/wait/fill/submit sequence (mirrors + # the change-password retry above). wait_for_health is a diagnostic pre-gate. + wait_for_health(BASE, timeout = 30.0, info = info) + relogin_err: Exception | None = None + for _relogin_attempt in range(3): + try: + try: + page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000) + except Exception as exc: + if not any(t in str(exc) for t in _tolerated_nav): + raise + info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login") + pw_field = page.locator("#password") + pw_field.wait_for(state = "visible", timeout = 60_000) + pw_field.fill(NEW2) + # Wait on the login POST so a transient 4xx/5xx is caught and retried + # here, not swallowed until the out-of-loop composer wait. + status, _ = click_and_wait_for_response( + page, + url_substr = "/api/auth/login", + method = "POST", + do_click = lambda: page.locator('button[type="submit"]').click(), + timeout_ms = 30_000, + info = lambda m: print(f"[ui] {m}", flush = True), + ) + if status is not None and status >= 400: + raise AssertionError( + f"login POST returned {status}; see console_errors={console_errors[:1]!r}" + ) + relogin_err = None + break + except Exception as e: + relogin_err = e + try: + cur_url = page.url + except Exception: + cur_url = "" + print( + f"[ui] re-login attempt {_relogin_attempt + 1} failed: " + f"{type(e).__name__}: {str(e)[:200]}; page.url={cur_url}; " + f"page_errors={len(page_errors)} console_errors={len(console_errors)}", + flush = True, + ) + if console_errors: + print( + f"[ui] first console.error: {console_errors[0][:200]!r}", + flush = True, + ) + if page_errors: + print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True) + try: + shoot(f"18-relogin-attempt-{_relogin_attempt + 1}-fail") + except Exception: + pass + if _relogin_attempt < 2: + # ERR_NO_BUFFER_SPACE needs the OS to recover socket + # buffers; back off 5s then 15s before retrying. + if "ERR_NO_BUFFER_SPACE" in str(e): + backoff_s = 5 if _relogin_attempt == 0 else 15 + print( + f"[ui] ENOBUFS detected; sleeping {backoff_s}s " + f"before retry to let OS recover socket buffers...", + flush = True, + ) + time.sleep(backoff_s) + # Replace the page if it died; otherwise next iteration's + # page.goto() handles the reload. + old_page = page + page = recover_or_replace_page( + page, + ctx, + default_timeout_ms = 60_000, + info = lambda m: print(f"[ui] recovery: {m}", flush = True), + ) + # A freshly created replacement page loses the pageerror/console + # listeners; re-attach so error tracking survives recovery. + if page is not old_page: + page.on("pageerror", lambda e: page_errors.append(str(e))) + page.on("console", _on_console) + if relogin_err is not None: + raise relogin_err + # Composer mount confirms the rotated session is authenticated. Kept OUTSIDE the + # retry: the loop breaks right after submit, so we never re-goto /login once login + # has set tokens -- that would hit the guest guard, redirect to /chat, and make a + # merely-slow composer look like a broken login. composer = page.locator('textarea[aria-label="Message input"]') composer.wait_for(state = "visible", timeout = 60_000) shoot("18-relogin-with-NEW2") From 5e43c623b98affc23efbf9dbe71061de7c1706a2 Mon Sep 17 00:00:00 2001 From: Etherl <61019402+Etherll@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:46:22 +0300 Subject: [PATCH 086/113] Fix FastSentenceTransformer Qwen embedding preprocessing (#6939) * Fix FastSentenceTransformer Qwen embedding preprocessing * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Document Transformer.load embedding modality fix for #6881 * Harden #6881 fix and add forwards/backwards-compatible regression tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fall back to Transformer constructor on legacy sentence-transformers without Hub-capable load * Mirror legacy sentence-transformers fallback in embedding-parity tripwire test * Tighten #6881 comments and docstrings * Skip embedding-parity test on CPU-only runners since FastSentenceTransformer requires CUDA * Honor the transformer module's saved subfolder when loading modules.json records a path for the Transformer module (root for decoder embedders like Qwen3-Embedding, 0_Transformer for the classic layout). Pooling/Normalize already load from their saved path; thread the same path into Transformer.load as subfolder so config and tokenizer resolve like stock ST. stays a no-op, so single-module models are unchanged. * Make embedding-parity test bf16-aware fp16 overflows to NaN on bf16-native embedders such as EmbeddingGemma (Gemma3), producing a false parity failure. Prefer bf16 when the GPU supports it so the tripwire can guard the full documented embedding matrix (Qwen3-Embedding, EmbeddingGemma, BGE-M3, all-MiniLM, GTE-ModernBERT), not just fp16-safe models. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- ...t_sentence_transformer_embedding_parity.py | 122 ++++++++++++++++++ ...st_sentence_transformers_pinned_symbols.py | 38 ++++++ unsloth/models/sentence_transformer.py | 65 +++++++++- 3 files changed, 222 insertions(+), 3 deletions(-) create mode 100644 tests/python/test_fast_sentence_transformer_embedding_parity.py diff --git a/tests/python/test_fast_sentence_transformer_embedding_parity.py b/tests/python/test_fast_sentence_transformer_embedding_parity.py new file mode 100644 index 0000000000..252d4486a5 --- /dev/null +++ b/tests/python/test_fast_sentence_transformer_embedding_parity.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text +like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a +"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building +via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings; +`_create_transformer_module` uses `Transformer.load(...)` instead. + +Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST +is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity, +opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected). +""" + +from __future__ import annotations + +import inspect +import os + +import pytest + + +def test_transformer_load_signature_supports_unsloth_kwargs(): + """Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs + the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back + to Transformer(...) there, so mirror that gate and skip.""" + models = pytest.importorskip("sentence_transformers.models") + load = getattr(models.Transformer, "load", None) + assert callable(load), ( + "sentence_transformers Transformer.load is missing; the #6881 fix in " + "unsloth.models.sentence_transformer._create_transformer_module depends on it." + ) + params = inspect.signature(load).parameters + accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) + # Mirror _create_transformer_module's hub_capable gate. + hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision")) + if not hub_capable: + pytest.skip( + "legacy Transformer.load(input_path); production path falls back to Transformer(...)" + ) + unsupported = [ + k + for k in ("token", "cache_folder", "revision", "trust_remote_code") + if not (accepts_var_kw or k in params) + ] + assert not unsupported, ( + f"installed sentence_transformers Transformer.load no longer accepts {unsupported} " + f"and has no **kwargs; update _create_transformer_module (#6881) before it silently " + f"falls back to Transformer(...)." + ) + + +def _probe_texts(): + return [ + "roasted chickpeas in 20 kg bags", + "The capital of France is Paris.", + "A fast brown fox jumps over the lazy dog.", + "recette de tarte aux pommes traditionnelle", + ] + + +def test_fast_sentence_transformer_matches_stock_st(): + """End-to-end: FastSentenceTransformer embeddings and tokenization must match a + stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and + GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners.""" + model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL") + if not model_id: + pytest.skip( + "set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model " + "(HF id or local path) to run the #6881 parity test" + ) + + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner") + np = pytest.importorskip("numpy") + pytest.importorskip("sentence_transformers") + from sentence_transformers import SentenceTransformer + + device = "cuda" + # Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native + # embedders such as EmbeddingGemma (Gemma3), which would mask real parity. + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + texts = _probe_texts() + max_seq_length = 256 + + # Control FIRST, before importing unsloth, so its global import patches never + # touch the stock reference (mirrors the issue's "restart runtime" repro). + ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype}) + ctrl.max_seq_length = max_seq_length + ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist() + ctrl_emb = np.asarray( + ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32 + ) + + import unsloth # noqa: F401 + from unsloth import FastSentenceTransformer + + fast = FastSentenceTransformer.from_pretrained( + model_id, + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = False, + load_in_16bit = True, + ) + fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist() + fast_emb = np.asarray( + fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32 + ) + + # Identical tokenization = no chat-template wrapping slipped in (the #6881 defect). + assert fast_ids == ctrl_ids, ( + f"tokenization diverged (chat-template wrapping regressed?):\n" + f" stock: {ctrl_ids}\n fast: {fast_ids}" + ) + + cos = (ctrl_emb * fast_emb).sum(1) / ( + np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1) + ) + assert float(cos.min()) > 0.99, ( + f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 " + f"(per-text {[round(float(c), 5) for c in cos]})" + ) diff --git a/tests/version_compat/test_sentence_transformers_pinned_symbols.py b/tests/version_compat/test_sentence_transformers_pinned_symbols.py index c0c35b9d5d..d7f9a54811 100644 --- a/tests/version_compat/test_sentence_transformers_pinned_symbols.py +++ b/tests/version_compat/test_sentence_transformers_pinned_symbols.py @@ -19,6 +19,8 @@ ST_TAGS = [ "v5.2.3", "v5.3.0", "v5.4.1", + "v5.5.1", + "v5.6.0", "master", ] @@ -120,6 +122,42 @@ def test_st_transformer_base_class_either_path(tag: str): ) +# Transformer.load classmethod: unsloth builds saved-ST modules through it (#6881). +@pytest.mark.parametrize("tag", ST_TAGS) +def test_st_transformer_load_accepts_unsloth_kwargs(tag: str): + """unsloth builds saved ST models via Transformer.load(...) so the saved + modality_config is honored (#6881). If .load stops accepting the hub kwargs it + passes (and has no **kwargs), update the fix before it silently regresses. Not + locating .load is a SKIP (may be inherited); the live test guards the install.""" + candidates = [ + "sentence_transformers/models/Transformer.py", + "sentence_transformers/models/transformer.py", + "sentence_transformers/base/modules/transformer.py", + "sentence_transformers/base/modules/module.py", + ] + for p in candidates: + src = fetch_text("UKPLab/sentence-transformers", tag, p) + if src is None or not has_def(src, "load", "func"): + continue + m = re.search(r"def\s+load\s*\((.*?)\)\s*(?:->[^:]*)?:", src, re.S) + if m is None: + continue + sig = m.group(1) + accepts_var_kw = "**" in sig + missing = [ + kw + for kw in ("token", "cache_folder", "revision", "trust_remote_code") + if not (accepts_var_kw or re.search(rf"\b{re.escape(kw)}\b", sig)) + ] + assert not missing, ( + f"{tag}: Transformer.load in {p} no longer accepts {missing} and has no " + f"**kwargs; update unsloth.models.sentence_transformer._create_transformer_module " + f"(#6881) before it silently falls back to Transformer(...)." + ) + return + pytest.skip(f"{tag}: Transformer.load not locatable in {candidates} (may be inherited)") + + # sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls. @pytest.mark.parametrize("tag", ST_TAGS) def test_st_util_helpers(tag: str): diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index c1172faa94..4a1a555bcf 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -990,7 +990,17 @@ class FastSentenceTransformer(FastModel): return None @staticmethod - def _create_transformer_module(model_name, model, tokenizer, max_seq_length, trust_remote_code): + def _create_transformer_module( + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + token = None, + cache_dir = None, + revision = None, + module_subfolder = "", + ): """Helper to create and configure a Transformer module.""" from sentence_transformers.models import Transformer @@ -1077,7 +1087,45 @@ class FastSentenceTransformer(FastModel): elif "tokenizer_args" in transformer_init_params: transformer_kwargs["tokenizer_args"] = trust_remote_code_kwargs.copy() - transformer_module = Transformer(model_name, **transformer_kwargs) + # Build via Transformer.load so the saved modality_config is honored: plain + # Transformer(...) makes ST 5.x infer a "message" modality for chat-template + # models (e.g. Qwen3-Embedding), chat-wrapping inputs and degrading embeddings + # (#6881). Only use .load when it resolves a Hub id (accepts the kwargs or + # **kwargs); legacy ST 3.x/4.x load(input_path) is local-only with no modality + # bug, so fall back to the constructor. + transformer_module = None + transformer_load = getattr(Transformer, "load", None) + has_modules_json = ( + FastSentenceTransformer._module_path( + model_name, token, cache_dir = cache_dir, revision = revision + ) + is not None + ) + if callable(transformer_load) and has_modules_json: + load_params = inspect.signature(transformer_load).parameters + accepts_var_kw = any( + p.kind is inspect.Parameter.VAR_KEYWORD for p in load_params.values() + ) + hub_capable = accepts_var_kw or any( + key in load_params for key in ("token", "cache_folder", "revision") + ) + if hub_capable: + load_kwargs = { + "token": token, + "cache_folder": cache_dir, + "revision": revision, + "trust_remote_code": trust_remote_code, + **transformer_kwargs, + } + # Resolve config/tokenizer from the module's saved subfolder + # (modules.json "path"), like stock ST; "" (root) is a no-op. + if module_subfolder: + load_kwargs["subfolder"] = module_subfolder + if not accepts_var_kw: + load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params} + transformer_module = Transformer.load(model_name, **load_kwargs) + if transformer_module is None: + transformer_module = Transformer(model_name, **transformer_kwargs) finally: # Restore original Auto* loading immediately AutoModel.from_pretrained = original_model_from_pretrained @@ -1191,6 +1239,10 @@ class FastSentenceTransformer(FastModel): tokenizer, max_seq_length, trust_remote_code, + token, + cache_dir, + revision, + module_subfolder = module_config.get("path") or "", ) modules[name] = transformer_module else: @@ -1226,7 +1278,14 @@ class FastSentenceTransformer(FastModel): ) transformer_module = FastSentenceTransformer._create_transformer_module( - model_name, model, tokenizer, max_seq_length, trust_remote_code + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + token, + cache_dir, + revision, ) modules["0"] = transformer_module From 6d674e5cc9aef396ce8aae45306b2b42beb76244 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:08:39 -0700 Subject: [PATCH 087/113] unsloth start: warn before running an agent's remote installer (#7024) When a coding agent is missing, `unsloth start ` offers to run the vendor's own installer (curl | bash, irm | iex, or npm) after an interactive confirm. Those installers execute with the user's privileges and there is no signature or hash check on the fetched content, so a blind "yes" is a supply-chain risk if the delivery path is compromised. Keep the auto-install convenience but make consent informed: before the prompt, name the exact remote source the installer fetches (or the command it runs for a package installer) and state that nothing verifies a signature or hash. Behavior is otherwise unchanged: non-interactive stdin still never executes anything, and the confirm still defaults to no. --- unsloth_cli/commands/start.py | 19 ++++++++++++++++++- unsloth_cli/tests/test_start.py | 26 ++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 764f5c7963..a8665b9be3 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -989,6 +989,12 @@ def _refresh_windows_path() -> None: os.environ["PATH"] = os.pathsep.join(entries) +def _install_source(install_hint: str) -> Optional[str]: + """The first http(s) URL an install hint fetches, or None (e.g. an npm install).""" + match = re.search(r"https?://[^\s'\")]+", install_hint) + return match.group(0) if match else None + + def _install_agent(name: str, install_hint: str) -> Optional[str]: # Missing agent under --launch: offer to run its documented install command, then # re-resolve it on PATH. Consent-based (we never auto-run a remote install script @@ -997,7 +1003,18 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]: if not sys.stdin.isatty(): return None typer.echo(f"`{name}` is not installed.") - if not typer.confirm(f"Install it now with `{install_hint}`?", default = False): + # Make the supply-chain risk explicit before the prompt: these are the vendors' + # own installers (curl | bash, irm | iex, npm), run with the user's privileges, + # and nothing checks a signature or hash on the fetched content. Naming the source + # turns a blind "yes" into informed consent. + source = _install_source(install_hint) + warning = ( + f"This will download and RUN a script from {source} with your privileges" + if source + else f"This will RUN `{install_hint}` with your privileges" + ) + typer.secho(f"{warning}; there is no signature or hash check.", fg = "yellow", err = True) + if not typer.confirm(f"Install `{name}` now with `{install_hint}`?", default = False): return None # Run each hint through the shell it is written for: PowerShell (irm | iex, or npm) # on Windows, /bin/sh (curl | bash, or npm) everywhere else. diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 18cb40f18d..87a295532e 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -128,6 +128,32 @@ def test_install_agent_uses_powershell_on_windows(monkeypatch): assert ran == [["powershell", "-NoProfile", "-Command", install_hint]] +def test_install_agent_warns_and_names_remote_source(monkeypatch, capsys): + # Before the confirm, a remote installer must name the URL it fetches so the + # user consents to a specific source rather than blindly accepting. + monkeypatch.setattr(start.os, "name", "nt") + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) # decline: nothing runs + hint = "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup" + assert start._install_agent("hermes", hint) is None + err = capsys.readouterr().err + assert "https://hermes-agent.nousresearch.com/install.ps1" in err + assert "download and RUN" in err + assert "signature or hash" in err + + +def test_install_agent_warns_for_package_installer(monkeypatch, capsys): + # An npm-style installer has no URL to fetch, but still runs with the user's + # privileges, so the warning names the command instead. + monkeypatch.setattr(start.os, "name", "posix") + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + err = capsys.readouterr().err + assert "npm install -g @openai/codex" in err + assert "with your privileges" in err + + def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch): monkeypatch.setattr(start.os, "name", "nt") From 0d4bd50768ca1d55009b51dfa097d7c71e819f4b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:26:24 -0700 Subject: [PATCH 088/113] Restore process-global torch.compile config on torch 2.12 so gradient checkpointing backward honors it (#7019) * Mirror dynamo/inductor config sets into defaults so torch 2.12 worker threads honor them torch 2.12 stores config user overrides in ContextVars, so direct assignments like torch._dynamo.config.recompile_limit = 1024 no longer reach the autograd engine worker threads. Gradient checkpointing recomputes fullgraph-compiled gpt-oss kernels inside backward on those threads, which then read the default recompile limit of 8 and raise FailOnRecompileLimitHit at step 0 of GRPO/SFT. Mirror direct config assignments into the process-global entry defaults on torch >= 2.12, restoring the torch <= 2.11 cross-thread semantics while leaving the context-scoped config.patch API untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep config.patch thread-local when mirroring dynamo/inductor sets config.patch(...) also assigns through ConfigModule.__setattr__, so the default-mirror was leaking its scoped, thread-local writes into the process-global entry default. Track patch enter/exit with a per-thread depth counter (wrapping ConfigModule.patch) and skip mirroring while inside a patch, so only genuine direct assignments restore the torch 2.11 cross-thread semantics and config.patch stays context-local. * Also keep config.load_config thread-local when mirroring config sets load_config restores a saved dynamo/inductor config by calling setattr per key, which the default-mirror would otherwise leak process-wide just like config.patch did. Wrap load_config with the same per-thread depth counter (renamed to _scoped_depth) so both scoped writers skip the mirror and stay context-local, while genuine direct assignments still restore the torch 2.11 cross-thread default. * Drop the pre-existing override replay from the config thread fix The replay was redundant: this runs from _gpu_init before unsloth sets any dynamo/inductor config, so the __setattr__ wrapper already mirrors every later assignment (recompile_limit included). It could also read a value that belonged to a config.patch context still active at import time and write that thread-local override into the global default. Removing it keeps the cross-thread fix and drops the now-unused _inductor.config import. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/_gpu_init.py | 6 ++ unsloth/import_fixes.py | 132 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 136 insertions(+), 2 deletions(-) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index e39b44488a..e6178e60f3 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -173,6 +173,7 @@ from .import_fixes import ( fix_vllm_guided_decoding_params, fix_vllm_pdl_blackwell, fix_triton_compiled_kernel_missing_attrs, + fix_dynamo_config_thread_visibility, patch_trunc_normal_precision_issue, ignore_logger_messages, patch_ipykernel_hf_xet, @@ -203,6 +204,10 @@ fix_vllm_guided_decoding_params() fix_trl_vllm_ascend() fix_vllm_pdl_blackwell() fix_triton_compiled_kernel_missing_attrs() +# Must run before unsloth_zoo's patch_torch_compile and the gpt-oss temporary +# patches raise the dynamo recompile limits, so those settings reach the +# autograd worker threads on torch >= 2.12. +fix_dynamo_config_thread_visibility() patch_trunc_normal_precision_issue() ignore_logger_messages() patch_ipykernel_hf_xet() @@ -233,6 +238,7 @@ del fix_vllm_guided_decoding_params del fix_trl_vllm_ascend del fix_vllm_pdl_blackwell del fix_triton_compiled_kernel_missing_attrs +del fix_dynamo_config_thread_visibility del patch_trunc_normal_precision_issue del ignore_logger_messages del patch_ipykernel_hf_xet diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index e5a5d01c2f..c5300ed4d0 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -1064,6 +1064,135 @@ def fix_triton_compiled_kernel_missing_attrs(): ) +def fix_dynamo_config_thread_visibility(): + """torch 2.12 made torch._dynamo/_inductor config overrides thread-local + (ContextVars), so `config.recompile_limit = 1024` set on the main thread is + invisible to the autograd worker threads that run backward. Gradient + checkpointing recompiles fullgraph gpt-oss kernels there against the default + limit of 8, raising FailOnRecompileLimitHit at step 0. Mirror direct config + assignments into the process-global entry default (torch <= 2.11 semantics). + config.patch(...) and config.load_config(...) also assign via __setattr__ but + are thread-local by design, so skip mirroring while inside one (tracked per + thread). No-op below torch 2.12 and on any torch without this internal layout. + """ + try: + import torch + + if Version(torch.__version__) < Version("2.12.0"): + return + import torch._dynamo.config as _dynamo_config + from torch.utils._config_module import ConfigModule + from contextvars import ContextVar + except Exception: + return + + try: + probe = getattr(_dynamo_config, "_config", {}).get("recompile_limit", None) + if probe is None or not isinstance(getattr(probe, "user_override", None), ContextVar): + # Overrides are not context-local on this torch; nothing to fix. + return + original_setattr = ConfigModule.__setattr__ + if getattr(original_setattr, "__unsloth_patched__", False): + return + except Exception: + return + + mirrored_modules = ("torch._dynamo.config", "torch._inductor.config") + + # config.patch(...) and config.load_config(...) also assign via __setattr__, but + # their writes are thread-local by design; a per-thread depth counter marks them + # so they are not mirrored into the process-global default. + import threading + + _scoped_depth = threading.local() + + def _in_scoped_write(): + return getattr(_scoped_depth, "n", 0) > 0 + + def _bump(delta): + _scoped_depth.n = getattr(_scoped_depth, "n", 0) + delta + + original_patch = ConfigModule.patch + if not getattr(original_patch, "__unsloth_patched__", False): + + @functools.wraps(original_patch) + def _patched_patch(self, *args, **kwargs): + ctx = original_patch(self, *args, **kwargs) + try: + cls = type(ctx) # patch() builds a fresh ConfigPatch class each call + if not getattr(cls, "__unsloth_patch_wrapped__", False): + _enter0, _exit0 = cls.__enter__, cls.__exit__ + + def _enter(s, _e = _enter0): + _bump(1) + try: + return _e(s) + finally: + _bump(-1) + + def _exit( + s, + *a, + _x = _exit0, + ): + _bump(1) + try: + return _x(s, *a) + finally: + _bump(-1) + + cls.__enter__, cls.__exit__ = _enter, _exit + cls.__unsloth_patch_wrapped__ = True + except Exception: + pass + return ctx + + _patched_patch.__unsloth_patched__ = True + ConfigModule.patch = _patched_patch + + # load_config restores a saved config by calling setattr per key (thread-local). + original_load_config = getattr(ConfigModule, "load_config", None) + if callable(original_load_config) and not getattr( + original_load_config, "__unsloth_patched__", False + ): + + @functools.wraps(original_load_config) + def _patched_load_config(self, *args, **kwargs): + _bump(1) + try: + return original_load_config(self, *args, **kwargs) + finally: + _bump(-1) + + _patched_load_config.__unsloth_patched__ = True + ConfigModule.load_config = _patched_load_config + + @functools.wraps(original_setattr) + def _patched_setattr(self, name, value): + original_setattr(self, name, value) + if _in_scoped_write(): + return # transient patch / load_config write: keep it thread-local + # Aliases (cache_size_limit -> recompile_limit) re-enter with the real name. + if self.__dict__.get("__name__", None) in mirrored_modules: + try: + entry = self.__dict__["_config"].get(name, None) + if entry is not None and entry.alias is None: + entry.default = value + except Exception: + pass + + _patched_setattr.__unsloth_patched__ = True + ConfigModule.__setattr__ = _patched_setattr + + # No replay of existing overrides: unsloth installs this before it sets any + # dynamo/inductor config, so the wrapper mirrors every later assignment. Replaying + # would also bake a still-active config.patch override into the global default. + logger.info( + "Unsloth: Patched torch config modules so dynamo/inductor settings " + "(e.g. recompile_limit) apply across threads on torch >= 2.12." + ) + + def patch_trunc_normal_precision_issue(): """ Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32. @@ -1323,8 +1452,7 @@ def fix_vllm_pdl_blackwell(): if patched: logger.info( - f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - " - f"patched: {', '.join(patched)}" + f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - patched: {', '.join(patched)}" ) else: # Just set the env var - vLLM might be an older version without supports_pdl From b509d47dd7427a1ba9ff1c80d1ca64fb9889bddf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:26:36 -0700 Subject: [PATCH 089/113] Silence torch._check_is_size FutureWarning and shim it if torch removes it (#7023) * Silence torch._check_is_size FutureWarning and shim it if torch removes it bitsandbytes 4-bit dequant calls torch._check_is_size, which torch deprecated with a FutureWarning ("Use _check(i >= 0) instead") that prints on every bnb-4bit load. Silence that warning in suppress_cuda_printf, and add fix_torch_check_is_size so a future torch that removes _check_is_size gets it shimmed to _check(i >= 0) (honoring the max bound) and bitsandbytes keeps working. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten fix_torch_check_is_size docstring Lead with what the shim does and drop the redundant line; two lines instead of three, same intent. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/_gpu_init.py | 3 +++ unsloth/import_fixes.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index e6178e60f3..984057e9f7 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -26,6 +26,7 @@ already_imported = [mod for mod in critical_modules if mod in sys.modules] # Fix some issues before importing other packages from .import_fixes import ( fix_message_factory_issue, + fix_torch_check_is_size, check_fbgemm_gpu_version, disable_broken_causal_conv1d, disable_broken_vllm, @@ -72,6 +73,7 @@ fix_bitsandbytes_rocm_arch_detection() disable_broken_causal_conv1d() disable_broken_vllm() fix_message_factory_issue() +fix_torch_check_is_size() check_fbgemm_gpu_version() torchvision_compatibility_check() fix_diffusers_warnings() @@ -81,6 +83,7 @@ del fix_bitsandbytes_rocm_arch_detection del disable_broken_causal_conv1d del disable_broken_vllm del fix_message_factory_issue +del fix_torch_check_is_size del check_fbgemm_gpu_version del torchvision_compatibility_check del fix_diffusers_warnings diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index c5300ed4d0..09de248c7b 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -172,6 +172,10 @@ if not UNSLOTH_ENABLE_LOGGING: # Deprecation warnings from torchao warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated") warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated") + # torch._check_is_size FutureWarning (called by bitsandbytes 4-bit dequant) + warnings.filterwarnings( + "ignore", message = r"_check_is_size will be removed", category = FutureWarning + ) # TorchAO deprecated import paths (https://github.com/pytorch/ao/issues/2752) warnings.filterwarnings( @@ -253,6 +257,30 @@ if not UNSLOTH_ENABLE_LOGGING: ) +def fix_torch_check_is_size(): + """Shim torch._check_is_size if a future torch removes it (bitsandbytes 4-bit + dequant calls it). The FutureWarning is silenced in suppress_cuda_printf.""" + try: + import torch + + if hasattr(torch, "_check_is_size"): + return + + def _check_is_size( + i, + message = None, + *, + max = None, + ): + torch._check(i >= 0, message) + if max is not None: + torch._check(i <= max, message) + + torch._check_is_size = _check_is_size + except Exception: + return + + # Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype' # MUST do this at the start primarily due to tensorflow causing issues def fix_message_factory_issue(): From c1e06e9ddfb53a9d40e1fb182030aded94f4b4bb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:47:59 -0700 Subject: [PATCH 090/113] unsloth start: add --persist to keep and reopen agent sessions (#7014) * unsloth start: add --resume to persist and reopen agent sessions `unsloth start ` launches a coding agent whose home is a throwaway temp dir wiped on exit, so codex/openclaw/hermes/pi (which relocate their whole home there) cannot resume a conversation after you quit. opencode and claude keep their session data in a fixed user dir, so they already resume. Add an opt-in --resume/--no-resume flag: it routes the launch to the stable Unsloth agents dir (the same one --no-launch already uses) so the session survives the exit, never touching the user's own ~/.. A bare --resume also reopens the last conversation via the agent's native flag (codex `resume --last`, opencode/claude/pi `--continue`). The default is unchanged: a plain launch still uses a temp dir and persists nothing. Add a dispatch-only `resume` job to the Local Agent Guides CI that drives the real launch path and asserts the split: codex/pi are wiped without --resume and persist with it, while opencode/claude persist either way. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * unsloth start: rename --resume to --persist The session flag collided with agents' own resume flags. `unsloth start claude --resume ` used to forward `--resume ` straight to Claude (which keeps its history in ~/.claude regardless), so a boolean --resume on unsloth start would have swallowed the session id and turned it into a stray prompt. Name the persistence flag --persist instead, so every agent's native resume flag (claude --resume , codex resume, opencode --continue, ...) still passes through untouched. Behavior is otherwise identical: --persist keeps a launched agent's session under the Unsloth agents dir, and a bare --persist reopens the last conversation. Add a regression test that `--resume ` passes through verbatim, and in the CI resume experiment skip the redundant second pass for opencode/claude (they persist either way, and a second CPU turn only risks a timeout). * unsloth start: correct --persist help and drop the buggy auto-resume Reword the --persist help to be accurate: claude and opencode keep sessions in the user's own stores and resume regardless, so --persist only stabilizes the otherwise-ephemeral relocated home of codex/openclaw/hermes/pi. Drop the bare-launch auto-append of native resume tokens: it errored on a first launch with no prior session, and was inconsistent between launch and no-launch. --persist now only keeps the session dir; resume via the agent's own command (e.g. `unsloth start codex --persist resume`), which now finds it. In the CI resume experiment, fail the pass when the launched turn exits non-zero, so a write-then-error is not misread as PERSISTED. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/scripts/agent-guides-drive.sh | 148 +++++++++++++++++ .github/workflows/local-agent-guides-ci.yml | 170 ++++++++++++++++++++ unsloth_cli/commands/start.py | 53 ++++-- unsloth_cli/tests/test_start.py | 142 ++++++++++++++++ 4 files changed, 502 insertions(+), 11 deletions(-) diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index defdb498c7..f4189a159e 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -527,6 +527,154 @@ case "$MODE" in echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" ;; + # ── resume: does a launched agent's session survive exit and resume? ──── + # Unlike the other modes, this drives the real LAUNCH path (`unsloth start + # ...`, the interactive default), not the --no-launch recipe. That + # path relocates each agent's home to a throwaway temp dir wiped on exit, so + # a session cannot be resumed -- unless --persist routes it to the stable + # Unsloth agents dir instead. We run one headless turn per pass and check + # whether the turn left a session in a persistent store (deterministic, no + # reliance on the model recalling anything), for a baseline pass and a + # --persist pass, and assert the expected split for this agent. + resume) + CODEWORD="PLATYPUS7" + T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK." + T2="What codeword did I ask you to remember? Reply with just that word." + WORK="$WORKDIR_BASE/${AGENT}-resume" + + # STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to. + # Read it from a --no-launch probe (which also writes the agent's config + # there). codex/pi relocate their whole home/HOME here; opencode/claude keep + # their session data in a fixed user dir, so STABLE_HOME stays empty for them. + parse_connect + case "$AGENT" in + codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;; + pi) STABLE_HOME="$(raw_env HOME)" ;; + *) STABLE_HOME="" ;; + esac + + # The persistent stores a session would land in if it were NOT wiped. We + # count files here before/after each turn; a positive delta means the + # session persisted (is resumable), zero means it went to a wiped temp dir. + resume_tracked_dirs() { + case "$AGENT" in + codex) printf '%s\n' "$HOME/.codex" ;; + opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;; + claude) printf '%s\n' "$HOME/.claude" ;; + pi) printf '%s\n' "$HOME/.pi" ;; + *) : ;; + esac + [ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME" + } + count_session_files() { + local total=0 d n + while IFS= read -r d; do + [ -n "$d" ] && [ -d "$d" ] || continue + n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n)) + done < <(resume_tracked_dirs) + echo "$total" + } + + # The headless first-turn subcommand per agent (mirrors file-edit's map), + # forwarded verbatim through the launch path as passthrough args. + set_t1_cmd() { + case "$AGENT" in + claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;; + codex) T1_CMD=(exec "$T1") ;; + opencode) T1_CMD=(run "$T1") ;; + pi) T1_CMD=(-p "$T1") ;; + *) guide_fail "resume mode does not cover agent '$AGENT'" ;; + esac + } + + # Run one headless turn through the launch path. $1=outfile, $2="" or + # "--persist", rest = the agent subcommand. --yolo auto-approves so no tool + # prompt can hang; --api-key attaches to the already-served CI model. + launch_turn() { + local out="$1" rflag="$2"; shift 2 + local flag=(); [ -n "$rflag" ] && flag=("$rflag") + run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \ + --api-key "$UNSLOTH_API_KEY" "$@" + local rc=$? + redact "$out" + return "$rc" + } + + # One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED + # from the session-store delta. Runs in the main shell (not a command + # substitution) so a hang's guide_fail actually fails the job and the + # progress lines reach the CI log. $1 = "" (baseline) or "--persist". + RESULT="" + run_pass() { + local rflag="$1" label="baseline" + [ -n "$rflag" ] && label="resume" + rm -rf "$WORK"; mkdir -p "$WORK" + set_t1_cmd + local out="$LOGS_DIR/${AGENT}-resume-${label}.txt" + local before after rc + before="$(count_session_files)" + pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK" + launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$? + popd >/dev/null || true + after="$(count_session_files)" + echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})" + # The turn must succeed for the delta to mean anything: an agent that writes a + # session file then errors would otherwise be misread as PERSISTED. Mirror the + # file-edit mode and fail the pass on a non-zero launch (the flagship codex recall + # below stays WARN-only, driven by its own launch_turn calls). + [ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \ + guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; } + if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi + } + + run_pass ""; BASELINE="$RESULT" + # Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix. + # opencode/claude persist either way, so the baseline already proves it and a + # second full CPU turn only risks a timeout; skip it for them. + case "$AGENT" in + codex|pi) run_pass "--persist"; RESUME="$RESULT" ;; + *) RESUME="n/a (persists either way)" ;; + esac + + # Expected: codex/pi relocate their whole home to the temp dir, so a plain + # launch is WIPED and only --persist PERSISTS. opencode/claude keep their + # session data in a fixed user dir, so the baseline already PERSISTS. + case "$AGENT" in + codex|pi) EXPECT_BASELINE="WIPED" ;; + opencode|claude) EXPECT_BASELINE="PERSISTED" ;; + esac + + echo "──────────────────────────────────────────────" + echo "[$AGENT] RESUME EXPERIMENT" + echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})" + echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}" + echo "──────────────────────────────────────────────" + + [ "$BASELINE" = "$EXPECT_BASELINE" ] \ + || guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}" + case "$AGENT" in + codex|pi) + [ "$RESUME" = "PERSISTED" ] \ + || guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;; + esac + + # Flagship behavioral proof (codex only, WARN-only): after a --persist plant, + # resume the session and check the model actually recalls the codeword. A + # miss is not a failure (the CI model is small); the mechanism gate above is + # the real assertion. + if [ "$AGENT" = "codex" ]; then + rm -rf "$WORK"; mkdir -p "$WORK" + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true + if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then + echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}" + else + echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed" + fi + fi + echo "[$AGENT] resume OK" + ;; + *) echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2 exit 2 diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 47f75dc1ba..25796bd5cf 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -471,6 +471,176 @@ jobs: redacted-configs/ retention-days: 7 + # ═════════════════════════════════════════════════════════════════════ + # Job: resume + # Does a conversation started with `unsloth start ` survive exit + # and resume? This drives the REAL launch path (not the --no-launch + # recipe the other jobs use). A plain launch relocates the agent home to + # a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the + # session to the stable Unsloth agents dir so it persists. opencode/claude + # keep their session data in a fixed user dir, so they persist either way. + # Dispatch-only: it is an end-to-end experiment, not a PR gate. + # ═════════════════════════════════════════════════════════════════════ + resume: + name: resume (${{ matrix.agent }}) + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + # codex/pi relocate their whole home (resume broken without --persist); + # opencode/claude keep session data in a fixed dir (resume already works). + # One agent from each class proves the split end to end; openclaw/hermes + # share codex's relocation mechanism and are covered by the unit tests. + agent: [codex, opencode, claude, pi] + env: + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18904' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: Resume experiment (launch path) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: resume-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + # ═════════════════════════════════════════════════════════════════════ # Job 3: prompt-cache # (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0 diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index a8665b9be3..48c0aca34b 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -133,6 +133,21 @@ _YOLO_OPTION = typer.Option( "flag/config. Any of the three spellings works for any agent." ), ) +_PERSIST_OPTION = typer.Option( + False, + "--persist/--no-persist", + help = ( + "Keep this agent's Unsloth-managed session dir so you can resume it later. " + "codex/openclaw/hermes/pi have their whole home relocated into an Unsloth dir " + "that is a throwaway temp dir (wiped on exit) by default; with --persist it " + "lives under the Unsloth agents dir and survives, so their own resume can reopen " + "it. claude and opencode keep sessions in your own stores (~/.claude, " + "~/.local/share/opencode), so they already resume regardless. To reopen a " + "session, pass the agent's own resume command through, e.g. " + "`unsloth start codex --persist resume` or `claude --resume `; those flow to " + "the agent unchanged." + ), +) # Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no # such flag (config only) and are handled in their config writers, so they are absent. @@ -1133,15 +1148,20 @@ def _agents_config_root() -> Path: @contextlib.contextmanager -def _session_config(agent: str, launch: bool): +def _session_config( + agent: str, + launch: bool, + persist: bool = False, +): """Yield a private directory for an agent's session config (never the user's own). - launch: an ephemeral temp dir removed after the agent process exits, so nothing - persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later - on this machine), reused across runs. Either way the user's real ~/. - config is left untouched. + launch (default): an ephemeral temp dir removed after the agent process exits, so + nothing persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run + later on this machine), reused across runs. persist (from --persist): use that same + stable dir even for a launch, so the agent's session survives the exit and can be + resumed next time. Either way the user's real ~/. config is left untouched. """ - if launch: + if launch and not persist: path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-")) try: yield path @@ -1453,6 +1473,7 @@ def claude( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point Claude Code at the running Studio server and start it.""" base, key, entry = _connect( @@ -1497,6 +1518,9 @@ def claude( # --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions. # IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a # sandbox is detected, and we don't want to falsely claim one on the user's host. + # claude keeps its history in ~/.claude/projects, which --settings/env never + # relocate, so a session already survives exit; resume it with `claude --continue` + # or `--resume ` passed through. command = [ "claude", "--model", @@ -1533,6 +1557,7 @@ def codex( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point OpenAI Codex at the running Studio server and start it.""" base, key, entry = _connect( @@ -1558,7 +1583,7 @@ def codex( *_yolo_command_flags("codex", yolo), *ctx.args, ] - with _session_config("codex", launch) as home: + with _session_config("codex", launch, persist = persist) as home: write_codex_config(base, entry, home) env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)} _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") @@ -1576,6 +1601,7 @@ def openclaw( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point OpenClaw at the running Studio server and start it.""" base, key, entry = _connect( @@ -1601,7 +1627,7 @@ def openclaw( if os.name == "nt" else "curl -fsSL https://openclaw.ai/install.sh | bash" ) - with _session_config("openclaw", launch) as cfg: + with _session_config("openclaw", launch, persist = persist) as cfg: config_path = cfg / "openclaw.json" # key lives in the config, not the env; --yolo writes the exec policy here too. write_openclaw_config(base, key, entry, config_path, yolo = yolo) @@ -1622,6 +1648,7 @@ def opencode( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point OpenCode at the running Studio server and start it.""" base, key, entry = _connect( @@ -1645,7 +1672,9 @@ def opencode( command = ["opencode", "--model", opencode_model] else: command = ["opencode"] - with _session_config("opencode", launch) as cfg: + # opencode keeps sessions in ~/.local/share/opencode (never relocated), so resume + # already survives exit; reopen the last one by passing `opencode --continue` through. + with _session_config("opencode", launch, persist = persist) as cfg: config_path = cfg / "opencode.json" # OPENCODE_CONFIG is an overlay (loaded between the user's global and project # configs), so this adds the Unsloth provider/model for the session without @@ -1697,6 +1726,7 @@ def hermes( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point Hermes (Nous Research) at the running Studio server and start it.""" base, key, entry = _connect( @@ -1708,7 +1738,7 @@ def hermes( ) command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] install_hint = _hermes_install_hint() - with _session_config("hermes", launch) as home: + with _session_config("hermes", launch, persist = persist) as home: # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session. write_hermes_config(base, entry, home / "config.yaml") @@ -1728,6 +1758,7 @@ def pi( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point Pi (coding agent) at the running Studio server and start it.""" base, key, entry = _connect( @@ -1752,7 +1783,7 @@ def pi( # --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs # no install scripts), so accepting the prompt skips dependency lifecycle scripts. install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent" - with _session_config("pi", launch) as home: + with _session_config("pi", launch, persist = persist) as home: # Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers # it over $HOME/.pi/agent), so pin it at the session dir: an inherited # PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 87a295532e..065972b275 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -2548,3 +2548,145 @@ def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path with start._session_config("codex", launch = False) as home2: assert home2 == home assert (home2 / "sessions" / "live.sqlite").read_text() == "state" + + +# ── --persist: persist the agent session so it can be resumed ──────────────── +def test_session_config_persist_uses_stable_dir_and_survives(monkeypatch, tmp_path): + # --persist routes a launch to the stable Unsloth agents dir (the one --no-launch + # already uses) instead of a throwaway temp dir, and never wipes it on exit. + monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents") + with start._session_config("codex", launch = True, persist = True) as home: + assert home == tmp_path / "agents" / "codex" + (home / "marker").write_text("kept") + assert home.exists() + assert (home / "marker").read_text() == "kept" + + +def test_session_config_default_launch_is_ephemeral(): + # Default launch (no --persist) still uses a throwaway temp dir wiped on exit. + with start._session_config("codex", launch = True) as home: + assert home.exists() + assert "unsloth-codex-" in home.name + assert not home.exists() + + +# The temp-dir agents: --persist points each one's home/state env at the stable dir; +# without it, at an ephemeral temp path. opencode is handled separately (only its +# config overlay is relocated; its session data was never in the temp dir). +_RESUME_ENV_VAR = { + "codex": "CODEX_HOME", + "openclaw": "OPENCLAW_STATE_DIR", + "hermes": "HERMES_HOME", + "pi": "HOME", +} + + +def _capture_launch(monkeypatch, argv): + captured = {} + + def run( + command, + env = None, + **kwargs, + ): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, argv) + assert result.exit_code == 0, result.output + return captured + + +@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR)) +def test_resume_persists_agent_home_to_stable_dir(agent, fake_studio, tmp_path, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}") + captured = _capture_launch(monkeypatch, [agent, "--persist"]) + stable = tmp_path / "agents" / agent + assert captured["env"][_RESUME_ENV_VAR[agent]] == str(stable) + # The stable dir survives the agent exit, so the session can be resumed. + assert stable.exists() + + +@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR)) +def test_default_launch_home_is_ephemeral(agent, fake_studio, tmp_path, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}") + captured = _capture_launch(monkeypatch, [agent]) + home = captured["env"][_RESUME_ENV_VAR[agent]] + assert f"unsloth-{agent}-" in home + assert str(tmp_path / "agents") not in home + + +def test_resume_opencode_config_in_stable_dir(fake_studio, tmp_path, monkeypatch): + # opencode's session data lives in ~/.local/share/opencode (never relocated), so + # resume already survives exit; --persist also stabilizes its config overlay dir. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode") + captured = _capture_launch(monkeypatch, ["opencode", "--persist"]) + stable = tmp_path / "agents" / "opencode" + assert captured["env"]["OPENCODE_CONFIG"] == str(stable / "opencode.json") + assert stable.exists() + + +def test_persist_bare_codex_launch_has_no_resume_token(fake_studio, monkeypatch): + # A bare `--persist` only persists the session dir; it must NOT auto-append a native + # resume token, or the very first launch (no session yet) would send codex down its + # no-session error path. The user resumes explicitly: `unsloth start codex --persist resume`. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + captured = _capture_launch(monkeypatch, ["codex", "--persist"]) + assert "resume" not in captured["command"] + # command[0] is the resolved executable path; assert the argv after it. + assert captured["command"][1:] == ["--oss", "--profile", start._CODEX_PROFILE] + + +def test_persist_bare_opencode_launch_has_no_resume_token(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode") + captured = _capture_launch(monkeypatch, ["opencode", "--persist"]) + assert "--continue" not in captured["command"] + assert captured["command"][1:] == ["--model", f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"] + + +def test_persist_bare_claude_launch_has_no_resume_token(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + captured = _capture_launch(monkeypatch, ["claude", "--persist"]) + assert "--continue" not in captured["command"] + assert captured["command"][1:] == ["--model", MODEL["id"]] + + +def test_resume_with_passthrough_does_not_auto_append(fake_studio, monkeypatch): + # When the caller drives their own subcommand, --persist only persists the dir; it + # must not inject a resume token that would collide with the user's command. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + captured = _capture_launch(monkeypatch, ["codex", "--persist", "exec", "hello"]) + assert "resume" not in captured["command"] + assert captured["command"][-2:] == ["exec", "hello"] + + +def test_default_launch_has_no_resume_token(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + captured = _capture_launch(monkeypatch, ["codex"]) + assert "resume" not in captured["command"] + + +def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatch): + # openclaw/hermes persist their session dir but have no non-interactive resume + # selector, so --persist must not append a token; their own picker resumes. + for agent in ("openclaw", "hermes"): + monkeypatch.setattr(start.shutil, "which", lambda _, a = agent: f"/usr/local/bin/{a}") + captured = _capture_launch(monkeypatch, [agent, "--persist"]) + assert "resume" not in captured["command"] + assert "--continue" not in captured["command"] + + +def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch): + # The persistence flag is --persist, NOT --resume, so an agent's own + # `--resume ` (e.g. `unsloth start claude --resume `) still flows + # through to the agent verbatim and is not swallowed as a Studio option. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + captured = _capture_launch(monkeypatch, ["claude", "--resume", "some-session-guid"]) + assert captured["command"][-2:] == ["--resume", "some-session-guid"] + # Studio never auto-appends its own resume token when the user drives resume. + assert captured["command"].count("--resume") == 1 + assert "--continue" not in captured["command"] From eb775d320778bb496378344c061bd538e4d39ad9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 03:20:02 -0700 Subject: [PATCH 091/113] Studio /v1/messages: accept thinking and unknown content blocks (#7017) * Studio /v1/messages: accept thinking and unknown content blocks The Anthropic-compatible /v1/messages endpoint modeled a message's content as Union[str, list[{text|image|tool_use|tool_result}]], so any other block type made Pydantic reject the whole request with `messages.N.content.str: Input should be a valid string`. Resuming a Claude session commonly replays assistant turns that carry `thinking` (extended thinking) blocks, and sometimes a null content for a tool-only turn, both of which tripped this and returned a 400. Accept them: - Add a permissive AnthropicUnknownBlock fallback (any block whose type is not one of the four known ones), so thinking/redacted_thinking/provider-specific/ future blocks validate. A validator keeps known types on their typed models, so a malformed known block (e.g. a tool_use without id) still fails cleanly. - Coerce a null message (and tool_result) content to "" so the converter's `for block in content` stays safe. The converter already drops block types it does not translate, so a thinking block is not forwarded to the model. * Studio /v1/messages: keep user content validation strict Make the thinking/null leniency role-aware so it never silently drops real user input. Assistant turns (replayed history) still accept unknown/thinking blocks and coerce a null tool-only turn to empty. User turns keep the strict boundary: a null user content is rejected, and a content block the converter cannot translate is rejected instead of being dropped into an empty prompt. Also remove an empty file committed by accident. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio /v1/messages: coalesce resumed user turns and tighten content checks - The /v1/messages count and generation paths now coalesce the adjacent user turns that dropping an empty or null assistant turn can leave behind, so a strict GGUF chat template no longer 400s on non-alternating roles. - A user content block with a non-string type (list / dict) is rejected as a clean 400 instead of raising TypeError and escaping as a 500. - The assistant null-to-empty coercion only applies to an explicit null; an assistant turn that omits content entirely still fails required-field validation instead of being silently coerced to an empty string. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio /v1/messages: tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/models/inference.py | 63 ++++++ studio/backend/routes/inference.py | 14 +- .../backend/tests/test_anthropic_messages.py | 184 ++++++++++++++++++ 3 files changed, 257 insertions(+), 4 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0f27b695fe..53b0f14b09 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1533,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel): tool_use_id: str content: Union[str, list] = "" + @field_validator("content", mode = "before") + @classmethod + def _coerce_null_content(cls, v): + # Some clients send null content for an empty tool result; the str|list + # union would 400 on it, so treat null as "". + return "" if v is None else v + + +# Block types the converter translates explicitly. Anything else (thinking / +# redacted_thinking, a provider block a resumed session replays, or a future type) +# is accepted as an unknown block and dropped by the converter, rather than 400-ing +# the whole request on strict validation. +_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"}) + + +class AnthropicUnknownBlock(BaseModel): + type: str + model_config = {"extra": "allow"} + + @field_validator("type") + @classmethod + def _only_unknown_types(cls, v): + # Known types parse as their typed models above (so a malformed known block + # still fails cleanly); this fallback only catches the rest. + if v in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError("known block type handled by its typed model") + return v + AnthropicContentBlock = Union[ AnthropicTextBlock, AnthropicImageBlock, AnthropicToolUseBlock, AnthropicToolResultBlock, + AnthropicUnknownBlock, ] @@ -1583,6 +1612,40 @@ class AnthropicMessage(BaseModel): role: Literal["user", "assistant"] content: Union[str, list[AnthropicContentBlock]] + @model_validator(mode = "before") + @classmethod + def _normalize_content(cls, data): + # Role-aware leniency that never silently drops real user input: + # - assistant: a resumed tool-only turn's null content -> "" (str|list would + # 400 on null; "" keeps the converter's `for block in content` safe). + # Unknown blocks (thinking / future types) validate via + # AnthropicUnknownBlock and are dropped by the converter. + # - user: keep strict. Null user content stays None so str|list rejects it + # (400) rather than forwarding an empty prompt; and reject block types the + # converter cannot translate, since it silently skips unknown user blocks + # -- a user turn made only of them would validate yet send no content + # (silent data loss). + if not isinstance(data, dict): + return data + content = data.get("content") + if data.get("role") == "assistant": + # Coerce only an explicit null (resumed tool-only turn). A missing + # content key stays malformed so the required-field check still 400s. + if "content" in data and content is None: + return {**data, "content": ""} + return data + if isinstance(content, list): + for block in content: + btype = ( + block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + ) + # Guard the value: a non-string type is unsupported too, and a + # membership test on an unhashable value would raise TypeError + # (escaping as a 500 instead of a clean 400). + if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError(f"unsupported content block type {btype!r} in a user message") + return data + class AnthropicTool(BaseModel): # Client tools have input_schema; server tools may only have type/name. diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d9901a5b2e..4bb9ce655e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10088,8 +10088,11 @@ async def anthropic_count_tokens( # Apply the same sanitization /messages does before generation, so the count # matches the prompt the real request would build (otherwise empty-assistant # sentinels / synthetic tool history inflate the count or hit the fallback). - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) openai_tools = anthropic_tools_to_openai(payload.tools or []) or None @@ -10217,8 +10220,11 @@ async def anthropic_messages( # builders apply the same strip; without it an Anthropic /v1/messages caller # replaying a prior provider-side tool_use forwards fake builtin tool # history to a backend with no matching function declarations. - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) # Enforce vision guard + re-encode embedded images to PNG so the Anthropic diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 0c6550a3bb..54454f6563 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1770,3 +1770,187 @@ class TestAnthropicMessagesToolRouting: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert backend.calls[0][0] == "plain" + + +def test_resumed_session_thinking_and_null_content_do_not_400(): + # A resumed session replays assistant turns with `thinking` (and sometimes null) + # content. Those must be accepted (thinking dropped by the converter), not 400ed. + from pydantic import ValidationError + + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "secret reasoning", "signature": "s"}, + {"type": "text", "text": "the answer"}, + {"type": "tool_use", "id": "t1", "name": "f", "input": {}}, + ], + }, + {"role": "assistant", "content": None}, # tool-only turn serialized as null + ], + ) + # Known blocks still parse as their typed models; only the unknown one is loose. + assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock" + assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock" + assert req.messages[2].content == "" # null coerced + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content")) + assert assistant["content"] == "the answer" + assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded + + # A malformed KNOWN block still fails cleanly instead of being swallowed. + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}], + ) + + +def test_user_null_content_rejected(): + # The null->"" leniency is assistant-only; a null user content must be rejected + # at the boundary, not coerced into an empty prompt and forwarded to the model. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": None}], + ) + + +def test_user_unknown_block_rejected_not_silently_dropped(): + # The converter skips user blocks it cannot translate, so a user turn whose only + # block is unknown would validate yet forward no content. Reject at the boundary + # to avoid that silent data loss (the assistant fallback is unaffected). + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "document", "source": {}}]}, + ], + ) + + +def test_user_translatable_blocks_still_accepted(): + # text / image / tool_result are translatable, so a real user message built from + # them must still pass; the unknown-block guard only trips on other types. + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "AA"}, + }, + {"type": "tool_result", "tool_use_id": "t1", "content": "ok"}, + ], + } + ], + ) + assert [type(b).__name__ for b in req.messages[0].content] == [ + "AnthropicTextBlock", + "AnthropicImageBlock", + "AnthropicToolResultBlock", + ] + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai) + + +def test_user_malformed_known_block_still_rejected(): + # The guard only allow-lists a user block's *type*; the union still validates its + # shape, so a known-but-malformed block (tool_result without tool_use_id) fails. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "tool_result", "content": "x"}]}, + ], + ) + + +def test_user_content_block_non_string_type_rejected_cleanly(): + # A user block whose `type` is a non-string (unhashable list / dict, or a stray + # int) must fail as a clean validation error, not raise TypeError from the + # frozenset membership test and escape as a 500. + from pydantic import ValidationError + for bad_type in ([], {}, 5): + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": [{"type": bad_type}]}], + ) + + +def test_assistant_missing_content_key_still_rejected(): + # The null -> "" leniency is only for an EXPLICIT null. An assistant message that + # omits content entirely stays malformed and must fail required-field validation. + from pydantic import ValidationError + + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant"}], + ) + # An explicit null is still accepted and coerced (regression guard). + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": None}, + ], + ) + assert req.messages[1].content == "" + + +def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch): + # user -> assistant(null) -> user is now accepted: the null assistant turn coerces + # to "" and is dropped. The route must then coalesce the two remaining user turns + # so a strict GGUF chat template does not 400 on non-alternating roles. + backend = _mock_backend(monkeypatch, context_length = 2048) + + class _Req: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/messages") + method = "POST" + + async def is_disconnected(self): + return False + + payload = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "please continue"}, + ], + ) + + response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t")) + assert response.status_code == 200 + + [(_path, kwargs)] = backend.calls + user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"] + assert len(user_turns) == 1 # the two user turns were merged, not left adjacent + merged = user_turns[0]["content"] + if isinstance(merged, list): + merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict)) + assert "first question" in merged and "please continue" in merged From 350233512092fc6847b42050b7768f5f5e9a4578 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Thu, 9 Jul 2026 07:39:48 -0300 Subject: [PATCH 092/113] Studio: add Vulkan llama.cpp support (#5819) * Studio: add Vulkan llama.cpp support * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address gemini's feedback * Studio: move the Vulkan VRAM probe into a standalone script * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Improve Vulkan probe error reporting * Resolve llama-server symlink so Vulkan build is detected * Drop unreachable Vulkan fallback in GPU free-memory dispatcher * Skip the Intel GPU probe when NVIDIA or ROCm is present * Reserve host RAM headroom for Vulkan integrated GPUs * Add a `UNSLOTH_FORCE_VULKAN` environment variable * [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 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honor GGML_VK_VISIBLE_DEVICES, reserve discrete Vulkan VRAM headroom, and clear Intel GPU on --cpu-fallback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Route Intel and forced-Vulkan hosts to the upstream Vulkan prebuilt, add arm64 Vulkan, keep Vulkan out of RAG auto-detect * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Clear the fork release pin when routing a Vulkan host to the upstream repo * Gate auto-Vulkan routing on no physical NVIDIA so hidden CUDA devices aren't used * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pin Vulkan launches with --device Vulkan instead of the raw GGML_VK_VISIBLE_DEVICES index space * Let user --device override the Vulkan pin, and gate direct Vulkan asset picks on no physical NVIDIA * Update RAG auto-backend test mocks for the _resolve_auto binary and Vulkan probes * Keep the add_dll_directory handle alive through the Vulkan probe DLL loads * Revert RAG auto Vulkan guard, guard multi-backend Vulkan detection, and preserve forced Vulkan across updates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use getattr for RTLD_GLOBAL in the Vulkan probe CDLL mode * Skip CUDA/ROCm APU and datacenter GPU tuning on Vulkan builds On a Vulkan llama.cpp build gpu_indices are ggml compact ordinals, not CUDA/ROCm physical ids, so _amd_apu_wants_unified_memory and _apply_datacenter_env were reading the wrong device. On a mixed AMD APU plus discrete GPU host that could raise a spurious system-RAM shortfall and block a valid discrete-GPU load. Gate all three call sites on not is_vulkan_backend; the Vulkan path already reserves iGPU host headroom and the backend ignores GGML_CUDA_* anyway. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten Vulkan-guard comment in load_model * Reduce comments in Vulkan support to be more succinct * Resolve shell-wrapper llama-server entrypoint to the real lib dir create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root when it cannot symlink into build/bin. _find_llama_server_binary returns that root entrypoint, but Path.resolve() does not follow a shell wrapper, so _llama_lib_dir returned the install root and _is_vulkan_backend missed libggml-vulkan.so -- silently skipping the Vulkan probe and --device pin on an otherwise valid Vulkan install. Follow the wrapper's exec target to build/bin. Regression test: test_shell_wrapper_entrypoint_resolves_to_real_lib_dir. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: danielhanchen --- .../backend/core/inference/_vulkan_probe.py | 110 ++++++++ studio/backend/core/inference/llama_cpp.py | 246 ++++++++++++++++-- .../tests/test_install_resolve_prebuilt.py | 170 ++++++++++++ studio/backend/tests/test_llama_cpp_update.py | 42 +++ .../tests/test_llama_cpp_vulkan_probe.py | 193 ++++++++++++++ studio/backend/utils/llama_cpp_update.py | 6 + studio/install_llama_prebuilt.py | 243 ++++++++++++++++- 7 files changed, 984 insertions(+), 26 deletions(-) create mode 100644 studio/backend/core/inference/_vulkan_probe.py create mode 100644 studio/backend/tests/test_llama_cpp_vulkan_probe.py diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py new file mode 100644 index 0000000000..706346daad --- /dev/null +++ b/studio/backend/core/inference/_vulkan_probe.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Standalone free-VRAM probe for the bundled ggml Vulkan backend. + +Run in a short-lived subprocess (``python _vulkan_probe.py ``) so the +Vulkan instance never lives in the long-running backend process. Loads the +bundled ggml Vulkan backend from ```` and prints one +``\\t\\t\\t`` line per device to stdout. +Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi +order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU +sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses +it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm +fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM. + +Uses only the standard library so it stays runnable as a bare script. +""" + +import ctypes +import os +import sys + +# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ... +_GGML_BACKEND_DEVICE_TYPE_IGPU = 2 + + +def _igpu_flags(base, lib, count: int) -> list[bool]: + """Per-device integrated-GPU flags via ggml's backend registry. + + The Vulkan reg enumerates devices in the same order as + ``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device = + i``), so reg index == device ordinal. Returns all-False on any failure so + the reader never over-caps a discrete card. + """ + flags = [False] * count + try: + lib.ggml_backend_vk_reg.restype = ctypes.c_void_p + lib.ggml_backend_vk_reg.argtypes = [] + base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t + base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p] + base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p + base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + base.ggml_backend_dev_type.restype = ctypes.c_int + base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p] + + reg = lib.ggml_backend_vk_reg() + if not reg: + return flags + dev_count = base.ggml_backend_reg_dev_count(reg) + for i in range(min(count, dev_count)): + dev = base.ggml_backend_reg_dev_get(reg, i) + if dev: + flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU + except Exception: + # Best-effort: any failure degrades to "discrete" so the memory + # readings still get through instead of crashing the probe. + pass + return flags + + +def main() -> int: + if len(sys.argv) < 2: + return 0 + bindir = sys.argv[1] + + # Hold add_dll_directory's handle for the rest of main() (the documented + # idiom) so bindir stays on the search path while the sibling ggml DLLs + # resolve below. + _dll_dir = None + if sys.platform == "win32": + base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll" + try: + _dll_dir = os.add_dll_directory(bindir) + except Exception: + pass + else: + base_name, vk_name = "libggml-base.so", "libggml-vulkan.so" + + # RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr + # falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode). + _rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0) + try: + base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global) + lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global) + except OSError as e: + print(f"ggml-vulkan load failed: {e}", file = sys.stderr) + return 1 + + lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int + lib.ggml_backend_vk_get_device_count.argtypes = [] + lib.ggml_backend_vk_get_device_memory.restype = None + lib.ggml_backend_vk_get_device_memory.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ] + + count = lib.ggml_backend_vk_get_device_count() + igpu = _igpu_flags(base, lib, count) + rows = [] + for i in range(count): + free, total = ctypes.c_size_t(0), ctypes.c_size_t(0) + lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total)) + rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value)) + sys.stdout.write("\n".join(rows)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index f61402aa5c..3ba9eff857 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1436,6 +1436,50 @@ def _backfill_usage_from_timings(usage, timings): return out +def _vulkan_lib_filename() -> str: + return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so" + + +# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit +# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared +# system RAM, so hold back the same margin rather than inventing a larger one. +_IGPU_HOST_RESERVE_MIB = 1024 + + +def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int: + """Reserve host headroom on an integrated (shared-memory) Vulkan GPU. + + An iGPU's reported free "VRAM" is really free system RAM, so sizing + context/offload against all of it would push the host into swap or the OOM + killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from + ggml's device type, so a discrete card is never touched; only ever reduces. + """ + if not is_igpu: + return free_mib + return max(0, free_mib - _IGPU_HOST_RESERVE_MIB) + + +def _llama_lib_dir(binary: str) -> Path: + # The installer exposes llama-server as a top-level entrypoint into build/bin/, + # where the ggml backend libs live, so callers looking for sibling libs (Vulkan + # detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a + # symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to + # a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it + # cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's + # exec target too, so a wrapper-based install still finds build/bin. + resolved = Path(binary).resolve() + try: + with open(resolved, "rb") as _f: + _head = _f.read(256) + if _head.startswith(b"#!"): + _m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore")) + if _m: + return (resolved.parent / _m.group(1)).resolve().parent + except OSError: + pass + return resolved.parent + + def _is_external_link(path: Path) -> bool: """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink or a Windows directory junction / reparse point. Such a link resolves into @@ -2278,6 +2322,30 @@ class LlamaCppBackend: return total + @staticmethod + def _is_vulkan_backend(binary: Optional[str] = None) -> bool: + """True if the installed llama.cpp build is Vulkan-only. + + The official prebuilts are single-backend, so the Vulkan ggml lib next + to llama-server identifies a Vulkan build. Keeps the free-memory probe + and GPU pin in ggml's Vulkan device-index space. For a custom + multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer + to that backend (torch-usable, better-understood probe/pin). + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return False + lib_dir = _llama_lib_dir(binary) + if not (lib_dir / _vulkan_lib_filename()).is_file(): + return False + for _backend in ("cuda", "hip"): + sibling = ( + f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so" + ) + if (lib_dir / sibling).is_file(): + return False + return True + @staticmethod def _resolve_visible_physical_ids() -> Optional[list[int]]: """Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on @@ -2440,11 +2508,42 @@ class LlamaCppBackend: return True @staticmethod - def _get_gpu_free_memory() -> list[tuple[int, int]]: + def _visible_devices_mask(env_name: str) -> Optional[set[int]]: + """Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset. + + ``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask + ("") yields an empty set (all devices hidden), distinct from an unset + var (None, no mask). Used by the nvidia-smi probe. + """ + raw = os.environ.get(env_name) + if raw is None: + return None + try: + return set(int(x.strip()) for x in raw.split(",") if x.strip()) + except ValueError: + return None + + @staticmethod + def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]: + """``--device Vulkan,...`` to pin a Vulkan launch to selected GPUs. + + The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory + reports and the registry names ``Vulkan``). Pin by that name, NOT via + GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw + vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices + and deduplicating ICDs), so a compact ordinal there could select a + different physical device or the CPU rasterizer. + """ + if not gpu_indices: + return [] + return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)] + + @staticmethod + def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]: """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by index; empty if no supported GPU is reachable. Thin wrapper over ``_get_gpu_memory`` for callers that only need free VRAM.""" - return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()] + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)] @staticmethod def _apple_metal_memory_budget_bytes() -> int: @@ -2475,7 +2574,7 @@ class LlamaCppBackend: return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION) @staticmethod - def _get_gpu_memory() -> list[tuple[int, int, int]]: + def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]: """Query free AND total memory per GPU. Order: @@ -2487,9 +2586,18 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. + On a Vulkan build the ggml Vulkan probe is authoritative, so the indices + are ggml's compact Vulkan ordinals (the space the pin selects via + ``--device Vulkan``). It reports ``total`` for discrete cards and 0 + for an iGPU (shared RAM) so the fit falls back to free*frac there. + Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no - supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. + supported GPU is reachable. """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if LlamaCppBackend._is_vulkan_backend(binary): + return LlamaCppBackend._get_gpu_free_memory_vulkan(binary) # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( @@ -2505,16 +2613,7 @@ class LlamaCppBackend: **_windows_hidden_subprocess_kwargs(), ) if result.returncode == 0: - allowed: Optional[set[int]] = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - # `if x.strip()` filters trailing-comma masks ("0,1,"). - # Empty mask (CVD="") yields an empty set -> all GPUs - # filtered out, per codebase convention. - allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) - except ValueError: - pass + allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES") gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): parts = [p.strip() for p in line.split(",")] @@ -2579,6 +2678,91 @@ class LlamaCppBackend: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]: + """Query free (and total) VRAM per device via the bundled ggml Vulkan backend. + + Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance + in this process) and returns (device_index, free_mib, total_mib) sorted + by index. The index is ggml's compact Vulkan ordinal -- the one the + registry names ``Vulkan`` and load_model pins with ``--device``, + NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set + ``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the + list already reflects it. iGPUs leave a host-RAM margin (see + ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass + their real total through. [] when no Vulkan build or device is reachable. + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return [] + binary_dir = _llama_lib_dir(binary) + if not (binary_dir / _vulkan_lib_filename()).is_file(): + return [] + + env = child_env_without_native_path_secret() + # Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so + # the probe enumerates the same device list the launch will, named + # Vulkan0..N in the compact order reported here and pinned by that name + # via --device -- probe, mask, and pin stay in one index space. Do NOT + # filter the mask in Python: ggml parses the env var in raw + # vkEnumeratePhysicalDevices space while this probe reports the compact + # post-filter ordinal, so a Python filter would compare mismatched spaces. + if sys.platform != "win32": + # Let the loader resolve sibling ggml libs next to the binary. + existing_ld = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ( + f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir) + ) + probe_script = Path(__file__).with_name("_vulkan_probe.py") + try: + result = subprocess.run( + [sys.executable, str(probe_script), str(binary_dir)], + capture_output = True, + text = True, + timeout = 15, + env = env, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + logger.debug( + f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}" + ) + return [] + except Exception as e: + logger.debug(f"vulkan GPU probe failed: {e}") + return [] + + gpus: list[tuple[int, int, int]] = [] + for line in result.stdout.strip().splitlines(): + parts = line.split("\t") + if len(parts) != 4: + continue + try: + idx = int(parts[0]) + free_mib = int(parts[1]) // (1024 * 1024) + is_igpu = parts[2] == "1" + # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the + # fit stays on free*frac (the host reserve below is its + # headroom); a discrete card passes its real total through. + total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024) + except ValueError: + continue + capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu) + if capped < free_mib: + logger.info( + f"Vulkan device VK{idx} is an integrated GPU sharing system " + f"RAM; reserving {free_mib - capped}MiB host headroom " + f"({free_mib}->{capped}MiB usable)" + ) + gpus.append((idx, capped, total_mib)) + gpus.sort(key = lambda g: g[0]) + if gpus: + logger.info( + "Vulkan GPU memory detected: " + + ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus) + ) + return gpus + @staticmethod def _available_system_memory_mib() -> Optional[int]: """Available system RAM in MiB (psutil, then /proc/meminfo), or None if @@ -2807,7 +2991,8 @@ class LlamaCppBackend: def _llama_server_env_for_binary(binary: str) -> dict[str, str]: """Build a subprocess env that lets llama-server resolve native libs.""" env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) + # _llama_lib_dir resolves the llama-server symlink to the real build/bin. + binary_dir = str(_llama_lib_dir(binary)) if sys.platform == "win32": # Ordering: see _build_windows_path_dirs. #5106. @@ -5210,6 +5395,7 @@ class LlamaCppBackend: # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() + is_vulkan_backend = self._is_vulkan_backend(binary) # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected @@ -5449,7 +5635,8 @@ class LlamaCppBackend: model_size = gguf_size + mmproj_size # 2-tuple gpus for existing logic + a total map for the absolute # per-GPU headroom (correct when the GPU is already partly used). - _gpu_mem = self._get_gpu_memory() + # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. + _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} @@ -6222,7 +6409,12 @@ class LlamaCppBackend: # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an # oversize load the OS would otherwise kill mid-flight. Base model # only: an optional MTP drafter is dropped by the MTP-drop fallback. - if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices): + # CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals. + if ( + model_size is not None + and not is_vulkan_backend + and self._amd_apu_wants_unified_memory(gpu_indices) + ): _ram_msg = self._apu_ram_shortfall_message( model_size, self._available_system_memory_mib() ) @@ -6485,6 +6677,12 @@ class LlamaCppBackend: ", ".join(unsupported_cache_flags), ) + # Vulkan pins via --device (a cmd arg, unlike the env-based + # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's + # last-wins parsing lets a user --device override Studio's pick. + if is_vulkan_backend and gpu_indices is not None: + cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Studio's auto-set flags. Already # validated by the route via validate_extra_args(). @@ -6536,23 +6734,25 @@ class LlamaCppBackend: env.setdefault("OMP_NUM_THREADS", "2") # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use - # shared system RAM. setdefault so a user value wins. - if self._amd_apu_wants_unified_memory(gpu_indices): + # shared system RAM. setdefault so a user value wins. Not on Vulkan + # (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids. + if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices): env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. - if self._apply_datacenter_env(env, gpu_indices): + if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices): multi_gpu = self._effective_gpu_count(gpu_indices) > 1 logger.info( f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) # Pin to selected GPU(s). On ROCm, narrowing only - # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full - # set, so set HIP_VISIBLE_DEVICES too. - if gpu_indices is not None: + # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so + # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device + # (above), not here. + if gpu_indices is not None and not is_vulkan_backend: pinned = ",".join(str(i) for i in gpu_indices) env["CUDA_VISIBLE_DEVICES"] = pinned try: diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index b825172a63..090d2932ea 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -55,6 +55,27 @@ def _host(**kw): return ilp.HostInfo(**base) +def test_force_cpu_clears_all_gpu_attributes_including_intel(): + # --cpu-fallback is the "select the CPU prebuilt even when a GPU is present" + # escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or + # the planner still prepends the Vulkan asset on an Intel-GPU host. + host = _host( + is_linux = True, + is_x86_64 = True, + has_usable_nvidia = True, + has_physical_nvidia = True, + has_rocm = True, + rocm_gfx_target = "gfx1100", + has_intel_gpu = True, + ) + forced = ilp._apply_host_overrides(host, force_cpu = True) + assert forced.has_usable_nvidia is False + assert forced.has_physical_nvidia is False + assert forced.has_rocm is False + assert forced.rocm_gfx_target is None + assert forced.has_intel_gpu is False + + def test_macos_upstream_pin_only_for_explicit_pre26_upstream(): pre26 = _host( system = "Darwin", @@ -313,3 +334,152 @@ def test_sm103_host_drops_cuda128_windows_build(): ) kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129]) assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name] + + +def _upstream_release(tag, asset_names): + return { + "tag_name": tag, + "assets": [ + {"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names + ], + } + + +def test_direct_upstream_arm64_intel_prefers_vulkan(): + # Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU + # second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset). + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-arm64" in kinds + assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz" + + +def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only(): + # A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical + # True, usable False) + an Intel iGPU must NOT get the Vulkan archive even + # when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES + # and could grab the reserved card. It falls through to the CPU asset. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-cpu"] + + +def test_direct_upstream_arm64_without_intel_is_cpu_only(): + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64") + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-arm64"] + + +def test_direct_upstream_x86_intel_prefers_vulkan(): + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-cpu" in kinds + + +def test_linux_vulkan_health_glob_matches_bare_cpu_lib(): + # The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU + # libs so a valid Vulkan install is not re-flagged unhealthy every check. + choice = ilp.AssetChoice( + repo = UPSTREAM, + tag = "b9925", + name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", + url = "https://example/x", + source_label = "upstream", + install_kind = "linux-vulkan", + ) + groups = ilp.runtime_payload_health_groups(choice) + assert ["libggml-cpu*.so*"] in groups + assert ["libggml-cpu-*.so*"] not in groups + + +def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin(): + # Routing fork -> upstream also drops the fork release pin, which is in a + # different tag namespace and would make the upstream resolver miss. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False) + assert repo == UPSTREAM + assert tag == "" + assert routed.has_intel_gpu is True + + +def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin(): + # A pin set WITH an explicit upstream repo is already on upstream -> kept. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + _routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False) + assert repo == UPSTREAM + assert tag == "b9596" + + +def test_route_to_vulkan_prebuilt_cpu_fallback_wins(): + # --cpu-fallback suppresses Vulkan routing even for an Intel host. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True) + assert repo == FORK + assert tag == "b9596-mix-abc" + assert routed is host + + +def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted(): + # A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1): + # physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or + # Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted(): + # An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_non_intel_unchanged(): + host = _host(is_linux = True, is_x86_64 = True) + routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + assert routed is host + + +def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys): + # The --resolve-prebuilt probe must agree with the install path: an + # auto-detected Intel host resolves against upstream (Vulkan), not the fork. + monkeypatch.setattr( + ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == UPSTREAM + assert out["repo"] == UPSTREAM diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 5138e90471..f405ebcbd1 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -448,6 +448,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path): assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5" +def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): + # A Vulkan install (marker asset carries 'vulkan') must re-assert + # UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to + # CUDA/ROCm and silently replaces the Vulkan build. + install_dir = tmp_path / "llama.cpp" + binary = _write_install( + install_dir, + "b9493", + repo = "ggml-org/llama.cpp", + asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz", + ) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + def _on_start(cmd): + _write_install( + install_dir, + "b9518", + repo = "ggml-org/llama.cpp", + asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz", + ) + + popen_kwargs: dict = {} + _patch_installer_popen( + monkeypatch, + lines = ["installed\n"], + on_start = _on_start, + captured_kwargs = popen_kwargs, + ) + + assert upd.start_update()["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1" + + def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") diff --git a/studio/backend/tests/test_llama_cpp_vulkan_probe.py b/studio/backend/tests/test_llama_cpp_vulkan_probe.py new file mode 100644 index 0000000000..92aaab4873 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_vulkan_probe.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Vulkan free-VRAM reader regression tests on a synthetic probe output. + +Covers the post-probe handling in +``LlamaCppBackend._get_gpu_free_memory_vulkan``: + + * integrated GPUs (probe reports is_igpu=1) leave a flat per-device host + margin matching llama.cpp's --fit-target, so context auto-sizing can't + over-commit shared RAM, and report total 0 (shared RAM is not a budget), + * discrete GPUs (is_igpu=0) keep their free untouched and pass their real + total through so the fit can reserve absolute headroom, + * an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged + (ggml applies it), not stripped or filtered in Python -- the probe reports + ggml's compact ordinal, which load_model pins with ``--device Vulkan``. + +The ggml Vulkan library is never loaded: subprocess.run is mocked to emit +the tab-separated lines the real ``_vulkan_probe.py`` would print. +""" + +from __future__ import annotations + +import subprocess +import sys +import types as _types +from pathlib import Path +from unittest import mock + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import importlib as _importlib # noqa: E402 + + +def _maybe_stub(name: str, builder): + try: + _importlib.import_module(name) + except ImportError: + sys.modules[name] = builder() + + +def _build_loggers_stub(): + m = _types.ModuleType("loggers") + m.get_logger = lambda name: __import__("logging").getLogger(name) + return m + + +_maybe_stub("loggers", _build_loggers_stub) +_maybe_stub("structlog", lambda: _types.ModuleType("structlog")) + +from core.inference import llama_cpp as _llama_mod # noqa: E402 +from core.inference.llama_cpp import ( # noqa: E402 + LlamaCppBackend, + _llama_lib_dir, + _vulkan_lib_filename, +) + +MIB = 1024 * 1024 +GIB = 1024 * MIB + + +def _make_vulkan_install(tmp_path: Path) -> str: + """A binary whose sibling dir holds the Vulkan ggml lib, so the + reader's ``is_vulkan_backend`` sibling-file check passes.""" + bindir = tmp_path / "build" / "bin" + bindir.mkdir(parents = True) + binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server") + binary.write_bytes(b"stub") + (bindir / _vulkan_lib_filename()).write_bytes(b"stub") + return str(binary) + + +def _mock_probe(rows: list[str], captured_env: dict | None = None): + """Patch subprocess.run so the _vulkan_probe.py call returns ``rows`` + (already tab-formatted), recording the env it was launched with.""" + real_run = subprocess.run + + def fake_run(cmd, *args, **kwargs): + if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd): + if captured_env is not None: + captured_env.clear() + captured_env.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess( + args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = "" + ) + return real_run(cmd, *args, **kwargs) + + return mock.patch("subprocess.run", side_effect = fake_run) + + +def _row( + idx: int, + free_bytes: int, + is_igpu: int, + total_bytes: int = 0, +) -> str: + return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}" + + +def test_integrated_gpu_leaves_host_margin(tmp_path): + binary = _make_vulkan_install(tmp_path) + # iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target). + # total stays 0: shared system RAM is not a VRAM budget for the fit. + rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus + + +def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path): + binary = _make_vulkan_install(tmp_path) + # 6 GiB free on a partially occupied 24 GiB card: free is untouched and the + # real total flows through so the fit reserves absolute headroom (CUDA/ROCm + # parity) instead of the looser free*frac budget. + rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus + + +def test_large_discrete_gpu_is_untouched(tmp_path): + binary = _make_vulkan_install(tmp_path) + # A 48 GiB discrete card stays untouched regardless of size; only the + # iGPU flag triggers the host margin, never a VRAM/RAM ratio. + rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus + + +def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch): + # The mask is NOT stripped or filtered in Python: ggml parses it in raw + # physical-device space while this probe reports the compact post-filter + # ordinal, so mixing spaces would be wrong. It is passed through unchanged + # so ggml applies it to the same device list the launch will enumerate. + binary = _make_vulkan_install(tmp_path) + monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1") + captured: dict = {} + rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows, captured_env = captured): + LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured + + +def test_vulkan_pin_args_uses_device_names_not_env_mask(): + # Pin by compact device name via --device (the space the probe reports and + # the registry names), never by writing a compact ordinal into the raw + # GGML_VK_VISIBLE_DEVICES index space. + assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"] + assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"] + assert LlamaCppBackend._vulkan_pin_args(None) == [] + assert LlamaCppBackend._vulkan_pin_args([]) == [] + + +def test_vulkan_only_build_is_detected(tmp_path): + binary = _make_vulkan_install(tmp_path) + assert LlamaCppBackend._is_vulkan_backend(binary) is True + + +def test_multi_backend_build_is_not_vulkan_only(tmp_path): + # A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be + # treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan + # device; defer to the CUDA/HIP path instead. + binary = _make_vulkan_install(tmp_path) + cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so" + (_llama_lib_dir(binary) / cuda).write_bytes(b"stub") + assert LlamaCppBackend._is_vulkan_backend(binary) is False + + +@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX") +def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path): + # create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root + # when it cannot symlink; _find_llama_server_binary returns that root entrypoint, + # so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else + # _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently + # never engage on a valid Vulkan install. + import os + + binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib + bindir = Path(binary).parent + wrapper = tmp_path / "llama-server" + wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n') + os.chmod(wrapper, 0o755) + assert _llama_lib_dir(str(wrapper)) == bindir + assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index c16ae91467..1bcbfbf95a 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path logger.info("llama update: installing", cmd = " ".join(cmd)) # Stream progress lines into job["progress"]. env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") + # Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm + # box would otherwise re-route and silently replace the Vulkan build. + # Re-assert it via the same env flag setup uses (mirrors + # _rocm_install_args). + if asset and "vulkan" in asset.lower(): + env["UNSLOTH_FORCE_VULKAN"] = "1" proc = subprocess.Popen( cmd, stdout = subprocess.PIPE, diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 6c75e6c394..856ba71478 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -10,6 +10,7 @@ import argparse import atexit import errno import fnmatch +import glob import hashlib import json import os @@ -265,6 +266,7 @@ class HostInfo: has_physical_nvidia: bool has_usable_nvidia: bool has_rocm: bool = False + has_intel_gpu: bool = False rocm_gfx_target: str | None = None # (major, minor) from platform.mac_ver(); None off macOS or if unparseable. # Skips a macos prebuilt whose minimum-OS exceeds this host. @@ -1482,6 +1484,24 @@ def direct_upstream_release_plan( install_kind = "windows-hip", ) ) + # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. Gate + # on no PHYSICAL NVIDIA (not just no usable one): a host that hid NVIDIA + # via CUDA_VISIBLE_DEVICES must not reach Vulkan, which ignores that mask + # and could enumerate the reserved card. Falls through to CPU below. + elif host.has_intel_gpu and not host.has_physical_nvidia: + vulkan_asset = f"llama-{release_tag}-bin-win-vulkan-x64.zip" + vulkan_url = assets.get(vulkan_asset) + if vulkan_url: + attempts.append( + AssetChoice( + repo = repo, + tag = release_tag, + name = vulkan_asset, + url = vulkan_url, + source_label = "upstream", + install_kind = "windows-vulkan", + ) + ) cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip" cpu_url = assets.get(cpu_asset) if cpu_url: @@ -1545,6 +1565,23 @@ def direct_upstream_release_plan( # ROCm hosts are excluded: this ggml-org path ships no per-gfx ROCm # asset, so they fall through to the empty-attempts raise (HIP source # build) rather than silently getting a CPU binary on a GPU host. + # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. The + # elif already excludes usable NVIDIA and ROCm; also require no PHYSICAL + # NVIDIA so a CUDA-hidden card isn't reached through Vulkan (CPU below). + if host.has_intel_gpu and not host.has_physical_nvidia: + vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-x64.tar.gz" + vulkan_url = assets.get(vulkan_asset) + if vulkan_url: + attempts.append( + AssetChoice( + repo = repo, + tag = release_tag, + name = vulkan_asset, + url = vulkan_url, + source_label = "upstream", + install_kind = "linux-vulkan", + ) + ) asset_name = f"llama-{release_tag}-bin-ubuntu-x64.tar.gz" asset_url = assets.get(asset_name) if asset_url: @@ -1564,6 +1601,23 @@ def direct_upstream_release_plan( # selector returned 0 attempts and the installer fell back to a # source build on every Linux ARM64 host (DGX Spark, Ampere # Altra, GitHub-hosted ubuntu-24.04-arm runners, etc.). + # Intel (or other non-NVIDIA/non-AMD) GPU: prefer the Vulkan prebuilt, + # mirroring the x86_64 branch. Upstream ships bin-ubuntu-vulkan-arm64. + # No physical NVIDIA: don't reach a CUDA-hidden card through Vulkan. + if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm: + vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-arm64.tar.gz" + vulkan_url = assets.get(vulkan_asset) + if vulkan_url: + attempts.append( + AssetChoice( + repo = repo, + tag = release_tag, + name = vulkan_asset, + url = vulkan_url, + source_label = "upstream", + install_kind = "linux-vulkan", + ) + ) asset_name = f"llama-{release_tag}-bin-ubuntu-arm64.tar.gz" asset_url = assets.get(asset_name) if asset_url: @@ -3075,6 +3129,40 @@ def detect_host() -> HostInfo: # Note: amdhip64.dll presence alone is NOT treated as GPU evidence # since the HIP SDK can be installed without an AMD GPU. + # Detect an Intel GPU; gates the Vulkan prebuilt. Linux reads the DRM sysfs + # vendor id (0x8086); Windows queries the WMI video controller list. Only + # probed with no usable NVIDIA and no ROCm (matching the Vulkan branches), + # keeping the probe (notably the Windows powershell call) off that path. + has_intel_gpu = False + if not has_usable_nvidia and not has_rocm: + if is_linux: + for _vendor_file in glob.glob("/sys/class/drm/card*/device/vendor"): + try: + with open(_vendor_file) as _vf: + if _vf.read().strip().lower() == "0x8086": + has_intel_gpu = True + break + except OSError: + continue + elif is_windows: + _ps = shutil.which("powershell") or shutil.which("pwsh") + if _ps: + try: + _result = run_capture( + [ + _ps, + "-NoProfile", + "-Command", + "Get-CimInstance Win32_VideoController | " + "Select-Object -ExpandProperty Name", + ], + timeout = 15, + ) + if _result.returncode == 0 and "intel" in _result.stdout.lower(): + has_intel_gpu = True + except Exception: + pass + return HostInfo( system = system, machine = machine, @@ -3090,6 +3178,7 @@ def detect_host() -> HostInfo: has_physical_nvidia = has_physical_nvidia, has_usable_nvidia = has_usable_nvidia, has_rocm = has_rocm, + has_intel_gpu = has_intel_gpu, rocm_gfx_target = rocm_gfx_target, macos_version = macos_version, ) @@ -3126,6 +3215,7 @@ def _apply_host_overrides( has_physical_nvidia = False, has_rocm = False, rocm_gfx_target = None, + has_intel_gpu = False, ) gfx = _normalize_forwarded_gfx(override_rocm_gfx) if gfx: @@ -3866,6 +3956,23 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice "falling back to source build with HIP support" ) + # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. No + # physical NVIDIA (not just no usable one): a CUDA-hidden card must not + # be reached through Vulkan, which ignores CUDA_VISIBLE_DEVICES. + if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm: + vulkan_name = f"llama-{llama_tag}-bin-ubuntu-vulkan-x64.tar.gz" + if vulkan_name in upstream_assets: + log(f"Intel GPU detected -- using upstream Vulkan prebuilt {vulkan_name}") + return AssetChoice( + repo = UPSTREAM_REPO, + tag = llama_tag, + name = vulkan_name, + url = upstream_assets[vulkan_name], + source_label = "upstream", + install_kind = "linux-vulkan", + ) + log("Intel GPU detected but no Vulkan prebuilt found -- falling back to CPU") + upstream_name = f"llama-{llama_tag}-bin-ubuntu-x64.tar.gz" if upstream_name not in upstream_assets: raise PrebuiltFallback("upstream Linux CPU asset was not found") @@ -3908,6 +4015,24 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice ) log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU") + # Intel (or other non-NVIDIA/non-AMD) GPU on Windows: use Vulkan. No + # physical NVIDIA so a CUDA-hidden card isn't reached through Vulkan. + if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm: + vulkan_name = f"llama-{llama_tag}-bin-win-vulkan-x64.zip" + if vulkan_name in upstream_assets: + log( + f"Intel GPU detected on Windows -- using upstream Vulkan prebuilt {vulkan_name}" + ) + return AssetChoice( + repo = UPSTREAM_REPO, + tag = llama_tag, + name = vulkan_name, + url = upstream_assets[vulkan_name], + source_label = "upstream", + install_kind = "windows-vulkan", + ) + log("Intel GPU detected on Windows but no Vulkan prebuilt found -- falling back to CPU") + upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip" if upstream_name not in upstream_assets: raise PrebuiltFallback("upstream Windows CPU asset was not found") @@ -4503,6 +4628,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: "linux-arm64-cuda", "linux-rocm", "linux-arm64", + "linux-vulkan", }: return ["llama-server", "llama-quantize", "llama-diffusion-gemma-visual-server", "lib*.so*"] if choice.install_kind in {"macos-arm64", "macos-x64"}: @@ -4516,6 +4642,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: "windows-cpu", "windows-cuda", "windows-hip", + "windows-vulkan", "windows-rocm", "windows-arm64", }: @@ -5731,8 +5858,10 @@ def validate_server( "linux-cuda", "linux-arm64-cuda", "linux-rocm", + "linux-vulkan", "windows-cuda", "windows-hip", + "windows-vulkan", "windows-rocm", "macos-arm64", } @@ -6354,6 +6483,20 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: ["libmtmd.so*"], ["libggml-hip.so*"], ] + if choice.install_kind == "linux-vulkan": + return [ + ["libllama-common.so*"], + ["libllama.so*"], + ["libggml.so*"], + ["libggml-base.so*"], + # Match the sibling globs (linux-cuda/-rocm): x64 bundles ship + # arch-suffixed libggml-cpu-.so, arm64 may ship a bare + # libggml-cpu.so; the '-' form missed the latter and re-flagged + # the install unhealthy on every check. + ["libggml-cpu*.so*"], + ["libmtmd.so*"], + ["libggml-vulkan.so*"], + ] if choice.install_kind in {"windows-cpu", "windows-arm64"}: return [["llama.dll"]] if choice.install_kind == "windows-cuda": @@ -6373,6 +6516,8 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: return groups if choice.install_kind in {"windows-hip", "windows-rocm"}: return [["llama.dll"], ["*hip*.dll"]] + if choice.install_kind == "windows-vulkan": + return [["llama.dll"], ["ggml-vulkan.dll"]] return [] @@ -6654,6 +6799,89 @@ def validate_prebuilt_attempts( raise PrebuiltFallback("no prebuilt bundle passed validation") +def force_vulkan_requested() -> bool: + """Whether UNSLOTH_FORCE_VULKAN opts this host into the Vulkan llama.cpp + prebuilt instead of its detected CUDA/ROCm backend (e.g. so an AMD user can + run the Vulkan build for inference). Scoped to the llama.cpp backend; the + torch/training stack installs separately and still sees the real GPU. + """ + return os.environ.get("UNSLOTH_FORCE_VULKAN", "").strip().lower() in ( + "1", + "true", + "yes", + ) + + +def _vulkan_only_host(host: HostInfo) -> HostInfo: + """Rewrite ``host`` so the asset selectors take their Vulkan branch. + + That branch fires on ``has_intel_gpu and not nvidia and not rocm``, so clear + the CUDA/ROCm flags and raise the integrated-GPU flag. The synthetic flag + never leaves install planning -- it only routes the llama.cpp prebuilt + choice, not the torch/training stack. + """ + return dataclasses_replace( + host, + has_usable_nvidia = False, + has_physical_nvidia = False, + has_rocm = False, + has_intel_gpu = True, + ) + + +def _route_to_vulkan_prebuilt( + host: HostInfo, published_repo: str, published_release_tag: str, *, force_cpu: bool +) -> tuple[HostInfo, str, str]: + """Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt. + + The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes + from UPSTREAM_REPO. Two triggers route here, both suppressed under + --cpu-fallback (the explicit "give me CPU" last resort wins): + * UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend; + * an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose + of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset. + Applied by BOTH the install path and the --resolve-prebuilt probe so the + "is a prebuilt available" answer matches what actually gets installed. + + Returns the (possibly rewritten) host, repo, and release tag. + """ + forced = force_vulkan_requested() + # Gate auto-routing on no PHYSICAL NVIDIA, not merely no usable one: a mixed + # NVIDIA+Intel host that hides NVIDIA with CUDA_VISIBLE_DEVICES=""/-1 keeps + # has_physical_nvidia=True while has_usable_nvidia goes False. Vulkan ignores + # CUDA_VISIBLE_DEVICES, so auto-routing such a host would let it grab the + # reserved NVIDIA GPU. An explicit UNSLOTH_FORCE_VULKAN still overrides. + auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm + if force_cpu or not (forced or auto_intel): + return host, published_repo, published_release_tag + if host.is_macos: + if forced: + log( + "UNSLOTH_FORCE_VULKAN is set but ignored on macOS " + "(Metal is used; there is no Vulkan prebuilt)" + ) + return host, published_repo, published_release_tag + if forced: + log( + "UNSLOTH_FORCE_VULKAN is set; installing the upstream Vulkan " + "llama.cpp prebuilt instead of the detected GPU backend" + ) + # Forcing may override a detected NVIDIA/ROCm host, so normalize it to + # Vulkan-only; an auto-detected Intel host already is. + host = _vulkan_only_host(host) + else: + log("Intel GPU detected; installing the upstream Vulkan llama.cpp prebuilt") + # Swapping the fork for upstream invalidates a fork release pin: the two use + # different tag namespaces (fork b9596-mix- vs upstream b9596), so a + # pinned fork tag would make the upstream resolver query a nonexistent + # release and fall back to source. Drop it and let the upstream resolver + # pick by the requested llama tag. A pin already on an explicit upstream repo + # (repo unchanged here) is preserved. + if published_repo != UPSTREAM_REPO: + published_release_tag = "" + return host, UPSTREAM_REPO, published_release_tag + + def diffusion_visual_server_backfill_needed( install_dir: Path, host: HostInfo, choice: AssetChoice ) -> bool: @@ -6696,6 +6924,9 @@ def install_prebuilt( override_rocm_gfx = override_rocm_gfx, force_cpu = force_cpu, ) + host, published_repo, published_release_tag = _route_to_vulkan_prebuilt( + host, published_repo, published_release_tag, force_cpu = force_cpu + ) choice: AssetChoice | None = None try: with install_lock(install_lock_path(install_dir)): @@ -6708,7 +6939,9 @@ def install_prebuilt( f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install" ) # Single resolver: every fork host selects from the release manifest; - # an explicit ggml-org override selects by asset filename instead. + # an explicit ggml-org override selects by asset filename instead. A + # forced-Vulkan host already has published_repo pointed at + # UPSTREAM_REPO above, so the resolver takes the Vulkan asset branch. requested_tag, release_plans = resolve_simple_install_release_plans( llama_tag, host, @@ -6994,10 +7227,14 @@ def main() -> int: override_rocm_gfx = args.rocm_gfx, force_cpu = args.cpu_fallback, ) - repo = args.published_repo + # Same Vulkan routing the install path applies, so the probe's answer + # matches what would install (an Intel/forced-Vulkan host -> upstream). + host, repo, release_tag = _route_to_vulkan_prebuilt( + host, args.published_repo, args.published_release_tag or "", force_cpu = args.cpu_fallback + ) try: _requested, plans = resolve_simple_install_release_plans( - args.resolve_prebuilt, host, repo, args.published_release_tag or "" + args.resolve_prebuilt, host, repo, release_tag ) choice = plans[0].attempts[0] if plans and plans[0].attempts else None if choice is None: From 216a1fad33561ee4fcf24fd47811fbf721b46f29 Mon Sep 17 00:00:00 2001 From: alkinun Date: Thu, 9 Jul 2026 13:46:47 +0300 Subject: [PATCH 093/113] Fix Windows installer torch index override (#6972) * Fix Windows installer torch index override * Clear inherited uv index env vars for pinned installs in studio/setup.ps1 (#6898) * Harden setup.ps1 index-var clearing to truly remove vars (#6898) * Apply UV_DEFAULT_INDEX torch index fix to Linux/Mac install.sh (#6898) * Neutralize all uv index env vars for pinned torch installs (#6898) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.ps1 | 24 ++++++--- install.sh | 28 ++++++---- studio/setup.ps1 | 13 ++++- .../test_tokenizers_and_torch_constraint.py | 51 +++++++++++++++++++ tests/sh/test_mac_intel_compat.sh | 2 +- 5 files changed, 99 insertions(+), 19 deletions(-) diff --git a/install.ps1 b/install.ps1 index 696f4e613a..0797cd3868 100644 --- a/install.ps1 +++ b/install.ps1 @@ -469,6 +469,17 @@ function Install-UnslothStudio { param( [Parameter(Mandatory = $true)][ScriptBlock]$Command ) + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when the command pins an index, clear every uv index env var so + # it wins, then restore in finally. Other installs keep the user's mirror. + $savedUvIndex = $null + if ($Command.ToString() -match '--default-index') { + $savedUvIndex = @{} + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue + } + } $prevEap = $ErrorActionPreference $ErrorActionPreference = "Continue" try { @@ -488,6 +499,7 @@ function Install-UnslothStudio { return [int]$LASTEXITCODE } finally { $ErrorActionPreference = $prevEap + if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } } } } @@ -2200,7 +2212,7 @@ exit 0 # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { # Transient AMD-index failure: fall back to a CPU base so the install # still completes; Studio setup retries ROCm afterwards. @@ -2209,7 +2221,7 @@ exit 0 # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU # torch>= range, so without it uv would keep the ROCm build and only swap # the companions -- a mismatched venv the flavor-repair block won't fix. - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2223,7 +2235,7 @@ exit 0 } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2306,7 +2318,7 @@ exit 0 # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx* - # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install + # is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops. if (-not $SkipTorch) { $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl @@ -2322,7 +2334,7 @@ exit 0 $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) @@ -2331,7 +2343,7 @@ exit 0 } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) diff --git a/install.sh b/install.sh index 0acc9ec0be..3f4ea92387 100755 --- a/install.sh +++ b/install.sh @@ -159,6 +159,12 @@ run_maybe_quiet() { run_install_cmd() { _label="$1" shift + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when we pass --default-index, neutralize every uv index env var so + # the pinned index wins. Other installs keep the user's mirror. + case " $* " in + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;; + esac if _is_verbose; then "$@" && return 0 _rc=$? @@ -2190,9 +2196,9 @@ _expected_torch_flavor_tag() { esac } -# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX / +# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX / # rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv -# resolves (torch + every transitive dep) via --index-url -- the same URLs the +# resolves (torch + every transitive dep) via --default-index -- the same URLs the # fresh-install paths above already use -- so a stale wheel is auto-repairable. # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. _torch_index_repairable() { @@ -2744,7 +2750,7 @@ if [ "$_MIGRATED" = true ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall fi ;; @@ -2870,7 +2876,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." # Pass explicit wheel URLs so the matched trio is @@ -2893,18 +2899,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm @@ -2964,7 +2970,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall fi ;; @@ -2999,14 +3005,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") - # Repair when flavor is wrong AND the index is plain --index-url reinstallable + # Repair when flavor is wrong AND the index is plain --default-index reinstallable # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only. if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" @@ -3017,7 +3023,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" - substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 07dcb17335..db01a1ecad 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2621,7 +2621,18 @@ function Fast-Install { param([Parameter(ValueFromRemainingArguments=$true)]$Args_) if ($UseUv) { $VenvPy = (Get-Command python).Source - $result = & uv pip install --python $VenvPy @Args_ 2>&1 + # An explicit --index-url must win. Inherited uv index env vars otherwise + # override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop + # them only for index-pinned installs; mirrors still apply elsewhere. + $saved = @{} + if (@($Args_) -contains '--index-url') { + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + $saved[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue + } + } + try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 } + finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } } if ($LASTEXITCODE -eq 0) { return } } & python -m pip install @Args_ 2>&1 diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py index 7390d7be9b..4322f0c7d6 100644 --- a/tests/python/test_tokenizers_and_torch_constraint.py +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -14,6 +14,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/ _REPO_ROOT = _TESTS_DIR.parent # unsloth/ _INSTALL_SH = _REPO_ROOT / "install.sh" _INSTALL_PS1 = _REPO_ROOT / "install.ps1" +_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1" _NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt" @@ -109,6 +110,56 @@ class TestStructuralInstallPs1Unchanged: assert '"torch>=2.4,<2.11.0"' in self._ps1 +class TestInstallPs1UvDefaultIndex: + """Installer-managed torch indexes must override inherited uv defaults.""" + + _ps1 = _read(_INSTALL_PS1) + + def test_torch_installs_use_default_index(self): + assert "--default-index $TorchIndexUrl" in self._ps1 + assert "--default-index $ROCmIndexUrl" in self._ps1 + + def test_torch_installs_do_not_use_deprecated_index_url(self): + assert "--index-url $TorchIndexUrl" not in self._ps1 + assert "--index-url $ROCmIndexUrl" not in self._ps1 + + def test_torch_installs_neutralize_all_uv_index_env_vars(self): + # Extra-index vars outrank --default-index, so pinned installs must clear them. + for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"): + assert var in self._ps1 + assert 'Remove-Item "Env:$n"' in self._ps1 + + +class TestSetupPs1FastInstallIndex: + """setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning.""" + + _ps1 = _read(_SETUP_PS1) + + def test_fast_install_clears_all_uv_index_env_vars(self): + for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"): + assert var in self._ps1 + # Must truly remove the vars (child sees no value), not set them empty. + assert 'Remove-Item "Env:$n"' in self._ps1 + + +class TestInstallShUvDefaultIndex: + """Linux/Mac installer torch indexes must override inherited uv defaults.""" + + _sh = _read(_INSTALL_SH) + + def test_torch_installs_use_default_index(self): + assert '--default-index "$TORCH_INDEX_URL"' in self._sh + + def test_torch_installs_do_not_use_deprecated_index_url(self): + assert '--index-url "$TORCH_INDEX_URL"' not in self._sh + + def test_torch_installs_neutralize_all_uv_index_env_vars(self): + # --default-index installs run with all uv index env vars unset via `env -u`. + assert ( + "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh + ) + + # Group 2 -- Shell snippet tests (bash subprocess, mocked python) class TestTorchConstraintShell: """Test the TORCH_CONSTRAINT block via bash with mocked python minor versions.""" diff --git a/tests/sh/test_mac_intel_compat.sh b/tests/sh/test_mac_intel_compat.sh index 3c3bbfaa5f..8a0ff4b641 100644 --- a/tests/sh/test_mac_intel_compat.sh +++ b/tests/sh/test_mac_intel_compat.sh @@ -312,7 +312,7 @@ if [ "$SKIP_TORCH" = true ]; then else echo "==> Installing PyTorch ($TORCH_INDEX_URL)..." uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi TORCH_EOF From cd9d251f157bc8a014a68f4688b961344d5d02f6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 04:10:59 -0700 Subject: [PATCH 094/113] Fix fast inference crash on compressed-tensors FP8 models (#7025) * Fix fast_gemv crash on compressed-tensors FP8 models Loading a compressed-tensors FP8 checkpoint (for example unsloth/Llama-3.2-1B-Instruct-FP8-Block) with fast_inference=False and running a forward crashed with 'Parameter object has no attribute absmax' inside fast_gemv. A compressed-tensors CompressedLinear exposes an already dequantized bf16 weight at forward time while keeping a weight_scale Parameter. The quant state resolution in get_lora_parameters/get_lora_parameters_bias fell back to that weight_scale, so a bf16 weight was routed into the bitsandbytes fast_gemv/fast_dequantize path, which expects a bitsandbytes QuantState with an absmax attribute. Only fall back to weight_scale_inv/weight_scale when the weight is still fp8. A decompressed bf16 weight then resolves to no quant state and flows through the normal bf16 path, which already handles bias and the LoRA backward. Real fp8 and bitsandbytes 4bit weights are unchanged. * Skip the fast_gemv dispatch test before importing unsloth when bitsandbytes is absent --- tests/test_fast_gemv_dispatch.py | 63 ++++++++++++++++++++++++++++++++ unsloth/kernels/utils.py | 27 ++++++++++++-- 2 files changed, 86 insertions(+), 4 deletions(-) create mode 100644 tests/test_fast_gemv_dispatch.py diff --git a/tests/test_fast_gemv_dispatch.py b/tests/test_fast_gemv_dispatch.py new file mode 100644 index 0000000000..7758db2cd9 --- /dev/null +++ b/tests/test_fast_gemv_dispatch.py @@ -0,0 +1,63 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""`get_lora_parameters` must not treat a `weight_scale` as a quant state for a weight that is +already dequantized to bf16 (e.g. a compressed-tensors layer at forward time). Otherwise the +bnb fast_gemv / fast_dequantize path reads a missing `absmax` and crashes. +""" + +from types import SimpleNamespace + +import pytest +import torch + +# unsloth.kernels.utils imports bitsandbytes unconditionally, so skip the whole module up +# front on runners without it (e.g. CPU-only) before importing unsloth, otherwise collection +# errors instead of producing a skip. Any other import error still surfaces as a failure. +pytest.importorskip("bitsandbytes") + +import unsloth # noqa: F401 (sets UNSLOTH_IS_PRESENT before transformers) +from unsloth.kernels.utils import get_lora_parameters_bias, _FP8_WEIGHT_DTYPES + +_FP8 = _FP8_WEIGHT_DTYPES[0] if _FP8_WEIGHT_DTYPES else None + + +def _proj(weight, weight_scale = None): + proj = SimpleNamespace(weight = weight, bias = None, merged = False) + if weight_scale is not None: + proj.weight_scale = weight_scale + return proj + + +def test_bf16_weight_scale_not_used_as_quant_state(): + """A bf16 weight carrying a weight_scale (compressed-tensors) -> quant state must be None.""" + proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16), torch.rand(2, 2)) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is None + + +def test_fp8_weight_keeps_scale(): + """An actual fp8 weight still resolves its weight_scale as the quant state.""" + if _FP8 is None: + pytest.skip("no float8 dtype in this torch build") + scale = torch.rand(2, 2) + proj = _proj(torch.randn(4, 4).to(_FP8), scale) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is scale + + +def test_plain_bf16_has_no_quant_state(): + proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16)) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is None diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 43ed198a4a..1b0b5ce12e 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -282,6 +282,21 @@ def QUANT_STATE(W): return getattr(W, "quant_state", None) +# fp8 weight dtypes. A `weight_scale` / `weight_scale_inv` should only be treated as a +# quant state when the weight itself is still fp8. compressed-tensors layers expose an +# already-dequantized bf16 weight at forward time while keeping a `weight_scale` around; +# reading that as a quant state routes a bf16 weight into the bitsandbytes fast_gemv / +# fast_dequantize path, which then reads a missing `absmax` and crashes. +_FP8_WEIGHT_DTYPES = tuple( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + ) + if dtype is not None +) + + def get_lora_parameters(proj): """Return (weight, weight quant_state, lora A, lora B, lora scale). With QAT enabled, also fake-quantizes the base layer and lora weights. @@ -298,9 +313,11 @@ def get_lora_parameters(proj): if weight_fake_quantizer is not None: W = weight_fake_quantizer(W) - # Get quant state for 4bit or FP8 + # Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the + # weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer) + # must not carry a scale as its quant state or fast_gemv will crash on it. W_quant = getattr(W, "quant_state", None) - if W_quant is None: + if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES: W_quant = getattr(base_layer, "weight_scale_inv", None) if W_quant is None: W_quant = getattr(base_layer, "weight_scale", None) @@ -349,9 +366,11 @@ def get_lora_parameters_bias(proj): ) # (proj.base_layer if hasattr(proj, "base_layer") else proj) W = base_layer.weight - # Get quant state for 4bit or FP8 + # Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the + # weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer) + # must not carry a scale as its quant state or fast_gemv will crash on it. W_quant = getattr(W, "quant_state", None) - if W_quant is None: + if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES: W_quant = getattr(base_layer, "weight_scale_inv", None) if W_quant is None: W_quant = getattr(base_layer, "weight_scale", None) From 534c877d2136b47b0ceec25cc45900c7f7f15e6d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 04:20:41 -0700 Subject: [PATCH 095/113] Keep native RoPE scaling when extending context; carry rope_theta for linear (#7028) * Keep native RoPE scaling when extending context; carry rope_theta for linear When max_seq_length exceeds a model's native window, the loader overwrote the model's rope_scaling with linear scaling. For models that already ship a scaled RoPE (llama3/yarn/longrope) that is far worse for long context, and on transformers v5 the linear dict omitted rope_theta (v5 keeps it under rope_parameters), so the rotary base fell back to 10000 and broke past ~8K tokens. Keep the native scaling and just widen the window; only synthesize linear for plain-RoPE models, and carry rope_theta so v5 keeps the real base. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Only preserve native llama3 when extending context; keep linear fallback otherwise The patched attention constructor (patch_llama_rope_scaling) rebuilds only linear, llama3 and longrope and its longrope branch reads a top-level original_max_position_embeddings, so preserving yarn or a nested-only longrope config would raise during construction on transformers <= 4.47.1. Keep only llama3 native; yarn/longrope/other types fall back to the linear override, still carrying rope_theta. * Correct long-context extension comment to match llama3-only preservation --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/utils/test_rope_scaling_drift.py | 33 ++++++++++++++ unsloth/models/llama.py | 61 +++++++++++++++++--------- 2 files changed, 73 insertions(+), 21 deletions(-) diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index 7a738e236c..b2ec1e5a20 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -257,6 +257,39 @@ def test_recompute_helper_scales_on_cpu(): ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled." +def test_extended_rope_scaling_keeps_llama3_and_carries_theta(): + # Long-context extension keeps native llama3, but falls back to linear for every other + # type (the patched attention constructor only rebuilds linear/llama3/longrope), and the + # linear dict carries rope_theta so transformers v5 does not fall back to base 10000. + from types import SimpleNamespace + + from unsloth.models.llama import _extended_rope_scaling + + # llama3 model: keep native scaling, do not synthesize linear. + scaling, native = _extended_rope_scaling(_make_config(LLAMA3_ROPE_SCALING), 2.0) + assert ( + scaling is None and native == "llama3" + ), "must keep native llama3 scaling instead of overwriting it with linear." + + # yarn is not rebuildable by the patcher -> keep the safe linear fallback, not native. + yarn = SimpleNamespace(rope_scaling = {"rope_type": "yarn", "factor": 2.0}, rope_theta = 500000.0) + scaling, _ = _extended_rope_scaling(yarn, 2.0) + assert scaling == { + "type": "linear", + "factor": 2.0, + "rope_theta": 500000.0, + }, f"yarn must fall back to linear (patcher cannot rebuild it), got {scaling}." + + # plain RoPE with theta only under v5 rope_parameters: linear must carry rope_theta. + v5 = SimpleNamespace(rope_parameters = {"rope_type": "default", "rope_theta": 1000000.0}) + scaling, _ = _extended_rope_scaling(v5, 2.0) + assert scaling == { + "type": "linear", + "factor": 2.0, + "rope_theta": 1000000.0, + }, f"linear override dropped rope_theta on v5 (got {scaling}); base would fall back to 10000." + + def test_extended_rotary_reads_config_factor(): # LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8 # (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405). diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 1f43f61443..05523bc27b 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1651,6 +1651,26 @@ def _rope_scaling_as_dict(rope_scaling): return {} +def _extended_rope_scaling(config, factor): + """RoPE scaling to extend a model past its native window. Keeps native llama3 as-is + (linear extension is far worse for long context); everything else gets linear. Returns + (scaling_or_None, type): None keeps llama3. The linear dict carries rope_theta so + transformers v5 (which stores it under rope_parameters) keeps the real base, not 10000. + Only llama3 is preserved because patch_llama_rope_scaling can only rebuild linear/llama3/ + longrope and its longrope branch needs a top-level original_max_position_embeddings.""" + existing = _rope_scaling_as_dict( + getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {} + ) + existing_type = existing.get("rope_type") or existing.get("type") + if existing_type == "llama3": + return None, existing_type + return { + "type": "linear", + "factor": factor, + "rope_theta": _get_rope_theta(config), + }, existing_type + + def _llama3_inv_freq_from_config( config, rope_scaling, @@ -2518,34 +2538,33 @@ class FastLlamaModel: max_seq_length = model_max_seq_length if (rope_scaling is None) and (max_seq_length > model_max_seq_length): - rope_scaling = max_seq_length / model_max_seq_length + factor = max_seq_length / model_max_seq_length if fast_inference: raise NotImplementedError( "Unsloth: Fast inference does not yet work with RoPE Scaling." ) - logger.warning_once( - f"Unsloth: {model_name} can only handle sequence lengths of at most " - f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of " - f"{round(rope_scaling, 3)}, it can be magically be extended to " - f"{max_seq_length}!" - ) - - # Warn RoPE scaling isn't allowed - if not has_rope_scaling: - raise RuntimeError( - f"However, {model_name} doesn't support RoPE Scaling!\n" - "Please file a feature request at https://github.com/unslothai/unsloth." + linear_scaling, native_type = _extended_rope_scaling(model_config, factor) + if linear_scaling is not None: + logger.warning_once( + f"Unsloth: {model_name} can only handle sequence lengths of at most " + f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of " + f"{round(factor, 3)}, it can be magically be extended to " + f"{max_seq_length}!" + ) + if not has_rope_scaling: + raise RuntimeError( + f"However, {model_name} doesn't support RoPE Scaling!\n" + "Please file a feature request at https://github.com/unslothai/unsloth." + ) + kwargs["rope_scaling"] = linear_scaling + else: + # Native llama3 scaling already handles long context; just widen the window. + logger.warning_once( + f"Unsloth: extending {model_name} to {max_seq_length} using its native " + f"{native_type} RoPE scaling." ) - - rope_scaling = { - "type": "linear", - "factor": rope_scaling, - } - - # Add to kwargs - kwargs["rope_scaling"] = rope_scaling from .loader_utils import ( check_and_disable_bitsandbytes_loading, From b5dca66cb1480b36ef738a4e62680e02cca2a65f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 04:52:30 -0700 Subject: [PATCH 096/113] scripts: refresh scan_packages allowlist baseline (#7032) * scripts: refresh scan_packages allowlist baseline Regenerate scripts/scan_packages_baseline.json against the current resolved dependency set so the blocking pip scan-packages gate matches what the scanner now finds. Refreshes evidence hashes for benign findings whose code shifted lines (unsloth-zoo mlx loader, gguf/mlx test /tmp fixtures) and adds two mainstream-library entries that were newly surfaced (torch inductor codecache base64+subprocess compile cache, torch testing common_utils socket import). Stale entries whose matching code changed and no longer triggers are dropped. All entries remain CRITICAL/HIGH findings manually judged benign; matched on (package, file, check, evidence_hash). * ci(security-audit): re-run scan when the allowlist baseline changes The security-audit pull_request trigger listed the scanners but not their allowlist baselines, so a baseline-only edit never re-ran the scan that consumes it. A refreshed baseline could therefore merge without CI confirming its evidence hashes match what the scanner finds. Add scan_packages_baseline.json and scan_npm_packages_baseline.json to the paths filter so baseline changes are validated on their own PR. --- .github/workflows/security-audit.yml | 6 +- scripts/scan_packages_baseline.json | 304 ++++++++++++--------------- 2 files changed, 140 insertions(+), 170 deletions(-) diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 0ef2ad1e9d..1275d12216 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -2,8 +2,8 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # Multi-language supply-chain audit. Triggers: -# - PRs touching any dependency manifest (Python / npm / Cargo) or -# this workflow file, +# - PRs touching any dependency manifest (Python / npm / Cargo), a +# scanner or its allowlist baseline, or this workflow file, # - push to main / pip, # - nightly @ 04:13 UTC so newly-published advisories surface even # when no PR opens, @@ -57,7 +57,9 @@ on: - 'studio/src-tauri/Cargo.lock' - 'pyproject.toml' - 'scripts/scan_packages.py' + - 'scripts/scan_packages_baseline.json' - 'scripts/scan_npm_packages.py' + - 'scripts/scan_npm_packages_baseline.json' - '.github/workflows/security-audit.yml' push: branches: [main, pip] diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 1d34cfb66d..3582517d31 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -39,7 +39,7 @@ "file": "botocore/utils.py", "check": "Reads credential paths AND makes network calls", "severity": "CRITICAL", - "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3719: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", "evidence_hash": "2d691bc373ab872aad23c744104596ba6d0d9f3b35aa101c7edbff4429b174c1" }, { @@ -55,23 +55,23 @@ "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", - "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" + "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", + "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" }, { "package": "datasets", "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", - "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" + "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", + "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" }, { "package": "diffusers", "file": "diffusers/utils/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence": "L1052: return importlib.import_module(\".\" + module_name, self.__name__)", "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { @@ -79,7 +79,7 @@ "file": "diffusers/utils/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L233: value = os.environ[key]\nNetwork: L688: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L709: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L728: image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", + "evidence": "Env: L236: value = os.environ[key]\nNetwork: L691: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L712: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L731: image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", "evidence_hash": "671190a6106c6ee9674e5e5942dc0940e1d2f8c78d5faf674413c2345b783fd9" }, { @@ -90,12 +90,20 @@ "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()", "evidence_hash": "894862e547cf91b90cd6e4b495db3fb05b7490ef0d63de7e795a7e3d9447d850" }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", + "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" + }, { "package": "fastmcp-slim", "file": "fastmcp/cli/apps_dev.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence": "Archive: L1353: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", "evidence_hash": "73a7a72013e9f800627ea07e6dbc3beeb8c905a6a5480c8fd896f0063173d25c" }, { @@ -103,8 +111,8 @@ "file": "fastmcp/cli/apps_dev.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L624: history.replaceState(null, \"\", url); sha256:fd8dbfa8af4dea2ce43f4d441f3f81239de341b76a2eb0a33c446f6757ce5f43\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", - "evidence_hash": "6ada4a9111213bdee5ea24c70a72ec4acdc8ffe0de4a01fd9835bc261ccab8f8" + "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" }, { "package": "fonttools", @@ -132,19 +140,35 @@ }, { "package": "huggingface-hub", - "file": "huggingface_hub/hf_api.py", + "file": "huggingface_hub/_sandbox.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", - "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" + "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", + "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/_sandbox.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", + "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272", - "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602" + "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", + "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", + "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" }, { "package": "huggingface-hub", @@ -159,8 +183,8 @@ "file": "huggingface_hub/utils/_http.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L443: while True: sha256:0ab4fed32d3af10f361963371f681923481377508a405b5d8770cef75f859168", - "evidence_hash": "1484f6b92f41c427ba8cbc7c4695a94975fea683dfa83aa510b4b0e982be4721" + "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", + "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" }, { "package": "huggingface-hub", @@ -218,6 +242,22 @@ "evidence": "L5: import socket sha256:915068303029fa5806199f256fb74504c65f253f9aee8ea23d8e384bb772b1c7", "evidence_hash": "30be130f165f418dfd37b144c5ae333de184b95f828ab8bd4010a67b84a5f814" }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", + "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", + "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" + }, { "package": "numba", "file": "numba/pycc/decorators.py", @@ -231,7 +271,7 @@ "file": "numba/tests/support.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)", + "evidence": "L1016: os.dup2(w, fd) | L1021: os.dup2(save, fd)", "evidence_hash": "fea7aa03d48bf0f4386302fa444984c4f5dfc772cfec3f1df199fd33a52eec10" }, { @@ -495,16 +535,16 @@ "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", - "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" + "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", + "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" }, { "package": "scikit-learn", "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", - "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" + "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", + "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" }, { "package": "scikit-learn", @@ -642,6 +682,14 @@ "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run(\nL2693: cmd.split(), capture_output=True, text=True, check=True\nL2694: ) | L2995: cmd_output = subprocess.run(\nL2996: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL2997: ) | L3707: out = subprocess.check_output(\nL3708: [\"ldd\", os.path.join(search, file)]\nL3709: ) | L3791: jobs.append(functools.partial(subprocess.check_call, cmd)) | L3876: subprocess.check_call(\nL3877: shlex.split(halide_cmd_gen.get_command_line())\nL3878: ) | L4336: subprocess.check_output(\nL4337: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4338: ) | L4591: output = subprocess.check_output(\nL4592: cmd_parts,\nL4593: stderr=subprocess.STDOUT,\nL4594: text=True,\nL4595: env=os.environ,\nL4596: )", "evidence_hash": "c09774087b702a6c5d6e2e85d9239c7c241ec938fbe9c0153e8f0b5c0710389b" }, + { + "package": "torch", + "file": "torch/_inductor/codecache.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1727: content = base64.b64decode(data)\nSubprocess: L3270: subprocess.run(\nL3271: cmd, capture_output=True, text=True, check=True\nL3272: ) | L3583: cmd_output = subprocess.run(\nL3584: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL3585: ) | L4338: out = subprocess.check_output(\nL4339: [\"ldd\", os.path.join(search, file)]\nL4340: ) | L4422: jobs.append(functools.partial(subprocess.check_call, cmd)) | L4507: subprocess.check_call(\nL4508: shlex.split(halide_cmd_gen.get_command_line())\nL4509: ) | L4992: subprocess.check_output(\nL4993: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4994: ) | L5247: output = subprocess.check_output(\nL5248: cmd_parts,\nL5249: stderr=subprocess.STDOUT,\nL5250: text=True,\nL5251: env=os.environ,\nL5252: )", + "evidence_hash": "87f77b5f51cb84fe9950fdeeb90fe8710e1b863100e90b5e2cfb228a725bee06" + }, { "package": "torch", "file": "torch/ao/__init__.py", @@ -695,7 +743,7 @@ "file": "torch/testing/_internal/common_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", + "evidence": "Env: L4900: env = os.environ.copy()\nNetwork: L4962: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4980: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", "evidence_hash": "704a851b9d68c9b885b9e15538bd7e96f03875503b618fe6f126c4438edd7386" }, { @@ -706,6 +754,14 @@ "evidence": "L32: import socket sha256:89faaaa8bc908e02dad73fd59b2b481fa91189c84b39b556c2766e71d2783bf3", "evidence_hash": "3d23d77ace91812a07cb9508cf352185d154176e8e8c8b9b28fa92cdbcfe0d53" }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L32: import socket sha256:ba439cbf568b194872f1d974c02b0487e51f677b67e379400522d0992600bd2d", + "evidence_hash": "88e98b227573997f86eedea8e885a407b0dd549d46d4a3f0b840ec5aafe66865" + }, { "package": "torchvision", "file": "torchvision/datasets/utils.py", @@ -743,8 +799,8 @@ "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1663: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", - "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" + "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", + "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" }, { "package": "transformers", @@ -759,15 +815,15 @@ "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", - "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" + "evidence": "L1699: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", + "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L284: value = os.environ[key] | L300: value = os.environ[key] | L2129: env = os.environ.copy() | L2251: for k in list(os.environ.keys()):\nNetwork: L2561: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", + "evidence": "Env: L288: value = os.environ[key] | L304: value = os.environ[key] | L2165: env = os.environ.copy() | L2287: for k in list(os.environ.keys()):\nNetwork: L2597: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", "evidence_hash": "73ff16aee09cf163fb3a7a04dfa2cf610595bde2f19460a579397695f728e3f4" }, { @@ -799,16 +855,16 @@ "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", - "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" + "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", + "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" }, { "package": "trl", "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", - "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" + "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", + "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" }, { "package": "trl", @@ -866,6 +922,14 @@ "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", "evidence_hash": "278ff15b0b702d37d7f0b30a1e55a31bf2b11883685718a47478fbb5ce7f5212" }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", + "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" + }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", @@ -919,16 +983,16 @@ "file": "tests/test_mlx_save_export_regressions.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L164: temporary_location=\"/tmp/ignored\", sha256:78837e80d48e872ef191aaacfe5e1c621a98a20df486a70a41d1a932d074a5b3", - "evidence_hash": "dd11376e664d0d7e7f4cc4baf57eacd4b7ae7b03222dce3912ce68b63dbfca1e" + "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:9f8502377b19666288b28399633dfc6740a64d0cb70ad1615e38b1269f94bf37", + "evidence_hash": "b7262d6e58f2ebad961dd3e64ca6c32bba356b5044d7a642d7dbd36a58cb6c81" }, { "package": "unsloth-zoo", "file": "tests/test_quantize_gguf_q2_k_l.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:32532cadc357beee1009f4e86481bdbe60a0b7bf47f6bb022b05ec1b8e15aed0", - "evidence_hash": "49f5b67379de17178f21a9bc93b79d6b94a70ecbdd16de86574934aac30a071d" + "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:06789b55e8f31426c233f37ff7d3729cc9e1f61c0829abd2c00c39216c63c7ad", + "evidence_hash": "ad4913d9099eb9b70e09d6860b242eb5f48c67e46d9bf4ae35c1c38a267d753b" }, { "package": "unsloth-zoo", @@ -951,7 +1015,7 @@ "file": "unsloth_zoo/llama_cpp.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2862: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", "evidence_hash": "b9f3b1652349fa8ef9ac2d1715978aca1e1632165851a00a2698dd47189e410c" }, { @@ -959,7 +1023,7 @@ "file": "unsloth_zoo/llama_cpp.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2862: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", "evidence_hash": "9cd0b1bb59c7eb1d814d7636dfd167c34f265eb7c4521a9d88b2bdcfd535b926" }, { @@ -1002,6 +1066,14 @@ "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" }, + { + "package": "cffi", + "file": "cffi/_cffi_gen_src.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", + "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" + }, { "package": "cffi", "file": "cffi/setuptools_ext.py", @@ -1127,7 +1199,7 @@ "file": "numba/tests/support.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)", + "evidence": "Obfusc: L874: __import__(modname)\nExec: L808: eval(co, globs, ns)", "evidence_hash": "649a7d750f903478243b0bcb9e8020521b505fc7fedc5b696ec01f4efc096109" }, { @@ -1159,7 +1231,7 @@ "file": "numba/tests/test_np_functions.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)", + "evidence": "Obfusc: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)", "evidence_hash": "9e81164131d16056fb56ad3cd11b8d129d1ff4f5855031e8b501e0335d5c14ed" }, { @@ -1175,16 +1247,16 @@ "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", - "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" + "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", + "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" }, { "package": "numpy", "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", - "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" + "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", + "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" }, { "package": "numpy", @@ -1199,7 +1271,7 @@ "file": "PIL/Image.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:", + "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3776: def eval(image: Image, *args: Callable[[int], float]) -> Image:", "evidence_hash": "c2c1e7ae44e15862caf8de549d09db7b35e93282450f07ef61aaf5450a408c13" }, { @@ -1255,7 +1327,7 @@ "file": "setuptools/_distutils/compilers/C/base.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):", + "evidence": "Obfusc: L1287: __import__(module_name)\nExec: L1114: if lib_type not in eval(expected):", "evidence_hash": "368651e9818ed2d1bb009027d3bcfbf94ae30639c0882a6c2bddde97b8c4f1e5" }, { @@ -1271,7 +1343,7 @@ "file": "setuptools/tests/config/test_pyprojecttoml.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", + "evidence": "Obfusc: L387: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", "evidence_hash": "067d41014f72a61d8b4adf25f3659d1f66a0e909f732223f48837aa7684df4e6" }, { @@ -1279,7 +1351,7 @@ "file": "setuptools/tests/test_editable_install.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)", + "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L447: exec(finder, loc, loc)", "evidence_hash": "a78d7f5af7eb4ba92656cda258c195b92f6337c585c97d0823e47a9d4a2eb15d" }, { @@ -1322,12 +1394,20 @@ "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace) | L909: exec(ln, {}, namespace) | L920: exec(c, namespace, funclocals)", "evidence_hash": "ab4f5819576a70038301668b8f3e4a781c4b757b146117d5d93eab1896a5a6cd" }, + { + "package": "tensorboard", + "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", + "check": "Python wheel ships large JS bundle (uncommon; manually review)", + "severity": "HIGH", + "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", + "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" + }, { "package": "torch", "file": "torch/_dynamo/bytecode_debugger.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)\nExec: L683: result = eval(arg, frame_globals, eval_locals) | L708: result = eval(cmd, frame_globals, eval_locals) | L716: exec(cmd, frame_globals, eval_locals)", + "evidence": "Anti: L1052: self._old_trace = sys.gettrace() | L1053: sys.settrace(self._settrace_callback) | L1113: sys.settrace(self._old_trace)\nExec: L684: result = eval(arg, frame_globals, eval_locals) | L709: result = eval(cmd, frame_globals, eval_locals) | L717: exec(cmd, frame_globals, eval_locals)", "evidence_hash": "dc2afd1769d357c15b69802bd2799fafa059c0b1dcdd4937528fb5b601962f1b" }, { @@ -1343,7 +1423,7 @@ "file": "torch/fx/experimental/rewriter.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)", + "evidence": "Obfusc: L44: code = compile(dest_ast, \"\", \"exec\")\nExec: L47: exec(code, globals_dict)", "evidence_hash": "76374f96feed416eec390458843621f33524cfb8d93ef0f3eb4cb1b47d0ad748" }, { @@ -1359,7 +1439,7 @@ "file": "torch/package/package_importer.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", + "evidence": "Obfusc: L599: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", "evidence_hash": "c7c0650f0c74a086d224112f77ee76634b8f47afc047ce27fee8c7fc45560512" }, { @@ -1391,7 +1471,7 @@ "file": "tests/test_mlx_trainer_internals.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L430: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L408: def eval(self):", + "evidence": "Obfusc: L1158: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L1136: def eval(self):", "evidence_hash": "c409327ef6420cc0c7224506fcb82b11bbc9838a6f2f97c9c2cfc00a40c4cdbf" }, { @@ -1407,7 +1487,7 @@ "file": "unsloth_zoo/compiler.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4294: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4294: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3385: exec(f\"import {model_location}\", globals()) | L3388: modeling_file = eval(model_location) | L3401: exec(\nL3402: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3403: ) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3407: globals(),\nL3408: locals(),\nL3409: ) | L3560: source = eval(f\"modeling_file.{module}\") | L3574: source = eval(f\"modeling_file.{module}\") | L3675: source = eval(f\"modeling_file.{module}\") | L3713: source = eval(f\"{model_location}.{module}\") | L3784: source = eval(f\"{model_location}.{module}\") | L3832: source = eval(f\"{model_location}.{module}\") | L4054: source = eval(f\"{model_location}.{module}\") | L4065: exec(\nL4066: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4067: globals(),\nL4068: ) | L4131: source = eval(f\"{model_location}.{module}\") | L4172: module_cls = eval(f\"{model_location}.{module}\") | L4209: module_cls = eval(f\"{model_location}.{module}\") | L4276: exec(\nL4277: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4278: globals(),\nL4279: ) | L4341: exec(inner_training_loop, globals()) | L4349: function = eval(f\"{model_location}.{module}\") | L4427: function = eval(f\"{model_location}.{module}\") | L4562: source = eval(f\"{model_location}.torch\") | L4569: function = eval(f\"source.nn.{module}\") | L4628: exec(\nL4629: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4630: globals(),\nL4631: locals(),\nL4632: ) | L4634: exec(\nL4635: f\"{model_location}.nn.{module}.forward = forward\",\nL4636: globals(),\nL4637: locals(),\nL4638: ) | L4642: exec(\nL4643: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4644: globals(),\nL4645: locals(),\nL4646: ) | L4648: exec(\nL4649: f\"combined_module.nn.{module}.forward = forward\",\nL4650: globals(),\nL4651: locals(),\nL4652: ) | L4669: exec(\nL4670: f\"{model_location}.{module} = combined_module.{module}\",\nL4671: globals(),\nL4672: locals(),\nL4673: ) | L4683: check_dicts = dir(eval(f\"{model_location}\")) | L4685: item = eval(f\"{model_location}.{check}\") | L4695: exec(\nL4696: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4697: globals(),\nL4698: locals(),\nL4699: )", + "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3389: exec(f\"import {model_location}\", globals()) | L3392: modeling_file = eval(model_location) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3407: ) | L3409: exec(\nL3410: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3411: globals(),\nL3412: locals(),\nL3413: ) | L3564: source = eval(f\"modeling_file.{module}\") | L3578: source = eval(f\"modeling_file.{module}\") | L3679: source = eval(f\"modeling_file.{module}\") | L3717: source = eval(f\"{model_location}.{module}\") | L3788: source = eval(f\"{model_location}.{module}\") | L3836: source = eval(f\"{model_location}.{module}\") | L4058: source = eval(f\"{model_location}.{module}\") | L4069: exec(\nL4070: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4071: globals(),\nL4072: ) | L4135: source = eval(f\"{model_location}.{module}\") | L4176: module_cls = eval(f\"{model_location}.{module}\") | L4213: module_cls = eval(f\"{model_location}.{module}\") | L4280: exec(\nL4281: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4282: globals(),\nL4283: ) | L4345: exec(inner_training_loop, globals()) | L4353: function = eval(f\"{model_location}.{module}\") | L4431: function = eval(f\"{model_location}.{module}\") | L4566: source = eval(f\"{model_location}.torch\") | L4573: function = eval(f\"source.nn.{module}\") | L4632: exec(\nL4633: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4634: globals(),\nL4635: locals(),\nL4636: ) | L4638: exec(\nL4639: f\"{model_location}.nn.{module}.forward = forward\",\nL4640: globals(),\nL4641: locals(),\nL4642: ) | L4646: exec(\nL4647: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4648: globals(),\nL4649: locals(),\nL4650: ) | L4652: exec(\nL4653: f\"combined_module.nn.{module}.forward = forward\",\nL4654: globals(),\nL4655: locals(),\nL4656: ) | L4673: exec(\nL4674: f\"{model_location}.{module} = combined_module.{module}\",\nL4675: globals(),\nL4676: locals(),\nL4677: ) | L4687: check_dicts = dir(eval(f\"{model_location}\")) | L4689: item = eval(f\"{model_location}.{check}\") | L4699: exec(\nL4700: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4701: globals(),\nL4702: locals(),\nL4703: )", "evidence_hash": "ec1875fd32d00fe885e566ebda75163e46e838ca31020abb57e0991892c2bdf7" }, { @@ -1423,8 +1503,8 @@ "file": "unsloth_zoo/mlx/loader.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L2218: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L140: mx.eval(model.parameters()) | L176: mx.eval(model.parameters()) | L2022: model.eval() | L2605: mx.eval(model.parameters()) | L2721: mx.eval(module.weight) | L4030: mx.eval(model.parameters()) | L4058: mx.eval(model.parameters()) | L4178: mx.eval(model.parameters())", - "evidence_hash": "9b29dade82912216c8b4808aa293b79749aa80ef1d2be35edd93bec7632810f1" + "evidence": "Obfusc: L2869: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L148: mx.eval(model.parameters()) | L180: mx.eval(model.parameters()) | L732: mx.eval(model.parameters()) | L733: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L799: mx.eval(model.parameters()) | L802: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L2673: model.eval() | L3256: mx.eval(model.parameters()) | L3372: mx.eval(module.weight) | L5666: mx.eval(model.parameters()) | L5716: mx.eval(model.parameters()) | L5859: mx.eval(model.parameters())", + "evidence_hash": "7b44760032c5df6d379ccfdd0bff3d23f857f64e08210fa0fba8d2881d457634" }, { "package": "unsloth-zoo", @@ -1439,7 +1519,7 @@ "file": "unsloth_zoo/saving_utils.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L3241: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3123: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3169: exec(save_pretrained, globals(), functions)", + "evidence": "Obfusc: L4015: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3897: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3943: exec(save_pretrained, globals(), functions)", "evidence_hash": "530b2383acd9fe8330aa65cd0bf86164aaacd47770e7c8d0752195bee36396ec" }, { @@ -1449,118 +1529,6 @@ "severity": "HIGH", "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)", "evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", - "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" - }, - { - "package": "unsloth-zoo", - "file": "scripts/scan_packages.py", - "check": "Writes to /tmp and executes (staged dropper)", - "severity": "CRITICAL", - "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", - "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", - "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" - }, - { - "package": "fastapi", - "file": "fastapi/routing.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L586: while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe", - "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb" - }, - { - "package": "tensorboard", - "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", - "check": "Python wheel ships large JS bundle (uncommon; manually review)", - "severity": "HIGH", - "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", - "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" - }, - { - "package": "fastapi", - "file": "fastapi/routing.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", - "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" - }, - { - "package": "fastmcp-slim", - "file": "fastmcp/cli/apps_dev.py", - "check": "Enumerates filesystem AND makes network calls", - "severity": "CRITICAL", - "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", - "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/_sandbox.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", - "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/_sandbox.py", - "check": "Writes to /tmp and executes (staged dropper)", - "severity": "CRITICAL", - "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", - "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/hf_api.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", - "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/utils/_http.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", - "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" - }, - { - "package": "cffi", - "file": "cffi/_cffi_gen_src.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", - "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" - }, - { - "package": "multiprocess", - "file": "multiprocess/forkserver.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L6: import socket sha256:6c707119169286c9a798e2c8d13a48614e481d8a503950916fd4ffb4c94d3182", - "evidence_hash": "50fec0f0522a8e4e636bf348b752002d7935d8455af31fb78c6f11e2eba19f6d" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3569: os.dup2(conn.fileno(), i) | L3601: \"test needs os.dup2()\") | L3619: os.dup2(fd, newfd) | L20: import socket sha256:c824dc0f409f242420c3fbb324790c53cb3078d2c8b07ee8f2a05694b01c2946", - "evidence_hash": "3878a2b430c175dbc5877a95195bfe52f9588ff73fb74e2261ed5e33087915ad" } ] } From fb5dc91bb4f33f8a4c5a41c89cbe88c93394e970 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 05:09:16 -0700 Subject: [PATCH 097/113] Studio: remove dead direct_linux_release_plan path (#7030) parse_direct_linux_release_bundle and direct_linux_release_plan are no longer reached by any live code path. Fork Linux installs resolve through _fork_manifest_release_plans -> _linux_published_attempts, and the upstream (ggml-org) path uses direct_upstream_release_plan. The dead parser also called _resolve_linux_bundle_profile, which no longer exists, so its CUDA branch would raise NameError if ever executed. Drop both functions and the obsolete TestDirectLinuxNvidiaCpuGate; its live equivalent TestLinuxPublishedAttemptsNvidiaCpuGate already covers the NVIDIA no-silent-CPU behaviour. --- studio/install_llama_prebuilt.py | 156 ------------------- tests/studio/install/test_selection_logic.py | 62 -------- 2 files changed, 218 deletions(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 856ba71478..40caebc040 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -1286,162 +1286,6 @@ def synthetic_checksums_for_release( ) -def parse_direct_linux_release_bundle( - repo: str, release: dict[str, Any] -) -> PublishedReleaseBundle | None: - release_tag = release.get("tag_name") - if not isinstance(release_tag, str) or not release_tag: - return None - - assets = release_asset_map(release) - artifacts: list[PublishedLlamaArtifact] = [] - inferred_labels: list[str] = [] - - linux_asset_re = re.compile( - r"^app-(?P