From 7f456352801664c7b60a978a5541449f000a0c14 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 01:27:27 -0700 Subject: [PATCH 01/49] Studio: auto-shut-down an exposed first-run instance if the admin password is never changed (#6651) * Studio: set the admin password before exposing it on the network On first run Studio seeds the default `unsloth` admin with a random bootstrap password and embeds it into index.html (window.__UNSLOTH_BOOTSTRAP__) so the local user can change it without typing it. A request with no Origin header counts as same-origin, which is what a normal top-level GET sends, so the page hands out the password to whoever loads it. That is harmless on the default 127.0.0.1 bind, but `--secure` (public Cloudflare tunnel) and `--host 0.0.0.0` (raw port reachable on the network) would serve the plaintext admin password to remote visitors during the bootstrap window. Fix this at the source: when launching a network-exposed web UI, prompt the operator in the terminal for a real admin password (with confirmation) before the socket binds or the tunnel opens, and persist it via update_password (which clears must_change_password and deletes the .bootstrap_password file). After that there is no bootstrap secret to leak. Non-interactive launches can supply it via UNSLOTH_STUDIO_ADMIN_PASSWORD. The masked reader echoes '*' per character and works on Linux, macOS, and Windows (PowerShell/cmd). Loopback binds, --api-only (no web UI), and Colab are unaffected. As defense in depth, the index handler now embeds the bootstrap object only for a direct local navigation: same-origin AND a loopback TCP peer with no proxy/tunnel forwarding headers (cf-ray, cf-connecting-ip, x-forwarded-for, x-forwarded-host, x-real-ip, forwarded). Colab stays exempt. This keeps the password off the wire even when the prompt is skipped (no TTY and no env var). Adds unit coverage for the prompt/confirm/decision logic, an integration test that provisioning clears the bootstrap state, and regression tests for the local-direct gate (loopback/IPv6/mapped/localhost peers, LAN/public peers, missing client, each forwarding header, spoofed XFF, and the Colab exemption). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fail fast on an explicitly empty admin-password env var resolve_admin_password_source treated UNSLOTH_STUDIO_ADMIN_PASSWORD="" like the var was unset and fell back to the bootstrap backstop. Treat any set value (including empty) as the env source so it reaches the minimum-length guard and refuses to expose the server instead of silently keeping the seeded password. * Studio: apply repo kwarg-spacing format to the secure-admin-password files * Studio: drop the pre-exposure password prompt; keep the local-direct gate Per review, the blocking prompt added friction for --secure / 0.0.0.0 first-run launches without extra security: the local-direct injection gate in main.py already keeps the bootstrap password off the network for any remote request. Remove the prompt module and its tests; the gate plus the existing must_change_password first-login flow are the fix. * Studio: shut down an exposed first-run instance if the admin password is never changed The local-direct gate keeps the seeded bootstrap password off the network, but it stays a valid credential until first login changes it. For an exposed web UI (--secure / 0.0.0.0, not --api-only, not Colab), arm a daemon timer: if the password is still the seeded one after the deadline (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT, default 3600s, 0 disables), print a message and shut Studio down via the existing graceful-shutdown path; if it was changed, leave Studio running. * Studio: revert the local-direct injection gate; keep the 1-hour auto-shutdown Per maintainer decision, keep the first-run auto-fill behavior unchanged (the bootstrap password still seeds the login form for convenience) and rely on the exposed-instance auto-shutdown to bound the window: an exposed web UI that never changes the seeded admin password is torn down after UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT (default 1h). Restores studio/backend/main.py and its origin test to upstream. * Studio: render the bootstrap-timeout shutdown message with a human duration The message hardcoded 'minute(s)' via timeout//60, so a sub-minute timeout (e.g. a 30s test value) printed 'within 1 minute(s)'. Add _format_duration so it reads '30 seconds' / '1 minute 30 seconds' / '60 minutes' as appropriate. The default 3600s still renders '60 minutes'. * Studio: drop stale local-direct gate reference from bootstrap_timeout docstring The gate was reverted (timer-only), so the module docstring should not describe a main.py gate that no longer exists. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/auth/bootstrap_timeout.py | 145 ++++++++++++++ studio/backend/run.py | 37 ++++ .../backend/tests/test_bootstrap_timeout.py | 185 ++++++++++++++++++ 3 files changed, 367 insertions(+) create mode 100644 studio/backend/auth/bootstrap_timeout.py create mode 100644 studio/backend/tests/test_bootstrap_timeout.py diff --git a/studio/backend/auth/bootstrap_timeout.py b/studio/backend/auth/bootstrap_timeout.py new file mode 100644 index 0000000000..728433dc54 --- /dev/null +++ b/studio/backend/auth/bootstrap_timeout.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Auto-shutdown for an exposed first-run Studio whose admin password is unchanged. + +On a fresh install the seeded bootstrap admin password stays a valid login +credential until first login changes it. When the web UI is put on the network +(``--secure`` / ``0.0.0.0``) and nobody completes that first-login change within +a deadline, tear Studio down so a fresh, unconfigured instance does not stay +publicly reachable indefinitely. If the password was changed, Studio keeps +running. + +Scope: web UI launches only (never ``--api-only``, which authenticates by API +key rather than the admin password, and never Colab). Configurable via +``UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT`` (seconds; default 3600; ``0`` disables). +""" + +import os +import sys +import threading + +BOOTSTRAP_TIMEOUT_ENV_VAR = "UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT" +DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS = 3600 + + +def bootstrap_timeout_seconds(env = None) -> int: + """Resolve the deadline in seconds. ``0`` (or invalid/negative) disables it. + + A malformed value falls back to the default rather than disabling, so a typo + cannot silently remove the protection. + """ + env = os.environ if env is None else env + raw = env.get(BOOTSTRAP_TIMEOUT_ENV_VAR) + if raw is None or raw.strip() == "": + return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + try: + value = int(raw) + except ValueError: + return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + return value if value > 0 else 0 + + +def _is_exposed_bind(host: str, secure: bool) -> bool: + """True when this launch puts the web UI on the network (tunnel or non-loopback).""" + if secure: + return True + if host in ("0.0.0.0", "::"): + return True + try: + from utils.host_policy import is_external_host + except Exception: + return False + return bool(is_external_host(host)) + + +def should_arm_bootstrap_timeout( + *, + host: str, + secure: bool, + api_only: bool, + frontend_served: bool, + is_colab: bool, + requires_change: bool, + timeout_seconds: int, +) -> bool: + """Whether to arm the deadline: only for an exposed web UI whose seeded admin + password is still unchanged. Pure decision (no I/O) for cheap unit testing.""" + if timeout_seconds <= 0: + return False + if api_only or not frontend_served or is_colab: + return False + if not requires_change: + return False + return _is_exposed_bind(host, secure) + + +def _format_duration(seconds: int) -> str: + """Human-friendly duration for the shutdown message (seconds under a minute).""" + + def _plural(n: int, unit: str) -> str: + return f"{n} {unit}{'' if n == 1 else 's'}" + + if seconds < 60: + return _plural(seconds, "second") + minutes, rem = divmod(seconds, 60) + label = _plural(minutes, "minute") + if rem: + label += f" {_plural(rem, 'second')}" + return label + + +def enforce_bootstrap_password_deadline( + storage, + trigger_shutdown, + *, + timeout_seconds: int, + logger = None, +) -> bool: + """Deadline handler: shut down iff the seeded admin password is still unchanged. + + Returns True if it shut Studio down, False if it left it running (the + password was changed in time). + """ + try: + still_default = storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME) + except Exception: + return False + if not still_default: + return False # password changed in time -> leave Studio running + + message = ( + "\nUnsloth Studio was exposed on the network but its default admin " + f"password was not changed within {_format_duration(timeout_seconds)}. " + "Shutting down to avoid leaving an unsecured public instance running.\n" + "Next time, sign in and change the password on first login, or set " + f"{BOOTSTRAP_TIMEOUT_ENV_VAR}=0 to disable this timeout." + ) + if logger is not None: + logger.warning(message) + print(message, file = sys.stderr, flush = True) + try: + trigger_shutdown() + except Exception as e: # shutdown is best-effort; never raise from the timer + if logger is not None: + logger.warning("Bootstrap-timeout shutdown failed: %s", e) + return True + + +def arm_bootstrap_timeout( + storage, + trigger_shutdown, + *, + timeout_seconds: int, + logger = None, +) -> "threading.Timer": + """Start a daemon timer that enforces the deadline. Returns the Timer.""" + timer = threading.Timer( + timeout_seconds, + enforce_bootstrap_password_deadline, + args = (storage, trigger_shutdown), + kwargs = {"timeout_seconds": timeout_seconds, "logger": logger}, + ) + timer.daemon = True + timer.start() + return timer diff --git a/studio/backend/run.py b/studio/backend/run.py index 9cb7868949..d4cbc26b41 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1199,6 +1199,43 @@ def run_server( _graceful_shutdown(_server) sys.exit(1) + # Time-box a freshly-exposed web UI: if nobody changes the seeded admin + # password within the deadline (default 1h), shut down rather than leave an + # unsecured public instance running. No-op for loopback, --api-only, Colab, + # an already-changed password, or UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT=0. + try: + from auth import storage as _auth_storage + from auth.bootstrap_timeout import ( + arm_bootstrap_timeout, + bootstrap_timeout_seconds, + should_arm_bootstrap_timeout, + ) + + _bootstrap_timeout = bootstrap_timeout_seconds() + if should_arm_bootstrap_timeout( + host = host, + secure = secure, + api_only = api_only, + frontend_served = bool(frontend_path) and not api_only, + is_colab = _IS_COLAB, + requires_change = _auth_storage.requires_password_change( + _auth_storage.DEFAULT_ADMIN_USERNAME + ), + timeout_seconds = _bootstrap_timeout, + ): + arm_bootstrap_timeout( + _auth_storage, + _trigger_shutdown, + timeout_seconds = _bootstrap_timeout, + logger = logger, + ) + logger.info( + "Studio will shut down in %ds unless the default admin password is changed.", + _bootstrap_timeout, + ) + except Exception as e: # best-effort: never block startup on the timeout + logger.warning("Bootstrap timeout not armed: %s", e) + if not silent: _emit_startup_output(host, port, display_host, secure = secure, enable_tools = enable_tools) diff --git a/studio/backend/tests/test_bootstrap_timeout.py b/studio/backend/tests/test_bootstrap_timeout.py new file mode 100644 index 0000000000..58d4829215 --- /dev/null +++ b/studio/backend/tests/test_bootstrap_timeout.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Coverage for the exposed-first-run auto-shutdown deadline. + +Tests the env parsing, the pure arm/no-arm decision matrix, and the deadline +handler (shut down iff the seeded admin password is still unchanged). The +threading.Timer itself is not exercised; the handler is invoked directly. +""" + +from types import SimpleNamespace + +from auth.bootstrap_timeout import ( + DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS, + _format_duration, + bootstrap_timeout_seconds, + enforce_bootstrap_password_deadline, + should_arm_bootstrap_timeout, +) + + +# ── bootstrap_timeout_seconds ─────────────────────────────────────── + + +def test_default_when_unset(): + assert bootstrap_timeout_seconds(env = {}) == DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + + +def test_default_when_empty(): + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": " "}) == ( + DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + ) + + +def test_explicit_value_parsed(): + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "1800"}) == 1800 + + +def test_zero_disables(): + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "0"}) == 0 + + +def test_negative_disables(): + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "-5"}) == 0 + + +def test_invalid_falls_back_to_default(): + # A typo must keep the protection, not silently disable it. + assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "abc"}) == ( + DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS + ) + + +# ── should_arm_bootstrap_timeout matrix ───────────────────────────── + + +def _arm_kwargs(**overrides): + kwargs = dict( + host = "0.0.0.0", + secure = False, + api_only = False, + frontend_served = True, + is_colab = False, + requires_change = True, + timeout_seconds = 3600, + ) + kwargs.update(overrides) + return kwargs + + +def test_arm_exposed_wildcard_web_ui(): + assert should_arm_bootstrap_timeout(**_arm_kwargs()) is True + + +def test_arm_secure_loopback_bind(): + # --secure forces a loopback bind but exposes a public tunnel. + assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = True)) is True + + +def test_no_arm_loopback_bind(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = False)) is False + + +def test_no_arm_api_only(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(api_only = True)) is False + + +def test_no_arm_no_frontend(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(frontend_served = False)) is False + + +def test_no_arm_colab(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(is_colab = True)) is False + + +def test_no_arm_password_already_changed(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(requires_change = False)) is False + + +def test_no_arm_timeout_disabled(): + assert should_arm_bootstrap_timeout(**_arm_kwargs(timeout_seconds = 0)) is False + + +# ── enforce_bootstrap_password_deadline ───────────────────────────── + + +def _fake_storage(requires_change: bool): + return SimpleNamespace( + DEFAULT_ADMIN_USERNAME = "unsloth", + requires_password_change = lambda _username: requires_change, + ) + + +def test_deadline_shuts_down_when_password_unchanged(): + calls = [] + result = enforce_bootstrap_password_deadline( + _fake_storage(requires_change = True), + lambda: calls.append("shutdown"), + timeout_seconds = 3600, + ) + assert result is True + assert calls == ["shutdown"] + + +def test_deadline_keeps_running_when_password_changed(): + calls = [] + result = enforce_bootstrap_password_deadline( + _fake_storage(requires_change = False), + lambda: calls.append("shutdown"), + timeout_seconds = 3600, + ) + assert result is False + assert calls == [] + + +def test_deadline_swallows_shutdown_errors(): + def _boom(): + raise RuntimeError("shutdown failed") + + # A failing shutdown must not propagate out of the timer thread. + result = enforce_bootstrap_password_deadline( + _fake_storage(requires_change = True), + _boom, + timeout_seconds = 3600, + ) + assert result is True + + +# ── _format_duration ──────────────────────────────────────────────── + + +def test_format_duration_sub_minute_uses_seconds(): + assert _format_duration(30) == "30 seconds" + + +def test_format_duration_singular_second(): + assert _format_duration(1) == "1 second" + + +def test_format_duration_exact_minutes(): + assert _format_duration(60) == "1 minute" + assert _format_duration(3600) == "60 minutes" + + +def test_format_duration_minutes_and_seconds(): + assert _format_duration(90) == "1 minute 30 seconds" + + +def test_shutdown_message_uses_formatted_duration(): + # The deadline message must reflect the real timeout, not a rounded + # "minute(s)" placeholder. Capture the warning via a fake logger. + logged = [] + + class _Logger: + def warning(self, msg, *args): + logged.append(msg) + + enforce_bootstrap_password_deadline( + _fake_storage(requires_change = True), + lambda: None, + timeout_seconds = 3600, + logger = _Logger(), + ) + assert any("60 minutes" in m for m in logged) + assert not any("minute(s)" in m for m in logged) From a9c8bcf0e1b517071060b456e2cb5d472cf96fd0 Mon Sep 17 00:00:00 2001 From: Abdul Moiz Date: Fri, 26 Jun 2026 13:44:41 +0500 Subject: [PATCH 02/49] Fix DDP crash from CPU-resident rotary inv_freq buffer (#6662) * Fix DDP crash from CPU-resident rotary inv_freq buffer DistributedDataParallel broadcasts all named buffers regardless of persistence or device, but Unsloth's RoPE inv_freq buffer is kept on CPU on purpose (per-GPU cos/sin caches are precomputed instead). That mismatch crashed multi-GPU DDP training with "No backend type associated with device type cpu" during _sync_module_states. Mark inv_freq/short_inv_freq/long_inv_freq buffers as DDP-ignored instead of moving them to GPU, so they're skipped during the buffer broadcast without disabling broadcast_buffers for the rest of the model. Fixes #6656 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: harden DDP-ignore against private API drift, re-apply after PEFT wrap - Wrap the private DistributedDataParallel._set_params_and_buffers_to_ignore_for_model call in try/except, falling back to setting _ddp_params_and_buffers_to_ignore directly so a future PyTorch API change can't block model loading. - Move _exclude_rope_inv_freq_from_ddp to loader_utils.py (shared by loader.py, llama.py, vision.py without circular imports) and call it again after get_peft_model wraps the model in a PeftModel, since the rotary buffers' fully qualified names change once nested under "base_model.model...". * [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: imagineer99 --- unsloth/models/llama.py | 4 +++- unsloth/models/loader.py | 3 +++ unsloth/models/loader_utils.py | 34 ++++++++++++++++++++++++++++++++++ unsloth/models/vision.py | 3 ++- 4 files changed, 42 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 2399334983..d5eca8d6df 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -28,7 +28,7 @@ from ._utils import ( is_bfloat16_supported, get_quant_type, ) -from .loader_utils import _get_fp8_mode_and_check_settings +from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings from ..utils.packing import ( get_packed_info_from_kwargs, mask_packed_sequence_boundaries, @@ -3049,6 +3049,7 @@ class FastLlamaModel: # Pre-wrapped PEFT model passes through here; still arm the detector so an RL # trainer can reset a compile cache poisoned by a pre-train forward. _unsloth_install_pretrain_detector(model) + model = _exclude_rope_inv_freq_from_ddp(model) return model else: raise TypeError( @@ -3404,6 +3405,7 @@ class FastLlamaModel: # Detect a stray pre-train forward so train() can drop the torch.compile # graph cache it would otherwise poison (see prepare_for_training_mode). _unsloth_install_pretrain_detector(model) + model = _exclude_rope_inv_freq_from_ddp(model) return model @staticmethod diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 75eba2d9f8..562afdd645 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -33,6 +33,7 @@ from transformers import AutoConfig from transformers import __version__ as transformers_version from peft import PeftConfig, PeftModel from .loader_utils import ( + _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings, _offline_quantize_to_fp8, _tag_model_with_fp8_torchao_config, @@ -885,6 +886,7 @@ class FastLanguageModel(FastLlamaModel): patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice) model = _fix_rope_inv_freq(model) + model = _exclude_rope_inv_freq_from_ddp(model) return model, tokenizer @@ -1822,6 +1824,7 @@ class FastModel(FastBaseModel): patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice) model = _fix_rope_inv_freq(model) + model = _exclude_rope_inv_freq_from_ddp(model) return model, tokenizer diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 7000c1587d..d6d6ce877e 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -499,6 +499,40 @@ def _get_fp8_mode_and_check_settings( return fp8_mode +# Rotary inv_freq buffers are deliberately kept on CPU - Unsloth pre-builds a +# cos/sin cache per GPU instead (see LlamaRotaryEmbedding.multi_gpu_cos_cached) +# so the GPU-resident lookup never needs to move the tiny inv_freq tensor itself. +# torch.nn.parallel.DistributedDataParallel ignores device entirely when it +# broadcasts buffers across ranks, so a CPU buffer crashes NCCL's +# _broadcast_coalesced with "No backend type associated with device type cpu". +# Telling DDP to skip these specific buffers avoids that crash without moving +# inv_freq to GPU (which would break the per-GPU cache design) and without +# disabling buffer broadcast for every other module (the user's workaround). +# Re-run this after wrapping with PEFT too - the buffers' fully qualified +# names change once they sit under a PeftModel (eg "base_model.model..."). +# https://github.com/unslothai/unsloth/issues/6656 +_ROTARY_INV_FREQ_BUFFER_NAMES = ("inv_freq", "short_inv_freq", "long_inv_freq") + + +def _exclude_rope_inv_freq_from_ddp(model): + ignored = list(getattr(model, "_ddp_params_and_buffers_to_ignore", None) or []) + for module_name, module in model.named_modules(): + for buffer_name, _ in module.named_buffers(recurse = False): + if buffer_name in _ROTARY_INV_FREQ_BUFFER_NAMES: + fqn = f"{module_name}.{buffer_name}" if module_name else buffer_name + if fqn not in ignored: + ignored.append(fqn) + if ignored: + try: + from torch.nn.parallel import DistributedDataParallel + DistributedDataParallel._set_params_and_buffers_to_ignore_for_model(model, ignored) + except Exception: + # Private PyTorch API - fall back to setting the attribute DDP reads + # directly if it ever moves or changes signature. + model._ddp_params_and_buffers_to_ignore = ignored + return model + + # ============================================================================= # Offline loading - single source of truth (shared by vision.py, loader.py and # the Studio exporter). Decide offline ONCE at the load boundary and force it diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index c72c1af4b1..39004b4d45 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -40,7 +40,7 @@ from ._utils import ( set_task_config_attr, ) from ._utils import * -from .loader_utils import _get_fp8_mode_and_check_settings +from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings from ..save import patch_saving_functions from ..models.loader_utils import is_distributed from unsloth_zoo.gradient_checkpointing import ( @@ -1741,6 +1741,7 @@ class FastBaseModel: # Detect a stray pre-train forward so train() can drop the torch.compile # graph cache it would otherwise poison (see prepare_for_training_mode). _unsloth_install_pretrain_detector(model) + model = _exclude_rope_inv_freq_from_ddp(model) return model @staticmethod From 3e43ed7b4afc42fe2afa295dd64c97ea9de64d77 Mon Sep 17 00:00:00 2001 From: Anmol Mishra Date: Fri, 26 Jun 2026 14:15:36 +0530 Subject: [PATCH 03/49] Patch FalconH1RMSNorm to fix float64 compilation crash on Intel Arc DG2 (#6691) * Patch FalconH1RMSNorm to fix float64 compilation on Intel Arc DG2 Fixes unslothai/unsloth#6555 Root cause: FalconH1RMSNorm.forward() does hidden_states.pow(2).mean().rsqrt() with self.variance_epsilon being a Python float64. When torch.compile fuses this pattern into the auto-generated Triton kernel 'triton_per_fused__to_copy_mean_mul_pow_rsqrt_*', the float64 epsilon causes type promotion to double. Intel Arc DG2 GPUs do not support double precision (Double type is not supported on this platform). The existing patch_rms_layernorm() only patches LlamaRMSNorm, not the separate FalconH1RMSNorm class in transformers.models.falcon_h1. Fix: add Unsloth_FalconH1RMSNorm that delegates to fast_rms_layernorm (@torch.compiler.disable, handles epsilon as tl.float32), and call the patch in FastFalconH1Model.pre_patch() before model creation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Condense FalconH1RMSNorm patch 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: Daniel Han --- unsloth/models/falcon_h1.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index b940c95c68..05bfd2ebb3 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -37,6 +37,8 @@ try: FalconH1DecoderLayer, FalconH1Model, FalconH1ForCausalLM, + FalconH1RMSNorm, + FalconH1RMSNormGated, FalconHybridMambaAttentionDynamicCache, ) except: @@ -677,6 +679,18 @@ def fix_prepare_inputs_for_generation(module): module.prepare_inputs_for_generation = _fast_prepare_inputs_for_generation +class Unsloth_FalconH1RMSNorm(FalconH1RMSNorm): + """fast_rms_layernorm (compiler-disabled, fp32 eps) avoids the float64 torch.compile RMSNorm kernel that fails on Intel Arc DG2 (issue #6555).""" + + def forward(self, hidden_states): + return fast_rms_layernorm(self, hidden_states, gemma = False) + + +def patch_falcon_h1_rms_layernorm(): + import transformers.models.falcon_h1.modeling_falcon_h1 + transformers.models.falcon_h1.modeling_falcon_h1.FalconH1RMSNorm = Unsloth_FalconH1RMSNorm + + class FastFalconH1Model(FastLlamaModel): @staticmethod def pre_patch(): @@ -708,6 +722,8 @@ class FastFalconH1Model(FastLlamaModel): transformers.models.falcon_h1.modeling_falcon_h1.FalconH1RotaryEmbedding = ( LlamaRotaryEmbedding ) + # Avoids the float64 RMSNorm compile kernel that fails on Intel Arc DG2 (issue #6555). + patch_falcon_h1_rms_layernorm() return @staticmethod From e9c6364e1ed8e16c2c3ccbf7a8ec7ecc139559b5 Mon Sep 17 00:00:00 2001 From: Matt Van Horn Date: Fri, 26 Jun 2026 01:55:59 -0700 Subject: [PATCH 04/49] feat: improve Unsloth Studio chat title generation quality (#6697) * feat: improve Unsloth Studio chat title generation quality * fix: address self-review (guard echoed role labels before punctuation stripping) * Address title generation review feedback Consolidate the echo guard into a single leading-label check (now also covering base and lora) and drop the post-punctuation duplicate that could never match a colon once punctuation is stripped. Swap the slice-based first-assistant lookup for an indexed find to avoid copying the messages array, and note the brace counter's assumptions in the test helper. --------- Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Co-authored-by: Daniel Han --- .../src/features/chat/runtime-provider.tsx | 28 +++-- tests/studio/test_chat_title_generation.py | 115 ++++++++++++++++++ 2 files changed, 136 insertions(+), 7 deletions(-) create mode 100644 tests/studio/test_chat_title_generation.py diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index fbaa4400ed..360e081f0f 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -427,26 +427,32 @@ function extractTextParts(m: ThreadMessage | undefined): string { async function generateTitleWithModel(payload: { userText: string; + assistantText?: string; }): Promise { const params = useChatRuntimeStore.getState().params; if (!params.checkpoint) return null; const user = clip(payload.userText, 256); - const parts: string[] = [user]; + const assistant = clip(payload.assistantText ?? "", 384); + const parts: string[] = [`User: ${user}`]; + if (assistant) { + parts.push(`Assistant: ${assistant}`); + } function normalizeTitle(raw: string): string | null { let title = raw.split(/\r?\n/, 1)[0] ?? ""; title = title.replace(/^\s*title\s*:\s*/i, ""); title = title.replace(/[^\x20-\x7E]+/g, " "); title = title.replace(/["'`]+/g, ""); - title = title.replace(/[.!?:;,]+/g, " "); - title = title.replace(/\s+/g, " ").trim(); - // Model echo fail-safe. - if (/\b(user|base|lora|assistant)\s*:/i.test(title)) { + // Echo fail-safe: reject leading role labels before punctuation strips the ":". + if (/^\s*(user|assistant|base|lora)\s*:/i.test(title)) { return null; } + title = title.replace(/[.!?:;,]+/g, " "); + title = title.replace(/\s+/g, " ").trim(); + const words = title.split(" ").filter(Boolean).slice(0, 6); const joined = words.join(" ").trim(); if (!joined) return null; @@ -468,7 +474,7 @@ async function generateTitleWithModel(payload: { { role: "system", content: - "Write 1 concise chat title for the user's message. Rules: 2-6 words, no quotes, no punctuation, ASCII only, do not echo input. Output title only.", + "Write 1 concise chat title summarizing the conversation topic, not the user's exact wording. Use the assistant reply as context when provided. Rules: 2-6 words, no quotes, no punctuation, ASCII only, do not echo input. Output title only.", }, { role: "user", content: parts.join("\n") }, ], @@ -730,8 +736,15 @@ function createStudioDbAdapter( return streamTitle(thread.title); } - const firstUser = messages.find((m) => m.role === "user"); + const firstUserIndex = messages.findIndex((m) => m.role === "user"); + const firstUser = + firstUserIndex === -1 ? undefined : messages[firstUserIndex]; + const firstAssistant = + firstUserIndex === -1 + ? undefined + : messages.find((m, i) => m.role === "assistant" && i > firstUserIndex); const userText = extractTextParts(firstUser) || defaultTitle; + const assistantText = extractTextParts(firstAssistant); if (!autoTitle) { const title = fallbackTitleFromUserText(userText); @@ -769,6 +782,7 @@ function createStudioDbAdapter( const title = (await generateTitleWithModel({ userText, + assistantText, })) || fallbackTitleFromUserText(userText); await persistTitle(title); diff --git a/tests/studio/test_chat_title_generation.py b/tests/studio/test_chat_title_generation.py new file mode 100644 index 0000000000..4f82b4ec03 --- /dev/null +++ b/tests/studio/test_chat_title_generation.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression checks for Studio chat title generation context.""" + +from __future__ import annotations + +from pathlib import Path + + +REPO = Path(__file__).resolve().parents[2] +RUNTIME_TSX = REPO / "studio/frontend/src/features/chat/runtime-provider.tsx" + + +def _source_until(src: str, anchor: str, end_anchor: str) -> str: + start = src.find(anchor) + assert start != -1, f"anchor {anchor!r} not found" + end = src.find(end_anchor, start) + assert end != -1, f"end anchor {end_anchor!r} not found" + return src[start:end] + + +def _balanced_block(src: str, anchor: str) -> str: + # Brace-counting only; assumes no unbalanced braces in strings, regexes, or comments. + start = src.find(anchor) + assert start != -1, f"anchor {anchor!r} not found" + body_start = src.find("{", start) + assert body_start != -1, f"body opener after {anchor!r} not found" + + depth = 0 + for index in range(body_start, len(src)): + char = src[index] + if char == "{": + depth += 1 + elif char == "}": + depth -= 1 + if depth == 0: + return src[start : index + 1] + raise AssertionError(f"unbalanced block after {anchor!r}") + + +def test_title_model_prompt_targets_conversation_topic(): + block = _source_until( + RUNTIME_TSX.read_text(), + "async function generateTitleWithModel", + "\nconst inflightTitleByKey", + ) + + assert "conversation topic" in block + assert "not the user's exact wording" in block + assert "Use the assistant reply as context when provided" in block + assert "Rules: 2-6 words" in block + + +def test_title_model_payload_includes_optional_assistant_reply(): + block = _source_until( + RUNTIME_TSX.read_text(), + "async function generateTitleWithModel", + "\nconst inflightTitleByKey", + ) + + assert "assistantText?: string;" in block + assert 'const assistant = clip(payload.assistantText ?? "", 384);' in block + assert "const parts: string[] = [`User: ${user}`];" in block + assert "if (assistant)" in block + assert "parts.push(`Assistant: ${assistant}`);" in block + assert 'parts.join("\\n")' in block + + +def test_generate_title_passes_first_assistant_reply_after_first_user(): + block = _balanced_block( + RUNTIME_TSX.read_text(), + "async generateTitle(remoteId", + ) + + assert 'const firstUserIndex = messages.findIndex((m) => m.role === "user");' in block + assert '.find((m, i) => m.role === "assistant" && i > firstUserIndex)' in block + assert "const assistantText = extractTextParts(firstAssistant);" in block + assert "generateTitleWithModel({" in block + assert "userText," in block + assert "assistantText," in block + + +def test_auto_title_disabled_uses_deterministic_user_text_fallback(): + block = _balanced_block( + RUNTIME_TSX.read_text(), + "async generateTitle(remoteId", + ) + auto_title_off = _balanced_block(block, "if (!autoTitle)") + + assert "fallbackTitleFromUserText(userText)" in auto_title_off + assert "generateTitleWithModel" not in auto_title_off + + +def test_model_failure_still_falls_back_to_user_text(): + block = _balanced_block( + RUNTIME_TSX.read_text(), + "async generateTitle(remoteId", + ) + + assert "})) || fallbackTitleFromUserText(userText);" in block + + +def test_title_normalizer_still_enforces_output_constraints(): + block = _source_until( + RUNTIME_TSX.read_text(), + "async function generateTitleWithModel", + "\nconst inflightTitleByKey", + ) + + assert r'replace(/[^\x20-\x7E]+/g, " ")' in block + assert 'replace(/["\'`]+/g, "")' in block + assert 'replace(/[.!?:;,]+/g, " ")' in block + assert 'title.split(" ").filter(Boolean).slice(0, 6)' in block + assert "joined.length > 60" in block From 2ef394137aef657f9578aab57813361366e72c17 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 03:31:33 -0700 Subject: [PATCH 05/49] Studio: harden background consumer loops and streaming paths against silent UI freezes (#6653) * Studio: harden the data-recipe and inference consumer loops against pump death Follow-up to #6643. The same single-unsupervised-consumer pattern the training pump had lives in two sibling loops, with the same failure mode: one bad event kills the only thread that updates the in-memory state every UI surface reads, while the worker subprocess keeps running. - data_recipe JobManager._pump_loop: a malformed worker log line that makes parse_log_message raise no longer kills the pump. Guard _handle_event, the queue read, and the worker-exit finalize, and broaden _drain_queue so a drain error still finalizes the job instead of leaving it wedged "active" (which also leaked the workflow-scoped API key until its 24h expiry). - inference InferenceOrchestrator._dispatcher_loop: guard the routing body so a malformed response or a mailbox put error can't kill the dispatcher and hang every in-flight generation (callers key liveness on the subprocess, not on this thread). Adds regression tests for both. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: extend consumer-loop hardening to RAG, hub, auth, and stream-reader paths Continuation of the data-recipe and inference pump hardening: the same "background producer updates in-memory state that a single unsupervised consumer surfaces to the UI" pattern shows up in several more Studio paths, each able to silently freeze a UI surface while the worker keeps running. RAG ingestion SSE (core/rag/ingestion.py): - job_events polled the queue with a blocking get and never noticed client disconnect or a dead worker, so a closed tab or a producer that died without emitting a terminal event left the stream hanging. It now polls with a timeout, emits heartbeats, ends on terminal job status, caps idle time, and always pops the job registry in finally. - Added _reap_finished_jobs() and call it from start_ingestion so finished job state does not accumulate. Startup reconcile (storage/rag_db.py, main.py): - reconcile_orphaned_ingestion_jobs() marks ingestion jobs (and their documents) that were left non-terminal by a previous crash as failed, so the UI does not show jobs stuck "running" forever after a restart. Wired in at startup next to cleanup_orphaned_runs(). Hub download watcher (hub/services/download_lifecycle.py): - _watch() could leave a job pinned "running" if finalize raised. Body is now guarded: on failure it logs and sets the job to error, and always invalidates the hf cache scan in finally. External provider stream (core/inference/external_provider.py): - read timeout was None (no stall ceiling); set to 300s so a wedged upstream surfaces as an error instead of an indefinitely hung stream. Auth store (auth/storage.py): - Enable WAL + busy_timeout on the auth DB so token validation (read on every request) and login writes stop serialising on the rollback journal. Matches studio_db / rag_db / providers_db. Login rate limiter (routes/auth.py): - _LOGIN_IP_BUCKETS could grow unbounded under spoofed-IP traffic; cap it and prune stale buckets, mirroring the per-account bucket handling. Training progress SSE (routes/training.py): - Break promptly on client disconnect instead of waiting for the next yield to fail on a closed socket, matching the export / data-recipe SSE routes. llama-server stdout drain (core/inference/llama_cpp.py): - Broaden the drain guard so an unexpected decode/read error logs at debug and stops the drainer cleanly instead of escaping the thread. Frontend stream readers (chat-api.ts, rag-api.ts): - Wrap the SSE read loops in try/finally + reader.cancel() so early return ([DONE]), thrown errors, and consumer aborts release the reader lock instead of holding it until GC. Tests: - test_training_progress_stream_nan: fake request now implements the async is_disconnected() the route polls, matching the other SSE route fakes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: address Codex review feedback on the consumer-loop hardening Four follow-ups from the automated review, all on code this PR introduced: - Data-recipe pump (manager.py): a queue read that keeps raising an error outside the read's narrow catch set (e.g. a broken queue pipe after the child died) hit the `continue` guard and skipped the dead-worker finalize below, spinning forever and leaving the job wedged "active" with its workflow key unretired. On a read failure, fall through to finalize when the worker is no longer alive. Added a regression test. - RAG ingestion SSE (ingestion.py): the 5-minute idle cap could end the stream while the job was still pending/running (a large document spends minutes in embedding/storing with no per-batch progress event). The route then sends [DONE], and the client treats a no-terminal-frame end as completion, marking the document indexed mid-ingestion. Drop the idle cap: while the worker is alive and non-terminal we keep heartbeating; the stream ends only on terminal DB status, the None sentinel, or client disconnect. - Login rate limiter (auth.py): the per-IP path pruned but then added the new IP unconditionally, so a spoofed-source-IP spray kept _LOGIN_IP_BUCKETS unbounded and made every new IP pay a full-dict prune scan. Gate the add on the cap, mirroring the account path. - Hub download watcher (download_lifecycle.py): if finalize raised before it reaped (proc.wait) and dropped the worker (e.g. an I/O error draining stderr), the crash path published a terminal state while the live Popen stayed registered and kept writing the cache, and the terminal set_job let claim() admit a retry on the same repo. Terminate + drop the worker before setting the terminal state. * Studio: keep login throttling working when the per-IP bucket dict saturates Review follow-up. The previous cap fix skipped creating a bucket for a new IP once _LOGIN_IP_BUCKETS was full, returning ip_fails=0. Under a sustained spray that also fills the account dict, every failure from such an IP then looked first-seen and _login_blocked had no bucket to enforce, so the cap effectively disabled throttling once saturated. Bound the dict with a FIFO eviction instead: if the IP is new and the dict is full, reclaim expired buckets (rate-limited so a burst of distinct IPs can't make each failure an O(n) sweep) and, if still full, evict the oldest-inserted IP. The new IP always gets a real bucket, so a saturating (e.g. spoofed X-Forwarded-For) spray stays throttled while memory stays bounded. Added a regression test that saturates the dict and asserts a later IP is still blocked. * Studio: address Codex review (RAG queue lifecycle, stream error, orphan chunks) Three follow-ups on the Phase 6 changes: - RAG ingestion SSE (ingestion.py): job_events removed the per-job queue in its finally on ANY exit, including an early client disconnect while the worker is still running. That dropped the worker's later events (the queue is the only one _emit writes to) and made a reconnect find no queue and receive only [DONE], which the client treats as completion. Only drop the queue on a terminal exit (None sentinel / terminal DB status); leftover terminal queues are still swept by _reap_finished_jobs. Added queue-lifecycle tests. - External provider stream (routes/inference.py): once the 300s read timeout can fire, the stream's except path failed the monitor but ended without an error frame or [DONE], so the chat client saw a bare EOF and saved the timed-out answer as a successful partial with no error. Emit an SSE error frame (and [DONE]) on stream failure so the client surfaces it. - RAG startup reconcile (storage/rag_db.py): marking a half-ingested document failed left its chunks/fts/vec rows intact, and retrieval filters by scope not status, so a failed document could still be retrieved and cited. Purge the document's chunks when reconciling it to failed (the doc row stays for re-ingest). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: release the remaining SSE stream readers (training, data-recipe, export) reviewer.py follow-up. The chat and RAG SSE readers were wrapped in try/finally + reader.cancel(), but the other three readers built on the same response.body.getReader() pattern were left without it: streamTrainingProgress, streamRecipeJobEvents, and streamExportLogs leak the ReadableStreamDefaultReader lock (held until GC) when the consumer aborts, returns early, or a parse/callback throws. Wrap each in try/finally + reader.cancel() (export already had a try/catch, so it only needed the finally). All five frontend SSE readers now release the reader symmetrically. * Tighten resilience comments and docstrings Condense the verbose explanatory comments and internal-helper docstrings added in this branch to shorter, clearer forms. Comment/whitespace only; verified no code changed via AST diff. No behaviour change. * Studio: keep chunks for completed docs during ingestion reconcile Startup reconciliation flips orphaned (non-terminal) ingestion jobs to failed and purges the document's chunks so a failed source can't be retrieved. But it dropped the chunks unconditionally, so a document the worker had already committed as 'completed' before the crash (only its job row left non-terminal) lost every chunk while still reporting 'completed'. That leaves an empty source that retrieval can't return and dedup (status != 'failed') blocks from re-ingest. Only purge chunks when the document UPDATE actually transitions it to failed; an already-completed document keeps its chunks. Adds reconcile regression tests for both the completed-doc and genuine in-flight-orphan cases. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: drop a finished RAG job's queue when the client disconnects job_events kept the per-job queue until it consumed the None sentinel, so a UI that stops on the terminal event (its reader.cancel aborts the stream before [DONE]) left the queue registered until the next _reap_finished_jobs sweep; a batch of uploads followed by idling retained them all. _run writes the terminal DB status before emitting the terminal event, so on generator exit, drop the queue when the job's DB row is already terminal (worker done, nothing to resume) and keep it only while the worker is still running. Adds a disconnect-after-terminal-event regression test. * Remove stray async task output files committed by mistake * Studio: harden login IP throttle and end progress stream on disconnect Two Codex review items: Login per-IP throttle: when the per-IP bucket dict saturated, FIFO eviction could drop a still-hot (blocked) bucket, so an IP could flood the dict with distinct (or spoofed) source IPs to push out its own bucket and retry as first-seen. Stop evicting hot buckets; a new IP that can't fit now shares a bounded overflow counter that still trips the per-IP threshold, so a saturating spray stays throttled and no live counter is reset. Progress SSE: on client disconnect the polling loop only broke and fell through to the unconditional final 'complete' frame, so a buffered or proxying consumer could read a still-active run as completed. Return from the generator instead. Adds regression tests for both (spray cannot reset a hot bucket; disconnect while active emits no complete frame). * Studio: shard the login overflow counter and stop cancelling chat stream after [DONE] Two Codex review items: Login throttle overflow: the single shared overflow counter meant that once a saturating spray pushed it past the per-IP threshold, _login_blocked returned 429 for every new unbucketed source IP, before credentials were checked -- a global login denial. Shard the overflow into a fixed array of counters keyed by hash(ip), so a hot shard only throttles the IPs that map to it while a single source's repeated failures still concentrate in one shard and stay throttled. Memory stays bounded and no live bucket is evicted. Adds a regression test that a hot overflow shard does not block an unrelated IP. Chat stream: the reader.cancel() in the SSE finally fired even after a natural [DONE]/EOF. The backend finalizes its api-monitor entry right after yielding the sentinel (the local pass-through finishes after the last yield), so a client cancel there can be observed as a disconnect and mark a completed request as cancelled. Track natural completion and only cancel on an early/abnormal exit. (No frontend unit test: the Studio frontend has no test harness.) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: give prep-timeout test fakes an is_disconnected method The progress stream now ends on client disconnect (await request.is_disconnected() before falling through to the terminal frame). After merging that into the prep-timeout tests added later on main, their _FakeRequest/_ReconnectRequest must provide is_disconnected or the generator raises AttributeError under CI. * Studio: keep the login overflow throttle when bucket capacity frees up _login_blocked only consulted the per-IP overflow shard while the bucket dict was still at capacity. If a slot freed before the 60s window expired (e.g. another IP's successful login calls _clear_login_bucket), a source counted in a hot shard stopped being blocked and its next failure got a fresh per-IP bucket, resetting the throttle the overflow path exists to preserve. Always max in the IP's shard (shards are empty outside saturation, so it is a no-op in the common case). Adds a regression test that a hot source stays throttled after a bucket frees. * Studio: clear a login IP's overflow throttle on successful login _clear_login_bucket reset the per-IP and per-account buckets on a successful login but not the overflow shard, so after the dict saturated and an IP was counted in overflow, a later successful login left those entries behind and the next failed attempt could immediately return 429. Store overflow entries as (timestamp, ip) so a source is throttled by its own count within the shard (also removing cross-IP collateral within a shard), and drop just that IP's entries in _clear_login_bucket. Adds a regression test that a successful login clears the overflow throttle. * Studio: bound the login overflow shard memory under high-cardinality spray The per-IP overflow tracked failures in a time-pruned deque of (timestamp, ip) tuples, so a spoofed-X-Forwarded-For spray of distinct one-off IPs grew memory and the per-check scan with request cardinality for the whole window -- undermining the bucket cap that exists to bound memory. Replace each shard with a fixed- capacity dict (ip -> [count, window_start]): O(1) lookups, and when a shard is full a one-off IP evicts the lowest-count entry (Space-Saving) so memory is hard- bounded while a persistent attacker keeps a high count and is never evicted. Adds a regression test that shards stay within the per-shard cap under a 5000-IP spray. * Studio: purge chunks for already-failed docs during ingestion reconcile The reconcile chunk-purge was gated on the documents UPDATE actually flipping a non-terminal doc to failed. A doc the worker had already marked 'failed' before the crash (job row left non-terminal) was not re-flipped, so its committed chunks were kept and stayed retrievable/citable, since retrieval filters by scope not status. Purge chunks whenever the document is not 'completed' (failed, in-flight, or gone), preserving the completed-doc carve-out. Adds a regression test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: don't inherit an evicted IP's count onto a new overflow source When a full overflow shard evicted the lowest-count entry, the new source inherited that count (Space-Saving base + 1). If a shard was saturated with hot entries, an unrelated new IP could land at/over the threshold and be 429'd after a single attempt -- cross-IP collateral despite the per-source-isolation intent. New entries now start clean at count 1; the only cost is that a heavy hitter that is the lowest-count entry in a fully saturated shard can briefly reset, which is preferable to blocking a bystander. Adds a regression test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: carry overflow failures into a new IP bucket on transition _login_blocked took max(per-IP bucket, overflow shard) rather than combining them, so a source could log (threshold-1) failures in overflow during saturation and, once a bucket slot freed, another (threshold-1) in a fresh bucket within the same window -- roughly doubling the per-IP limit. When a saturated-era IP first gets a real bucket, migrate its windowed overflow count into that bucket (and drop the overflow entry) so the combined failures throttle at the intended limit. Adds a regression test. * Studio: reconcile a completed doc's orphaned job to completed, not failed When a crash left an ingestion job non-terminal after its document was already committed as completed, reconcile marked the job failed. After restart the upload UI has no in-memory SSE queue and falls back to getJob(), which treats a failed job as an indexing failure and removes/toasts a document that is actually searchable. Mark the job completed (keeping its chunks) when its document is completed. Extends the completed-doc reconcile test to assert the job status. * Studio: clamp the overflow failure count migrated into a login bucket A saturated source could accrue an unbounded overflow count, then materialize one deque entry per recorded failure when a bucket slot freed, allocating an arbitrarily large deque under the login lock. Only at-or-above the per-IP threshold matters for blocking, so cap the count there at the record and take sites; the migration is now bounded without weakening the limit. * Studio: keep the RAG job stream alive on a transient status read The heartbeat poll read the job row unguarded; a momentarily-locked DB would raise out of job_events, which the SSE route turns into a terminal error frame, and the UI drops a document whose worker is still running. Treat a failed status read as non-terminal: heartbeat and retry, and keep the queue so a reconnect can resume. * Studio: set busy_timeout before journal_mode on the auth DB Switching journal_mode needs a lock, so if a refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY and the shared try leaves the connection on SQLite's default zero lock wait. Set busy_timeout first so the switch waits instead of failing. * [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/auth/storage.py | 11 + .../backend/core/data_recipe/jobs/manager.py | 99 ++++--- .../core/inference/external_provider.py | 8 +- studio/backend/core/inference/llama_cpp.py | 7 +- studio/backend/core/inference/orchestrator.py | 45 ++-- studio/backend/core/rag/ingestion.py | 91 ++++++- .../hub/services/download_lifecycle.py | 61 +++-- studio/backend/main.py | 10 + studio/backend/routes/auth.py | 137 +++++++++- studio/backend/routes/inference.py | 8 + studio/backend/routes/training.py | 5 + studio/backend/storage/rag_db.py | 68 +++++ .../tests/test_data_recipe_pump_resilience.py | 152 +++++++++++ .../test_inference_dispatcher_resilience.py | 91 +++++++ studio/backend/tests/test_login_rate_limit.py | 245 ++++++++++++++++++ .../test_rag_job_events_queue_lifecycle.py | 114 ++++++++ .../tests/test_rag_reconcile_orphaned.py | 114 ++++++++ .../test_training_progress_prep_timeout.py | 6 + .../test_training_progress_stream_nan.py | 24 ++ .../src/features/chat/api/chat-api.ts | 157 ++++++----- .../src/features/export/api/export-api.ts | 7 + .../frontend/src/features/rag/api/rag-api.ts | 51 ++-- .../src/features/recipe-studio/api/index.ts | 47 ++-- .../src/features/training/api/train-api.ts | 61 +++-- 24 files changed, 1388 insertions(+), 231 deletions(-) create mode 100644 studio/backend/tests/test_data_recipe_pump_resilience.py create mode 100644 studio/backend/tests/test_inference_dispatcher_resilience.py create mode 100644 studio/backend/tests/test_rag_job_events_queue_lifecycle.py create mode 100644 studio/backend/tests/test_rag_reconcile_orphaned.py diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 1f153699d7..a0da2b2096 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -110,6 +110,17 @@ def get_connection() -> sqlite3.Connection: except OSError: pass conn.row_factory = sqlite3.Row + # WAL lets token reads run concurrently with refresh-token writes; + # busy_timeout bounds lock waits. Matches the other Studio SQLite stores. + # Set busy_timeout first: switching journal_mode needs a lock, so if a + # refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY; + # with busy_timeout already in effect it waits instead of failing and leaving + # this connection on SQLite's default zero lock wait. + try: + conn.execute("PRAGMA busy_timeout=5000") + conn.execute("PRAGMA journal_mode=WAL") + except sqlite3.Error: + pass conn.execute( """ CREATE TABLE IF NOT EXISTS auth_user ( diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index c238c250bd..0e0044702e 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -28,6 +28,9 @@ from .constants import ( from .parse import apply_update, coerce_event, parse_log_message from .types import Job from .worker import run_job_process +from loggers import get_logger + +logger = get_logger(__name__) _CTX = mp.get_context("spawn") @@ -445,54 +448,86 @@ class JobManager: events.append(coerce_event(q.get_nowait())) except queue.Empty: return events - except (EOFError, OSError, ValueError): + except Exception: + # Return what we have so the run still finalizes rather than wedging "active". + logger.exception( + "Data-recipe job pump: queue drain failed; finalizing with drained events" + ) return events + def _safe_handle_event(self, job: Job, event: dict) -> None: + """Apply one event, swallowing any handler error so the pump can't die.""" + try: + self._handle_event(job, event) + except Exception: + etype = event.get("type") if isinstance(event, dict) else type(event).__name__ + logger.exception("Data-recipe job pump: failed to handle %s event; skipping", etype) + def _pump_loop(self) -> None: - """Background thread: consumes worker events + updates job snapshot.""" + """Background thread: consume worker events and update the job snapshot. + + Guarded so no single event can end the loop; it is the sole writer of the + snapshot the UI polls, so its death would freeze status/SSE. + """ while True: snap = self._snapshot() if snap is None: return job, proc, mp_q = snap - event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25) + try: + event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25) + except Exception: + # If a read keeps raising after the worker died, finalize instead + # of spinning forever; only retry while the worker is still alive. + logger.exception("Data-recipe job pump: queue read failed; continuing") + if proc.is_alive(): + time.sleep(0.1) + continue + event = None + if event is not None: - self._handle_event(job, event) + self._safe_handle_event(job, event) continue if proc.is_alive(): continue - for e in self._drain_queue(mp_q): - self._handle_event(job, e) + # Worker exited: drain + finalize, guarded so an error can't strand the run "active". + try: + for e in self._drain_queue(mp_q): + self._safe_handle_event(job, e) - retired_job: Job | None = None - with self._lock: - if self._job and self._job.status in { - "pending", - "active", - "cancelling", - }: - if self._job.status == "cancelling": - self._job.status = "cancelled" - else: - self._job.status = "error" - self._job.error = self._job.error or "process exited" - self._job.finished_at = time.time() - event_type = ( - EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR - ) - self._emit( - { - "type": event_type, - "ts": time.time(), - "job_id": self._job.job_id, - } - ) - retired_job = self._job - if retired_job is not None: - self._retire_workflow_key(retired_job) + retired_job: Job | None = None + with self._lock: + if self._job and self._job.status in { + "pending", + "active", + "cancelling", + }: + if self._job.status == "cancelling": + self._job.status = "cancelled" + else: + self._job.status = "error" + self._job.error = self._job.error or "process exited" + self._job.finished_at = time.time() + event_type = ( + EVENT_JOB_CANCELLED + if self._job.status == "cancelled" + else EVENT_JOB_ERROR + ) + self._emit( + { + "type": event_type, + "ts": time.time(), + "job_id": self._job.job_id, + } + ) + retired_job = self._job + if retired_job is not None: + self._retire_workflow_key(retired_job) + except Exception: + logger.exception("Data-recipe job pump: finalization after worker exit failed") return def _handle_event(self, job: Job, event: dict) -> None: diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index cae001c34d..20312e067c 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -771,11 +771,9 @@ class ExternalProviderClient: self.base_url = self.base_url[: -len("/openai")] self.api_key = api_key self._timeout = httpx.Timeout(timeout, connect = 10.0) - # Disable read timeout on SSE streams: reasoning-heavy models pause - # tens of seconds between bytes while thinking, and httpx's read - # timeout is the per-byte gap, not wall clock. connect/write bounds - # still surface real network failures. - self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = None) + # Generous per-byte read timeout: reasoning models pause tens of seconds + # between bytes, but a dead upstream must eventually error, not hang forever. + self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = 300.0) def _auth_headers(self) -> dict[str, str]: """Build authentication headers using the provider's registry config.""" diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 152a3f19b2..255fc4140f 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -3147,9 +3147,10 @@ class LlamaCppBackend: except (ValueError, OSError): # Log file closed under us; tee silently. pass - except (ValueError, OSError): - # Pipe closed -- process terminating. - pass + except Exception: + # Never let the drain thread die: a full stdout pipe can deadlock + # llama-server (Windows). Pipe-closed on exit is the common case. + logger.debug("llama-server stdout drain stopped", exc_info = True) # GGUF KV type sizes for fast skipping _GGUF_TYPE_SIZE = { diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index c980dbde2d..5dbd5fb479 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -534,29 +534,34 @@ class InferenceOrchestrator: except (EOFError, OSError, ValueError): break - rid = resp.get("request_id") - rtype = resp.get("type", "") + # Sole consumer of the response queue; if it died every in-flight + # stream would hang, so never let routing kill the dispatcher. + try: + rid = resp.get("request_id") + rtype = resp.get("type", "") - # Status messages — log and skip - if rtype == "status": - logger.info("Subprocess status: %s", resp.get("message", "")) - continue - - # Route to mailbox if a matching request_id exists - if rid: - with self._mailbox_lock: - mbox = self._mailboxes.get(rid) - if mbox is not None: - mbox.put(resp) + # Status messages: log and skip + if rtype == "status": + logger.info("Subprocess status: %s", resp.get("message", "")) continue - # No matching mailbox (a _gen_lock reader or orphaned). Can't - # un-get from mp.Queue, so just log. (status was handled above.) - logger.debug( - "Dispatcher: no mailbox for request_id=%s type=%s, dropping", - rid, - rtype, - ) + # Route to mailbox if a matching request_id exists + if rid: + with self._mailbox_lock: + mbox = self._mailboxes.get(rid) + if mbox is not None: + mbox.put(resp) + continue + + # No matching mailbox; can't un-get from mp.Queue, so just log. + logger.debug( + "Dispatcher: no mailbox for request_id=%s type=%s, dropping", + rid, + rtype, + ) + except Exception: + logger.exception("Inference dispatcher: failed to route a response; continuing") + continue def _generate_dispatched( self, diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index c0c9a9f656..77bbdc92f5 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -26,6 +26,11 @@ _jobs_lock = threading.Lock() _EMBED_BATCH = 64 # bounds peak memory +# Poll with a timeout so the generator wakes periodically to detect a gone +# client or a terminal job whose worker died without the None sentinel. +_SSE_POLL_SECONDS = 1.0 +_TERMINAL_JOB_STATUSES = {"completed", "failed"} + def _sha256_file(path: str) -> str: h = hashlib.sha256() @@ -178,6 +183,9 @@ def start_ingestion( if ext not in config.UPLOAD_EXTS: raise ValueError(f"unsupported file type: {ext}") + # Reclaim queues for finished jobs so the registry stays bounded. + _reap_finished_jobs() + sha = _sha256_file(stored_path) conn = rag_db.get_connection() try: @@ -248,19 +256,86 @@ def _new_job( return job_id +def _reap_finished_jobs() -> None: + """Drop per-job queues whose DB row already reached a terminal status. + + Otherwise removed only by ``job_events`` after the ``None`` sentinel, so a + caller that polls ``/jobs/{id}`` instead of streaming would grow ``_jobs`` + forever. Safe while streaming: ``job_events`` holds its queue reference. + """ + with _jobs_lock: + job_ids = list(_jobs.keys()) + for jid in job_ids: + row = get_job_status(jid) + if row is not None and row.get("status") in _TERMINAL_JOB_STATUSES: + with _jobs_lock: + _jobs.pop(jid, None) + + def job_events(job_id: str): - """Yield job events for SSE; ends when the worker signals completion.""" + """Yield job events for SSE; ends when the worker signals completion. + + Timed ``get`` so the generator can't block forever: it wakes to heartbeat, + to notice a disconnected client, and to stop on a terminal DB status (a hard + worker death that skipped the ``None`` sentinel). Drops the queue only on a + terminal exit, never on an early client disconnect. + + It deliberately does *not* end on idle alone: a long silent stage (e.g. + embedding a large doc) is not a failure, and ending there would send + ``[DONE]`` with the row still pending, which the client treats as completion. + The stream ends only on a terminal status, the ``None`` sentinel, or disconnect. + """ with _jobs_lock: q = _jobs.get(job_id) if q is None: return - while True: - event = q.get() - if event is None: - break - yield event - with _jobs_lock: - _jobs.pop(job_id, None) + terminal = False + try: + while True: + try: + event = q.get(timeout = _SSE_POLL_SECONDS) + except queue.Empty: + try: + row = get_job_status(job_id) + except Exception: # noqa: BLE001 + # A transient status read (e.g. the DB momentarily locked) must + # not abort the stream: routes/rag.py would turn the raised + # exception into a terminal {type: error} frame and the UI would + # drop a document whose worker is still running. Heartbeat and + # retry on the next poll instead. + logger.warning( + "job_events status read failed for %s; continuing", job_id, exc_info = True + ) + yield {"type": "heartbeat"} + continue + if row is None or row.get("status") in _TERMINAL_JOB_STATUSES: + # Worker finished (or row gone); stop and let the client reconcile via getJob. + terminal = True + break + yield {"type": "heartbeat"} + continue + if event is None: + terminal = True + break + yield event + finally: + # Drop the queue once nothing more will be emitted into it: either a + # terminal exit, or a disconnect after the job already finished (the UI + # stops on the terminal event, before [DONE], so terminal is still False + # here -- _run writes the terminal DB status before emitting it). Keep it + # only while the worker is still running, so an early disconnect can + # reconnect and resume its events. + if not terminal: + try: + row = get_job_status(job_id) + terminal = row is None or row.get("status") in _TERMINAL_JOB_STATUSES + except Exception: # noqa: BLE001 + # Can't confirm terminality (transient DB error) -- keep the queue so + # a reconnect can resume rather than orphaning a live worker's events. + terminal = False + if terminal: + with _jobs_lock: + _jobs.pop(job_id, None) def get_job_status(job_id: str) -> dict | None: diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py index c2b99c0f18..44c39337fb 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -314,25 +314,50 @@ def register_worker( worker_token = hf_token def _watch() -> None: - finalize_worker_exit( - registry, - key, - proc, - hf_token = worker_token, - label = label, - log_prefix = log_prefix, - logger = logger, - repo_type = repo_type, - repo_id = repo_id, - transport = transport, - ) - if registry.get_job(key).state in ("error", "cancelled"): - download_registry.purge_empty_marker_dir( - repo_type, - repo_id, - download_registry.variant_from_key(key), + try: + finalize_worker_exit( + registry, + key, + proc, + hf_token = worker_token, + label = label, + log_prefix = log_prefix, + logger = logger, + repo_type = repo_type, + repo_id = repo_id, + transport = transport, ) - hf_cache_scan.invalidate_hf_cache_scans() + except Exception: + # finalize_worker_exit is the only thing that clears running/cancelling; + # if it raises, force a terminal state so claim() isn't blocked until restart. + logger.exception("download watcher crashed for %s", key) + # finalize may have raised before reaping the worker; terminate the + # still-registered Popen first, else the terminal set_job clears the + # repo guard and a live worker would race a retry on the same repo. + try: + kill_and_reap_process(proc, label = label, logger = logger) + except Exception: + logger.exception("failed to reap worker after watcher crash for %s", key) + try: + registry.drop_process(key, proc) + except Exception: + logger.exception("failed to drop worker after watcher crash for %s", key) + try: + registry.set_job(key, "error", "download watcher crashed") + except Exception: + logger.exception("failed to mark %s errored after watcher crash", key) + finally: + try: + if registry.get_job(key).state in ("error", "cancelled"): + download_registry.purge_empty_marker_dir( + repo_type, + repo_id, + download_registry.variant_from_key(key), + ) + except Exception: + logger.exception("post-finalize marker cleanup failed for %s", key) + finally: + hf_cache_scan.invalidate_hf_cache_scans() threading.Thread(target = _watch, name = watch_name, daemon = True).start() return True diff --git a/studio/backend/main.py b/studio/backend/main.py index a8e81b68e6..731625ca74 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -487,6 +487,16 @@ async def lifespan(app: FastAPI): import structlog structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc) + # Same for RAG: fail ingestion jobs stranded mid-ingest by a crash. + try: + from storage.rag_db import reconcile_orphaned_ingestion_jobs + reconcile_orphaned_ingestion_jobs() + except Exception as exc: + import structlog + structlog.get_logger(__name__).warning( + "reconcile_orphaned_ingestion_jobs failed at startup: %s", exc + ) + _start_helper_precache_if_enabled() # Warm the RAG embedder so the first upload skips the cold load. Non-fatal. diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index d2b3bf94e9..92ecdbfb5b 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -74,9 +74,81 @@ _LOGIN_WINDOW_SECONDS = 60.0 _LOGIN_MAX_FAILS = 5 _LOGIN_IP_MAX_FAILS = 30 _LOGIN_LOCKOUT_SECONDS = 60 -# Bucket-dict cap. On overflow, prune stale entries; if still full the failure -# folds into the per-IP aggregate only. +# Bucket-dict cap. On overflow, reclaim expired buckets; a new IP that still can't +# fit falls back to a sharded overflow rather than evicting a hot bucket. _LOGIN_MAX_BUCKETS = 4096 +# Last full stale-sweep time; rate-limits the O(n) sweep under a burst of new IPs. +_LAST_IP_PRUNE = 0.0 +# Sharded overflow for per-IP failures that can't get their own bucket while the +# dict is saturated. Each shard is a small fixed-capacity dict ``ip -> [count, +# window_start]``: a per-IP count (so a source is throttled, and cleared on +# success, by its own failures -- no cross-IP collateral) with hard-bounded +# memory and O(1) lookups. When a shard is full a new IP evicts the lowest-count +# entry (and starts clean, never inheriting its count) rather than growing without +# bound, so a high-cardinality spray can't blow memory/CPU the way a per-failure +# deque could; a persistent attacker keeps a high count and is never the one +# evicted. +_LOGIN_IP_OVERFLOW_SHARDS = 256 +_LOGIN_IP_OVERFLOW_MAX = 64 # distinct IPs tracked per shard +_LOGIN_IP_OVERFLOW: list[dict] = [dict() for _ in range(_LOGIN_IP_OVERFLOW_SHARDS)] + + +def _overflow_shard(ip: str) -> dict: + return _LOGIN_IP_OVERFLOW[hash(ip) % _LOGIN_IP_OVERFLOW_SHARDS] + + +def _overflow_record(ip: str, now: float) -> int: + """Record an overflow failure for ``ip`` and return its windowed count.""" + shard = _overflow_shard(ip) + entry = shard.get(ip) + if entry is not None: + if now - entry[1] > _LOGIN_WINDOW_SECONDS: + entry[0], entry[1] = 1, now + else: + # Only "at or above the per-IP threshold" matters for blocking, so cap + # the count there. This also keeps the migration into a per-IP bucket + # bounded -- without the cap a saturated source could accrue an + # unbounded count, then materialize one deque entry per failure + # (``[start] * carried``) on the next attempt, allocating an arbitrarily + # large deque while holding the login lock. + entry[0] = min(entry[0] + 1, _LOGIN_IP_MAX_FAILS) + return entry[0] + if len(shard) >= _LOGIN_IP_OVERFLOW_MAX: + # Make room by dropping the lowest-count entry, but the new source starts + # clean -- never inherit the evicted IP's failures, or an unrelated source + # could be 429'd after one attempt. Worst case under a saturated shard is + # that a heavy hitter briefly resets, not that a bystander is blocked. + del shard[min(shard, key = lambda k: shard[k][0])] + shard[ip] = [1, now] + return 1 + + +def _overflow_blocked(ip: str, now: float) -> int: + """Seconds this IP is throttled by its own overflow count, or 0.""" + shard = _overflow_shard(ip) + entry = shard.get(ip) + if entry is None: + return 0 + if now - entry[1] > _LOGIN_WINDOW_SECONDS: + del shard[ip] + return 0 + if entry[0] >= _LOGIN_IP_MAX_FAILS: + return max(1, int(_LOGIN_WINDOW_SECONDS - (now - entry[1]))) + return 0 + + +def _overflow_take(ip: str, now: float) -> tuple[int, float]: + """Pop ip's overflow entry, returning its ``(count, window_start)`` so the + count can migrate into a fresh per-IP bucket. ``(0, now)`` if none/expired.""" + entry = _overflow_shard(ip).pop(ip, None) + if entry is None or now - entry[1] > _LOGIN_WINDOW_SECONDS: + return 0, now + # Cap the carried count so the bucket migration never allocates more than the + # per-IP threshold worth of deque entries (defensive; _overflow_record already + # clamps, but keep the bound at the consumption site too). + return min(entry[0], _LOGIN_IP_MAX_FAILS), entry[1] + + # Unrepresentable as a real username (leading NUL); folds unknown-user attempts # into one slot so attacker cardinality can't blow the bucket dict. _UNKNOWN_LOGIN_USER = "\x00unknown-user" @@ -169,13 +241,50 @@ def _prune_stale_buckets(now: float) -> None: _LOGIN_BUCKETS.pop(key, None) +def _prune_stale_ip_buckets(now: float) -> None: + """Drop empty / expired per-IP buckets to bound memory under spray. + + The dict is otherwise reclaimed only on a successful login, so a failure-only + spray from many (or spoofed) IPs would grow it without bound. + """ + stale: list[str] = [] + for bucket_ip, bucket in _LOGIN_IP_BUCKETS.items(): + _prune_bucket(bucket, now) + if not bucket: + stale.append(bucket_ip) + for bucket_ip in stale: + _LOGIN_IP_BUCKETS.pop(bucket_ip, None) + + def _record_login_failure(key: tuple[str, str]) -> int: + global _LAST_IP_PRUNE now = time.monotonic() ip, _username = key with _LOGIN_BUCKETS_LOCK: - ip_bucket = _LOGIN_IP_BUCKETS.setdefault(ip, deque()) - _prune_bucket(ip_bucket, now) - ip_bucket.append(now) + # Keep the dict bounded without disabling throttling and without letting a + # spray reset a hot bucket: for a new IP at the cap, reclaim expired buckets + # (rate-limited) to make room. + ip_bucket = _LOGIN_IP_BUCKETS.get(ip) + if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS: + if now - _LAST_IP_PRUNE >= 1.0: + _prune_stale_ip_buckets(now) + _LAST_IP_PRUNE = now + if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS: + # Still full -- every bucket is hot. Count this failure in the IP's + # bounded overflow shard instead of evicting a live one, so the spray + # stays throttled but can't push out (and reset) any IP's own counter. + ip_fails = _overflow_record(ip, now) + else: + if ip_bucket is None: + ip_bucket = _LOGIN_IP_BUCKETS[ip] = deque() + # Carry over any overflow failures this IP accrued while the dict + # was saturated, so straddling the overflow -> bucket transition + # can't double the effective per-IP limit. + carried, start = _overflow_take(ip, now) + ip_bucket.extend([start] * carried) + _prune_bucket(ip_bucket, now) + ip_bucket.append(now) + ip_fails = len(ip_bucket) if key not in _LOGIN_BUCKETS and len(_LOGIN_BUCKETS) >= _LOGIN_MAX_BUCKETS: _prune_stale_buckets(now) @@ -184,8 +293,8 @@ def _record_login_failure(key: tuple[str, str]) -> int: _prune_bucket(account_bucket, now) account_bucket.append(now) return len(account_bucket) - # Bucket dict at cap; per-IP cap still applies via ip_bucket. - return len(ip_bucket) + # Both dicts at cap (sustained spray): fall back to the per-IP count. + return ip_fails def _blocked_for(bucket: deque | None, now: float, max_fails: int) -> int: @@ -202,10 +311,16 @@ def _login_blocked(key: tuple[str, str]) -> int: now = time.monotonic() ip, _username = key with _LOGIN_BUCKETS_LOCK: - return max( - _blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), + # Honor the IP's overflow shard regardless of current dict capacity: a + # source counted there during saturation must stay throttled until those + # failures age out, even if a bucket later frees up -- otherwise a fresh + # bucket would reset it. Shards are empty outside saturation, so this is a + # no-op in the common case. + ip_blocked = max( _blocked_for(_LOGIN_IP_BUCKETS.get(ip), now, _LOGIN_IP_MAX_FAILS), + _overflow_blocked(ip, now), ) + return max(_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), ip_blocked) def _clear_login_bucket(key: tuple[str, str]) -> None: @@ -213,6 +328,10 @@ def _clear_login_bucket(key: tuple[str, str]) -> None: with _LOGIN_BUCKETS_LOCK: _LOGIN_BUCKETS.pop(key, None) _LOGIN_IP_BUCKETS.pop(ip, None) + # A successful login resets the IP's throttle, including any overflow it + # accumulated during saturation (drop only this IP's entry, so a + # shard-mate's throttle is untouched). + _overflow_shard(ip).pop(ip, None) # Sync def (not async): compute_identity_proof touches SQLite on the first call, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e77616fb9f..a9c85fa99a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4486,6 +4486,14 @@ async def _proxy_to_external_provider( except Exception as exc: logger.error("external_provider.stream_error", error = str(exc)) api_monitor.fail(monitor_id, _friendly_error(exc)) + # Surface the failure: a bare EOF (e.g. after a read timeout) is treated + # by the chat client as success, saving a partial answer with no error. + yield ( + "data: " + + json.dumps({"error": {"message": _friendly_error(exc), "type": "server_error"}}) + + "\n\n" + ) + yield "data: [DONE]\n\n" finally: try: await gen.aclose() diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 38a2cab389..3818fe9f73 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -847,6 +847,11 @@ async def stream_training_progress( ) while backend.is_training_active(): + # Client gone: end the generator without falling through to the final + # "complete" frame, which a buffered/proxy consumer could otherwise read + # as a finished run while training is still active. + if await request.is_disconnected(): + return try: tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None) live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0 diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index 564e3284f8..4da600d768 100644 --- a/studio/backend/storage/rag_db.py +++ b/studio/backend/storage/rag_db.py @@ -156,3 +156,71 @@ def vec_table_exists(conn: sqlite3.Connection) -> bool: "SELECT 1 FROM sqlite_master WHERE type='table' AND name='chunks_vec'" ).fetchone() return row is not None + + +def _delete_document_chunks(conn, document_id: str) -> None: + """Delete a document's chunk rows (chunks/chunks_fts/chunks_vec), keeping the + documents row. Used when reconciling a half-ingested doc to failed: retrieval + filters by scope not status, so leftover chunks would stay citable.""" + chunk_ids = [ + r["id"] + for r in conn.execute( + "SELECT id FROM chunks WHERE document_id=?", (document_id,) + ).fetchall() + ] + if not chunk_ids: + return + has_vec = vec_table_exists(conn) + for chunk_id in chunk_ids: + conn.execute("DELETE FROM chunks_fts WHERE chunk_id=?", (chunk_id,)) + if has_vec: + conn.execute("DELETE FROM chunks_vec WHERE chunk_id=?", (chunk_id,)) + conn.execute("DELETE FROM chunks WHERE document_id=?", (document_id,)) + + +def reconcile_orphaned_ingestion_jobs() -> int: + """Fail ingestion jobs/documents left mid-flight by a crash so they stop + showing as stuck "processing" and become re-ingestible. Run at startup. + No-op without RAG. Returns the number of jobs reset. + """ + if not RAG_AVAILABLE: + return 0 + conn = get_connection() + try: + rows = conn.execute( + "SELECT id, document_id FROM ingestion_jobs " + "WHERE status NOT IN ('completed', 'failed')" + ).fetchall() + for row in rows: + doc = conn.execute( + "SELECT status FROM documents WHERE id=?", (row["document_id"],) + ).fetchone() + if doc is not None and doc["status"] == "completed": + # Worker finished indexing before the crash but didn't retire the + # job row. Mark the job completed (not failed) and keep its chunks, + # so the UI's getJob fallback after restart doesn't flag a + # searchable document as a failed ingestion. + conn.execute( + "UPDATE ingestion_jobs SET status='completed', stage='done', " + "progress=1.0, error=NULL WHERE id=?", + (row["id"],), + ) + continue + conn.execute( + "UPDATE ingestion_jobs SET status='failed', stage='error', " + "error='Server restarted during ingestion' WHERE id=?", + (row["id"],), + ) + conn.execute( + "UPDATE documents SET status='failed' " + "WHERE id=? AND status NOT IN ('completed', 'failed')", + (row["document_id"],), + ) + # A failed or still-in-flight doc must not leave citable chunks + # (retrieval filters by scope, not status); also drops any chunks of a + # doc already 'failed' before the crash. + _delete_document_chunks(conn, row["document_id"]) + conn.commit() + return len(rows) + finally: + conn.close() diff --git a/studio/backend/tests/test_data_recipe_pump_resilience.py b/studio/backend/tests/test_data_recipe_pump_resilience.py new file mode 100644 index 0000000000..e702be7811 --- /dev/null +++ b/studio/backend/tests/test_data_recipe_pump_resilience.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Data-recipe job pump resilience. + +The pump is the sole consumer of worker events and sole writer of the job +snapshot the status/SSE endpoints read; a handler error must not kill it, or the +job stays wedged "active" and the workflow key is never retired. Fakes only. +""" + +from __future__ import annotations + +import queue +import sys +import threading +import time +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 core.data_recipe.jobs.manager import JobManager # noqa: E402 +from core.data_recipe.jobs.types import Job # noqa: E402 + + +class _FakeProc: + def __init__(self, alive: bool = True): + self._alive = alive + + def is_alive(self): + return self._alive + + +class _ScriptedQueue: + def __init__(self, events): + self._events = list(events) + + def get(self, timeout = None): + if self._events: + return self._events.pop(0) + raise queue.Empty + + def get_nowait(self): + if self._events: + return self._events.pop(0) + raise queue.Empty + + +def _wait_until(predicate, timeout = 5.0): + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + +def _manager_with_active_job(): + m = JobManager.__new__(JobManager) + m._lock = threading.Lock() + job = Job(job_id = "job-test") + job.status = "active" + m._job = job + m._proc = _FakeProc(alive = True) + m._mp_q = _ScriptedQueue([]) + return m + + +def test_pump_survives_handler_exception_and_still_finalizes(monkeypatch): + m = _manager_with_active_job() + handled: list = [] + + def fake_handle(job, event): + if event.get("type") == "boom": + raise RuntimeError("malformed log line") + handled.append(event.get("type")) + + emitted: list = [] + retired: list = [] + monkeypatch.setattr(m, "_handle_event", fake_handle) + monkeypatch.setattr(m, "_emit", lambda e: emitted.append(e)) + monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j)) + + m._mp_q = _ScriptedQueue( + [{"type": "boom"}, {"type": "log"}, {"type": "boom"}, {"type": "progress"}] + ) + + pump = threading.Thread(target = m._pump_loop, daemon = True) + pump.start() + try: + assert _wait_until( + lambda: handled == ["log", "progress"] + ), "pump must keep processing events after a handler raises" + assert pump.is_alive() + finally: + m._proc._alive = False # worker exits -> pump should finalize and stop + pump.join(timeout = 5) + + assert not pump.is_alive() + # The exited worker is finalized as error (not left wedged "active") and the + # workflow key is retired despite the earlier handler exceptions. + assert m._job.status == "error" + assert retired and retired[0] is m._job + + +def test_pump_finalizes_when_drain_raises(monkeypatch): + m = _manager_with_active_job() + monkeypatch.setattr(m, "_emit", lambda e: None) + retired: list = [] + monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j)) + + class _BadDrainQueue: + def get(self, timeout = None): + raise queue.Empty + + def get_nowait(self): + raise RuntimeError("corrupt drain payload") + + m._proc = _FakeProc(alive = False) + m._mp_q = _BadDrainQueue() + + m._pump_loop() # returns once it sees the dead worker + + assert m._job.status == "error" + assert retired and retired[0] is m._job + + +def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch): + # A read that keeps raising after the child died must not spin the pump + # forever: once the worker is gone it falls through to finalize. + m = _manager_with_active_job() + monkeypatch.setattr(m, "_emit", lambda e: None) + retired: list = [] + monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j)) + + class _BrokenReadQueue: + def get(self, timeout = None): + raise RuntimeError("broken queue pipe") + + def get_nowait(self): + raise queue.Empty + + m._proc = _FakeProc(alive = False) + m._mp_q = _BrokenReadQueue() + + pump = threading.Thread(target = m._pump_loop, daemon = True) + pump.start() + pump.join(timeout = 5) + assert not pump.is_alive(), "pump must finalize a dead worker even when reads keep raising" + assert m._job.status == "error" + assert retired and retired[0] is m._job diff --git a/studio/backend/tests/test_inference_dispatcher_resilience.py b/studio/backend/tests/test_inference_dispatcher_resilience.py new file mode 100644 index 0000000000..2de833db0e --- /dev/null +++ b/studio/backend/tests/test_inference_dispatcher_resilience.py @@ -0,0 +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 + +"""Inference dispatcher resilience. + +The dispatcher thread is the sole consumer of the response queue; if a malformed +response killed it, every in-flight generation would hang forever. A bad response +must be logged and skipped, not fatal. Fakes only. +""" + +from __future__ import annotations + +import queue +import sys +import threading +import time +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 core.inference.orchestrator import InferenceOrchestrator # noqa: E402 + + +class _ScriptedQueue: + def __init__(self, items): + self._items = list(items) + + def get(self, timeout = None): + if self._items: + return self._items.pop(0) + raise queue.Empty + + +def _dispatcher(): + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._dispatcher_stop = threading.Event() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + return o + + +def test_dispatcher_survives_malformed_response_and_routes_next(): + o = _dispatcher() + rid = "req-1" + mbox = queue.Queue() + o._mailboxes = {rid: mbox} + # A non-dict response (resp.get -> AttributeError) must not kill the loop; + # the following valid response must still reach its mailbox. + o._resp_queue = _ScriptedQueue([12345, {"request_id": rid, "type": "token", "text": "hi"}]) + + t = threading.Thread(target = o._dispatcher_loop, daemon = True) + t.start() + try: + got = mbox.get(timeout = 5) + assert got["text"] == "hi", "valid response must route despite the prior bad one" + assert t.is_alive(), "dispatcher must survive a malformed response" + finally: + o._dispatcher_stop.set() + t.join(timeout = 5) + assert not t.is_alive() + + +def test_dispatcher_survives_mailbox_put_error(): + o = _dispatcher() + rid = "req-2" + + class _BadMailbox: + def put(self, _resp): + raise RuntimeError("mailbox is broken") + + good = queue.Queue() + o._mailboxes = {rid: _BadMailbox(), "req-3": good} + o._resp_queue = _ScriptedQueue( + [ + {"request_id": rid, "type": "token", "text": "boom"}, + {"request_id": "req-3", "type": "token", "text": "ok"}, + ] + ) + + t = threading.Thread(target = o._dispatcher_loop, daemon = True) + t.start() + try: + got = good.get(timeout = 5) + assert got["text"] == "ok" + assert t.is_alive() + finally: + o._dispatcher_stop.set() + t.join(timeout = 5) + assert not t.is_alive() diff --git a/studio/backend/tests/test_login_rate_limit.py b/studio/backend/tests/test_login_rate_limit.py index 14b10576da..6f9635e41e 100644 --- a/studio/backend/tests/test_login_rate_limit.py +++ b/studio/backend/tests/test_login_rate_limit.py @@ -29,9 +29,15 @@ def _reset_buckets(): auth_routes._LOGIN_BUCKETS.clear() auth_routes._LOGIN_IP_BUCKETS.clear() + for _shard in auth_routes._LOGIN_IP_OVERFLOW: + _shard.clear() + auth_routes._LAST_IP_PRUNE = 0.0 yield auth_routes._LOGIN_BUCKETS.clear() auth_routes._LOGIN_IP_BUCKETS.clear() + for _shard in auth_routes._LOGIN_IP_OVERFLOW: + _shard.clear() + auth_routes._LAST_IP_PRUNE = 0.0 @pytest.fixture @@ -215,6 +221,245 @@ class TestBucketKeyAndBlocking: # Hard cap respected; further keys don't allocate. assert len(auth_routes._LOGIN_BUCKETS) <= 10 + def test_ip_bucket_cap_bounds_without_disabling_throttling(self, env_no_proxy, monkeypatch): + """The per-IP dict is bounded, but saturating it must NOT disable + throttling: a new IP that keeps failing after the cap is hit is still + blocked (now via the shared overflow counter).""" + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + # Saturate the per-IP dict with distinct source IPs. + for idx in range(50): + auth_routes._record_login_failure((f"198.51.100.{idx}", "admin")) + assert len(auth_routes._LOGIN_IP_BUCKETS) <= 10 # bounded + + # A brand-new IP arriving after saturation is still throttled: it can't get + # its own bucket, so its failures land in the shared overflow counter. + victim = ("203.0.113.99", "admin") + for _ in range(5): + auth_routes._record_login_failure(victim) + assert auth_routes._login_blocked(victim) > 0 + + def test_saturating_spray_cannot_reset_a_hot_ip_bucket(self, env_no_proxy, monkeypatch): + """An IP flooding the dict must not evict (and reset) its own hot bucket. + + With FIFO eviction the oldest-inserted bucket -- the attacker's own, now + blocked -- was popped once enough fresh IPs arrived, letting the attacker + retry as first-seen. The overflow counter must keep it throttled. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + # Neutralize account-bucket blocking so this isolates the per-IP path. + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + attacker = ("203.0.113.7", "admin") + for _ in range(5): + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) > 0 # attacker is throttled + + # Attacker sprays many distinct IPs to try to push its own bucket out. + for idx in range(100): + auth_routes._record_login_failure((f"198.51.100.{idx}", "admin")) + + # Still throttled: its hot bucket survived rather than being evicted. + assert auth_routes._login_blocked(attacker) > 0 + + def test_overflow_is_sharded_so_a_hot_ip_does_not_block_unrelated_ips( + self, env_no_proxy, monkeypatch + ): + """A saturating spray must not globally deny login: a hot overflow shard + throttles only the IPs that hash to it, not every new unbucketed client. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + # Neutralize account-bucket blocking so this isolates the per-IP path. + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + # Saturate the bucket dict so further new IPs fall through to overflow. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + + # Drive one IP's real overflow shard hot. + attacker_ip = "198.51.100.7" + for _ in range(5): + auth_routes._record_login_failure((attacker_ip, "admin")) + assert auth_routes._login_blocked((attacker_ip, "admin")) > 0 + + # A new IP in a *different* shard must not be denied (a single global + # counter would block it; a sharded one preserves per-source isolation). + attacker_shard = auth_routes._overflow_shard(attacker_ip) + victim_ip = next( + f"203.0.113.{i}" + for i in range(256) + if auth_routes._overflow_shard(f"203.0.113.{i}") is not attacker_shard + ) + assert auth_routes._login_blocked((victim_ip, "admin")) == 0 + + def test_overflow_throttle_survives_capacity_freeing(self, env_no_proxy, monkeypatch): + """A source throttled via overflow must stay throttled even if a bucket + frees up before the window expires; otherwise a fresh bucket resets it. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + # Neutralize account-bucket blocking so this isolates the per-IP path. + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + # Saturate the dict, then drive a source's overflow shard hot. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + attacker = ("198.51.100.7", "admin") + for _ in range(5): + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) > 0 + + # A successful login from another IP frees a bucket slot. + auth_routes._clear_login_bucket(("10.0.0.0", "admin")) + assert len(auth_routes._LOGIN_IP_BUCKETS) < auth_routes._LOGIN_MAX_BUCKETS + + # Still throttled (overflow shard still hot), and a new failure that now + # gets a fresh per-IP bucket must not reset the throttle. + assert auth_routes._login_blocked(attacker) > 0 + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) > 0 + + def test_overflow_shard_is_memory_bounded_under_cardinality_spray( + self, env_no_proxy, monkeypatch + ): + """A high-cardinality spray must not grow overflow memory without bound: + each shard tracks at most _LOGIN_IP_OVERFLOW_MAX distinct IPs. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_OVERFLOW_MAX", 8) + + # Saturate the dict, then spray thousands of distinct one-off IPs. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + for idx in range(5000): + auth_routes._record_login_failure((f"198.51.{idx // 256}.{idx % 256}", "admin")) + + assert all(len(shard) <= 8 for shard in auth_routes._LOGIN_IP_OVERFLOW) + + def test_overflow_eviction_does_not_inherit_count_onto_new_ip(self, env_no_proxy, monkeypatch): + """Evicting a hot entry to make room must not hand its failure count to the + new source; one attempt from an unrelated IP must not 429 it. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_OVERFLOW_MAX", 2) + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + # Force every overflow IP into one shard so we can saturate it. + shard0 = auth_routes._LOGIN_IP_OVERFLOW[0] + monkeypatch.setattr(auth_routes, "_overflow_shard", lambda _ip: shard0) + + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + # Fill the shard (cap 2) with two hot IPs at/over the threshold. + for _ in range(5): + auth_routes._record_login_failure(("198.51.100.1", "admin")) + for _ in range(5): + auth_routes._record_login_failure(("198.51.100.2", "admin")) + assert len(shard0) == 2 + + # A new IP evicts the lowest-count entry; it must start clean, so one + # failure leaves it below the threshold and unblocked. + new_ip = ("203.0.113.50", "admin") + auth_routes._record_login_failure(new_ip) + assert auth_routes._login_blocked(new_ip) == 0 + + def test_overflow_count_migrates_into_new_bucket(self, env_no_proxy, monkeypatch): + """Straddling the overflow -> bucket transition must not double the per-IP + limit: the overflow count carries into the freshly created bucket. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + # Saturate, then push one IP to 4 overflow failures (one below threshold). + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + attacker = ("198.51.100.7", "admin") + for _ in range(4): + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) == 0 # 4 < 5 + + # Free a slot so the next failure lands in a fresh per-IP bucket. + auth_routes._clear_login_bucket(("10.0.0.0", "admin")) + # One more failure must throttle (4 carried + 1 = 5), not reset to 1. + auth_routes._record_login_failure(attacker) + assert auth_routes._login_blocked(attacker) > 0 + + def test_overflow_migration_is_bounded_not_one_entry_per_failure( + self, env_no_proxy, monkeypatch + ): + """A saturated IP can rack up many overflow failures; migrating them into a + fresh bucket must allocate at most the per-IP threshold worth of entries, + not one deque entry per recorded failure (which would let a single later + attempt allocate an arbitrarily large deque under the login lock). + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100000) + + # Saturate the dict, then hammer one IP far past the threshold in overflow. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + attacker_ip = "198.51.100.7" + attacker = (attacker_ip, "admin") + for _ in range(5000): + auth_routes._record_login_failure(attacker) + # The stored overflow count is clamped at the threshold, not 5000. + entry = auth_routes._overflow_shard(attacker_ip).get(attacker_ip) + assert entry is not None and entry[0] <= auth_routes._LOGIN_IP_MAX_FAILS + + # Free a slot so the next failure migrates the overflow count into a bucket. + auth_routes._clear_login_bucket(("10.0.0.0", "admin")) + auth_routes._record_login_failure(attacker) + bucket = auth_routes._LOGIN_IP_BUCKETS[attacker_ip] + # Bounded by the threshold (+1 for the triggering failure), not ~5000. + assert len(bucket) <= auth_routes._LOGIN_IP_MAX_FAILS + 1 + # Still throttled -- bounding the migration must not weaken the limit. + assert auth_routes._login_blocked(attacker) > 0 + + def test_successful_login_clears_overflow_throttle(self, env_no_proxy, monkeypatch): + """A successful login resets the IP's throttle, including overflow, so a + single later typo is not immediately blocked. + """ + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_FAILS", 100) + + # Saturate the dict, then push one IP into overflow until it is throttled. + for idx in range(10): + auth_routes._record_login_failure((f"10.0.0.{idx}", "admin")) + ip = ("198.51.100.7", "admin") + for _ in range(5): + auth_routes._record_login_failure(ip) + assert auth_routes._login_blocked(ip) > 0 + + # A successful login from that IP clears its overflow entries... + auth_routes._clear_login_bucket(ip) + assert auth_routes._login_blocked(ip) == 0 + # ...and a single subsequent failure does not immediately re-block it. + auth_routes._record_login_failure(ip) + assert auth_routes._login_blocked(ip) == 0 + # ---------- /login 429 body ---------- diff --git a/studio/backend/tests/test_rag_job_events_queue_lifecycle.py b/studio/backend/tests/test_rag_job_events_queue_lifecycle.py new file mode 100644 index 0000000000..0eb115c562 --- /dev/null +++ b/studio/backend/tests/test_rag_job_events_queue_lifecycle.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""job_events keeps the per-job queue registered only while the worker runs. + +``_emit()`` writes to ``_jobs[job_id]`` while the worker runs; if an early SSE +disconnect removed that queue, later events would be dropped and a reconnect +would see only ``[DONE]`` and mark a running job complete. So keep it on an early +disconnect of a running job, but drop it on a terminal exit or a disconnect after +the job already finished; ``_reap_finished_jobs`` sweeps any leftovers. +""" + +import queue +import sqlite3 +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) + +import core.rag.ingestion as ing + + +def test_early_disconnect_keeps_queue_registered(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + # Job is still running; nothing terminal has happened. + monkeypatch.setattr(ing, "get_job_status", lambda _jid: {"status": "running"}) + jid = "job-early-disconnect" + ing._jobs[jid] = queue.Queue() + try: + gen = ing.job_events(jid) + next(gen) # enter loop: Empty -> non-terminal -> heartbeat + gen.close() # client disconnects before the job finishes + assert ( + jid in ing._jobs + ), "queue must survive an early disconnect so the worker can still emit" + finally: + ing._jobs.pop(jid, None) + + +def test_terminal_sentinel_removes_queue(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + jid = "job-terminal-sentinel" + q = queue.Queue() + q.put({"type": "progress", "stage": "embedding", "progress": 0.5}) + q.put(None) # worker finished -> sentinel + ing._jobs[jid] = q + try: + events = list(ing.job_events(jid)) # drains progress, then None -> terminal + assert any(e.get("type") == "progress" for e in events) + assert jid not in ing._jobs, "queue must be removed once the job is terminal" + finally: + ing._jobs.pop(jid, None) + + +def test_disconnect_after_terminal_event_removes_queue(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + # Worker finished: the DB row is terminal and a complete event is queued. The + # UI reads that event and disconnects (reader.cancel) before the None sentinel, + # so the queue must still drop rather than linger until the next reap. + monkeypatch.setattr(ing, "get_job_status", lambda _jid: {"status": "completed"}) + jid = "job-disconnect-after-complete" + q = queue.Queue() + q.put({"type": "complete", "num_chunks": 3}) + q.put(None) + ing._jobs[jid] = q + try: + gen = ing.job_events(jid) + assert next(gen)["type"] == "complete" # client receives the terminal event + gen.close() # disconnects before draining the sentinel + assert jid not in ing._jobs, "a finished job's queue must drop on disconnect" + finally: + ing._jobs.pop(jid, None) + + +def test_transient_status_read_failure_does_not_end_stream(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + # The heartbeat poll hits a momentarily-locked DB. That must not propagate: the + # SSE route would turn the raised error into a terminal {type: error} frame and + # the UI would drop a document whose worker is still running. The stream should + # heartbeat and keep the queue so the worker can finish / a reconnect can resume. + calls = {"n": 0} + + def flaky_status(_jid): + calls["n"] += 1 + if calls["n"] == 1: + raise sqlite3.OperationalError("database is locked") + return {"status": "running"} + + monkeypatch.setattr(ing, "get_job_status", flaky_status) + jid = "job-transient-read-failure" + ing._jobs[jid] = queue.Queue() + try: + gen = ing.job_events(jid) + assert next(gen) == {"type": "heartbeat"} # transient error -> heartbeat, no raise + gen.close() + assert jid in ing._jobs, "an unconfirmed (transient-error) status must keep the queue" + finally: + ing._jobs.pop(jid, None) + + +def test_terminal_db_status_removes_queue(monkeypatch): + monkeypatch.setattr(ing, "_SSE_POLL_SECONDS", 0.01) + # No events arrive, but the DB row reports the job finished (hard worker death + # that skipped the sentinel): the stream ends and the queue is reaped. + monkeypatch.setattr(ing, "get_job_status", lambda _jid: {"status": "completed"}) + jid = "job-terminal-db" + ing._jobs[jid] = queue.Queue() + try: + list(ing.job_events(jid)) + assert jid not in ing._jobs, "a terminal DB status must remove the queue" + finally: + ing._jobs.pop(jid, None) diff --git a/studio/backend/tests/test_rag_reconcile_orphaned.py b/studio/backend/tests/test_rag_reconcile_orphaned.py new file mode 100644 index 0000000000..c6932e4588 --- /dev/null +++ b/studio/backend/tests/test_rag_reconcile_orphaned.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Startup reconciliation must not strip chunks from already-completed docs. + +A crash can leave an ingestion_jobs row non-terminal after the worker already +committed the document as ``completed`` with all its chunks. Reconciliation flips +the orphaned job to ``failed`` but must touch the document (and its chunks) only +when it actually transitions the document to ``failed`` -- otherwise a completed +source loses every chunk yet still reports ``completed``, so retrieval finds +nothing and dedup (``status != 'failed'``) blocks re-ingest. +""" + +import math + +from core.rag import store +from core.rag.chunking import Chunk +from storage import rag_db + +VOCAB = ["alpha", "bravo", "charlie", "delta"] + + +def _embed(text): + v = [float(text.lower().count(w)) for w in VOCAB] + n = math.sqrt(sum(x * x for x in v)) or 1.0 + return [x / n for x in v] + + +def _chunk(text, index = 0): + return Chunk( + text = text, + token_count = len(text.split()), + page_number = None, + source_page_index = 0, + chunk_index = index, + page_char_start = 0, + page_char_end = len(text), + ) + + +def _add_doc(conn, scope, doc_id, status, texts): + store.create_document( + conn, scope = scope, filename = f"{doc_id}.txt", sha256 = doc_id, document_id = doc_id + ) + store.add_chunks( + conn, scope, doc_id, [_chunk(t, i) for i, t in enumerate(texts)], [_embed(t) for t in texts] + ) + store.set_document_status(conn, doc_id, status, num_chunks = len(texts)) + + +def _orphan_job( + conn, + doc_id, + scope, + status = "running", +): + conn.execute( + "INSERT INTO ingestion_jobs(id, document_id, scope, status, stage, progress, created_at) " + "VALUES(?,?,?,?,?,?,datetime('now'))", + (f"job-{doc_id}", doc_id, scope, status, "embedding", 0.5), + ) + conn.commit() + + +def _chunk_count(conn, doc_id): + return conn.execute("SELECT COUNT(*) FROM chunks WHERE document_id=?", (doc_id,)).fetchone()[0] + + +def _job_status(conn, doc_id): + return conn.execute( + "SELECT status FROM ingestion_jobs WHERE id=?", (f"job-{doc_id}",) + ).fetchone()["status"] + + +def test_completed_doc_keeps_chunks_when_its_job_is_orphaned(rag_conn): + # Worker finished the document but crashed before retiring the job row. + _add_doc(rag_conn, "kb_a", "done", "completed", ["alpha bravo", "charlie delta"]) + _orphan_job(rag_conn, "done", "kb_a") + + assert rag_db.reconcile_orphaned_ingestion_jobs() == 1 + + # Document stays completed with all chunks; dedup still finds it. + assert store.get_document(rag_conn, "done")["status"] == "completed" + assert _chunk_count(rag_conn, "done") == 2 + assert store.document_by_hash(rag_conn, "kb_a", "done") == "done" + # The orphaned job is reconciled to completed (not failed), so the UI's getJob + # fallback doesn't flag a searchable document as a failed ingestion. + assert _job_status(rag_conn, "done") == "completed" + + +def test_in_flight_doc_is_failed_and_its_chunks_dropped(rag_conn): + # Partial chunks committed, document never marked terminal -> genuine orphan. + _add_doc(rag_conn, "kb_a", "partial", "processing", ["alpha bravo"]) + _orphan_job(rag_conn, "partial", "kb_a") + + assert rag_db.reconcile_orphaned_ingestion_jobs() == 1 + + assert store.get_document(rag_conn, "partial")["status"] == "failed" + assert _chunk_count(rag_conn, "partial") == 0 + # Failed doc is re-ingestible (not deduped). + assert store.document_by_hash(rag_conn, "kb_a", "partial") is None + + +def test_already_failed_doc_has_its_chunks_dropped(rag_conn): + # Worker committed chunks then marked the doc 'failed', but crashed before + # retiring the job row. Reconcile won't re-flip the doc (already failed), but + # its chunks must still be purged so they aren't retrievable/citable. + _add_doc(rag_conn, "kb_a", "failed_doc", "failed", ["alpha bravo"]) + _orphan_job(rag_conn, "failed_doc", "kb_a") + + assert rag_db.reconcile_orphaned_ingestion_jobs() == 1 + + assert store.get_document(rag_conn, "failed_doc")["status"] == "failed" + assert _chunk_count(rag_conn, "failed_doc") == 0 diff --git a/studio/backend/tests/test_training_progress_prep_timeout.py b/studio/backend/tests/test_training_progress_prep_timeout.py index a7e6d4f839..28e2ee37b9 100644 --- a/studio/backend/tests/test_training_progress_prep_timeout.py +++ b/studio/backend/tests/test_training_progress_prep_timeout.py @@ -71,11 +71,17 @@ class _Backend: class _FakeRequest: headers = {} + async def is_disconnected(self): + return False + class _ReconnectRequest: # Reconnect carrying the last step the client already received. headers = {"last-event-id": "10"} + async def is_disconnected(self): + return False + def _raw(response): async def _drain(): diff --git a/studio/backend/tests/test_training_progress_stream_nan.py b/studio/backend/tests/test_training_progress_stream_nan.py index 899527a04d..5cd84bbca5 100644 --- a/studio/backend/tests/test_training_progress_stream_nan.py +++ b/studio/backend/tests/test_training_progress_stream_nan.py @@ -62,6 +62,16 @@ class _FakeBackend: class _FakeRequest: headers = {} + async def is_disconnected(self): + return False + + +class _DisconnectedRequest: + headers = {} + + async def is_disconnected(self): + return True + def _collect_events(response, timeout = 15): async def _drain(): @@ -116,6 +126,20 @@ def test_inactive_stream_completes_with_live_step_and_null_loss(monkeypatch): assert final["loss"] is None +def test_disconnect_while_active_does_not_emit_complete(monkeypatch): + # Client drops mid-run: the stream must end without a terminal "complete" + # frame, which a buffered/proxy consumer could otherwise read as a finished + # run while training is still active. + backend = _FakeBackend(active_polls = 5) + monkeypatch.setattr(rt, "get_training_backend", lambda: backend) + + response = asyncio.run( + rt.stream_training_progress(_DisconnectedRequest(), current_subject = "tester") + ) + raw = _collect_events(response) + assert "event: complete" not in raw + + def test_stream_uses_finite_history_when_progress_in_sync(monkeypatch): backend = _FakeBackend(active_polls = 2) # Live progress agrees with the history tail: normal finite behavior. diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 7112a7877c..15066caa8f 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -844,79 +844,96 @@ export async function* streamChatCompletions( const reader = response.body.getReader(); const decoder = new TextDecoder(); let buffer = ""; + let completed = false; - while (true) { - const { done, value } = await reader.read(); - if (done) { - break; + try { + while (true) { + const { done, value } = await reader.read(); + if (done) { + completed = true; + break; + } + + buffer += decoder.decode(value, { stream: true }); + + let separatorIndex = buffer.search(/\r?\n\r?\n/); + while (separatorIndex >= 0) { + const rawEvent = buffer.slice(0, separatorIndex); + const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; + buffer = buffer.slice(separatorIndex + separatorLength); + + const dataLines = parseSseEvent(rawEvent); + if (dataLines.length === 0) { + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + + const dataText = dataLines.join("\n"); + if (dataText === "[DONE]") { + completed = true; + return; + } + + const parsed = JSON.parse(dataText) as + | OpenAIChatChunk + | { type?: string; content?: string; error?: { message?: string } }; + if ("error" in parsed && parsed.error) { + throw new Error(parsed.error.message || "Stream error"); + } + // Tool status events are custom SSE payloads, not OpenAI chunks + if ("type" in parsed && parsed.type === "tool_status") { + yield { + _toolStatus: parsed.content ?? "", + } as unknown as OpenAIChatChunk; + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + // Diffusion frame: a per-step canvas snapshot. Custom SSE payload (not an OpenAI chunk) with + // no assistant text, surfaced as a transient marker for the in-place renderer, never the transcript. + if ("type" in parsed && parsed.type === "diffusion_frame") { + yield { + _diffusionFrame: parsed, + } as unknown as OpenAIChatChunk; + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + // Tool start/end events carry full input/output for the tool outputs panel + if ( + "type" in parsed && + (parsed.type === "tool_start" || parsed.type === "tool_end") + ) { + yield { _toolEvent: parsed } as unknown as OpenAIChatChunk; + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + // Relay server-side reasoning duration. + if ( + parsed && + typeof parsed === "object" && + "type" in parsed && + parsed.type === "reasoning_summary" + ) { + yield { + _reasoningDurationMs: (parsed as { duration_ms?: number }).duration_ms, + } as unknown as OpenAIChatChunk; + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + yield parsed as OpenAIChatChunk; + separatorIndex = buffer.search(/\r?\n\r?\n/); + } } - - buffer += decoder.decode(value, { stream: true }); - - let separatorIndex = buffer.search(/\r?\n\r?\n/); - while (separatorIndex >= 0) { - const rawEvent = buffer.slice(0, separatorIndex); - const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; - buffer = buffer.slice(separatorIndex + separatorLength); - - const dataLines = parseSseEvent(rawEvent); - if (dataLines.length === 0) { - separatorIndex = buffer.search(/\r?\n\r?\n/); - continue; + } finally { + // Only abort on an early/abnormal exit. After a natural [DONE] (or server + // EOF) the request is logically complete and the backend finalizes its + // api-monitor entry right after the sentinel; cancelling here can be seen as + // a disconnect and mark a successful request as cancelled. + if (!completed) { + try { + await reader.cancel(); + } catch { + // already closed } - - const dataText = dataLines.join("\n"); - if (dataText === "[DONE]") { - return; - } - - const parsed = JSON.parse(dataText) as - | OpenAIChatChunk - | { type?: string; content?: string; error?: { message?: string } }; - if ("error" in parsed && parsed.error) { - throw new Error(parsed.error.message || "Stream error"); - } - // Tool status events are custom SSE payloads, not OpenAI chunks - if ("type" in parsed && parsed.type === "tool_status") { - yield { - _toolStatus: parsed.content ?? "", - } as unknown as OpenAIChatChunk; - separatorIndex = buffer.search(/\r?\n\r?\n/); - continue; - } - // Diffusion frame: a per-step canvas snapshot. Custom SSE payload (not an OpenAI chunk) with - // no assistant text, surfaced as a transient marker for the in-place renderer, never the transcript. - if ("type" in parsed && parsed.type === "diffusion_frame") { - yield { - _diffusionFrame: parsed, - } as unknown as OpenAIChatChunk; - separatorIndex = buffer.search(/\r?\n\r?\n/); - continue; - } - // Tool start/end events carry full input/output for the tool outputs panel - if ( - "type" in parsed && - (parsed.type === "tool_start" || parsed.type === "tool_end") - ) { - yield { _toolEvent: parsed } as unknown as OpenAIChatChunk; - separatorIndex = buffer.search(/\r?\n\r?\n/); - continue; - } - // Relay server-side reasoning duration. - if ( - parsed && - typeof parsed === "object" && - "type" in parsed && - parsed.type === "reasoning_summary" - ) { - yield { - _reasoningDurationMs: (parsed as { duration_ms?: number }).duration_ms, - } as unknown as OpenAIChatChunk; - separatorIndex = buffer.search(/\r?\n\r?\n/); - continue; - } - yield parsed as OpenAIChatChunk; - separatorIndex = buffer.search(/\r?\n\r?\n/); } } } diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index 659dd44f88..f055f1bd61 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -425,5 +425,12 @@ export async function streamExportLogs(options: { } catch (err) { if (isAbortError(err)) return; throw err; + } finally { + // Release the stream lock now instead of leaking the reader until GC. + try { + await reader.cancel(); + } catch { + // already closed + } } } diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index 34ded24ac3..b1215d2d46 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -197,31 +197,40 @@ export async function* streamJobEvents( const decoder = new TextDecoder(); let buffer = ""; - while (true) { - const { done, value } = await reader.read(); - if (done) break; - buffer += decoder.decode(value, { stream: true }); + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); - let separatorIndex = buffer.search(/\r?\n\r?\n/); - while (separatorIndex >= 0) { - const rawEvent = buffer.slice(0, separatorIndex); - const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; - buffer = buffer.slice(separatorIndex + separatorLength); + let separatorIndex = buffer.search(/\r?\n\r?\n/); + while (separatorIndex >= 0) { + const rawEvent = buffer.slice(0, separatorIndex); + const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; + buffer = buffer.slice(separatorIndex + separatorLength); - const dataLines: string[] = []; - for (const line of rawEvent.split(/\r?\n/)) { - if (line.startsWith("data:")) dataLines.push(line.slice(5).trimStart()); - } - if (dataLines.length > 0) { - const dataText = dataLines.join("\n"); - if (dataText === "[DONE]") return; - try { - yield JSON.parse(dataText) as JobEvent; - } catch { - // Ignore unparseable frames; [DONE] still ends the loop. + const dataLines: string[] = []; + for (const line of rawEvent.split(/\r?\n/)) { + if (line.startsWith("data:")) dataLines.push(line.slice(5).trimStart()); } + if (dataLines.length > 0) { + const dataText = dataLines.join("\n"); + if (dataText === "[DONE]") return; + try { + yield JSON.parse(dataText) as JobEvent; + } catch { + // Ignore unparseable frames; [DONE] still ends the loop. + } + } + separatorIndex = buffer.search(/\r?\n\r?\n/); } - separatorIndex = buffer.search(/\r?\n\r?\n/); + } + } finally { + // Release the stream lock now instead of leaking the reader until GC. + try { + await reader.cancel(); + } catch { + // already closed } } } diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index b227b960d1..4fe24b1f6d 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -400,28 +400,37 @@ export async function streamRecipeJobEvents(options: { const decoder = new TextDecoder(); let buffer = ""; - while (true) { - const { value, done } = await reader.read(); - if (done) { - break; - } - buffer += decoder.decode(value, { stream: true }); - let separatorIndex = buffer.search(/\r?\n\r?\n/); - while (separatorIndex >= 0) { - const rawEvent = buffer.slice(0, separatorIndex); - const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; - buffer = buffer.slice(separatorIndex + separatorLength); + try { + while (true) { + const { value, done } = await reader.read(); + if (done) { + break; + } + buffer += decoder.decode(value, { stream: true }); + let separatorIndex = buffer.search(/\r?\n\r?\n/); + while (separatorIndex >= 0) { + const rawEvent = buffer.slice(0, separatorIndex); + const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; + buffer = buffer.slice(separatorIndex + separatorLength); - if (rawEvent.startsWith("retry:")) { + if (rawEvent.startsWith("retry:")) { + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + + const parsed = parseJobEvent(rawEvent); + if (parsed) { + options.onEvent(parsed); + } separatorIndex = buffer.search(/\r?\n\r?\n/); - continue; } - - const parsed = parseJobEvent(rawEvent); - if (parsed) { - options.onEvent(parsed); - } - separatorIndex = buffer.search(/\r?\n\r?\n/); + } + } finally { + // Release the stream lock now instead of leaking the reader until GC. + try { + await reader.cancel(); + } catch { + // already closed } } } diff --git a/studio/frontend/src/features/training/api/train-api.ts b/studio/frontend/src/features/training/api/train-api.ts index af3ea347c2..609e4f8591 100644 --- a/studio/frontend/src/features/training/api/train-api.ts +++ b/studio/frontend/src/features/training/api/train-api.ts @@ -143,37 +143,46 @@ export async function streamTrainingProgress(options: { const decoder = new TextDecoder(); let buffer = ""; - while (true) { - const { value, done } = await reader.read(); - if (done) { - break; - } + try { + while (true) { + const { value, done } = await reader.read(); + if (done) { + break; + } - buffer += decoder.decode(value, { stream: true }); + buffer += decoder.decode(value, { stream: true }); - let separatorIndex = buffer.search(/\r?\n\r?\n/); - while (separatorIndex >= 0) { - const rawEvent = buffer.slice(0, separatorIndex); - const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; - buffer = buffer.slice(separatorIndex + separatorLength); + let separatorIndex = buffer.search(/\r?\n\r?\n/); + while (separatorIndex >= 0) { + const rawEvent = buffer.slice(0, separatorIndex); + const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2; + buffer = buffer.slice(separatorIndex + separatorLength); + + if (rawEvent.startsWith("retry:")) { + separatorIndex = buffer.search(/\r?\n\r?\n/); + continue; + } + + try { + const event = parseSseEvent(rawEvent); + if (event) { + options.onEvent(event); + } + } catch (error) { + if (!isAbortError(error)) { + throw error; + } + } - if (rawEvent.startsWith("retry:")) { separatorIndex = buffer.search(/\r?\n\r?\n/); - continue; } - - try { - const event = parseSseEvent(rawEvent); - if (event) { - options.onEvent(event); - } - } catch (error) { - if (!isAbortError(error)) { - throw error; - } - } - - separatorIndex = buffer.search(/\r?\n\r?\n/); + } + } finally { + // Release the stream lock now instead of leaking the reader until GC. + try { + await reader.cancel(); + } catch { + // already closed } } } From b693ed7c913e97331728079df520d3d7126f27f3 Mon Sep 17 00:00:00 2001 From: James Dawdy Date: Fri, 26 Jun 2026 10:56:52 -0500 Subject: [PATCH 06/49] fix: wrap unprotected evaluate() calls with robust_evaluate() to handle navigation context loss (#6677) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: wrap unprotected evaluate() calls with robust_evaluate() to handle navigation context loss Fixes PR #5911 - Playwright UI test error: 'Execution context was destroyed' The test had several direct page.evaluate() and locator.evaluate() calls that weren't wrapped with robust_evaluate(), which retries when navigation destroys the execution context mid-operation. Changes: - Wrap picker_visible_text() evaluate in robust_evaluate() - Wrap _bubble_count() evaluate in robust_evaluate() - Wrap assistant text query in robust_evaluate() - Wrap theme_item click evaluation in robust_evaluate() - Wrap background color/theme query in robust_evaluate() This ensures all execution context losses from concurrent navigation are properly caught and retried with exponential backoff, preventing transient failures in the UI test suite. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: revert robust_evaluate on theme_item.evaluate per Codex review The theme_item.evaluate('el => el.click()') is side-effecting — retrying after a context loss could double-toggle the theme. It's already inside a 3-attempt try/except loop that handles click failures gracefully. The other 4 changes (all read-only queries) remain wrapped in robust_evaluate() since retrying them is safe. * fix: wrap remaining chat UI evaluate --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 --- tests/studio/playwright_chat_ui.py | 32 ++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index f9147054ee..a892d3414d 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -494,12 +494,15 @@ with sync_playwright() as p: # typeahead actually filters (else an ignored-input regression # would silently pass). def picker_visible_text(): - return page.evaluate("""() => { + return robust_evaluate( + page, + """() => { const el = document.querySelector( '[role="dialog"], [role="listbox"], [role="menu"]' ); return el ? (el.innerText || '').trim() : ''; - }""") + }""", + ) search.fill("qwen") page.wait_for_timeout(800) @@ -535,9 +538,12 @@ with sync_playwright() as p: def _bubble_count(): """Total [data-role='assistant'] elements (empty or not).""" - return page.evaluate("""() => { + return robust_evaluate( + page, + """() => { return document.querySelectorAll('[data-role="assistant"]').length; - }""") + }""", + ) def send_and_wait(prompt, idx): # 1. Wait until the previous turn fully stopped: Send attached @@ -626,8 +632,11 @@ with sync_playwright() as p: send_and_wait(p_, i) shoot("04-after-five-turns") - texts = page.evaluate("""() => Array.from(document.querySelectorAll('[data-role="assistant"]')) - .map(e => (e.innerText || '').trim())""") + texts = robust_evaluate( + page, + """() => Array.from(document.querySelectorAll('[data-role="assistant"]')) + .map(e => (e.innerText || '').trim())""", + ) if len(texts) < len(prompts): fail(f"expected >= {len(prompts)} assistant bubbles, got {len(texts)}") info(f"five turn lengths = {[len(t) for t in texts[:5]]}") @@ -840,7 +849,9 @@ with sync_playwright() as p: # Settle. The ".dark" class on is the ground truth # (theme-store toggles only that); don't gate on ".light". page.wait_for_timeout(700) - bg = page.evaluate("""() => { + bg = robust_evaluate( + page, + """() => { const root = document.documentElement; return { cls: root.className, @@ -848,7 +859,8 @@ with sync_playwright() as p: bg: getComputedStyle(document.body).backgroundColor, rbg: getComputedStyle(root).backgroundColor, }; - }""") + }""", + ) observed.append(bg) shoot(f"10-theme-cycle-{cycle + 1}") info(f" cycle {cycle + 1}: dark={bg['isDark']} body bg={bg['bg']!r}") @@ -1050,7 +1062,8 @@ with sync_playwright() as p: shoot("15d-recent-clicked") info(f"OK clicked recent entry: {t[:60]!r}") # The landed thread must include at least one of our prompts. - turns_text = page.evaluate( + turns_text = robust_evaluate( + page, """() => { const els = document.querySelectorAll( '[data-role="user"], [data-role="assistant"]' @@ -1058,7 +1071,6 @@ with sync_playwright() as p: return Array.from(els).map(e => (e.innerText || '') .toLowerCase()).join(' '); }""", - None, ) clicked_recent = True if any(k in turns_text for k in PROMPT_KEYWORDS): From cb274484a6835daa64b00c84dadf45f82f4ffd61 Mon Sep 17 00:00:00 2001 From: Avaya Aggarwal <119044997+OnePunchMonk@users.noreply.github.com> Date: Sat, 27 Jun 2026 00:00:10 +0530 Subject: [PATCH 07/49] Add GGUF --tensor-parallel CLI option (#6561) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/models/inference.py | 3 +- unsloth_cli/_inference.py | 108 +++++-- unsloth_cli/commands/chat.py | 26 +- unsloth_cli/commands/inference.py | 26 +- unsloth_cli/tests/test_inference_chat.py | 355 +++++++++++++++++++++++ 5 files changed, 494 insertions(+), 24 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 26825a472e..4a3162b09e 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -106,8 +106,7 @@ class LoadRequest(BaseModel): "Extra arguments forwarded verbatim to llama-server for GGUF models. " "One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. " "Studio-managed flags (model identity, port, context length, GPU placement, " - "auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for " - "non-GGUF models." + "auth, UI/server mode) are rejected. Ignored for non-GGUF models." ), ) diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index 0baeb2ffe5..e8a3414d73 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -3,11 +3,12 @@ """Model loading and streaming shared by `inference` and `chat`.""" +import asyncio import os import re import sys from pathlib import Path -from typing import Optional +from typing import List, Optional import typer @@ -211,28 +212,65 @@ def resolve_model_config(model: str, *, hf_token: Optional[str]): return model_config -def _load_gguf_backend(model_config, *, hf_token, max_seq_length): +def _validate_llama_extra_args_or_exit(llama_extra_args: Optional[List[str]]) -> list[str]: + from core.inference.llama_server_args import validate_extra_args + try: + return validate_extra_args(llama_extra_args) + except ValueError as exc: + typer.echo(f"Error: {exc}", err = True) + raise typer.Exit(code = 1) + + +def _load_gguf_backend( + model_config, + *, + hf_token, + max_seq_length, + tensor_parallel: bool = False, + llama_extra_args: Optional[List[str]] = None, +): ensure_studio_backend_path() from core.inference.llama_cpp import LlamaCppBackend + from core.inference.tensor_fallback import load_with_tensor_fallback llama_backend = LlamaCppBackend() + extra_args = _validate_llama_extra_args_or_exit(llama_extra_args) common = dict( hf_variant = model_config.gguf_variant, model_identifier = model_config.identifier, is_vision = model_config.is_vision, n_ctx = max_seq_length, ) - if model_config.gguf_hf_repo: - loaded = llama_backend.load_model( - hf_repo = model_config.gguf_hf_repo, hf_token = hf_token, **common + + async def _attempt_gguf_load( + requested_tensor_parallel: bool, attempt_extra_args: Optional[List[str]] + ) -> bool: + attempt_common = dict( + common, + tensor_parallel = requested_tensor_parallel, + extra_args = attempt_extra_args, ) - else: - loaded = llama_backend.load_model( + if model_config.gguf_hf_repo: + return llama_backend.load_model( + hf_repo = model_config.gguf_hf_repo, + hf_token = hf_token, + **attempt_common, + ) + return llama_backend.load_model( gguf_path = model_config.gguf_file, mmproj_path = model_config.gguf_mmproj_file, mtp_draft_path = model_config.gguf_mtp_file, - **common, + **attempt_common, ) + + loaded = asyncio.run( + load_with_tensor_fallback( + _attempt_gguf_load, + requested_tensor = tensor_parallel, + extra_args = extra_args, + label = model_config.identifier, + ) + ) if not loaded: typer.echo("Model load failed", err = True) raise typer.Exit(code = 1) @@ -245,6 +283,8 @@ def load_chat_backend( hf_token: Optional[str], max_seq_length: int, load_in_4bit: bool, + tensor_parallel: bool = False, + llama_extra_args: Optional[List[str]] = None, model_config = None, fresh_backend: bool = False, ): @@ -259,7 +299,13 @@ def load_chat_backend( typer.echo(f"Loading {model}", err = True) if model_config.is_gguf: - return _load_gguf_backend(model_config, hf_token = hf_token, max_seq_length = max_seq_length) + 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() @@ -447,18 +493,31 @@ class HttpChatBackend: # No redirects: this carries a bearer token (see urlopen_no_redirect). return urlopen_no_redirect(request, timeout = timeout) - def ensure_loaded(self, model: str, *, hf_token, max_seq_length, load_in_4bit) -> None: + def ensure_loaded( + self, + model: str, + *, + hf_token, + max_seq_length, + load_in_4bit, + tensor_parallel: bool = False, + llama_extra_args: Optional[List[str]] = None, + ) -> None: typer.echo(f"Loading {model} on the Studio server", err = True) + payload = { + "model_path": model, + "hf_token": hf_token, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "tensor_parallel": tensor_parallel, + } + if llama_extra_args: + payload["llama_extra_args"] = llama_extra_args try: self._request( "POST", "/api/inference/load", - { - "model_path": model, - "hf_token": hf_token, - "max_seq_length": max_seq_length, - "load_in_4bit": load_in_4bit, - }, + payload, ).close() except Exception as exc: typer.echo(f"Model load failed: {exc}", err = True) @@ -538,7 +597,15 @@ class HttpChatBackend: pass -def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit): +def connect_studio_server( + model: str, + *, + hf_token, + max_seq_length, + load_in_4bit, + tensor_parallel: bool = False, + llama_extra_args: Optional[List[str]] = None, +): """Backend on a running Studio server, or None (caller loads locally).""" base_url = find_studio_server() if not base_url: @@ -576,6 +643,11 @@ def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit) return _refuse("couldn't self-issue a Studio token (is Studio set up here?).") backend = HttpChatBackend(base_url, token) backend.ensure_loaded( - model, hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit + model, + hf_token = hf_token, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + tensor_parallel = tensor_parallel, + llama_extra_args = llama_extra_args, ) return backend diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py index c62916bc75..a483aeeb6c 100644 --- a/unsloth_cli/commands/chat.py +++ b/unsloth_cli/commands/chat.py @@ -1,7 +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 -from typing import Optional +from typing import List, Optional import typer from rich.console import Console @@ -153,6 +153,22 @@ def chat( ), max_seq_length: int = typer.Option(4096, "--max-seq-length"), load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"), + tensor_parallel: bool = typer.Option( + False, + "--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." + ), + ), + llama_extra_args: Optional[List[str]] = typer.Option( + None, + "--llama-extra-arg", + help = ( + "Extra llama-server arg for GGUF models. Repeat for multiple " + "tokens, e.g. --llama-extra-arg=--top-k --llama-extra-arg 20." + ), + ), think: bool = typer.Option( False, "--think/--no-think", @@ -190,7 +206,13 @@ def chat( err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False) raise typer.Exit(code = 1) - load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit) + load_opts = dict( + hf_token = hf_token, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + tensor_parallel = tensor_parallel, + llama_extra_args = llama_extra_args, + ) # 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) diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py index 5dbc32c7d2..1401fa9e9e 100644 --- a/unsloth_cli/commands/inference.py +++ b/unsloth_cli/commands/inference.py @@ -1,7 +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 -from typing import Optional +from typing import List, Optional import typer @@ -31,6 +31,22 @@ def inference( ), max_seq_length: int = typer.Option(2048, "--max-seq-length"), load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"), + tensor_parallel: bool = typer.Option( + False, + "--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." + ), + ), + llama_extra_args: Optional[List[str]] = typer.Option( + None, + "--llama-extra-arg", + help = ( + "Extra llama-server arg for GGUF models. Repeat for multiple " + "tokens, e.g. --llama-extra-arg=--top-k --llama-extra-arg 20." + ), + ), think: bool = typer.Option( False, "--think/--no-think", @@ -55,7 +71,13 @@ def inference( # A running Studio server keeps the model warm between runs, which is # exactly what a one-shot command wants. - load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit) + load_opts = dict( + hf_token = hf_token, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + tensor_parallel = tensor_parallel, + llama_extra_args = llama_extra_args, + ) chat_backend = None if no_server else connect_studio_server(model, **load_opts) if chat_backend is None: chat_backend = load_chat_backend(model, **load_opts) diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py index 770bd0db29..013b7c5a81 100644 --- a/unsloth_cli/tests/test_inference_chat.py +++ b/unsloth_cli/tests/test_inference_chat.py @@ -9,6 +9,7 @@ import inspect import sys import types from pathlib import Path +from types import SimpleNamespace _REPO_ROOT = Path(__file__).resolve().parents[2] if str(_REPO_ROOT) not in sys.path: @@ -16,6 +17,7 @@ if str(_REPO_ROOT) not in sys.path: import typer +import pytest from rich.console import Console from typer.testing import CliRunner @@ -43,6 +45,14 @@ def _chat_app(): return cli +def _inference_app(): + from unsloth_cli.commands.inference import inference + + cli = typer.Typer() + cli.command()(inference) + return cli + + def test_visible_text_passthrough_when_shown(): text = "reasoninganswer" assert visible_text(text, show_thinking = True) == text @@ -81,6 +91,16 @@ def test_inference_think_defaults_off(): assert "--think/--no-think" in (getattr(opt, "param_decls", None) or []) +def test_inference_exposes_gguf_runtime_options(): + from unsloth_cli.commands.inference import inference + + tensor = _option(inference, "tensor_parallel") + assert "--tensor-parallel/--no-tensor-parallel" in (getattr(tensor, "param_decls", None) or []) + + extra = _option(inference, "llama_extra_args") + assert "--llama-extra-arg" in (getattr(extra, "param_decls", None) or []) + + def test_chat_command_is_registered_with_options(): params = inspect.signature(chatmod.chat).parameters assert "model" in params @@ -94,6 +114,12 @@ def test_chat_command_is_registered_with_options(): verbose = _option(chatmod.chat, "verbose") assert {"--verbose", "-v"} <= set(getattr(verbose, "param_decls", None) or []) + tensor = _option(chatmod.chat, "tensor_parallel") + assert "--tensor-parallel/--no-tensor-parallel" in (getattr(tensor, "param_decls", None) or []) + + extra = _option(chatmod.chat, "llama_extra_args") + assert "--llama-extra-arg" in (getattr(extra, "param_decls", None) or []) + class _FakeBackend: def __init__(self): @@ -324,6 +350,243 @@ def test_http_backend_streams_cumulative_text(monkeypatch): assert out == ["He", "Hello"] +def test_http_backend_load_forwards_gguf_runtime_options(monkeypatch): + backend = HttpChatBackend("http://localhost:8888", "token") + requests = [] + + class _OK: + def close(self): + pass + + def fake_request( + method, + path, + payload = None, + timeout = None, + ): + requests.append((method, path, payload, timeout)) + return _OK() + + monkeypatch.setattr(backend, "_request", fake_request) + + backend.ensure_loaded( + "org/model-GGUF", + hf_token = "hf_x", + max_seq_length = 8192, + load_in_4bit = False, + tensor_parallel = True, + llama_extra_args = ["--top-k", "20"], + ) + + assert requests == [ + ( + "POST", + "/api/inference/load", + { + "model_path": "org/model-GGUF", + "hf_token": "hf_x", + "max_seq_length": 8192, + "load_in_4bit": False, + "tensor_parallel": True, + "llama_extra_args": ["--top-k", "20"], + }, + None, + ) + ] + + +def test_http_backend_load_sends_explicit_false_tensor_parallel(monkeypatch): + backend = HttpChatBackend("http://localhost:8888", "token") + requests = [] + + class _OK: + def close(self): + pass + + monkeypatch.setattr( + backend, + "_request", + lambda method, path, payload = None, timeout = None: ( + requests.append((method, path, payload, timeout)), + _OK(), + )[1], + ) + + backend.ensure_loaded( + "org/model-GGUF", + hf_token = None, + max_seq_length = 4096, + load_in_4bit = True, + tensor_parallel = False, + ) + + assert requests[0][2]["tensor_parallel"] is False + + +def test_load_gguf_backend_forwards_local_runtime_options(monkeypatch): + import unsloth_cli._inference as inference + + calls = [] + + class _FakeLlamaCppBackend: + def load_model(self, **kwargs): + calls.append(kwargs) + return True + + fake_llama_cpp = types.ModuleType("core.inference.llama_cpp") + fake_llama_cpp.LlamaCppBackend = _FakeLlamaCppBackend + fake_args = types.ModuleType("core.inference.llama_server_args") + fake_args.validate_extra_args = lambda args: list(args or []) + fake_tensor_fallback = types.ModuleType("core.inference.tensor_fallback") + + async def _passthrough( + attempt_load, + *, + requested_tensor, + extra_args, + label = "", + cancelled = None, + ): + return await attempt_load(requested_tensor, extra_args) + + fake_tensor_fallback.load_with_tensor_fallback = _passthrough + + monkeypatch.setitem(sys.modules, "core", types.ModuleType("core")) + monkeypatch.setitem(sys.modules, "core.inference", types.ModuleType("core.inference")) + monkeypatch.setitem(sys.modules, "core.inference.llama_cpp", fake_llama_cpp) + monkeypatch.setitem(sys.modules, "core.inference.llama_server_args", fake_args) + monkeypatch.setitem(sys.modules, "core.inference.tensor_fallback", fake_tensor_fallback) + monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None) + + config = SimpleNamespace( + gguf_variant = "Q4_K_M", + identifier = "org/model-GGUF", + is_vision = False, + gguf_hf_repo = "org/model-GGUF", + ) + + backend = inference._load_gguf_backend( + config, + hf_token = "hf_x", + max_seq_length = 8192, + tensor_parallel = True, + llama_extra_args = ["--top-k", "20"], + ) + + assert isinstance(backend, ChatBackend) + assert calls == [ + { + "hf_repo": "org/model-GGUF", + "hf_token": "hf_x", + "hf_variant": "Q4_K_M", + "model_identifier": "org/model-GGUF", + "is_vision": False, + "n_ctx": 8192, + "tensor_parallel": True, + "extra_args": ["--top-k", "20"], + } + ] + + +def test_load_gguf_backend_exits_cleanly_on_invalid_extra_args(monkeypatch): + import unsloth_cli._inference as inference + + fake_llama_cpp = types.ModuleType("core.inference.llama_cpp") + fake_llama_cpp.LlamaCppBackend = object + fake_args = types.ModuleType("core.inference.llama_server_args") + + def _raise(_args): + raise ValueError("llama-server flag '--model' is managed by Unsloth Studio") + + fake_args.validate_extra_args = _raise + fake_tensor_fallback = types.ModuleType("core.inference.tensor_fallback") + fake_tensor_fallback.load_with_tensor_fallback = None + + monkeypatch.setitem(sys.modules, "core", types.ModuleType("core")) + monkeypatch.setitem(sys.modules, "core.inference", types.ModuleType("core.inference")) + monkeypatch.setitem(sys.modules, "core.inference.llama_cpp", fake_llama_cpp) + monkeypatch.setitem(sys.modules, "core.inference.llama_server_args", fake_args) + monkeypatch.setitem(sys.modules, "core.inference.tensor_fallback", fake_tensor_fallback) + monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None) + + config = SimpleNamespace( + gguf_variant = "Q4_K_M", + identifier = "org/model-GGUF", + is_vision = False, + gguf_hf_repo = "org/model-GGUF", + ) + + with pytest.raises(typer.Exit) as excinfo: + inference._load_gguf_backend( + config, + hf_token = "hf_x", + max_seq_length = 8192, + llama_extra_args = ["--model"], + ) + + assert excinfo.value.exit_code == 1 + + +def test_load_gguf_backend_uses_tensor_fallback(monkeypatch): + import unsloth_cli._inference as inference + + calls = [] + fallback_calls = [] + + class _FakeLlamaCppBackend: + def load_model(self, **kwargs): + calls.append(kwargs) + return kwargs["tensor_parallel"] is False + + fake_llama_cpp = types.ModuleType("core.inference.llama_cpp") + fake_llama_cpp.LlamaCppBackend = _FakeLlamaCppBackend + fake_args = types.ModuleType("core.inference.llama_server_args") + fake_args.validate_extra_args = lambda args: list(args or []) + fake_tensor_fallback = types.ModuleType("core.inference.tensor_fallback") + + async def _fallback( + attempt_load, + *, + requested_tensor, + extra_args, + label = "", + cancelled = None, + ): + fallback_calls.append((requested_tensor, extra_args, label)) + ok = await attempt_load(requested_tensor, extra_args) + if ok: + return True + return await attempt_load(False, ["--split-mode", "layer"]) + + fake_tensor_fallback.load_with_tensor_fallback = _fallback + + monkeypatch.setitem(sys.modules, "core", types.ModuleType("core")) + monkeypatch.setitem(sys.modules, "core.inference", types.ModuleType("core.inference")) + monkeypatch.setitem(sys.modules, "core.inference.llama_cpp", fake_llama_cpp) + monkeypatch.setitem(sys.modules, "core.inference.llama_server_args", fake_args) + monkeypatch.setitem(sys.modules, "core.inference.tensor_fallback", fake_tensor_fallback) + monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None) + + config = SimpleNamespace( + gguf_variant = "Q4_K_M", + identifier = "org/model-GGUF", + is_vision = False, + gguf_hf_repo = "org/model-GGUF", + ) + + backend = inference._load_gguf_backend( + config, + hf_token = "hf_x", + max_seq_length = 8192, + tensor_parallel = True, + ) + + assert isinstance(backend, ChatBackend) + assert fallback_calls == [(True, [], "org/model-GGUF")] + assert [call["tensor_parallel"] for call in calls] == [True, False] + assert calls[1]["extra_args"] == ["--split-mode", "layer"] + + def test_http_backend_merges_emoji_split_across_deltas(monkeypatch): backend = HttpChatBackend("http://localhost:8888", "token") response = _FakeSSEResponse( @@ -365,6 +628,98 @@ def test_chat_prefers_running_studio_server(monkeypatch): assert closed == ["http"] +def test_chat_forwards_gguf_runtime_options_to_loader(monkeypatch): + loads = [] + + class _FakeHttpBackend: + def close(self): + pass + + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr( + chatmod, + "connect_studio_server", + lambda model, **kwargs: (loads.append((model, kwargs)), _FakeHttpBackend())[1], + ) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: None) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + + result = CliRunner().invoke( + _chat_app(), + [ + "fake-model", + "--tensor-parallel", + "--llama-extra-arg=--top-k", + "--llama-extra-arg", + "20", + ], + input = "/exit\n", + ) + + assert result.exit_code == 0, result.output + assert loads == [ + ( + "fake-model", + { + "hf_token": None, + "max_seq_length": 4096, + "load_in_4bit": True, + "tensor_parallel": True, + "llama_extra_args": ["--top-k", "20"], + }, + ) + ] + + +def test_inference_forwards_gguf_runtime_options_to_loader(monkeypatch): + from unsloth_cli.commands import inference as infermod + + loads, streams, closed = [], [], [] + + class _FakeBackend: + def stream(self, messages, **kwargs): + streams.append((messages, kwargs)) + return iter(["answer"]) + + def close(self): + closed.append(True) + + monkeypatch.setattr( + infermod, + "connect_studio_server", + lambda model, **kwargs: (loads.append((model, kwargs)), _FakeBackend())[1], + ) + monkeypatch.setattr(infermod, "load_chat_backend", lambda *a, **k: None) + + result = CliRunner().invoke( + _inference_app(), + [ + "fake-model", + "hello", + "--tensor-parallel", + "--llama-extra-arg=--top-k", + "--llama-extra-arg", + "20", + ], + ) + + assert result.exit_code == 0, result.output + assert loads == [ + ( + "fake-model", + { + "hf_token": None, + "max_seq_length": 2048, + "load_in_4bit": True, + "tensor_parallel": True, + "llama_extra_args": ["--top-k", "20"], + }, + ) + ] + assert streams[0][0] == [{"role": "user", "content": "hello"}] + assert closed == [True] + + def test_chat_server_mode_compare_loads_base_locally(monkeypatch): streamed, closed, base_loads = [], [], [] From 9451aef51e2639e4003a1b490048afed1471e51d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 12:07:53 -0700 Subject: [PATCH 08/49] studio: return a clean model id from the OpenAI API instead of the local .gguf path (#6518) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 9 +++ .../core/inference/llama_server_args.py | 5 ++ studio/backend/core/inference/model_ids.py | 57 +++++++++++++++++ studio/backend/routes/inference.py | 62 +++++++++++++++---- studio/backend/tests/test_model_ids.py | 53 ++++++++++++++++ .../tests/test_openai_models_path_leak.py | 45 ++++++++++++++ 6 files changed, 218 insertions(+), 13 deletions(-) create mode 100644 studio/backend/core/inference/model_ids.py create mode 100644 studio/backend/tests/test_model_ids.py create mode 100644 studio/backend/tests/test_openai_models_path_leak.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 255fc4140f..31969afb14 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -5476,6 +5476,15 @@ class LlamaCppBackend: "--no-context-shift", ] + # Report a clean public model id (matching GET /v1/models) rather + # than the raw -m path in llama-server's own /v1/models and the + # "model" field of its chat/completions responses. + from core.inference.model_ids import public_model_id + + _alias = public_model_id(self._model_identifier or model_path) + if _alias: + cmd.extend(["--alias", _alias]) + fully_gpu_offloaded = False if use_fit: cmd.extend(["--fit", "on"]) diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index b42be5ee0d..f400d2ae40 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -25,6 +25,11 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = ( # Model identity: Studio resolves it from LoadRequest; a second -m would # load a different model than Studio thinks it loaded. frozenset({"-m", "--model"}), + # Public model id: Studio sets a sanitized --alias so the OpenAI API never + # exposes the local .gguf path. A user-supplied alias is appended after + # Studio's and, with llama.cpp's last-wins parsing, would reintroduce the + # path leak this is meant to prevent. + frozenset({"-a", "--alias"}), frozenset({"-mu", "--model-url"}), frozenset({"-dr", "--docker-repo"}), frozenset({"-hf", "-hfr", "--hf-repo"}), diff --git a/studio/backend/core/inference/model_ids.py b/studio/backend/core/inference/model_ids.py new file mode 100644 index 0000000000..4d409cb4b5 --- /dev/null +++ b/studio/backend/core/inference/model_ids.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Public model identifiers for the OpenAI-compatible API. + +The exposed API must report a stable, clean model id rather than the absolute +on-disk path of a local GGUF. The internal identifier for a direct local load is +the absolute ``.gguf`` path, which leaks the host filesystem layout and is +awkward for clients to round-trip. ``public_model_id`` maps such an internal +identifier to a clean name while leaving Hugging Face repo ids (``org/model``) +and already-clean names untouched. +""" + +from __future__ import annotations + +import os +from typing import Optional + +_GGUF_SUFFIX = ".gguf" + + +def _looks_like_path(identifier: str) -> bool: + """True when *identifier* is a local filesystem path, not a HF repo id. + + A repo id is ``org/model`` (a single forward slash, no leading separator, no + drive, no ``.gguf``). Anything ending in ``.gguf``, starting with a path + separator or a relative/home prefix (``./``, ``../``, ``~``), carrying a + Windows drive, or with three or more ``/`` segments is treated as a local + path. + """ + if identifier.lower().endswith(_GGUF_SUFFIX): + return True + if identifier.startswith(("/", "\\", "./", "../", ".\\", "..\\", "~")): + return True + if len(identifier) >= 2 and identifier[1] == ":": # Windows drive, e.g. C:\ + return True + if identifier.count("/") >= 2 or "\\" in identifier: + return True + return False + + +def public_model_id(identifier: Optional[str]) -> Optional[str]: + """Return a clean, path-free public id for *identifier*. + + - Local GGUF path -> the file stem with ``.gguf`` stripped, e.g. + ``/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf`` -> ``Qwen3-30B-A3B-Q4_K_M``. + - HF repo id (``org/model``) and already-clean names -> returned unchanged. + - ``None`` / empty -> returned unchanged. + """ + if not identifier: + return identifier + if not _looks_like_path(identifier): + return identifier + name = os.path.basename(identifier.replace("\\", "/").rstrip("/")) + if name.lower().endswith(_GGUF_SUFFIX): + name = name[: -len(_GGUF_SUFFIX)] + return name or identifier diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a9c85fa99a..1061542ca6 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1107,6 +1107,7 @@ from auth.authentication import get_current_subject from state.tool_approvals import resolve_tool_decision 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.providers import get_base_url @@ -3703,7 +3704,7 @@ async def generate_audio( # Pick backend — both return (wav_bytes, sample_rate) llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False): - model_name = llama_backend.model_identifier + model_name = public_model_id(llama_backend.model_identifier) gen = lambda: llama_backend.generate_audio_response( text = text, audio_type = llama_backend._audio_type, @@ -3721,7 +3722,7 @@ async def generate_audio( model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_audio"): raise HTTPException(status_code = 400, detail = "Active model is not an audio model.") - model_name = backend.active_model_name + model_name = public_model_id(backend.active_model_name) gen = lambda: backend.generate_audio_response( text = text, temperature = payload.temperature, @@ -4837,7 +4838,8 @@ async def openai_chat_completions( return response if using_gguf: - model_name = llama_backend.model_identifier or payload.model + # Echo a clean public id in the response, never the absolute .gguf path. + model_name = public_model_id(llama_backend.model_identifier) or payload.model if getattr(llama_backend, "_is_audio", False): if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF audio chat completions") @@ -4852,7 +4854,9 @@ async def openai_chat_completions( status_code = 400, detail = "No model loaded. Call POST /inference/load first.", ) - model_name = backend.active_model_name or payload.model + # Clean public id so the response never echoes a local path; the audio + # branch below receives this sanitized label too. + model_name = public_model_id(backend.active_model_name) or payload.model if _wants_multiple_choices(payload): _raise_unsupported_n("non-GGUF chat completions") @@ -6401,7 +6405,9 @@ def _openai_model_objects() -> list[dict]: llama_backend = get_llama_cpp_backend() if llama_backend.is_loaded: entry = { - "id": llama_backend.model_identifier, + # Public id, never the absolute .gguf path (which leaks the host + # filesystem layout); see core.inference.model_ids.public_model_id. + "id": public_model_id(llama_backend.model_identifier), "object": "model", "created": _created, "owned_by": "local", @@ -6422,7 +6428,7 @@ def _openai_model_objects() -> list[dict]: if backend.active_model_name: model_info = backend.models.get(backend.active_model_name, {}) entry = { - "id": backend.active_model_name, + "id": public_model_id(backend.active_model_name), "object": "model", "created": _created, "owned_by": "local", @@ -6463,9 +6469,24 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge model, or 404 model_not_found otherwise. Defined after the LIST route so it does not shadow it; ``{model_id:path}`` keeps ids with slashes intact. """ - for model in _openai_model_objects(): + objects = _openai_model_objects() + for model in objects: if model["id"] == model_id: return model + # Backward compatibility: a client may still send the legacy raw identifier + # (e.g. an absolute .gguf path cached from an older /v1/models). Resolve it to + # the clean object so it keeps working, without ever echoing the path back. + llama_backend = get_llama_cpp_backend() + backend = get_inference_backend() + for raw in ( + llama_backend.model_identifier if llama_backend.is_loaded else None, + backend.active_model_name or None, + ): + if raw and model_id == raw: + clean = public_model_id(raw) + for model in objects: + if model["id"] == clean: + return model raise HTTPException( status_code = 404, detail = openai_error_body( @@ -7402,6 +7423,15 @@ async def _responses_stream( target_url = f"{llama_backend.base_url}/v1/chat/completions" async def event_generator(): + # Clean public id for every response envelope. Prefer the loaded model's + # id so the stream agrees with /v1/models, chat/completions and the + # non-streaming twin; fall back to a sanitized payload.model (a legacy + # raw .gguf path is stripped, never echoed back). + _clean_model = ( + public_model_id(getattr(llama_backend, "model_identifier", None)) + or public_model_id(payload.model) + or payload.model + ) full_text = "" full_reasoning = "" input_tokens = 0 @@ -7563,7 +7593,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "failed", - "model": payload.model, + "model": _clean_model, "output": _snapshot_output(), "usage": { "input_tokens": input_tokens, @@ -7587,7 +7617,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "in_progress", - "model": payload.model, + "model": _clean_model, "output": [], "usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, }, @@ -7627,7 +7657,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "failed", - "model": payload.model, + "model": _clean_model, "output": [], "error": {"code": 502, "message": _friendly_error(e)}, }, @@ -7653,7 +7683,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "failed", - "model": payload.model, + "model": _clean_model, "output": [], "error": { "code": resp.status_code, @@ -8002,7 +8032,7 @@ async def _responses_stream( "object": "response", "created_at": created_at, "status": "completed", - "model": payload.model, + "model": _clean_model, "output": _snapshot_output(), "usage": { "input_tokens": input_tokens, @@ -8274,7 +8304,13 @@ async def anthropic_messages( ), ) - model_name = getattr(llama_backend, "model_identifier", None) or payload.model + # Clean public id so /v1/messages never echoes the local .gguf path (and a + # legacy raw path sent as payload.model is sanitized rather than returned). + model_name = ( + public_model_id(getattr(llama_backend, "model_identifier", None)) + or public_model_id(payload.model) + or payload.model + ) message_id = f"msg_{uuid.uuid4().hex[:24]}" # ── Translate Anthropic → OpenAI ────────────────────────── diff --git a/studio/backend/tests/test_model_ids.py b/studio/backend/tests/test_model_ids.py new file mode 100644 index 0000000000..1b0cd927d8 --- /dev/null +++ b/studio/backend/tests/test_model_ids.py @@ -0,0 +1,53 @@ +# 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 pathlib import Path + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.model_ids import public_model_id # noqa: E402 + + +def test_local_gguf_path_becomes_clean_stem(): + assert public_model_id("/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf") == "Qwen3-30B-A3B-Q4_K_M" + assert public_model_id("/home/u/.cache/models/llama.gguf") == "llama" + + +def test_hf_repo_id_unchanged(): + assert public_model_id("unsloth/Qwen3-30B-A3B-GGUF") == "unsloth/Qwen3-30B-A3B-GGUF" + assert public_model_id("Qwen3-30B-A3B") == "Qwen3-30B-A3B" + + +def test_none_and_empty_passthrough(): + assert public_model_id(None) is None + assert public_model_id("") == "" + + +def test_windows_path(): + assert public_model_id("C:\\models\\foo.gguf") == "foo" + assert public_model_id("models\\sub\\bar.gguf") == "bar" + + +def test_directory_path_uses_basename(): + assert public_model_id("/opt/models/MyModelDir") == "MyModelDir" + # A 3+ segment relative path is a local path, not an org/model repo id. + assert public_model_id("a/b/c") == "c" + + +def test_relative_and_home_paths_are_sanitized(): + # ./ ../ ~ prefixed paths are local and must not be echoed raw. + assert public_model_id("./model.gguf") == "model" + assert public_model_id("../models/foo.gguf") == "foo" + assert public_model_id("~/models/baz.gguf") == "baz" + assert public_model_id("./mistral") == "mistral" + assert public_model_id("~/mistral") == "mistral" + assert public_model_id(".\\models\\foo.gguf") == "foo" + + +def test_dotted_repo_id_not_mistaken_for_relative_path(): + # A leading dot that is not ./ or ../ is an ordinary clean name. + assert public_model_id(".hidden-model") == ".hidden-model" + assert public_model_id("org/.config") == "org/.config" diff --git a/studio/backend/tests/test_openai_models_path_leak.py b/studio/backend/tests/test_openai_models_path_leak.py new file mode 100644 index 0000000000..a84a33f840 --- /dev/null +++ b/studio/backend/tests/test_openai_models_path_leak.py @@ -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 + +"""GET /v1/models must report a clean public id, never the on-disk .gguf path.""" + +import json +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import routes.inference as inf # noqa: E402 + + +class _FakeLlama: + is_loaded = True + model_identifier = "/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf" + context_length = 4096 + max_context_length = None + native_context_length = None + + +class _FakeUnsloth: + active_model_name = None + models: dict = {} + context_length = None + max_seq_length = None + + +def test_openai_models_returns_clean_id_without_path(monkeypatch): + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth()) + + objs = inf._openai_model_objects() + + assert len(objs) == 1 + assert objs[0]["id"] == "Qwen3-30B-A3B-Q4_K_M" + # The serialized payload must not leak the absolute path or the .gguf suffix. + blob = json.dumps(objs) + assert "/srv/models" not in blob + assert ".gguf" not in blob + # Context fields still flow through. + assert objs[0]["context_length"] == 4096 From 4a5d41eb3d3731f41347c6adec0f6dde95cf3253 Mon Sep 17 00:00:00 2001 From: Luca Cesarano Date: Fri, 26 Jun 2026 21:48:30 +0200 Subject: [PATCH 09/49] fix(install): enable UV_NATIVE_TLS on macOS for corporate TLS-inspection proxies (#6671) --------- Co-authored-by: Luca Cesarano Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- README.md | 5 +++++ install.sh | 15 +++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/README.md b/README.md index 9162d29b1c..e3fd4e6980 100644 --- a/README.md +++ b/README.md @@ -246,6 +246,11 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh $env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex ``` +On macOS, the installer defaults to the system certificate store (`UV_SYSTEM_CERTS=1`) so uv trusts the CAs in your Keychain, needed behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.). Opt out with: +```bash +curl -fsSL https://unsloth.ai/install.sh | UV_SYSTEM_CERTS=0 sh +``` + Point the frontend build at a corporate npm mirror/proxy with `UNSLOTH_NPM_REGISTRY` (for the developer install behind a firewall that blocks `registry.npmjs.org`): ```bash UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh --local diff --git a/install.sh b/install.sh index 548e6f702a..7a9f0be87f 100755 --- a/install.sh +++ b/install.sh @@ -1636,6 +1636,21 @@ export UV_HTTP_RETRIES : "${UV_HTTP_TIMEOUT:=180}" export UV_HTTP_TIMEOUT +# macOS: trust the system Keychain so uv uses SecureTransport instead of rustls. +# Required behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.) which +# present their own CA certificate. rustls (uv's default) ignores the Keychain +# and rejects intercepted connections with "invalid peer certificate: UnknownIssuer". +# Set both vars: UV_SYSTEM_CERTS is the modern one (uv >= 0.11), UV_NATIVE_TLS the +# legacy one understood by uv 0.8.16-0.10.x, which the installer keeps if already +# present (UV_MIN_VERSION) and which ignores UV_SYSTEM_CERTS. Mirror the choice onto +# both so it works on either uv. Opt out with UV_SYSTEM_CERTS=0. +if [ "$OS" = "macos" ]; then + : "${UV_SYSTEM_CERTS:=1}" + : "${UV_NATIVE_TLS:=$UV_SYSTEM_CERTS}" +fi +[ -n "${UV_SYSTEM_CERTS:-}" ] && export UV_SYSTEM_CERTS +[ -n "${UV_NATIVE_TLS:-}" ] && export UV_NATIVE_TLS + version_ge() { # returns 0 if $1 >= $2 _a=$1 From e594e5d2018618c665009135335562ec35846178 Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Sat, 27 Jun 2026 06:54:28 +0800 Subject: [PATCH 10/49] fix: stop faking 8bit load flag (#6708) Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- unsloth/models/_utils.py | 81 ++++++++++++++++++++++++++++++++++++++++ unsloth/models/llama.py | 12 ++---- unsloth/models/vision.py | 12 ++---- 3 files changed, 89 insertions(+), 16 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 7a056cef82..7ad6e8ea33 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -86,6 +86,8 @@ __all__ = [ "is_moe_model", "get_moe_target_parameters", "make_fast_generate_wrapper", + "_mark_unsloth_disable_data_parallel", + "_patch_transformers_trainer_data_parallel", ] import torch @@ -160,6 +162,85 @@ from unsloth_zoo.training_utils import ( ) +def _iter_wrapped_models(model): + seen = set() + current = model + while current is not None and id(current) not in seen: + yield current + seen.add(id(current)) + next_model = getattr(current, "model", None) + if next_model is None: + next_model = getattr(current, "base_model", None) + if next_model is None: + next_model = getattr(current, "module", None) + current = next_model + + +def _patch_transformers_trainer_data_parallel(): + try: + from transformers.trainer import Trainer + except (ImportError, ModuleNotFoundError): + return False + + original_wrap_model = getattr(Trainer, "_wrap_model", None) + if original_wrap_model is None: + return False + if getattr(original_wrap_model, "_unsloth_data_parallel_patched", False): + return True + try: + supports_dataloader = "dataloader" in inspect.signature(original_wrap_model).parameters + except (TypeError, ValueError): + supports_dataloader = True + + def _call_original_wrap_model(self, model, wrap_args, wrap_kwargs): + if supports_dataloader: + return original_wrap_model(self, model, *wrap_args, **wrap_kwargs) + + if "dataloader" in wrap_kwargs: + wrap_kwargs = {k: v for k, v in wrap_kwargs.items() if k != "dataloader"} + return original_wrap_model(self, model, *wrap_args, **wrap_kwargs) + + @functools.wraps(original_wrap_model) + def _unsloth_wrap_model(self, model, *wrap_args, **wrap_kwargs): + args = getattr(self, "args", None) + disable_data_parallel = getattr(model, "_unsloth_disable_data_parallel", False) + is_real_8bit = getattr(model, "is_loaded_in_8bit", False) + if ( + args is None + or not disable_data_parallel + or is_real_8bit + or getattr(args, "n_gpu", 0) <= 1 + ): + return _call_original_wrap_model(self, model, wrap_args, wrap_kwargs) + + had_n_gpu = hasattr(args, "_n_gpu") + old_n_gpu = getattr(args, "_n_gpu", None) + args._n_gpu = 1 + try: + return _call_original_wrap_model(self, model, wrap_args, wrap_kwargs) + finally: + if had_n_gpu: + args._n_gpu = old_n_gpu + else: + try: + delattr(args, "_n_gpu") + except AttributeError: + pass + + _unsloth_wrap_model._unsloth_data_parallel_patched = True + _unsloth_wrap_model._unsloth_original_wrap_model = original_wrap_model + Trainer._wrap_model = _unsloth_wrap_model + return True + + +def _mark_unsloth_disable_data_parallel(model, disable = True): + if disable: + _patch_transformers_trainer_data_parallel() + for module in _iter_wrapped_models(model): + setattr(module, "_unsloth_disable_data_parallel", bool(disable)) + return model + + def resolve_hip_gpu_stats_name(gpu_stats): name = str(getattr(gpu_stats, "name", "") or "").strip() name = re.sub(r"\s*\([^)]*\)\s*$", "", name).strip() diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d5eca8d6df..14ee5ee24e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2832,13 +2832,11 @@ class FastLlamaModel: internal_model = model while hasattr(internal_model, "model"): internal_model._saved_temp_tokenizer = tokenizer - # Also set is_loaded_in_8bit to disable incorrect DDP - internal_model.is_loaded_in_8bit = True internal_model = internal_model.model internal_model._saved_temp_tokenizer = tokenizer - # Also set is_loaded_in_8bit to disable incorrect DDP - internal_model.is_loaded_in_8bit = True + # Prevent Transformers Trainer from auto-wrapping Unsloth LoRA models in DP. + _mark_unsloth_disable_data_parallel(model) # For transformers > 4.47.1, we need to add rotary_emb to all attention layers if IS_ATTENTION_REFACTOR or hasattr(model.model, "rotary_emb"): @@ -3379,13 +3377,11 @@ class FastLlamaModel: while hasattr(internal_model, "model"): if hasattr(internal_model, "_saved_temp_tokenizer"): internal_model._saved_temp_tokenizer.padding_side = "right" - # Also set is_loaded_in_8bit to disable incorrect DDP - internal_model.is_loaded_in_8bit = True internal_model = internal_model.model if hasattr(internal_model, "_saved_temp_tokenizer"): internal_model._saved_temp_tokenizer.padding_side = "right" - # Also set is_loaded_in_8bit to disable incorrect DDP - internal_model.is_loaded_in_8bit = True + # Prevent Transformers Trainer from auto-wrapping Unsloth LoRA models in DP. + _mark_unsloth_disable_data_parallel(model) # Clear deleted GPU items for _ in range(3): diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 39004b4d45..bdc2bd9ef6 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1440,16 +1440,14 @@ class FastBaseModel: while hasattr(m, "model"): m.max_seq_length = max_seq_length m._saved_temp_tokenizer = tokenizer - # Also set is_loaded_in_8bit to disable incorrect DDP - m.is_loaded_in_8bit = True if not full_finetuning else False m = m.model m.max_seq_length = max_seq_length # Save to modules as well for module in model.modules(): module.max_seq_length = max_seq_length m._saved_temp_tokenizer = tokenizer - # Also set is_loaded_in_8bit to disable incorrect DDP - m.is_loaded_in_8bit = True if not full_finetuning else False + # Prevent Transformers Trainer from auto-wrapping Unsloth LoRA models in DP. + _mark_unsloth_disable_data_parallel(model, disable = not full_finetuning) # Patch generate if os.environ.get("UNSLOTH_DISABLE_FAST_GENERATION", "0") == "0" and hasattr( @@ -1826,14 +1824,12 @@ class FastBaseModel: if hasattr(m, "_saved_temp_tokenizer"): if hasattr(m._saved_temp_tokenizer, "tokenizer"): m._saved_temp_tokenizer.tokenizer.padding_side = "left" - # Also set is_loaded_in_8bit to disable incorrect DDP - m.is_loaded_in_8bit = True if not full_finetuning else False m = m.model if hasattr(m, "_saved_temp_tokenizer"): if hasattr(m._saved_temp_tokenizer, "tokenizer"): m._saved_temp_tokenizer.tokenizer.padding_side = "left" - # Also set is_loaded_in_8bit to disable incorrect DDP - m.is_loaded_in_8bit = True if not full_finetuning else False + # Prevent Transformers Trainer from auto-wrapping Unsloth LoRA models in DP. + _mark_unsloth_disable_data_parallel(model, disable = not full_finetuning) # Clear deleted GPU items for _ in range(3): From b11966b2db759f933342364c825211f740a9002a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 16:42:06 -0700 Subject: [PATCH 11/49] studio: list the full local model catalog from /v1/models (#6519) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/core/inference/model_ids.py | 14 ++ studio/backend/routes/inference.py | 105 ++++++++++-- studio/backend/routes/models.py | 161 +++++++++-------- studio/backend/tests/test_model_ids.py | 11 +- studio/backend/tests/test_openai_catalog.py | 181 ++++++++++++++++++++ 5 files changed, 388 insertions(+), 84 deletions(-) create mode 100644 studio/backend/tests/test_openai_catalog.py diff --git a/studio/backend/core/inference/model_ids.py b/studio/backend/core/inference/model_ids.py index 4d409cb4b5..548cc60f94 100644 --- a/studio/backend/core/inference/model_ids.py +++ b/studio/backend/core/inference/model_ids.py @@ -55,3 +55,17 @@ def public_model_id(identifier: Optional[str]) -> Optional[str]: if name.lower().endswith(_GGUF_SUFFIX): name = name[: -len(_GGUF_SUFFIX)] return name or identifier + + +def model_id_matches(requested: Optional[str], internal: Optional[str]) -> bool: + """Whether a client-supplied *requested* id refers to *internal*. + + Accepts the clean public id (preferred) and, for backward compatibility, the + raw internal identifier (e.g. a legacy absolute path a client cached from an + older ``/v1/models`` response). + """ + if requested is None or internal is None: + return False + if requested == internal: + return True + return public_model_id(internal) == requested diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1061542ca6..caf2a262a3 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6391,6 +6391,9 @@ async def serve_sandbox_file( # OpenAI-Compatible Models Listing (/models → /v1/models) # ===================================================================== +# `owned_by` marker on every /v1/models entry (loaded and available alike). +_OWNED_BY = "unsloth-studio" + def _openai_model_objects() -> list[dict]: """The model objects GET /v1/models exposes (one per loaded local backend). @@ -6410,7 +6413,7 @@ def _openai_model_objects() -> list[dict]: "id": public_model_id(llama_backend.model_identifier), "object": "model", "created": _created, - "owned_by": "local", + "owned_by": _OWNED_BY, } _ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None)) if _ctx is not None: @@ -6431,7 +6434,7 @@ def _openai_model_objects() -> list[dict]: "id": public_model_id(backend.active_model_name), "object": "model", "created": _created, - "owned_by": "local", + "owned_by": _OWNED_BY, } _ctx = _positive_int_or_none(model_info.get("context_length")) if _ctx is None: @@ -6449,15 +6452,86 @@ def _openai_model_objects() -> list[dict]: return models +# Brief cache for the local-model filesystem scan so repeated /v1/models calls +# don't rescan the HF cache and models dirs on every request. +_CATALOG_CACHE: dict = {"at": 0.0, "models": []} +_CATALOG_TTL_S = 30.0 +_CATALOG_LOCK = asyncio.Lock() + + +async def _cached_local_catalog() -> list: + """Locally available models (models dir + HF caches + LM Studio + scan + folders), cached for a few seconds. Returns a list of LocalModelInfo. + + The scan walks several directories and stats many files, so it runs in a + worker thread (asyncio.to_thread) -- calling it inline would block the event + loop and stall every concurrent request and in-flight inference stream. A + lock with a double-check collapses a burst of simultaneous /v1/models calls + into a single scan instead of one per request.""" + # Validity is keyed on "at" (set only after a scan), not on list contents, so + # an empty/errored scan is still cached instead of rescanning on every poll. + now = time.monotonic() + if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: + return _CATALOG_CACHE["models"] + async with _CATALOG_LOCK: + now = time.monotonic() + if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S: + return _CATALOG_CACHE["models"] + try: + from routes.models import collect_local_models + _CATALOG_CACHE["models"] = await asyncio.to_thread( + collect_local_models, Path("./models").resolve() + ) + except Exception as exc: + logger.debug("model catalog scan failed: %s", exc) + _CATALOG_CACHE["models"] = [] + # Stamp after the scan, not the pre-scan "now": a scan slower than the TTL + # would otherwise leave the cache already expired, so every waiter rescans. + _CATALOG_CACHE["at"] = time.monotonic() + return _CATALOG_CACHE["models"] + + +async def _openai_catalog_objects() -> list[dict]: + """Every model the server knows about for ``GET /v1/models``: the loaded + model(s) plus locally available (downloaded/cached) models discovered by + scanning. Loaded entries keep their context fields and are marked + ``loaded: true``. All ids are clean public ids (never absolute paths).""" + _created = int(time.time()) + # Loaded models first (clean ids + context fields), marked loaded. + by_id: dict[str, dict] = {} + for entry in _openai_model_objects(): + by_id[entry["id"]] = {**entry, "loaded": True} + + # Locally available (downloaded/cached) models that are not already loaded. + for info in await _cached_local_catalog(): + cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None)) + if not cid or cid in by_id: + continue + obj = { + "id": cid, + "object": "model", + "created": _created, + "owned_by": _OWNED_BY, + "loaded": False, + } + display = getattr(info, "display_name", None) + if display: + obj["display_name"] = display + by_id[cid] = obj + + return list(by_id.values()) + + @router.get("/models") async def openai_list_models(current_subject: str = Depends(get_current_subject)): """ - OpenAI-compatible model listing endpoint. + OpenAI-compatible model listing endpoint (``GET /v1/models``). - Returns the currently loaded model in the format expected by - OpenAI-compatible clients (``GET /v1/models``). + Lists every model available on this server -- the loaded model(s) plus + locally available (downloaded/cached) models -- not only what is resident in + memory. Each entry carries a clean public id and a ``loaded`` flag. """ - return {"object": "list", "data": _openai_model_objects()} + return {"object": "list", "data": await _openai_catalog_objects()} @router.get("/models/{model_id:path}") @@ -6465,11 +6539,20 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge """ OpenAI-compatible single-model retrieval endpoint (``GET /v1/models/{id}``). - Returns the bare model object when ``model_id`` matches a loaded local - model, or 404 model_not_found otherwise. Defined after the LIST route so - it does not shadow it; ``{model_id:path}`` keeps ids with slashes intact. + Returns the bare model object when ``model_id`` matches a known model + (loaded or locally available), or 404 model_not_found otherwise. Defined + after the LIST route so it does not shadow it; ``{model_id:path}`` keeps ids + with slashes intact. """ - objects = _openai_model_objects() + from core.inference.model_ids import model_id_matches + + # Loaded models resolve without a catalog scan (the common case); only build + # the full catalog -- which may hit the filesystem -- for unloaded ids. + for entry in _openai_model_objects(): + if entry["id"] == model_id: + return {**entry, "loaded": True} + + objects = await _openai_catalog_objects() for model in objects: if model["id"] == model_id: return model @@ -6482,7 +6565,7 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge llama_backend.model_identifier if llama_backend.is_loaded else None, backend.active_model_name or None, ): - if raw and model_id == raw: + if raw and model_id_matches(model_id, raw): clean = public_model_id(raw) for model in objects: if model["id"] == clean: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 951c2960f3..e22f65751c 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -722,6 +722,94 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca return found +def collect_local_models(models_root: Path) -> List[LocalModelInfo]: + """Scan ``models_root``, the HF caches, LM Studio dirs, and user scan folders, + returning a deduplicated, hidden-filtered list of discovered local models. + + Shared by ``GET /models/local`` (the model picker) and the OpenAI-compatible + catalog (``GET /v1/models``) so the UI and the API never drift. ``models_root`` + must already be validated/trusted by the caller. + """ + from storage.studio_db import list_scan_folders + from utils.paths import ( + hf_default_cache_dir, + legacy_hf_cache_dir, + lmstudio_model_dirs, + ) + + hf_cache_dir = _resolve_hf_cache_dir() + legacy_hf = legacy_hf_cache_dir() + hf_default = hf_default_cache_dir() + lm_dirs = lmstudio_model_dirs() + + local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) + + # Resolve once; an inaccessible aux cache must skip that scan, not 500. + hf_cache_real = _safe_resolve(hf_cache_dir) + legacy_real = _safe_resolve(legacy_hf) + default_real = _safe_resolve(hf_default) + + # Scan legacy Unsloth HF cache for backward compatibility. + if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real: + local_models += _scan_hf_cache(legacy_hf) + + # Scan HF system default cache (may differ under env overrides). + if _safe_is_dir(hf_default) and default_real != hf_cache_real and default_real != legacy_real: + local_models += _scan_hf_cache(hf_default) + + # Scan LM Studio directories. + for lm_dir in lm_dirs: + local_models += _scan_lmstudio_dir(lm_dir) + + # Scan user-added custom folders (per-folder cap). + _MAX_MODELS_PER_FOLDER = 200 + try: + custom_folders = list_scan_folders() + except Exception as e: + logger.warning("Could not load custom scan folders: %s", e) + custom_folders = [] + for folder in custom_folders: + folder_path = Path(folder["path"]) + try: + # Filter Ollama .studio_links/ from generic scanners to + # avoid duplicates and leaking internal paths into the UI. + _generic = [ + m + for m in ( + _scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER) + + _scan_hf_cache(folder_path) + + _scan_lmstudio_dir(folder_path) + ) + if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts) + ] + custom_models = _generic + if len(custom_models) < _MAX_MODELS_PER_FOLDER: + custom_models += _scan_ollama_dir( + folder_path, + limit = _MAX_MODELS_PER_FOLDER - len(custom_models), + ) + except OSError as e: + logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e) + continue + local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models] + + # Deduplicate, but always keep custom folder entries (keyed by + # (id, source)) so they show in the "Custom Folders" UI section + # even when the model is also in the HF cache. + deduped: dict[str, LocalModelInfo] = {} + for model in local_models: + key = f"{model.id}\x00custom" if model.source == "custom" else model.id + if key not in deduped: + deduped[key] = model + + models = sorted( + deduped.values(), + key = lambda item: (item.updated_at or 0), + reverse = True, + ) + return [m for m in models if not _is_hidden_model(m.id, m.path)] + + @router.get("/local", response_model = LocalModelListResponse) async def list_local_models( models_dir: str = Query( @@ -770,78 +858,7 @@ async def list_local_models( ) try: - local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir) - - # Resolve once; an inaccessible aux cache must skip that scan, not 500. - hf_cache_real = _safe_resolve(hf_cache_dir) - legacy_real = _safe_resolve(legacy_hf) - default_real = _safe_resolve(hf_default) - - # Scan legacy Unsloth HF cache for backward compatibility. - if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real: - local_models += _scan_hf_cache(legacy_hf) - - # Scan HF system default cache (may differ under env overrides). - if ( - _safe_is_dir(hf_default) - and default_real != hf_cache_real - and default_real != legacy_real - ): - local_models += _scan_hf_cache(hf_default) - - # Scan LM Studio directories. - for lm_dir in lm_dirs: - local_models += _scan_lmstudio_dir(lm_dir) - - # Scan user-added custom folders (per-folder cap). - from storage.studio_db import list_scan_folders - - _MAX_MODELS_PER_FOLDER = 200 - try: - custom_folders = list_scan_folders() - except Exception as e: - logger.warning("Could not load custom scan folders: %s", e) - custom_folders = [] - for folder in custom_folders: - folder_path = Path(folder["path"]) - try: - # Filter Ollama .studio_links/ from generic scanners to - # avoid duplicates and leaking internal paths into the UI. - _generic = [ - m - for m in ( - _scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER) - + _scan_hf_cache(folder_path) - + _scan_lmstudio_dir(folder_path) - ) - if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts) - ] - custom_models = _generic - if len(custom_models) < _MAX_MODELS_PER_FOLDER: - custom_models += _scan_ollama_dir( - folder_path, - limit = _MAX_MODELS_PER_FOLDER - len(custom_models), - ) - except OSError as e: - logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e) - continue - local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models] - - # Deduplicate, but always keep custom folder entries (keyed by - # (id, source)) so they show in the "Custom Folders" UI section - # even when the model is also in the HF cache. - deduped: dict[str, LocalModelInfo] = {} - for model in local_models: - key = f"{model.id}\x00custom" if model.source == "custom" else model.id - if key not in deduped: - deduped[key] = model - - models = sorted( - deduped.values(), - key = lambda item: (item.updated_at or 0), - reverse = True, - ) - models = [m for m in models if not _is_hidden_model(m.id, m.path)] + models = collect_local_models(models_root) return LocalModelListResponse( models_dir = str(models_root), diff --git a/studio/backend/tests/test_model_ids.py b/studio/backend/tests/test_model_ids.py index 1b0cd927d8..f9116afec3 100644 --- a/studio/backend/tests/test_model_ids.py +++ b/studio/backend/tests/test_model_ids.py @@ -8,7 +8,7 @@ _BACKEND = Path(__file__).resolve().parents[1] if str(_BACKEND) not in sys.path: sys.path.insert(0, str(_BACKEND)) -from core.inference.model_ids import public_model_id # noqa: E402 +from core.inference.model_ids import model_id_matches, public_model_id # noqa: E402 def test_local_gguf_path_becomes_clean_stem(): @@ -51,3 +51,12 @@ def test_dotted_repo_id_not_mistaken_for_relative_path(): # A leading dot that is not ./ or ../ is an ordinary clean name. assert public_model_id(".hidden-model") == ".hidden-model" assert public_model_id("org/.config") == "org/.config" + + +def test_matches_clean_and_legacy(): + path = "/srv/models/Qwen3-Q4.gguf" + assert model_id_matches("Qwen3-Q4", path) # clean public id + assert model_id_matches(path, path) # legacy raw path + assert not model_id_matches("other", path) + assert not model_id_matches(None, path) + assert not model_id_matches("x", None) diff --git a/studio/backend/tests/test_openai_catalog.py b/studio/backend/tests/test_openai_catalog.py new file mode 100644 index 0000000000..f9baf20a66 --- /dev/null +++ b/studio/backend/tests/test_openai_catalog.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 + +"""GET /v1/models lists the full server catalog (loaded + locally available).""" + +import asyncio +import json +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parents[1] +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import routes.inference as inf # noqa: E402 + + +class _Info: + def __init__( + self, + id, + display_name, + model_id = None, + ): + self.id = id + self.display_name = display_name + self.model_id = model_id + + +class _FakeLlama: + is_loaded = True + model_identifier = "/srv/models/Qwen3-Q4.gguf" + context_length = 4096 + max_context_length = None + native_context_length = None + + def __init__(self, loaded = True): + self.is_loaded = loaded + + +class _FakeUnsloth: + active_model_name = None + models: dict = {} + context_length = None + max_seq_length = None + + +def test_catalog_lists_loaded_and_available(monkeypatch): + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth()) + + async def _fake_catalog(): + return [ + _Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup + _Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded + _Info("models--org--Foo", "Foo", model_id = "org/Foo"), # hf cache repo id + ] + + monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog) + + data = asyncio.run(inf._openai_catalog_objects()) + ids = {m["id"]: m for m in data} + + # Loaded model is present, marked loaded, and keeps context fields. + assert ids["Qwen3-Q4"]["loaded"] is True + assert ids["Qwen3-Q4"]["context_length"] == 4096 + # Available-but-not-loaded models are listed too. + assert ids["Llama-8B-Q8"]["loaded"] is False + assert ids["org/Foo"]["loaded"] is False + # The loaded gguf and the on-disk copy collapse to one clean id. + assert [m["id"] for m in data].count("Qwen3-Q4") == 1 + # No absolute paths or .gguf suffixes leak anywhere. + blob = json.dumps(data) + assert ".gguf" not in blob + assert "/srv/" not in blob + assert "/data/" not in blob + + +def test_empty_and_errored_scans_are_cached(monkeypatch): + # Cache validity is keyed on the timestamp, not list contents, so an empty + # (fresh install / no local models) or errored scan is still cached for the + # TTL instead of rescanning the filesystem on every /v1/models poll. + import routes.models as models_mod + for outcome in ("empty", "error"): + calls = {"n": 0} + + def _scan(_root, _outcome = outcome): + calls["n"] += 1 + if _outcome == "error": + raise RuntimeError("scan blew up") + return [] + + monkeypatch.setattr(models_mod, "collect_local_models", _scan) + monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []}) + + async def _run(): + return [await inf._cached_local_catalog() for _ in range(3)] + + results = asyncio.run(_run()) + assert results == [[], [], []], outcome + assert calls["n"] == 1, f"{outcome} scan ran {calls['n']}x (TTL not honored)" + + +def test_catalog_ttl_starts_after_scan_completes(monkeypatch): + # The cache timestamp must be taken AFTER the scan, not before it. A scan that + # outlives the TTL would otherwise leave the cache born-expired, so the next + # caller rescans instead of reusing the just-computed catalog. + import routes.models as models_mod + + clock = {"t": 1000.0} + monkeypatch.setattr(inf.time, "monotonic", lambda: clock["t"]) + monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []}) + + calls = {"n": 0} + + def _slow_scan(_root): + calls["n"] += 1 + clock["t"] += inf._CATALOG_TTL_S + 10 # the scan itself outlives the TTL + return [_Info("/m/A.gguf", "A")] + + monkeypatch.setattr(models_mod, "collect_local_models", _slow_scan) + + async def _run(): + first = await inf._cached_local_catalog() + second = await inf._cached_local_catalog() # clock unchanged since scan end + return first, second + + first, second = asyncio.run(_run()) + assert [i.id for i in first] == ["/m/A.gguf"] + assert calls["n"] == 1, "TTL started before the scan -> cache born expired, rescanned" + + +def test_retrieve_loaded_model_skips_catalog_scan(monkeypatch): + # Retrieving a loaded id must resolve from the loaded set alone, never paying + # for the filesystem scan that _cached_local_catalog drives. + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth()) + + async def _boom(): + raise AssertionError("catalog scan must not run for a loaded id") + + monkeypatch.setattr(inf, "_cached_local_catalog", _boom) + + model = asyncio.run(inf.openai_retrieve_model("Qwen3-Q4", current_subject = "t")) + assert model["id"] == "Qwen3-Q4" + assert model["loaded"] is True + + +def test_cached_local_catalog_offloads_and_caches(monkeypatch): + # The filesystem scan must run off the event loop (asyncio.to_thread) and be + # cached, so a burst of /v1/models calls does not re-scan or block. + calls = {"scan": 0, "threaded": 0} + + def _fake_collect(_root): + calls["scan"] += 1 + return [_Info("/data/models/A.gguf", "A")] + + import routes.models as models_mod + + monkeypatch.setattr(models_mod, "collect_local_models", _fake_collect) + + real_to_thread = inf.asyncio.to_thread + + async def _counting_to_thread(fn, *a, **k): + calls["threaded"] += 1 + return await real_to_thread(fn, *a, **k) + + monkeypatch.setattr(inf.asyncio, "to_thread", _counting_to_thread) + # Fresh cache for a deterministic count. + monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []}) + + async def _run(): + first = await inf._cached_local_catalog() + second = await inf._cached_local_catalog() # within TTL -> cached + return first, second + + first, second = asyncio.run(_run()) + assert [i.id for i in first] == ["/data/models/A.gguf"] + assert second is first or [i.id for i in second] == [i.id for i in first] + assert calls["scan"] == 1 # cached: scanned once for two calls + assert calls["threaded"] == 1 # offloaded to a worker thread From 101de1927a60c2a40446ffb47648b092a8a64005 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 19:45:32 -0700 Subject: [PATCH 12/49] Silence torchao _C*.so load-failure WARNING on torch >= 2.11 (#6712) On torch >= 2.11 torchao tries to dlopen each prebuilt _C*.so and logs a per-file "Failed to load .../_C*.so" WARNING via the torchao logger when one cannot load. This happens on an ABI tag mismatch in the prebuilt wheel (for example a cp310 .so under a cp312 runtime, as on Colab) or when the kernel targets an arch the GPU does not have (mxfp8 needs FP8 hardware, _C_cutlass_90a is Hopper/SM90 only). torchao falls back to its non-cpp paths and Unsloth's bnb-4bit / Triton kernels do not use these, so the warning is cosmetic. Add a HideLoggingMessage filter on the same torchao logger that already filters the torch < 2.11 "Skipping import of cpp extensions" message, so only these records are dropped rather than raising the whole logger to ERROR. --- unsloth/import_fixes.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 5afd6f4b37..bff55b4e7b 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -158,6 +158,11 @@ if not UNSLOTH_ENABLE_LOGGING: logging.getLogger("torchao").addFilter( HideLoggingMessage("Skipping import of cpp extensions due to incompatible torch version") ) + # torch >= 2.11 path: torchao dlopens each prebuilt _C*.so and logs "Failed to load + # .../_C*.so" when one can't (ABI tag mismatch in the wheel, e.g. a cp310 .so under a + # cp312 runtime on Colab, or an arch-specific kernel the GPU lacks). It falls back to + # non-cpp paths and Unsloth doesn't use these kernels, so drop the cosmetic record. + logging.getLogger("torchao").addFilter(HideLoggingMessage("Failed to load ")) # SyntaxWarning: invalid escape sequence '\.' warnings.filterwarnings("ignore", message = "invalid escape sequence", category = SyntaxWarning) # PYTORCH_CUDA_ALLOC_CONF is deprecated warning from torch From 1fcd69e662f233896e075fc03472f45ac60e8819 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 19:45:46 -0700 Subject: [PATCH 13/49] Harden flaky Studio CI: retry VS-hide rename and tolerate same-URL nav interrupt (#6713) Two intermittent Studio CI failures, both runner-environment flakes unrelated to test logic: Windows 'Studio install + inference without Visual Studio': the 'Hide Visual Studio + CMake' step renames C:\Program Files\Microsoft Visual Studio to simulate a host with no build tools. A background handle on a Program Files directory (Defender scan or an MSBuild node) makes Rename-Item intermittently fail with 'Access is denied', and $ErrorActionPreference = Stop turns that into a hard job failure. Wrap the VS and cmake renames in both Hide steps in a short Rename-WithRetry (6 tries, 3s apart) to ride out the transient lock. macOS 'Chat UI Tests': the re-login goto to /login can be interrupted by the SPA auth guard redirecting to the same /login URL, which Playwright reports as 'Navigation to .../login is interrupted by another navigation to .../login'. The goto already tolerated ERR_ABORTED; broaden it to also tolerate the same-URL interrupt (the password-field wait right after confirms we landed on /login), and add the same signature to the two Playwright flake-retry harnesses as a safety net for any other navigation. Validated: playwright_chat_ui.py parses + byte-compiles, both workflow YAMLs parse, bash -n on the retry harnesses, PowerShell AST parse on all pwsh steps, and a functional check of Rename-WithRetry (succeeds, and rethrows after exhausting retries). --- .github/workflows/studio-mac-ui-smoke.yml | 29 ++++++++++--------- .../studio-windows-inference-smoke.yml | 22 ++++++++++++-- tests/studio/playwright_chat_ui.py | 12 ++++---- 3 files changed, 42 insertions(+), 21 deletions(-) diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index 512af54d53..20ca247b9f 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -185,13 +185,14 @@ jobs: # Retry up to 3 times to absorb known macos-14 free-runner # flakes: (1) Playwright Node 24 pipeTransport.js 'Unexpected # end of JSON input' crash when the Chromium browser process - # dies mid-test, and (2) Chromium net::ERR_NO_BUFFER_SPACE - # when the runner's kernel briefly runs out of socket buffers. - # The retry FULLY resets Studio (kill, reset-password, reboot, - # wait /api/health, re-export bootstrap pw) before re-running - # the script. A real test failure (assertion / timeout) does - # NOT match either pattern so it bypasses retry and surfaces - # immediately. + # dies mid-test, (2) Chromium net::ERR_NO_BUFFER_SPACE when the + # runner's kernel briefly runs out of socket buffers, and (3) a + # goto 'interrupted by another navigation' when the SPA auth + # guard redirects mid-navigation. The retry FULLY resets Studio + # (kill, reset-password, reboot, wait /api/health, re-export + # bootstrap pw) before re-running the script. A real test failure + # (assertion / timeout) does NOT match any pattern so it bypasses + # retry and surfaces immediately. run: | mkdir -p logs/playwright attempt=1 @@ -204,8 +205,9 @@ jobs: if [ "$rc" -eq 0 ]; then break fi - if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \ - || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log; } \ + if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \ + || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log \ + || grep -q "interrupted by another navigation" logs/playwright_attempt_${attempt}.log; } \ && [ "$attempt" -lt "$max_attempts" ]; then echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..." kill "${STUDIO_PID}" 2>/dev/null || true @@ -280,8 +282,8 @@ jobs: STUDIO_UI_TURN_TIMEOUT_MS: '540000' GGUF_REPO: ${{ env.GGUF_REPO }} GGUF_VARIANT: ${{ env.GGUF_VARIANT }} - # Same flake-retry shape as "Drive the chat UI with Playwright" - # -- catches pipeTransport JSON crash and ERR_NO_BUFFER_SPACE. + # Same flake-retry shape as "Drive the chat UI with Playwright" -- catches + # pipeTransport JSON crash, ERR_NO_BUFFER_SPACE, and nav interrupts. run: | mkdir -p logs/playwright_extra attempt=1 @@ -294,8 +296,9 @@ jobs: if [ "$rc" -eq 0 ]; then break fi - if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \ - || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log; } \ + if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \ + || grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log \ + || grep -q "interrupted by another navigation" logs/playwright_extra_attempt_${attempt}.log; } \ && [ "$attempt" -lt "$max_attempts" ]; then echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..." kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index c44c68278d..08a0ee782d 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -1338,11 +1338,19 @@ jobs: 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 } + } + } # 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-Item -LiteralPath $d -NewName ((Split-Path $d -Leaf) + '.vsoff') + Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff') Write-Host "Hid VS: $d" } } @@ -1351,7 +1359,7 @@ jobs: $hidden = @() foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) { if ($c.Source -and (Test-Path -LiteralPath $c.Source)) { - Rename-Item -LiteralPath $c.Source -NewName ((Split-Path $c.Source -Leaf) + '.off') + Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off') $hidden += $c.Source Write-Host "Hid cmake: $($c.Source)" } @@ -1536,8 +1544,16 @@ jobs: 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 } + } + } foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - if (Test-Path -LiteralPath $d) { Rename-Item -LiteralPath $d -NewName ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" } + if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" } } - name: Windows CUDA and ROCm prebuilts exist in unslothai/llama.cpp (what GPU users download, no VS) diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index a892d3414d..a53534acc0 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -1225,15 +1225,17 @@ with sync_playwright() as p: # ───────────────────────────────────────────────────── 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; resolve on - # domcontentloaded and tolerate it -- the pw-field wait confirms /login. + # 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. + _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 "ERR_ABORTED" not in str(exc): + if not any(t in str(exc) for t in _tolerated_nav): raise - info(f"goto /login aborted ({exc!r}); password-field wait will confirm /login") + 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) From c8bcacc3fea27e97bea0ea09c9ad7554c729724f Mon Sep 17 00:00:00 2001 From: oobabooga Date: Sat, 27 Jun 2026 02:43:36 -0300 Subject: [PATCH 14/49] Fix fast_inference crash on ABI-broken vLLM: probe compiled extensions, not just import vllm (#6621) * Fix fast_inference crash on ABI-broken vLLM: force-load compiled extensions in the broken-vLLM probe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Broaden broken-vLLM probe: catch non-libcudart .so failures and _moe_C_stable_libtorch * Revert stray reformat of the PDL fix log line * Trim verbose comments in the broken-vLLM probe * Drop non-existent vllm._moe_C_stable_libtorch from the broken-vLLM probe * Shorten comments in broken vLLM extension detection Condense the docstrings and inline comments for the lazy-loaded vLLM probe and the new regression test while keeping the rationale. Comments only, no code changes (verified with an AST signature check and the existing tests). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- tests/test_vllm_broken_detection.py | 166 ++++++++++++++++++++++++++++ unsloth/_gpu_init.py | 19 ++-- unsloth/import_fixes.py | 27 ++++- 3 files changed, 197 insertions(+), 15 deletions(-) create mode 100644 tests/test_vllm_broken_detection.py diff --git a/tests/test_vllm_broken_detection.py b/tests/test_vllm_broken_detection.py new file mode 100644 index 0000000000..ee89ddbedf --- /dev/null +++ b/tests/test_vllm_broken_detection.py @@ -0,0 +1,166 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Regression test for #6590: modern vLLM lazy-loads its compiled extensions, so +a bare ``import vllm`` succeeds even when ``vllm._C`` (or a sibling) is ABI-broken +and ``disable_broken_vllm`` missed it. GPU-free, via a synthetic vLLM.""" + +from __future__ import annotations + +import contextlib +import importlib.abc +import importlib.machinery +import importlib.util +import sys +import types + +import pytest + + +_LIBCUDART_ERROR = "libcudart.so.13: cannot open shared object file: No such file or directory" + + +class _ExtensionLoader(importlib.abc.Loader): + """A compiled extension that loads cleanly or fails on dlopen.""" + + def __init__(self, broken, error): + self.broken = broken + self.error = error + + def create_module(self, spec): + return None + + def exec_module(self, module): + if self.broken: + raise ImportError(self.error) + + +class _FakeVllmFinder(importlib.abc.MetaPathFinder): + """Lazy vLLM: ``import vllm`` succeeds; each ``vllm._*`` ext is healthy, + ABI-broken, or absent, as real vLLM only loads ``_C`` & friends on use.""" + + def __init__(self, present, broken, error): + self.present = present + self.broken = broken + self.error = error + + def find_spec( + self, + fullname, + path = None, + target = None, + ): + if fullname in self.present: + return importlib.machinery.ModuleSpec( + name = fullname, + loader = _ExtensionLoader(broken = fullname in self.broken, error = self.error), + is_package = False, + ) + return None # absent -> ModuleNotFoundError, which the guard ignores + + +@contextlib.contextmanager +def _fake_vllm( + present, + broken, + error = _LIBCUDART_ERROR, +): + """Install a synthetic lazy vLLM, restoring VLLM_BROKEN, find_spec, + meta_path, and the vllm* sys.modules entries on exit.""" + from unsloth import import_fixes + + submodules = import_fixes._VLLM_COMPILED_EXTENSIONS + saved_meta_path = list(sys.meta_path) + saved_find_spec = importlib.util.find_spec + saved_broken = import_fixes.VLLM_BROKEN + saved_modules = {n: sys.modules.get(n) for n in ("vllm", *submodules)} + try: + import_fixes.VLLM_BROKEN = False + fake_vllm = types.ModuleType("vllm") + fake_vllm.__path__ = [] + fake_vllm.__spec__ = importlib.machinery.ModuleSpec("vllm", loader = None, is_package = True) + sys.modules["vllm"] = fake_vllm + for name in submodules: + sys.modules.pop(name, None) + sys.meta_path.insert(0, _FakeVllmFinder(present, broken, error)) + yield import_fixes + finally: + import_fixes.VLLM_BROKEN = saved_broken + sys.meta_path[:] = saved_meta_path + importlib.util.find_spec = saved_find_spec + for name, module in saved_modules.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + +@pytest.mark.parametrize( + "broken_ext", + ["vllm._C", "vllm._C_stable_libtorch"], + ids = ["core_C", "sibling_C_stable_libtorch"], +) +def test_disable_broken_vllm_detects_lazy_loaded_broken_extension(broken_ext): + # A CUDA-major mismatch breaks every ext; whichever one loads first must trip detection. + present = {"vllm._C", "vllm._C_stable_libtorch"} + with _fake_vllm(present = present, broken = {broken_ext}) as import_fixes: + detected = import_fixes.disable_broken_vllm() + + assert detected is True, ( + f"disable_broken_vllm missed an ABI-broken {broken_ext} behind a " + "lazily-importable vllm package — issue #6590 would resurface." + ) + assert import_fixes.VLLM_BROKEN is True + # Once disabled, vLLM must look absent so callers fall back cleanly. + assert importlib.util.find_spec("vllm") is None + + +@pytest.mark.parametrize( + "error", + [ + "libnccl.so.2: cannot open shared object file: No such file or directory", + "libcuda.so.1: cannot open shared object file: No such file or directory", + ], + ids = ["libnccl", "libcuda"], +) +def test_disable_broken_vllm_detects_non_cudart_so_failure(error): + # A CUDA mismatch can surface through a non-libcudart .so (libnccl, libcuda), + # which the old libcudart/libcublas/libnvrtc allow-list let slip through. + with _fake_vllm(present = {"vllm._C"}, broken = {"vllm._C"}, error = error) as import_fixes: + detected = import_fixes.disable_broken_vllm() + + assert detected is True, ( + f"disable_broken_vllm missed a present-but-broken vllm._C raising " + f"{error!r} — vLLM would be left enabled and crash later." + ) + assert import_fixes.VLLM_BROKEN is True + + +@pytest.mark.parametrize( + "present", + [{"vllm._C"}, {"vllm._C", "vllm._C_stable_libtorch", "vllm._moe_C"}], + ids = ["core_only", "all_present"], +) +def test_disable_broken_vllm_keeps_healthy_vllm_enabled(present): + # Healthy install: an absent sibling (ModuleNotFoundError) or an extra present + # ext that loads cleanly must NOT be mistaken for an ABI break. + with _fake_vllm(present = present, broken = set()) as import_fixes: + detected = import_fixes.disable_broken_vllm() + + assert detected is False + assert import_fixes.VLLM_BROKEN is False + assert importlib.util.find_spec("vllm") is not None + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index 917717e08e..7f080336aa 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -36,16 +36,15 @@ from .import_fixes import ( fix_huggingface_hub, ) -# Redirect a read-only Hugging Face cache before anything below can import -# huggingface_hub / transformers / vllm (disable_broken_vllm probes -# `import vllm`, check_fbgemm_gpu_version imports transformers, and -# fix_huggingface_hub imports huggingface_hub itself), all of which can -# freeze Hub's cache constants with the un-redirected paths. unsloth_zoo -# runs the same redirect at import, but that happens after these probes. -# hf_cache.py is stdlib-only, so load it straight from its file without -# triggering the full unsloth_zoo package init this early; the zoo's own -# call later is an idempotent no-op. Older unsloth_zoo without hf_cache.py -# is skipped silently. +# Redirect a read-only Hugging Face cache before anything below imports +# huggingface_hub / transformers / vllm (disable_broken_vllm probes `import vllm` +# and its compiled extensions, check_fbgemm_gpu_version imports transformers, +# fix_huggingface_hub imports huggingface_hub) -- any of which would freeze Hub's +# cache constants with the un-redirected paths. unsloth_zoo runs the same redirect +# at import, but only after these probes. hf_cache.py is stdlib-only, so load it +# straight from its file without triggering the full unsloth_zoo init this early; +# the zoo's later call is an idempotent no-op. Older unsloth_zoo without it is +# skipped silently. try: import importlib.util as _importlib_util from pathlib import Path as _Path diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index bff55b4e7b..d979e5f22f 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -2354,11 +2354,9 @@ def _is_broken_vllm_error(error) -> bool: ) ) or ("vllm" in message and "undefined symbol" in message): return True - # Also catch CUDA shared library mismatches during vllm import - # e.g. "libcudart.so.12: cannot open shared object file" - if ( - "libcudart" in message or "libcublas" in message or "libnvrtc" in message - ) and "cannot open shared object file" in message: + # Forced extension load raises the bare loader error (no "vllm._C" + # wrapper); match any .so failure as callers feed only vLLM imports. + if "cannot open shared object file" in message: return True current = getattr(current, "__cause__", None) or getattr(current, "__context__", None) return False @@ -2550,6 +2548,16 @@ def _clear_vllm_modules(): sys.modules.pop(module_name, None) +# vLLM's compiled extensions. A CUDA-major ABI break hits all of them, so +# probing the eagerly-loaded _C and its siblings reliably trips it. +_VLLM_COMPILED_EXTENSIONS = ( + "vllm._C", + "vllm._C_stable_libtorch", + "vllm._moe_C", + "vllm._rocm_C", +) + + def disable_broken_vllm(error = None): """Disable vLLM dynamically when its shared library is ABI-broken.""" global VLLM_BROKEN @@ -2567,6 +2575,15 @@ def disable_broken_vllm(error = None): try: import vllm # noqa: F401 + + # Lazy vLLM lets a bare `import vllm` succeed even when an extension + # is ABI-broken; force-load each to surface the .so failure here. + # A missing one raises ModuleNotFoundError (skipped below). + for _ext in _VLLM_COMPILED_EXTENSIONS: + try: + importlib.import_module(_ext) + except ModuleNotFoundError: + pass return False except Exception as import_error: failure = import_error From 98a01e70cd8a4fac4be61b0bb50aac1b95dfc648 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 27 Jun 2026 01:52:18 -0700 Subject: [PATCH 15/49] Studio: restore tensor parallelism for vision/mmproj GGUFs (#6659) * Studio: restore tensor parallelism for vision/mmproj GGUFs #6416 disabled --split-mode tensor for any GGUF that ships an mmproj projector to dodge a GGML_ASSERT crash (#6415) seen on an older llama.cpp build with consumer Blackwell (sm_120). The blanket skip silently dropped tensor_parallel=true for every multimodal/MTP GGUF (e.g. Qwen3.6-35B-A3B-MTP); on hardware where the model fits on one GPU the load then collapsed to a single GPU. mmproj + --split-mode tensor works on current builds (verified end to end on B200/sm_100), so the skip was disabling a working configuration. Make the vision skip self-healing per binary: - attempt tensor for vision models by default - skip upfront only on a binary already seen to abort on tensor + mmproj this session (_vision_tensor_split_aborts), recorded when such a launch crashes at startup (_record_vision_tensor_split_abort). Process scoped, so a studio update re-probes the new build. The route-level layer-split fallback stays the net. - add _select_gpus(min_gpus=...) so a downgraded tensor request can keep multiple GPUs instead of collapsing to one (default 1, no behavior change). Add tests/test_tp_vision_regression.py: an AST allowlist guard over the tensor_parallel drop sites (which would have flagged #6416), plus cache and _select_gpus coverage. No GPU required. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: address review on vision tensor-parallel self-healing Three fixes from the PR review: - Record a vision-tensor abort only after every startup retry fails. The first version cached the binary on the first spawn crash, which on every build (including capable ones) is the benign --fit step abort that the existing --fit off retry resolves. That poisoned the cache so the next vision load in the same process skipped tensor. Recording now happens at the post-retry failure block (after fit-off, flash-attn-off and MTP-drop), so a binary that actually works is never cached. - Gate the record on the tensor/mmproj crash signature: a hard signal fault (_is_signal_crash) with no non-tensor cause (_output_has_nonprojector_diagnostic excludes OOM and unknown-arch), so an OOM, bad extra args, or MTP/flash-attn crash no longer marks an otherwise capable binary incompatible. - Preserve the multi-GPU request on the cached downgrade. The vision gate now raises _layer_min_gpus to the visible GPU count and threads it through the layer-split GPU selection (_select_gpus min_gpus and the subset loops), so a downgraded tensor request still spreads across GPUs instead of collapsing to a single card the model happens to fit. Verified two vision+tensor loads in one backend process both tensor-split across 4 GPUs (the benign fit abort no longer poisons the cache). Tests updated. * Studio: harden vision tensor-parallel self-healing (review round 2) Address the second review round on the vision/mmproj tensor-parallel fix: - Preserve vision on the first load: a --split-mode tensor + --mmproj GGML_ASSERT now raises so the route-level tensor->layer fallback retries layer split with the projector intact, instead of stripping --mmproj and silently loading text-only (which returned success and skipped the fallback, losing vision on the first load until the next cached load). - Symmetric multi-GPU preservation: the pooled-VRAM tensor downgrade now raises _layer_min_gpus from the usable tensor GPUs like the vision downgrade, so it no longer collapses a multi-GPU request to a single card. - Base the layer fallback minimum on usable GPUs: _select_gpus caps min_gpus to the count of cards with usable VRAM, so a downgrade never forces a nearly-full card in (or trips --fit) just to hit the count. - Re-probe after in-app updates: key the per-binary abort cache on (path, mtime) like _capability_cache, so POST /api/llama/update swapping the binary in place (no backend restart) re-probes the new build instead of inheriting the old build's abort. - Bump _layer_min_gpus for a known-bad vision binary independent of the tensor drop, so the route fallback's layer retry (tensor already off) still spreads across GPUs. Adds deterministic non-GPU regression tests for each. * Studio: gate cached-vision layer minimum on the current tensor request The cached-vision _layer_min_gpus bump fired for every later vision load on a binary recorded as tensor+mmproj-incompatible, including loads that did not request tensor parallelism. A plain non-tensor vision load that fits on one card would then grab every GPU just because an earlier TP attempt aborted in the same backend process. Re-tie the bump to the current tensor request (back inside the tensor-drop guard), so only a downgraded tensor request preserves the multi-GPU spread; a non-tensor vision load minimizes device count as before. * Studio: preserve GPU count + confirm assert on vision tensor fallback Third review round on the vision/mmproj tensor-parallel fix: - Preserve multi-GPU on the first tensor->layer fallback. The route-level retry runs tensor-off, so the in-function downgrades can't see the original tensor request and a fits-on-one-card model loaded the first successful fallback on a single GPU. The GGUF load closure now passes preserve_multi_gpu_on_layer (the toggle asked for tensor, this attempt is layer) and load_model raises _layer_min_gpus for it, so the downgrade still spreads across GPUs. - Cap the auto-context layer loops to usable GPUs. They bypass _select_gpus, so a raised _layer_min_gpus could force a nearly-full card into the subset (or trip --fit). They now start from _auto_min_gpus, capped to the GPUs with usable VRAM. - Confirm the tensor/mmproj assert before caching. Recording (and the layer-retry raise) now require the ggml assert marker via _is_tensor_split_assert, not the bare-signal predicate shared with the projector-incompat branch, so a corrupt or too-new projector that SIGSEGVs independent of split mode is no longer cached as tensor/mmproj-incompatible. Adds deterministic non-GPU regression tests for each. * Studio: extend multi-GPU fallback to extra/env tensor + overhead-aware cap Fourth review round on the vision/mmproj tensor-parallel fix: - Preserve multi-GPU fallback for all tensor requests, not just the UI toggle. Tensor can also be requested via --split-mode tensor in extra args or an inherited LLAMA_ARG_SPLIT_MODE=tensor env; the fallback retries those too, so the preserve_multi_gpu_on_layer hint now keys off _effective_tensor_parallel (the same check the fallback uses), comparing the overall request against the current attempt instead of only request.tensor_parallel. - Cap the auto-context layer fallback to GPUs that can pay the per-device layer overhead. The cap counted any card with positive usable VRAM, so a nearly-full GPU with a few MiB free stayed eligible and could be exposed to llama.cpp and OOM. It now mirrors _select_gpus: a card counts only if usable VRAM exceeds the per-device pipeline overhead. Adds deterministic non-GPU regression tests for both. * Studio: match the #6415 split-axis assert + replay layer-preserve hint Fifth review round on the vision/mmproj tensor-parallel fix: - Narrow the tensor/mmproj crash signature. _is_tensor_split_assert matched any GGML_ASSERT/GGML_ABORT, so an unrelated invariant a corrupt GGUF or projector trips with --mmproj present could be cached as tensor/mmproj-incompatible. It now matches the specific #6415 warmup assertion (GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0) in ggml-backend-meta), whose split-axis signature is inherent to tensor splitting. A reworded future assert just re-crashes-then-falls-back (vision preserved via layer split) instead of poisoning the cache for other models. - Persist the layer-preserve hint for respawns. A successful tensor->layer fallback committed _last_load_kwargs without preserve_multi_gpu_on_layer, so _respawn_if_dead replayed only --split-mode layer + tensor_parallel=False and a mid-session respawn of a fits-on-one-card model came back single-GPU. The hint is now in the replay snapshot, so recovery keeps the multi-GPU placement. Adds deterministic non-GPU regression tests for both. * Studio: tighten comments on the vision tensor-parallel fix Make the comments and docstrings added by this PR succinct: collapse the multi-line block comments in llama_cpp.py / inference.py to one or two lines, trim the verbose test docstrings (the names and assert messages already carry the intent), and shorten the module docstring. No code changes; verified comment-only with scripts/comment_tools.py check --strip-docstrings. * Studio: cache vision tensor abort only on the split-axis token _is_tensor_split_assert also accepted any GGML_ASSERT/GGML_ABORT from ggml-backend-meta, but that file holds many asserts, so an unrelated scheduler/projector/model invariant on an --mmproj launch could cache the binary as tensor/mmproj-incompatible and make later compatible vision models skip tensor parallelism. Match the GGML_BACKEND_SPLIT_AXIS_* token itself (unique to the #6415 warmup assert), not the source file name. * Studio: don't leak the httpx test stub into later tests The regression module stubbed httpx via sys.modules.setdefault, which installs the lightweight stub even when real httpx is present but not yet imported. The stub then persists for the whole pytest process, so provider/HF tests collected later (importing httpx or huggingface_hub.errors) got a module missing HTTPError/Response. Mirror the neighboring llama_cpp helper tests: import real httpx first and only fall back to a stub on ImportError. * Studio: latch the #6415 tensor-split abort on the first spawn, key it per model The self-heal recorded the --split-mode tensor abort only in the post-retry failure block, after the flash-attn-off retry. But SPLIT_MODE_TENSOR requires flash_attn, so the flash-off retry can't run tensor and its output no longer carries the warmup split-axis assert (ggml-backend-meta :541). The record therefore never fired on the real reproducer and the crash loop repeated on every load (reported by oobabooga on #6659). Latch instead on the first spawn that shows the signal crash + split-axis marker: record it, kill the process, and raise straight to the route's layer fallback, skipping the futile flash-attn/MTP retry ladder for this crash. The crash is a tensor-split geometry limit (e.g. MQA n_head_kv=1 splitting to GGML_BACKEND_SPLIT_AXIS_0), not a vision/mmproj property: it reproduces without --mmproj and even single-GPU tensor. So drop the vision/mmproj scoping, rename _vision_tensor_* -> _tensor_split_*, and key the session cache on (binary, mtime, model) rather than (binary, mtime) so one model's abort no longer skips tensor for every other model on the same build. Regression tests updated to pin the early-spawn record, the per-model cache, and that an unrelated ggml-backend-meta assert is not treated as the marker. * Studio: reload on explicit tensor-off after a multi-GPU layer fallback When a tensor load is downgraded to layer but kept multi-GPU to honor the tensor request (preserve_multi_gpu_on_layer, the geometry-cache gate, or the budget downgrade), the server reports tensor_parallel=False with --split-mode layer stored. A later Apply that explicitly turns the tensor toggle off then matched the loaded state and deduped to already_loaded, so Studio kept the fallback's all-GPU CUDA_VISIBLE_DEVICES placement instead of re-selecting normal placement (a single GPU for a model that fits on one card). Latch a _layer_preserves_tensor_intent flag in load_model whenever a tensor request is downgraded to layer with the multi-GPU floor raised (_layer_min_gpus > 1), clear it when tensor stays on or on unload, and force a reload in _request_matches_loaded_settings when the user explicitly turns the tensor toggle off while that flag is set. An Apply that does not touch the toggle still dedupes, so a working multi-GPU layer server is not churned. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: address reviewer.py findings on the tensor-split self-heal P1 (dedup): tensor intent can be dropped via extras, not only the toggle. An explicit llama_extra_args=["--split-mode", "layer"] matches the stored fallback extras, so _request_matches_loaded_settings deduped to the preserved all-GPU placement instead of reloading. Now reload when layer_preserves_tensor_intent and the user explicitly drops tensor via the toggle OR via extras (_effective_tensor_parallel of the explicit extras is false). P1 (downgrade symmetry): the len(tp_gpus) < 2 compute-buffer downgrade cleared tensor_parallel without raising _layer_min_gpus, unlike the budget and geometry downgrades. GPUs below tensor's replicated compute-buffer reserve can still take layer split's lower overhead, so keep the multi-GPU request (len(gpus) >= 2) and let _select_gpus cap unusable cards. P2 (cache key): key the tensor-split abort cache on st_mtime_ns, so a binary replaced in place within the same second after an abort is re-probed instead of inheriting the stale entry. P2 (test hygiene): load routes/inference.py via importlib in the regression tests instead of importing the routes package, which runs routes/__init__.py and pulls in every router (e.g. python-multipart). Added regression coverage for the extras-off reload, the compute-buffer multi-GPU preservation, and the same-second nanosecond cache invalidation. * Studio: record the tensor-split abort on the Windows CRT abort exit too The first-spawn split-axis latch only recorded when _is_signal_crash matched (POSIX signal or 0xC0000000+ NTSTATUS). On MSVC builds GGML_ASSERT terminates through the CRT abort() path with exit code 3, which is neither, so the cache never filled on Windows and every later load of the same bad binary/model repeated the tensor crash before falling back to layer. The split-axis marker is definitive, so accept either a signal crash or the Windows abort() exit (3) when the marker is present. Add _is_abort_exit and a unit test, and assert the early latch honors it. * Studio: fix UnboundLocalError on --fit-on fallback, reload backend fast path Two follow-ups from review on the tensor-split self-heal: UnboundLocalError: _layer_min_gpus was initialized inside the GPU-selection try. If NVML probing or GGUF/mmproj sizing raised, the except path logged "using --fit on" and fell through to the command builder, where the new self._layer_preserves_tensor_intent = _layer_min_gpus > 1 then raised, turning a safe --fit-on layer fallback into a hard load failure. Bind _layer_min_gpus before the try so the except path always has it. Backend fast path: _request_matches_loaded_settings forces a reload when a preserved tensor->layer fallback gets an explicit tensor-off request, but load_model's own _already_in_target_state still matched the tensor-off/layer settings and short-circuited, so the placement re-selection never ran. Mirror the guard there: reload when layer_preserves_tensor_intent and the request drops tensor intent. The flag clears on that reload, so there's no loop. Added regression coverage for both. * Studio: testable tensor-split record decision; skip futile fit-off retry Follow-ups from a deeper review of the tensor-split self-heal: Extract the record decision into _should_record_tensor_split_abort(rc, output) (marker AND (signal crash OR Windows abort)) and call it from the early latch. The combined boolean was only covered by source-inspection substring checks, so an or->and typo would silently stop recording on Windows (CRT abort exit 3 is not a signal) with every test still green. Add a behavioral test over the POSIX / Windows / NTSTATUS / clean-exit / SIGKILL / no-marker matrix. Skip the --fit off retry inside _spawn_and_wait when the crash already shows the split-axis marker: that abort is fit-independent, so the retry just warms up and crashes a second time before the latch records it. Skipping it lets the caller latch immediately and corrects the latch comment. Also clarify the dedup-guard comments (toggle read from model_fields_set vs extras via _effective_tensor_parallel without env; the backend fast path is intentionally broader and only ever forces a reload). * Studio: don't reload-loop tensor-off requests under env tensor The preserved-fallback reload guard fired on the raw tensor toggle, ignoring LLAMA_ARG_SPLIT_MODE=tensor. For an env-driven tensor user, an explicit tensor_parallel=false request then forced a reload that re-engaged tensor via the env and re-created the same preserved layer fallback, so every /load reloaded -- bypassing the env-downgrade matching that exists to avoid exactly this loop. Gate the guard on the env-aware effective tensor state: reload only when an explicit toggle/extras change leaves _effective_tensor_parallel (which consults the env) off. If the env still forces tensor, fall through to the existing env-downgrade match, which dedupes instead of looping. Added a regression test with LLAMA_ARG_SPLIT_MODE=tensor set. * Studio: tighten comments and test docstrings on the TP self-heal Condense the verbose comments and test docstrings added across the review rounds into fewer, succinct lines without changing their intent: the early-latch and downgrade-site rationale, the cache/key and helper docstrings, the dedup-guard comments, and the per-test docstrings. No code changes (AST-verified comments and docstrings only); tests and lint unchanged. * Studio: clear preserved tensor flag on diffusion; carry it across non-drop reloads Two follow-ups on the preserved-fallback machinery: Diffusion: the DiffusionGemma path early-returns from load_model before the command builder that sets/clears _layer_preserves_tensor_intent, so the flag from a prior tensor->layer fallback leaked onto a later diffusion load and forced needless reloads of the diffusion server on tensor-off/extra Applies. Clear it when starting diffusion. Settings reload: the preserve hint was recomputed only from the new request, so a reload for an unrelated setting (e.g. max_seq_length) with the tensor toggle omitted dropped a preserved multi-GPU layer placement back to one GPU. Carry llama_backend.layer_preserves_tensor_intent into the hint when the request is not an explicit tensor-off/extras-off drop, so a fitting model stays multi-GPU. Added regression tests for the diffusion clear, the carry-forward, and the updated tensor-intent computation. * Studio: gate the preserve carry-forward on the same model being loaded The tensor-intent carry-forward read llama_backend.layer_preserves_tensor_intent without checking it belonged to the model being loaded. On a direct model switch (load B without an explicit /unload of A), the flag is still set from A's downgrade (it isn't reset until B's load_model reaches the command builder, after the route reads it), so a plain load of B got preserve_multi_gpu_on_layer=True and was spread across all GPUs even though it fits on one and the user never requested tensor for it. The backend dedup doesn't have this leak (it checks model_identifier first); the leak was only in the route hint. Extract the decision into _carry_preserved_tensor_intent(preserved, same_model, explicit_drop) and gate it on the backend still holding the same model. Add a behavioral truth-table test (catches a `not` inversion and a missing same-model guard) and tighten the compute-buffer downgrade test to bound its source window. * Studio: match the HF quant too when carrying preserved tensor intent The same-model guard on the preserve carry-forward compared only model_identifier, which is variant-agnostic for HF repos. A later load of the same repo with a different gguf_variant (which already bypassed dedupe on the variant mismatch) was treated as the same model, so a request that omits tensor settings inherited the prior variant's preserved intent and forced multi-GPU layer placement for a quant that never requested tensor. Also require the loaded hf_variant to match for HF repos (local direct-file loads already differ by model_identifier path). Added a regression test for the variant guard. * Studio: match the loaded GGUF by path too when carrying preserved tensor intent A local directory holding multiple GGUF variants keeps one variant-agnostic model_identifier (the directory) while config.gguf_file selects the file, so the same-model guard let variant B inherit variant A's preserved tensor->layer fallback and forced B onto multi-GPU. Mirror _already_in_target_state's identity logic: match by resolved path when both sides have a local file, else by HF variant. #6659 * Studio: let implicit same-settings reloads dedupe after a preserved fallback The backend _already_in_target_state mirror forced a reload on ANY effective tensor-off request once a tensor->layer fallback was preserved. In the HF auto-pick / local-directory flows the route-level dedup is skipped, so an identical /load with tensor omitted reached this guard and reloaded every time even without an explicit drop. Thread the route's preserve_multi_gpu_on_layer decision in so only an explicit drop reloads; implicit carry-forward dedupes. #6659 * Studio: only an explicit tensor/split-mode change drops preserved intent The explicit-drop test treated request.llama_extra_args is not None as a drop, so a same-model reload that merely added an unrelated pass-through arg (e.g. --top-k 20) without touching the tensor field or --split-mode disabled the carry-forward and collapsed a fitting model back to one GPU. A drop now requires an explicit tensor_parallel field change or a non-tensor --split-mode override, via a shared _is_explicit_tensor_drop helper used by both the already-loaded dedup and the load carry-forward so the two readers agree. #6659 * Studio: treat an explicit clear of extras as a tensor drop When tensor intent was extras-driven (--split-mode tensor) and fell back to a preserved layer split, a later request that explicitly clears extras (llama_extra_args=[]) but omits tensor_parallel left the empty list with no split-mode override, so the carry-forward kept the model pinned multi-GPU instead of returning to normal layer selection. _is_explicit_tensor_drop now also counts an explicit empty-list clear as a drop, while an unrelated extra (--top-k) or inherit (None) still carries the preserved intent. #6659 * Studio: don't treat the UI's tensor_parallel echo as a tensor drop The Studio frontend always sends tensor_parallel and copies the /load response's resolved value back into its state, so after a tensor->layer fallback every ctx/settings reload carries tensor_parallel=false even though the user never changed it. Keying the drop on the field (or on an empty extras clear) collapsed the preserved multi-GPU placement on the next reload. A fallback also always stores --split-mode layer, never a tensor split mode, so a clear never wipes tensor intent. _is_explicit_tensor_drop now drops only on an explicit non-tensor --split-mode override; the bare field echo, an empty clear, an unrelated extra, and inherit all keep the preserved placement, and --split-mode tensor / tensor_parallel=true re-engage tensor. #6659 * Studio: match the resolved config.identifier when carrying tensor intent The same-model guard for the carry-forward compared the raw request id, but ModelConfig.from_identifier normalizes it (adds the unsloth/ prefix for a shorthand, fixes repo-id case) before load_model stores config.identifier. So a ctx/settings reload using the shorthand id missed the match, dropped _carry_preserved_tensor_intent, and could collapse a preserved multi-GPU layer placement to one GPU. Compare against config.identifier (what the backend stores), keeping it symmetric with _already_in_target_state. #6659 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 203 ++++- studio/backend/routes/inference.py | 85 ++ .../tests/test_tp_vision_regression.py | 805 ++++++++++++++++++ 3 files changed, 1073 insertions(+), 20 deletions(-) create mode 100644 studio/backend/tests/test_tp_vision_regression.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 31969afb14..8ec11fa79f 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1271,6 +1271,9 @@ class LlamaCppBackend: self._cache_type_kv: Optional[str] = None # Whether --split-mode tensor was applied on the active load. self._tensor_parallel: bool = False + # Layer load kept multi-GPU only to honor a downgraded tensor request, so a + # later explicit tensor-off reloads instead of deduping to it (#6659). + self._layer_preserves_tensor_intent: bool = False self._reasoning_default: bool = True self._speculative_type: Optional[str] = None # Canonical UI-facing mode the user requested @@ -1643,6 +1646,11 @@ class LlamaCppBackend: """Whether --split-mode tensor is active on the loaded server.""" return self._tensor_parallel + @property + def layer_preserves_tensor_intent(self) -> bool: + """True when a downgraded tensor request kept this layer load multi-GPU.""" + return self._layer_preserves_tensor_intent + @property def speculative_type(self) -> Optional[str]: return self._speculative_type @@ -2430,6 +2438,37 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # (binary, mtime, model) that aborted on --split-mode tensor this process (#6415 + # geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't + # skip tensor for others; tensor is tried by default, recorded only on a real abort. + _tensor_split_abort_keys: set[tuple[str, int, str]] = set() + + @classmethod + def _tensor_split_cache_key( + cls, binary: Optional[str], model: Optional[str] + ) -> Optional[tuple[str, int, str]]: + """(path, mtime_ns, model) key; ns mtime re-probes a same-second binary swap.""" + if not binary or not model: + return None + try: + mtime = Path(binary).stat().st_mtime_ns + except OSError: + mtime = 0 + return (binary, mtime, model) + + @classmethod + def _tensor_split_aborts(cls, binary: Optional[str], model: Optional[str]) -> bool: + """True if (binary, model) aborted on --split-mode tensor this session.""" + key = cls._tensor_split_cache_key(binary, model) + return key is not None and key in cls._tensor_split_abort_keys + + @classmethod + def _record_tensor_split_abort(cls, binary: Optional[str], model: Optional[str]) -> None: + """Remember a (binary, model) that aborts on --split-mode tensor.""" + key = cls._tensor_split_cache_key(binary, model) + if key is not None: + cls._tensor_split_abort_keys.add(key) + @staticmethod def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: """Return DLL dirs from pip-installed CUDA wheels under @@ -2569,9 +2608,13 @@ class LlamaCppBackend: usable_fraction: Optional[float] = None, total_by_idx: Optional[dict[int, int]] = None, per_device_overhead_bytes: int = 0, + min_gpus: int = 1, ) -> tuple[Optional[list[int]], bool]: """Pick GPU(s) for a model from estimated VRAM and free memory. + ``min_gpus`` (default 1, capped at ``len(gpus)``) keeps a downgraded + tensor/multi-GPU request spread instead of collapsing to one card. + ``model_size_bytes`` should include weights and estimated KV cache. ``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides headroom for compute buffers, CUDA context, and other runtime @@ -2590,9 +2633,11 @@ class LlamaCppBackend: if not gpus: return None, True + min_gpus = max(1, min(min_gpus, len(gpus))) model_size_mib = model_size_bytes / (1024 * 1024) if usable_fraction is None: usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION + overhead_mib = per_device_overhead_bytes / (1024 * 1024) # Per-GPU usable budget: free - (1-frac)*total when total is known, else # the legacy free*frac (also covers a total-0 two-column probe). @@ -2606,19 +2651,26 @@ class LlamaCppBackend: # card can have less usable room than a less-used small one. ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True) - # Try 1 GPU at the usable-VRAM threshold. - if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: + # Cap a downgraded multi-GPU request to the usable count so it doesn't pull + # in a near-full card to hit min_gpus. No-op for the default min_gpus == 1. + usable_count = sum(1 for idx, free_mib in ranked if _usable(idx, free_mib) > overhead_mib) + min_gpus = max(1, min(min_gpus, usable_count or 1)) + + # Try 1 GPU at the usable-VRAM threshold (only when one device is allowed). + if min_gpus <= 1 and _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: return [ranked[0][0]], False - # Try N GPUs (accumulate usable memory from most-free). Each GPU past the - # first adds a fixed per-device overhead the pool must hold. - overhead_mib = per_device_overhead_bytes / (1024 * 1024) + # Try N GPUs (most-free first); each past the first adds per-device overhead. + # Require at least min_gpus devices before accepting a fit. cumulative = 0.0 selected = [] for idx, free_mib in ranked: selected.append(idx) cumulative += _usable(idx, free_mib) - if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib: + if ( + len(selected) >= min_gpus + and cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib + ): return sorted(selected), False # Too large even for all GPUs; let --fit handle it @@ -3868,7 +3920,7 @@ class LlamaCppBackend: logger.debug(f"Could not list repo files for {label}: {e}") break logger.debug( - f"Could not list repo files for {label} " f"(attempt {attempt + 1}/3): {e}" + f"Could not list repo files for {label} (attempt {attempt + 1}/3): {e}" ) if attempt < 2: self._cancel_event.wait(2**attempt) @@ -4332,6 +4384,17 @@ class LlamaCppBackend: ) ) + @staticmethod + def _is_tensor_split_assert(output: str) -> bool: + """True only for the #6415 split-axis warmup assert (GGML_BACKEND_SPLIT_AXIS_*), + not any ggml assert/abort, so an unrelated invariant isn't cached. stderr is + merged into output.""" + text = (output or "").lower() + if "ggml_assert" not in text and "ggml_abort" not in text: + return False + # the split-axis enum token, unique to this assert (not the source file). + return "split_axis" in text + @staticmethod def _is_signal_crash(returncode: Optional[int]) -> bool: """True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a @@ -4344,6 +4407,20 @@ class LlamaCppBackend: return True return -returncode in (4, 6, 7, 8, 11) # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV + @staticmethod + def _is_abort_exit(returncode: Optional[int]) -> bool: + """Windows CRT abort() exit code (3) from GGML_ASSERT on MSVC -- not a POSIX + signal or 0xC0000000+ NTSTATUS.""" + return returncode == 3 + + @classmethod + def _should_record_tensor_split_abort(cls, returncode: Optional[int], output: str) -> bool: + """The #6415 split-axis abort: the marker plus a hard crash (POSIX signal or + Windows abort exit). Marker required so a generic crash isn't cached.""" + return cls._is_tensor_split_assert(output) and ( + cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode) + ) + @staticmethod def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: """Return cmd with flash attention forced off, or None when its effective @@ -4488,6 +4565,8 @@ class LlamaCppBackend: n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, extra_args: Optional[List[str]] = None, + # Route-level tensor->layer fallback retry: keep the layer split multi-GPU. + preserve_multi_gpu_on_layer: bool = False, ) -> bool: """Start llama-server with a GGUF model. @@ -4518,6 +4597,8 @@ class LlamaCppBackend: "n_gpu_layers": n_gpu_layers, "n_parallel": n_parallel, "extra_args": list(extra_args) if extra_args is not None else None, + # Replayed by _respawn_if_dead so a downgraded model stays multi-GPU. + "preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer, } # Serialise the whole load so concurrent /load calls never leave two # llama-server processes alive (#5401 / #5161). Doesn't block /unload. @@ -4541,6 +4622,7 @@ class LlamaCppBackend: chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, + preserve_multi_gpu_on_layer = preserve_multi_gpu_on_layer, ): logger.info( f"load_model: backend already in target state for " @@ -4626,6 +4708,9 @@ class LlamaCppBackend: # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; # serve them with the diffusion runner (same OpenAI-compat interface). if self._is_diffusion: + # Not a tensor/layer GGUF: clear any preserved-fallback flag from a + # prior load (this path skips the command builder that clears it). + self._layer_preserves_tensor_intent = False with self._lock: if self._cancel_event.is_set(): logger.info("Load cancelled before diffusion server start") @@ -4780,6 +4865,9 @@ class LlamaCppBackend: "image input will be disabled for this session" ) model_size = None # set in the fit try; used by the APU RAM guard + # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound + # before the try so the --fit-on except path still has it (no UnboundLocal). + _layer_min_gpus = 1 try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -5064,10 +5152,8 @@ class LlamaCppBackend: _apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024) def _restore_after_tensor_downgrade(): - # Tensor mode dropped a quantized KV and stripped the cache - # extras (it rejects quantized); layer split supports them, so - # restore the original type + extras (minus --split-mode) and - # clear the env flag so the layer launch re-emits them. + # Restore the quantized KV + extras tensor dropped (layer + # split supports them), minus --split-mode. nonlocal cache_type_kv, _cache_type_from_env, extra_args if _tensor_dropped_cache_type_kv is not None: cache_type_kv = _tensor_dropped_cache_type_kv @@ -5078,13 +5164,22 @@ class LlamaCppBackend: else extra_args ) - if tensor_parallel and effective_is_vision: + # The route fallback retry is tensor-off; keep it multi-GPU. + if preserve_multi_gpu_on_layer: + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) + + if tensor_parallel and self._tensor_split_aborts(binary, model_identifier): + # Aborted on tensor for this model this session (#6415); skip + # tensor upfront, layer split serves it. logger.info( - "Tensor parallelism skipped for vision model: " - "--split-mode tensor is incompatible with --mmproj " - "in the current llama.cpp build; using layer split." + "Tensor parallelism skipped: this llama.cpp build aborted " + "on --split-mode tensor for this model earlier this " + "session; using layer split across %d GPU(s).", + len(gpus), ) tensor_parallel = False + # Keep the multi-GPU request (gated on it, not the cache). + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) _restore_after_tensor_downgrade() # Tensor mode replicates a compute buffer on every GPU, so drop @@ -5124,6 +5219,11 @@ class LlamaCppBackend: len(gpus), ) tensor_parallel = False + # GPUs below tensor's compute-buffer reserve can still do layer + # split, so keep multi-GPU (mirrors the budget/geometry drops); + # _select_gpus caps unusable cards. + if len(gpus) >= 2: + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) # Layer split supports a quantized KV the tensor attempt # dropped; restore the original cache type + extras (minus # --split-mode) so the layer launch re-emits them. @@ -5160,8 +5260,12 @@ class LlamaCppBackend: "per-device compute buffers; falling back to layer split." ) tensor_parallel = False - # Restore the dropped quantized KV + original cache extras - # (minus --split-mode); layer split supports them. + # Weights needed >1 card, so keep multi-GPU across the + # usable tensor GPUs. + if len(tp_gpus) >= 2: + _layer_min_gpus = max(_layer_min_gpus, len(tp_gpus)) + # Restore the dropped quantized KV + cache extras (minus + # --split-mode); layer split supports them. _restore_after_tensor_downgrade() if tensor_parallel and tp_gpus: @@ -5263,6 +5367,7 @@ class LlamaCppBackend: usable_fraction = _pin_fraction, total_by_idx = total_by_idx, per_device_overhead_bytes = _pipeline_overhead_bytes, + min_gpus = _layer_min_gpus, ) # No silent shrink: effective_ctx stays == requested_ctx. else: @@ -5273,7 +5378,22 @@ class LlamaCppBackend: ranked = sorted( gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True ) - for n_gpus in range(1, len(ranked) + 1): + # Skips _select_gpus, so apply its cap: count only cards + # whose usable VRAM clears the per-device layer overhead. + _pipeline_overhead_mib = _pipeline_overhead_bytes / (1024 * 1024) + _auto_min_gpus = max( + 1, + min( + _layer_min_gpus, + sum( + 1 + for g in ranked + if _gpu_usable(g, pin_fraction) > _pipeline_overhead_mib + ) + or 1, + ), + ) + for n_gpus in range(_auto_min_gpus, len(ranked) + 1): subset = ranked[:n_gpus] pool_budget = _pool_budget_mib(subset, pin_fraction) _ms = _subset_model_size(n_gpus) @@ -5303,7 +5423,7 @@ class LlamaCppBackend: # at 131k may pin fine with a 4096 KV (#5106). effective_ctx = min(4096, effective_ctx) if effective_ctx > 0: - for n_gpus in range(1, len(ranked) + 1): + for n_gpus in range(_auto_min_gpus, len(ranked) + 1): subset = ranked[:n_gpus] kv = self._estimate_kv_cache_bytes( effective_ctx, @@ -5339,6 +5459,7 @@ class LlamaCppBackend: usable_fraction = _pin_fraction, total_by_idx = total_by_idx, per_device_overhead_bytes = _pipeline_overhead_bytes, + min_gpus = _layer_min_gpus, ) if use_fit and not explicit_ctx: # Weights don't fit on any subset; default UI to 4096 @@ -5578,12 +5699,15 @@ class LlamaCppBackend: ] ) self._tensor_parallel = True + self._layer_preserves_tensor_intent = False logger.info( "Tensor parallelism: --split-mode tensor, --tensor-split %s", tp_tensor_split, ) else: self._tensor_parallel = False + # > 1 only when a tensor request was downgraded but kept multi-GPU. + self._layer_preserves_tensor_intent = _layer_min_gpus > 1 # Speculative decoding. See _build_speculative_flags for the # mode resolution, benchmarks, and llama.cpp references. @@ -5867,7 +5991,17 @@ class LlamaCppBackend: _startup_crashed = ( self._process.poll() is not None and self._process.returncode != 0 ) - if _spawn_attempt == 0 and _fit_retry_allowed and _startup_crashed: + # A split-axis abort (#6415) is fit-independent: skip the + # --fit off retry and let the caller latch it. + _split_axis_crash = self._is_tensor_split_assert( + "\n".join(self._stdout_lines[-50:]) + ) + if ( + _spawn_attempt == 0 + and _fit_retry_allowed + and _startup_crashed + and not _split_axis_crash + ): logger.warning( "llama-server crashed during startup (exit code %s) " "with the default memory-fit step enabled; Studio " @@ -5913,6 +6047,21 @@ class LlamaCppBackend: ) healthy = _spawn_and_wait(cmd) + # #6415 split-mode tensor warmup abort. Latch it on THIS first spawn: + # the flash-attn-off retry below can't run tensor (needs flash_attn), + # so its output drops the marker and recording later would miss it, + # looping every load. Record and raise to the route's layer fallback, + # skipping the futile flash-attn/MTP retries. + if not healthy and self._tensor_parallel and not self._cancel_event.is_set(): + _ts_out = "\n".join(self._stdout_lines[-50:]) + _ts_rc = self._process.poll() if self._process is not None else None + if self._should_record_tensor_split_abort(_ts_rc, _ts_out): + LlamaCppBackend._record_tensor_split_abort(binary, model_identifier) + self._kill_process() + raise RuntimeError( + "llama-server aborted on --split-mode tensor " + "(split-axis geometry); retrying with layer split." + ) # Flash-attention kernels hard-crash at startup on some ROCm/GPU # builds (frequently inside the vision tower). Disabling FA keeps # both vision and MTP, so retry that way before dropping either. @@ -6057,6 +6206,7 @@ class LlamaCppBackend: # Read the crash code before _kill_process() clears _process. _crash_rc = self._process.poll() if self._process is not None else None self._kill_process() + # The #6415 split-axis abort is latched earlier (first spawn). # Skip if a cancel/unload is pending (mirrors the MTP guard). if ( launched_with_mmproj @@ -6488,6 +6638,7 @@ class LlamaCppBackend: spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, mtp_draft_path: Optional[str] = None, + preserve_multi_gpu_on_layer: bool = False, ) -> bool: """True iff the live server already satisfies these load kwargs. @@ -6530,6 +6681,17 @@ class LlamaCppBackend: # server. An identical request would downgrade the same way. if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel): return False + # Preserved tensor->layer fallback + an EXPLICIT tensor drop: reload so + # placement re-selects instead of keeping the all-GPU mask (mirrors the route, + # #6659). preserve_multi_gpu_on_layer carries the route's carry-forward decision + # (True for an implicit same-settings reload), so those still dedupe -- the HF + # auto-pick / local-dir flows skip the route guard and only reach here. + if ( + self._layer_preserves_tensor_intent + and not _effective_tensor_parallel(extra_args, tensor_parallel) + and not preserve_multi_gpu_on_layer + ): + return False # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. @@ -6641,6 +6803,7 @@ class LlamaCppBackend: self._supports_tools = False self._cache_type_kv = None self._tensor_parallel = False + self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index caf2a262a3..d3f56edec5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -683,7 +683,9 @@ try: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _effective_tensor_parallel, _tensor_parallel_matches_loaded, + parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -718,7 +720,9 @@ except ImportError: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _effective_tensor_parallel, _tensor_parallel_matches_loaded, + parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -2078,6 +2082,32 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[ ) +def _carry_preserved_tensor_intent( + *, preserved: bool, same_model: bool, explicit_drop: bool +) -> bool: + """Carry a preserved multi-GPU layer fallback forward only for a reload of the + SAME loaded model that doesn't explicitly drop tensor intent, so a fitting model + isn't collapsed to one GPU on a ctx-only change -- but an unrelated model switch + (without /unload) or an explicit tensor-off doesn't inherit it (#6659).""" + return preserved and same_model and not explicit_drop + + +def _is_explicit_tensor_drop(request: LoadRequest) -> bool: + """True only when the request explicitly selects a non-tensor --split-mode (e.g. + layer/row/none), a deliberate departure from a preserved tensor->layer fallback. + + A bare tensor_parallel field is NOT a drop: the Studio UI always sends it and echoes + the /load response's resolved value back, so after a fallback every reload carries + tensor_parallel=false even though the user never changed it -- treating that as a drop + would collapse the preserved multi-GPU placement on the next ctx/settings reload. An + empty clear is not a drop either (a fallback always stores --split-mode layer, never a + tensor split mode, so a clear never wipes tensor intent), nor is an unrelated extra + (--top-k) or inherit (None). tensor_parallel=true / --split-mode tensor re-engage + tensor. Shared by the already-loaded dedup and the load carry-forward (#6659).""" + override = parse_split_mode_override(request.llama_extra_args) + return override is not None and override.strip().lower() != "tensor" + + def _request_matches_loaded_settings( request: LoadRequest, llama_backend: LlamaCppBackend, @@ -2116,6 +2146,13 @@ def _request_matches_loaded_settings( effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False + # Preserved tensor->layer fallback (both report tensor=off, so the check above + # matches): if the user now explicitly drops tensor intent, reload so placement + # re-selects instead of keeping the all-GPU mask (#6659). The effective check + # includes the env, so an env-only tensor (LLAMA_ARG_SPLIT_MODE=tensor) that + # can't actually be dropped falls through to the env-downgrade match, not a loop. + if llama_backend.layer_preserves_tensor_intent and _is_explicit_tensor_drop(request): + return False # Spec decoding works on vision models too (MTP is mmproj-compatible, # llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare # the real requested mode -- coercing vision to ``off`` here used to @@ -2810,6 +2847,48 @@ async def load_model( hf_variant = config.gguf_variant, ) + # Tensor intent for this load: the request itself, or a preserved + # multi-GPU layer fallback carried across a reload of the SAME model that + # doesn't drop it (e.g. a ctx-only change), so a fitting model doesn't + # silently collapse to one GPU. Only an explicit non-tensor --split-mode + # override counts as the drop -- the tensor field echo / unrelated extras keep + # the preserved placement; the same-model guard stops a switch-without-unload + # inheriting the prior model's intent. + _explicit_tensor_drop = _is_explicit_tensor_drop(request) + # Compare the resolved config.identifier (what load_model stores), not the + # raw request id: from_identifier normalizes shorthands (adds unsloth/, fixes + # case), so a reload with the shorthand would otherwise miss the match and + # drop the carry-forward. #6659 + _same_model_loaded = ( + llama_backend.is_loaded + and (llama_backend.model_identifier or "").lower() + == (config.identifier or "").lower() + ) + # model_identifier is variant-agnostic for HF repos and dir-level for a + # local multi-variant directory, so also require the loaded quant to match + # (path else variant, mirroring _already_in_target_state) -- otherwise a + # different variant inherits the prior one's preserved intent. #6659 + if _same_model_loaded: + if config.gguf_file and llama_backend.gguf_path: + try: + _same_model_loaded = ( + Path(llama_backend.gguf_path).resolve() + == Path(config.gguf_file).resolve() + ) + except OSError: + _same_model_loaded = False + else: + _same_model_loaded = (llama_backend.hf_variant or "").lower() == ( + config.gguf_variant or "" + ).lower() + _tensor_intent_overall = _effective_tensor_parallel( + extra_llama_args, request.tensor_parallel + ) or _carry_preserved_tensor_intent( + preserved = llama_backend.layer_preserves_tensor_intent, + same_model = _same_model_loaded, + explicit_drop = _explicit_tensor_drop, + ) + # Run a single load attempt with the given tensor flag + extras. async def _attempt_gguf_load( tensor_parallel: bool, attempt_extra_args: Optional[list[str]] @@ -2823,6 +2902,12 @@ async def load_model( **_source_load_kwargs, **attempt_kwargs, tensor_parallel = tensor_parallel, + # True on the layer fallback retry (tensor wanted overall but not on + # this attempt): keep multi-GPU. Mirrors the fallback's key. + preserve_multi_gpu_on_layer = bool( + _tensor_intent_overall + and not _effective_tensor_parallel(attempt_extra_args, tensor_parallel) + ), ) # Tensor parallelism is arch-gated in llama.cpp and crashes some loads diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py new file mode 100644 index 0000000000..09af876da6 --- /dev/null +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -0,0 +1,805 @@ +# 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 guards for silent tensor-parallel downgrades in load_model. + +PR #6416 blanket-disabled tensor parallelism for vision models to dodge a +--split-mode tensor + --mmproj GGML_ASSERT (#6415), which silently single-GPU'd +any mmproj/MTP GGUF that fit on one card. The fix makes the skip self-healing: +tensor is tried by default and recorded per (binary, model) only on a real abort. + +load_model is too entangled to drive end-to-end, so these tests inspect the +source / drive the pure helpers. The headline test pins the set of TP-drop +conditions, so a new silent drop fails CI. No GPU; fully deterministic. +""" + +from __future__ import annotations + +import ast +import importlib.util +import inspect +import os +import sys +import textwrap +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# External-dep stubs so importing the backend doesn't require structlog / httpx / +# loggers -- but only when the real module is missing, so a lightweight stub never +# shadows the real package (or `loggers.handlers` submodule) for tests collected +# later in the same pytest process. +try: + import structlog # noqa: F401 +except ImportError: + _structlog_stub = _types.ModuleType("structlog") + _structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + sys.modules["structlog"] = _structlog_stub +try: + import loggers # noqa: F401 +except ImportError: + _loggers_stub = _types.ModuleType("loggers") + _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) + sys.modules["loggers"] = _loggers_stub +try: + import httpx as _httpx_real # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None}) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Client = type( + "C", + (), + { + "__init__": lambda s, **kw: None, + "__enter__": lambda s: s, + "__exit__": lambda s, *a: None, + }, + ) + sys.modules["httpx"] = _httpx_stub + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + +_GB = 1024**3 + + +def _load_inference_routes_module(): + """Load routes/inference.py directly, bypassing routes/__init__.py (which imports + every router, dragging in unrelated deps like python-multipart) (Codex #6659).""" + route_path = Path(_BACKEND_DIR) / "routes" / "inference.py" + spec = importlib.util.spec_from_file_location( + "tp_vision_regression_inference_routes", route_path + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_model_ast() -> ast.FunctionDef: + """Parse load_model into an AST FunctionDef (no import side effects).""" + src = textwrap.dedent(inspect.getsource(LlamaCppBackend.load_model)) + return ast.parse(src).body[0] + + +def _tensor_parallel_false_drop_guards() -> list[str]: + """Source of the guard expression for every `if ...: tensor_parallel = False` + (the LOCAL variable, not self._tensor_parallel) inside load_model.""" + fn = _load_model_ast() + + def _body_drops_tp(body) -> bool: + for n in body: + if ( + isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets) + and isinstance(n.value, ast.Constant) + and n.value.value is False + ): + return True + return False + + return [ + ast.unparse(node.test) + for node in ast.walk(fn) + if isinstance(node, ast.If) and _body_drops_tp(node.body) + ] + + +# Every condition that may flip a requested tensor_parallel back to False. Adding +# one must be conscious: update this allowlist and keep multi-GPU where possible. +_ALLOWED_TP_DROP_GUARDS = { + # Capability: --split-mode tensor aborted for this (binary, model) (#6415). + # Self-healing -- tried by default, skipped only after a real abort (vs #6416). + "tensor_parallel and self._tensor_split_aborts(binary, model_identifier)", + # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. + "tensor_parallel and len(tp_gpus) < 2", + # Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split. + "_tp_weight_budget_mib <= _tp_required_mib", +} + + +def test_tensor_parallel_drop_sites_match_allowlist(): + """The set of reasons a requested TP can be dropped is fixed and reviewed: a new + drop site fails this set-equality until consciously allowlisted (would catch #6416).""" + found = set(_tensor_parallel_false_drop_guards()) + assert found == _ALLOWED_TP_DROP_GUARDS, ( + "tensor_parallel drop sites changed.\n" + f" unexpected (new) : {sorted(found - _ALLOWED_TP_DROP_GUARDS)}\n" + f" missing (removed): {sorted(_ALLOWED_TP_DROP_GUARDS - found)}\n" + "A new drop means a user's TP request is ignored for a new reason -- " + "review it, keep multi-GPU where possible, surface it, then update " + "_ALLOWED_TP_DROP_GUARDS." + ) + + +def test_every_tp_drop_is_logged_not_silent(): + """Each tensor_parallel downgrade must log why, so it never disappears silently.""" + fn = _load_model_ast() + + def _body_drops_tp(body): + return any( + isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets) + and isinstance(n.value, ast.Constant) + and n.value.value is False + for n in body + ) + + def _body_logs(body) -> bool: + for n in ast.walk(ast.Module(body = list(body), type_ignores = [])): + if ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "logger" + ): + return True + return False + + for node in ast.walk(fn): + if isinstance(node, ast.If) and _body_drops_tp(node.body): + assert _body_logs(node.body), ( + f"TP drop under `{ast.unparse(node.test)}` has no logger call -- " + "downgrades must explain themselves." + ) + + +def test_tensor_split_gate_is_self_healing_not_blanket(): + """Skip is conditional on a recorded (binary, model) abort, not a blanket + is_vision disable (the #6416 regression).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert "self._tensor_split_aborts(binary, model_identifier)" in src + assert "if tensor_parallel and is_vision:" not in src + assert "if tensor_parallel and effective_is_vision:" not in src + + +def test_tensor_split_skip_documents_layer_split_fallback(): + """When the skip fires (known-bad binary+model), it states the fallback.""" + src = inspect.getsource(LlamaCppBackend.load_model) + gate = src.find("self._tensor_split_aborts(binary, model_identifier)") + assert gate != -1 + block = src[gate : gate + 600] + assert "layer split" in block, "the skip should state it falls back to layer split" + + +def test_tensor_split_abort_recorded_early_on_first_spawn(): + """Recorded on the first spawn showing the marker, before the flash-attn-off + retry (which can't run tensor so drops the marker) -- else it loops (oobabooga, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + idx = src.find("_record_tensor_split_abort(binary, model_identifier)") + assert idx != -1, "load_model must record a (binary, model) tensor-split abort" + guard = src[max(0, idx - 600) : idx] + assert "self._tensor_parallel" in guard + assert ( + "_should_record_tensor_split_abort" in guard + ), "record must be gated on the marker-plus-hard-crash decision helper" + # Recorded before the flash-attn-off retry, not after the full ladder. + fa_off = src.find("_with_flash_attn_off") + assert 0 <= idx < fa_off, "recording must latch on the first spawn, before flash-off" + + +def test_vision_downgrade_preserves_multi_gpu_intent(): + """The vision downgrade raises _layer_min_gpus and threads it into both the + _select_gpus and auto-context layer paths, so a fitting model still spreads.""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert "_layer_min_gpus = max(_layer_min_gpus, len(gpus))" in src + assert src.count("min_gpus = _layer_min_gpus") >= 2 + assert "range(_auto_min_gpus, len(ranked) + 1)" in src + auto = src.find("_auto_min_gpus = max(") + assert auto != -1 and "_layer_min_gpus" in src[auto : auto + 200] + + +# ── per-binary capability cache (pure) ─────────────────────────────── + + +def test_tensor_attempted_by_default_for_unknown_binary(): + """A (binary, model) not seen to abort -> tensor is attempted (not skipped).""" + assert LlamaCppBackend._tensor_split_aborts("/never/seen/llama-server", "m") is False + assert LlamaCppBackend._tensor_split_aborts(None, "m") is False + assert LlamaCppBackend._tensor_split_aborts("/x", None) is False + + +def test_recorded_tensor_abort_is_per_model(): + """A recorded (binary, model) abort trips the gate for that model only -- a + different model on the same binary still attempts tensor (oobabooga, #6659).""" + b = f"/tmp/llama-server-{id(object())}" + try: + assert LlamaCppBackend._tensor_split_aborts(b, "model-a") is False + LlamaCppBackend._record_tensor_split_abort(b, "model-a") + assert LlamaCppBackend._tensor_split_aborts(b, "model-a") is True + # a different model on the same binary is unaffected + assert LlamaCppBackend._tensor_split_aborts(b, "model-b") is False + finally: + LlamaCppBackend._tensor_split_abort_keys.discard( + LlamaCppBackend._tensor_split_cache_key(b, "model-a") + ) + + +# ── _select_gpus: single-GPU collapse vs honored multi-GPU intent (pure) ── + + +def test_select_gpus_collapses_to_single_gpu_when_model_fits(): + """Default (min_gpus=1): a 39 GB model on four 183 GB GPUs pins ONE GPU -- the + 'single GPU' symptom once TP drops, and why the downgrade needs min_gpus.""" + gpus = [(0, 180000), (1, 180000), (2, 180000), (3, 180000)] # (idx, free MiB) + gpu_indices, _use_fit = LlamaCppBackend._select_gpus(int(39 * _GB), gpus) + assert gpu_indices is not None and len(gpu_indices) == 1 + + +def test_select_gpus_min_gpus_keeps_multi_gpu_for_fitting_model(): + """min_gpus>=2 must NOT collapse to one GPU for a model that fits on one.""" + gpus = [(0, 180000), (1, 180000), (2, 180000), (3, 180000)] + gpu_indices, _ = LlamaCppBackend._select_gpus(int(39 * _GB), gpus, min_gpus = 2) + assert gpu_indices is not None and len(gpu_indices) >= 2 + + +def test_select_gpus_min_gpus_capped_to_available(): + """min_gpus larger than the GPU count is capped, not an error.""" + gpus = [(0, 180000), (1, 180000)] + gi, _ = LlamaCppBackend._select_gpus(int(10 * _GB), gpus, min_gpus = 8) + assert gi is not None and len(gi) == 2 + + +def test_select_gpus_uses_multiple_gpus_when_model_does_not_fit(): + """Sanity: selection spreads across GPUs when one card can't hold the model.""" + gpus = [(0, 40000), (1, 40000), (2, 40000), (3, 40000)] # 40 GB free each + gpu_indices, _use_fit = LlamaCppBackend._select_gpus(int(120 * _GB), gpus) + assert gpu_indices is not None and len(gpu_indices) >= 2 + + +def test_select_gpus_min_gpus_excludes_unusable_gpu(): + """min_gpus caps to usable cards: 2 free + 1 nearly-full -> 2-GPU split, not + forcing the full card (OOM) or tripping --fit (#6659).""" + gpus = [(0, 180000), (1, 180000), (2, 500)] # GPU 2 is nearly full + total = {0: 180000, 1: 180000, 2: 180000} + gi, _ = LlamaCppBackend._select_gpus( + int(39 * _GB), + gpus, + min_gpus = 3, + total_by_idx = total, + per_device_overhead_bytes = int(1 * _GB), + ) + assert gi is not None + assert 2 not in gi, "a nearly-full GPU must not be forced in to satisfy min_gpus" + assert len(gi) == 2 + + +def test_tensor_abort_cache_invalidated_on_binary_mtime_change(tmp_path): + """Cache keys on (path, mtime, model), so a binary swapped in place (in-app + update, no restart) is re-probed instead of inheriting the old abort (#6659).""" + binp = tmp_path / "llama-server" + binp.write_text("v1") + p = str(binp) + try: + LlamaCppBackend._record_tensor_split_abort(p, "m") + assert LlamaCppBackend._tensor_split_aborts(p, "m") is True + # Simulate an in-place update bumping the binary's mtime. + st = binp.stat() + os.utime(p, (st.st_atime, st.st_mtime + 10)) + assert ( + LlamaCppBackend._tensor_split_aborts(p, "m") is False + ), "a binary swapped in place (new mtime) must be re-probed" + # A same-second replacement (sub-second mtime bump) must also re-probe: + # second-resolution mtime would inherit the stale abort (reviewer.py P2). + sec_ns = (binp.stat().st_mtime_ns // 1_000_000_000) * 1_000_000_000 + os.utime(p, ns = (sec_ns, sec_ns)) + LlamaCppBackend._record_tensor_split_abort(p, "m") + binp.write_text("v2") + os.utime(p, ns = (sec_ns, sec_ns + 1)) + assert ( + LlamaCppBackend._tensor_split_aborts(p, "m") is False + ), "a same-second in-place swap (ns mtime bump) must be re-probed" + finally: + for key in list(LlamaCppBackend._tensor_split_abort_keys): + if key and key[0] == p: + LlamaCppBackend._tensor_split_abort_keys.discard(key) + + +def test_tensor_split_abort_raises_early_to_layer_fallback(): + """The first-spawn abort raises to the route's layer fallback (not the text-only + mmproj strip), before the flash-attn-off retry, preserving the projector (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + raise_idx = src.find("(split-axis geometry); retrying with layer split") + assert raise_idx != -1, "the split-axis abort must raise to trigger a layer retry" + # raises before both the flash-attn-off retry and the text-only mmproj strip + assert raise_idx < src.find("_with_flash_attn_off") + assert raise_idx < src.find("_strip_mmproj_args(_last_spawn_cmd)") + # gated on the marker-plus-crash helper, which also drives the record just above + guard = src[max(0, raise_idx - 600) : raise_idx] + assert "_should_record_tensor_split_abort" in guard + rec_idx = src.find("_record_tensor_split_abort(binary, model_identifier)") + assert rec_idx != -1 and rec_idx < raise_idx + + +def test_budget_downgrade_preserves_multi_gpu_intent(): + """The pooled-VRAM downgrade raises _layer_min_gpus from the usable tensor GPUs + too, symmetric with the vision downgrade (reviewer.py asymmetric fix, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + budget = src.find("_tp_weight_budget_mib <= _tp_required_mib") + assert budget != -1 + block = src[budget : budget + 1000] + assert "tensor_parallel = False" in block + assert ( + "_layer_min_gpus = max(_layer_min_gpus, len(tp_gpus))" in block + ), "the budget downgrade must preserve multi-GPU intent like the vision gate" + + +def test_compute_buffer_downgrade_preserves_multi_gpu_intent(): + """The len(tp_gpus) < 2 compute-buffer downgrade raises _layer_min_gpus from the + full GPU set too, so it is symmetric with the budget/geometry downgrades and + doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + gate = src.find("tensor_parallel and len(tp_gpus) < 2") + assert gate != -1 + # Bound to exactly this block: from its gate to the next (budget) downgrade. + nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate) + assert nxt != -1 + block = src[gate:nxt] + assert "tensor_parallel = False" in block + assert ( + "_layer_min_gpus = max(_layer_min_gpus, len(gpus))" in block + ), "the compute-buffer downgrade must preserve multi-GPU intent like the others" + + +def test_tensor_split_layer_min_gpus_bump_requires_tensor_request(): + """Every guard that bumps _layer_min_gpus off the abort cache also tests + tensor_parallel, so a non-tensor load on a known-bad binary doesn't grab every + GPU for a fitting model (#6659).""" + fn = _load_model_ast() + checked = 0 + for node in ast.walk(fn): + if isinstance(node, ast.If): + test_src = ast.unparse(node.test) + if "self._tensor_split_aborts(binary, model_identifier)" not in test_src: + continue + body = "\n".join(ast.unparse(n) for n in node.body) + if "_layer_min_gpus" in body: + checked += 1 + assert "tensor_parallel" in test_src, ( + "the cached _layer_min_gpus bump must require a current tensor " + f"request, but fires under `{test_src}`" + ) + assert checked >= 1, "expected an abort-cache guard that bumps _layer_min_gpus" + + +# ── round-2 follow-up: route-fallback retry + auto-context cap + assert marker ── + + +def test_layer_fallback_retry_preserves_multi_gpu_intent(): + """load_model takes a preserve_multi_gpu_on_layer hint and raises _layer_min_gpus + for it, so the tensor-off fallback retry still spreads a fitting model (#6659).""" + sig = inspect.signature(LlamaCppBackend.load_model) + assert "preserve_multi_gpu_on_layer" in sig.parameters + assert sig.parameters["preserve_multi_gpu_on_layer"].default is False + fn = _load_model_ast() + found = any( + isinstance(n, ast.If) + and "preserve_multi_gpu_on_layer" in ast.unparse(n.test) + and "_layer_min_gpus" in "\n".join(ast.unparse(b) for b in n.body) + for n in ast.walk(fn) + ) + assert found, "preserve_multi_gpu_on_layer must raise _layer_min_gpus" + + +def test_auto_context_layer_loops_capped_to_usable_gpus(): + """The auto-context loops bypass _select_gpus, so they apply its cap: a card + counts only if usable VRAM clears the per-device layer overhead (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert ( + "range(max(1, _layer_min_gpus), len(ranked) + 1)" not in src + ), "auto-context loops must cap _layer_min_gpus to usable GPUs, not use it raw" + assert "_auto_min_gpus" in src + assert "range(_auto_min_gpus, len(ranked) + 1)" in src + # the eligibility threshold is the per-device layer overhead, not bare > 0 + auto = src.find("_auto_min_gpus = max(") + assert auto != -1 + block = src[auto : auto + 400] + assert "_pipeline_overhead_mib" in block, ( + "a card must clear the per-device layer overhead to count, mirroring " + "_select_gpus, so a nearly-full GPU is not exposed and OOMs" + ) + + +def test_fallback_hint_uses_effective_tensor_request_not_just_toggle(): + """Tensor intent keys off _effective_tensor_parallel (toggle + extras + env), not + just the toggle, so extra/env-driven tensor users keep multi-GPU (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_tensor_intent_overall = _effective_tensor_parallel(") + assert idx != -1, "the GGUF load closure must compute tensor intent" + block = src[idx : idx + 300] + assert "extra_llama_args, request.tensor_parallel" in block + pres = src.find("preserve_multi_gpu_on_layer = bool(") + assert ( + "_effective_tensor_parallel(attempt_extra_args, tensor_parallel)" in src[pres : pres + 200] + ) + # not the toggle-only form this replaced + assert ( + "bool(\n request.tensor_parallel and not tensor_parallel" not in src + ) + + +def test_carry_preserved_tensor_intent_truth_table(): + """Behavioral check of the carry-forward decision: carried only for the SAME + model, preserved, and not an explicit drop. Catches a `not` inversion (ctx-only + collapse) and a missing same-model guard (cross-model leak) (#6659).""" + inference_routes = _load_inference_routes_module() + f = inference_routes._carry_preserved_tensor_intent + assert f(preserved = True, same_model = True, explicit_drop = False) is True + assert f(preserved = True, same_model = True, explicit_drop = True) is False # explicit drop + assert f(preserved = True, same_model = False, explicit_drop = False) is False # model switch + assert f(preserved = False, same_model = True, explicit_drop = False) is False # not a fallback + + +def test_preserved_fallback_carried_across_non_drop_reload(): + """The hint carries the preserved fallback via _carry_preserved_tensor_intent, + gated on the same model loaded, so a ctx-only reload keeps multi-GPU but a model + switch / explicit drop doesn't inherit it (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_tensor_intent_overall = _effective_tensor_parallel(") + assert idx != -1 + block = src[idx : idx + 400] + assert "_carry_preserved_tensor_intent(" in block + assert "preserved = llama_backend.layer_preserves_tensor_intent" in block + assert "same_model = _same_model_loaded" in block + assert "explicit_drop = _explicit_tensor_drop" in block + + +def test_same_model_guard_checks_path_and_variant(): + """The same-model guard matches the resolved config.identifier (what load_model + stores, after from_identifier normalizes shorthands) -- not the raw request id -- + and also matches the loaded quant by path (local multi-variant dir) else variant (HF + repo), so a reload keeps the carry-forward and a different variant doesn't inherit + the prior one's preserved tensor intent (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_same_model_loaded = (") + assert idx != -1 + block = src[idx : idx + 1300] + # Identity compares the normalized config.identifier, not the raw model_identifier. + head = src[idx : idx + 200] + assert "config.identifier" in head and "== (model_identifier" not in head + assert "llama_backend.gguf_path" in block and "config.gguf_file" in block + assert "llama_backend.hf_variant" in block and "config.gguf_variant" in block + + +def test_diffusion_load_clears_preserved_tensor_flag(): + """The diffusion early-return path (skips the command builder) clears the + preserved-fallback flag, so a prior tensor fallback doesn't churn it (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + diff = src.find("if self._is_diffusion:") + assert diff != -1 + start = src.find("return self._start_diffusion_server", diff) + assert start != -1 + assert "self._layer_preserves_tensor_intent = False" in src[diff:start] + + +def test_is_tensor_split_assert_marker(): + """Matches the specific #6415 split-axis assert, not any ggml assert/abort, so + an unrelated invariant a corrupt GGUF/projector trips isn't cached (#6659).""" + f = LlamaCppBackend._is_tensor_split_assert + # the real #6415 warmup assert (split-axis enum, in ggml-backend-meta) + assert ( + f( + "ggml-backend-meta.cpp:541: GGML_ASSERT(src_ss[0].axis != " + "GGML_BACKEND_SPLIT_AXIS_0) failed" + ) + is True + ) + # the split-axis token alone (file path elided / reworded) still matches + assert f("GGML_ASSERT(x.axis != GGML_BACKEND_SPLIT_AXIS_1) failed") is True + # UNRELATED asserts must NOT match -- including a different invariant from the + # same multi-assert source file (matched on the token, not the file name). + assert f("ggml-backend-meta.cpp:99: GGML_ASSERT(buf != NULL) failed") is False + assert f("/x/ggml.c:1234: GGML_ASSERT(ne == 1) failed") is False + assert f("ggml_abort: something else entirely") is False + assert f("Segmentation fault (core dumped)") is False + assert f("") is False + assert f(None) is False + + +def test_layer_preserve_hint_replayed_on_respawn(): + """The preserve hint is in the replay snapshot (_pending_load_kwargs), so a + respawn keeps the downgraded model multi-GPU (Codex review on #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + pend = src.find("_pending_load_kwargs = {") + assert pend != -1 + block = src[pend : src.find("}", pend) + 1] + assert '"preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer' in block, ( + "the layer-preserve hint must be in the replay snapshot so _respawn_if_dead " + "keeps the multi-GPU placement" + ) + + +def test_should_record_tensor_split_abort_decision(): + """Behavioral check of marker AND (signal crash OR Windows abort), so an + or->and typo or caching a generic crash fails here, not just the source pins.""" + f = LlamaCppBackend._should_record_tensor_split_abort + marker = "ggml-backend-meta.cpp:541: GGML_ASSERT(x.axis != GGML_BACKEND_SPLIT_AXIS_0) failed" + # marker + a hard crash records, across every platform's abort encoding + assert f(-6, marker) is True # POSIX SIGABRT + assert f(-11, marker) is True # POSIX SIGSEGV + assert f(3, marker) is True # Windows CRT abort() exit (not a signal) + assert f(0xC0000005, marker) is True # Windows NTSTATUS access violation + # marker present but no hard crash -> not recorded + assert f(0, marker) is False # clean exit + assert f(-9, marker) is False # SIGKILL (OOM / unload), not a fault + assert f(None, marker) is False # still running + # hard crash but not the split-axis marker -> not recorded (no over-caching) + assert f(3, "some other failure") is False + assert f(-6, "GGML_ASSERT(buf != NULL) failed") is False + assert f(0xC0000005, "") is False + + +def test_fit_off_retry_skipped_on_split_axis_abort(): + """The fit-independent --fit off retry is skipped on the split-axis marker, else + the model crashes a second time before the latch records it (reviewer.py, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + retry = src.find('run_cmd = [*run_cmd, "--fit", "off"]') + assert retry != -1 + guard = src[max(0, retry - 1000) : retry] + assert "_fit_retry_allowed" in guard and "_startup_crashed" in guard + assert ( + "not _split_axis_crash" in guard + ), "the fit-off retry must be skipped when the crash is a split-axis abort" + + +def test_is_abort_exit_recognizes_windows_crt_abort(): + """exit code 3 (MSVC abort()) counts as a crash; signals / clean exits do not.""" + f = LlamaCppBackend._is_abort_exit + assert f(3) is True + assert f(0) is False + assert f(-6) is False # POSIX SIGABRT is handled by _is_signal_crash, not here + assert f(None) is False + + +# ── tensor-off after a multi-GPU fallback forces a reload (route dedup) ─ + + +class _NoopProcess: + """Stand-in for Popen so is_loaded is True and atexit cleanup doesn't crash.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def _fallback_loaded_backend(layer_preserves_tensor_intent: bool) -> LlamaCppBackend: + """A loaded backend in the tensor->layer fallback state (tensor off, --split-mode + layer stored), differing only in the preserved-multi-GPU flag.""" + b = LlamaCppBackend() + b._model_identifier = "owner/repo" + b._requested_n_ctx = 0 + b._cache_type_kv = None + b._tensor_parallel = False + b._layer_preserves_tensor_intent = layer_preserves_tensor_intent + b._extra_args = ["--split-mode", "layer"] + b._requested_spec_mode = "auto" + b._chat_template_override = None + b._gguf_path = None + return b + + +def test_tensor_off_echo_preserves_multi_gpu_fallback(): + """The Studio UI always sends tensor_parallel and echoes the /load response's + resolved value, so after a fallback a ctx/settings reload carries tensor_parallel= + false even though the user never changed it. That echo must NOT collapse the + preserved multi-GPU placement -- it dedupes (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo", tensor_parallel = False) + assert "tensor_parallel" in req.model_fields_set, "the UI always sends the field" + + # Preserved fallback + bare tensor=false echo: dedupe, keep multi-GPU (no collapse). + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + # A genuine layer load (no preserved intent): tensor-off also dedupes, no churn. + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = False) + ) + is True + ) + + +def test_explicit_split_mode_layer_extras_reloads_after_multi_gpu_fallback(): + """Tensor intent can be dropped via extras too: an explicit --split-mode layer + matches the stored fallback extras but must still reload (reviewer.py P1, #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"]) + assert "llama_extra_args" in req.model_fields_set + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is False + ) + + +def test_tensor_off_reload_requires_explicit_toggle(): + """An Apply that doesn't touch the toggle (e.g. a context change) isn't churned + by the preserved-fallback reload -- the working server is kept (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo") # tensor_parallel left unset + assert "tensor_parallel" not in req.model_fields_set + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + + +def test_tensor_off_under_env_tensor_does_not_reload_loop(monkeypatch): + """With LLAMA_ARG_SPLIT_MODE=tensor set, a tensor-off request can't drop tensor + intent, so the env-aware guard dedupes instead of reload-looping (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + monkeypatch.setenv("LLAMA_ARG_SPLIT_MODE", "tensor") + + req = LoadRequest(model_path = "owner/repo", tensor_parallel = False) + assert "tensor_parallel" in req.model_fields_set + # env still forces tensor -> not a real drop -> dedupe (no reload loop). + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + + +def test_is_explicit_tensor_drop_truth_table(): + """Only an explicit non-tensor --split-mode override is a drop. A bare + tensor_parallel field (the UI always sends it and echoes the fallback's false), an + empty clear, an unrelated extra (--top-k), or inherit (None) must NOT collapse a + preserved fallback; --split-mode tensor / tensor_parallel=true re-engage (Codex + #6659).""" + from models.inference import LoadRequest + + f = _load_inference_routes_module()._is_explicit_tensor_drop + # A non-tensor split-mode override is the one deliberate departure -> drop. + assert ( + f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"])) is True + ) + # tensor / retry re-engages, never a drop. + assert ( + f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "tensor"])) + is False + ) + # A bare tensor_parallel field is the UI echo, not a drop (would collapse on reload). + assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = False)) is False + assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = True)) is False + # Unrelated extra / empty clear / inherit all keep the preserved placement. + assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--top-k", "20"])) is False + assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = [])) is False + assert f(LoadRequest(model_path = "owner/repo")) is False + + +def test_explicit_tensor_drop_uses_shared_helper_in_both_readers(): + """Both the already-loaded dedup and the load carry-forward derive the drop from + _is_explicit_tensor_drop, so they agree on what counts as a drop -- a reload for + an unrelated extra still carries the preserved intent rather than collapsing to one + GPU (Codex #6659).""" + src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() + # Dedup reader (the preserved-fallback reload guard). + assert "layer_preserves_tensor_intent and _is_explicit_tensor_drop(request)" in src + # Load carry-forward reader feeds the same decision into the carry-forward. + assert "_explicit_tensor_drop = _is_explicit_tensor_drop(request)" in src + + +def test_layer_preserves_tensor_intent_set_only_on_preserved_downgrade(): + """load_model latches the flag from _layer_min_gpus (raised only when a tensor + request is downgraded but kept multi-GPU), and clears it when tensor stays on.""" + src = inspect.getsource(LlamaCppBackend.load_model) + on = src.find("self._tensor_parallel = True") + off = src.find("self._tensor_parallel = False") + assert 0 <= on and 0 <= off + assert "self._layer_preserves_tensor_intent = False" in src[on : on + 120] + assert "self._layer_preserves_tensor_intent = _layer_min_gpus > 1" in src[off : off + 400] + + +def test_layer_min_gpus_bound_before_gpu_selection_try(): + """_layer_min_gpus is bound before the GPU-selection try, so the --fit-on except + path can't UnboundLocalError when the command builder reads it (Codex #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert src.count("_layer_min_gpus = 1") == 1, "exactly one init, before the try" + init = src.find("_layer_min_gpus = 1") + try_body = src.find("gguf_size = self._get_gguf_size_bytes") + fit_except = src.find("GPU selection failed") + use_after = src.find("self._layer_preserves_tensor_intent = _layer_min_gpus > 1") + assert ( + -1 < init < try_body < fit_except < use_after + ), "the init must precede the try body, the except, and the command-builder use" + + +def test_already_in_target_state_reloads_on_tensor_off_after_fallback(): + """The backend fast path mirrors the route dedup: a preserved fallback reloads on + an EXPLICIT tensor-off request, but an implicit same-settings reload (carry-forward + preserve_multi_gpu_on_layer=True) still dedupes (Codex #6659).""" + + def _backend(layer_preserves: bool) -> LlamaCppBackend: + b = _fallback_loaded_backend(layer_preserves_tensor_intent = layer_preserves) + b._process = _NoopProcess() + b._healthy = True + return b + + kwargs = dict( + gguf_path = None, + mtp_draft_path = None, + model_identifier = "owner/repo", + hf_variant = None, + n_ctx = 0, + cache_type_kv = None, + speculative_type = None, + spec_draft_n_max = None, + tensor_parallel = False, + chat_template_override = None, + extra_args = ["--split-mode", "layer"], + is_vision = False, + ) + # Preserved fallback + EXPLICIT tensor drop -> reload (not already in target state). + assert _backend(True)._already_in_target_state(**kwargs) is False + # Same preserved fallback but an implicit reload that carries the intent forward + # (HF auto-pick / local-dir flows skip the route guard and reach here) -> dedupe. + assert ( + _backend(True)._already_in_target_state(**kwargs, preserve_multi_gpu_on_layer = True) is True + ) + # A genuine layer load (no preserved intent) -> dedupe, no churn. + assert _backend(False)._already_in_target_state(**kwargs) is True From 4c72e09480d5f0c21d6739c62b23e7d501464ffc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 27 Jun 2026 05:21:05 -0700 Subject: [PATCH 16/49] Studio: stop handing CI/user secrets to downloaded llama.cpp binaries (#6696) * Studio: stop handing CI/user secrets to downloaded llama.cpp binaries The macOS prebuilt path installs llama.cpp from the unslothai/llama.cpp fork's latest (unpinned, mutable) release and then executes the downloaded llama-server / llama-quantize binaries during install-time validation. binary_env() built that child environment from a full os.environ.copy(), so a compromised or tampered prebuilt would inherit every secret in the process: HF_TOKEN and the workflow GitHub tokens in CI, and HF / cloud credentials for end users running install.sh / setup.sh. We publish prebuilts daily, so pinning a release tag is not workable. Instead, neutralise the impact: these binaries have no reason to read any token, so strip secret-bearing variables (exact names plus TOKEN/SECRET/PASSWORD/CREDENTIAL/PRIVATE_KEY/API_KEY markers) before handing the env to a downloaded binary. The installer's own GitHub and Hugging Face API calls read os.environ directly, so authentication and release-API rate limiting are unaffected; PATH, LD_LIBRARY_PATH, DYLD_LIBRARY_PATH and CUDA/ROCm vars are preserved. One change covers the install-time validation path for all six macOS workflows and end users. Follow-up (separate, sequenced): publish build-provenance attestations from the fork's prebuilt workflows and verify them in CI, so a forged release is rejected rather than merely starved of secrets. * Strip KUBECONFIG, SSH_AUTH_SOCK, and PASSPHRASE-marked vars from binary env Extend the deny-list per PR review: KUBECONFIG and SSH_AUTH_SOCK are credential pointers/capabilities a downloaded binary never needs, and a PASSPHRASE marker catches SSH_PASSPHRASE / GPG_PASSPHRASE. Tests updated. * Studio: also scrub proxy/index env vars and URL-embedded credentials before running prebuilt binaries * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope mlx-ci secrets to the install + download commands for PR #6696 Drop the ambient step-level env block and pass GH/GITHUB/HF tokens only on the installer and GGUF-download commands, so the directly invoked llama-quantize / llama-server smoke runs see no secrets. The installer still reads tokens from os.environ for the releases API and probe fetch. * Trim verbose comments around the secret-env scrubber for PR #6696 Comment-only: condense the block comments added across this PR. Logic unchanged (comment_tools.py check confirms code-only signature equal). * Redirect HOME / cache pointers to an empty dir for prebuilt binaries (PR #6696) Address Codex P2: stripping token env vars still let a tampered binary read on-disk token stores (~/.cache/huggingface/token, ~/.aws/credentials, ~/.config/gh) through $HOME and the cache/config pointers. Point HOME plus the HF / XDG / Windows home pointers at a single empty throwaway dir for the downloaded-binary env. Defense in depth: a binary resolving the real home via getpwuid is out of scope and needs OS sandboxing. * Close residual credential-probe gaps for PR #6696 Address the latest Codex review: - Strip token-only URL userinfo too (scheme://ghp_token@host), not just the user:pass form. - Redirect HOMEDRIVE/HOMEPATH alongside USERPROFILE so a Windows binary cannot reconstruct the real profile from %HOMEDRIVE%%HOMEPATH%. - Drop explicit credential-file pointers (NETRC, PIP_CONFIG_FILE, DOCKER_CONFIG, GIT_CONFIG_GLOBAL) that live outside HOME. - Probe ldd with a secret-free env: linux_runtime_dirs ran ldd on the untrusted prebuilt with the inherited os.environ, and ldd may execute the binary, so it could observe HF_TOKEN/GITHUB_TOKEN during the probe. Factored the shared scrub into secret_free_environ(). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Separate token-bearing install from binary smoke; drop CI command files (PR #6696) Address the two P1s in the latest review: - mlx-ci: GitHub bakes secrets into the run-script text, so inline token assignments in a step that later runs the prebuilt let a tampered binary read them from the script. Split into a token-bearing install + download step that never launches a binary, and a secret-free smoke step that runs llama-quantize / llama-server. - secret_free_environ now drops the GitHub Actions command files (GITHUB_ENV, GITHUB_PATH, GITHUB_OUTPUT, GITHUB_STEP_SUMMARY, BASH_ENV) and the smoke step unsets them, so a tampered prebuilt cannot inject PATH/env into the later token-bearing MLX steps. * Run the prebuilt smoke last, after all token-bearing steps (PR #6696) Address the P1 workspace-poisoning vector: even with no secrets in its env, a tampered prebuilt could edit the checkout or installed modules, and the later HF_TOKEN MLX steps would then execute that poisoned code on push builds. Move the prebuilt install + smoke to the end of the job so the untrusted binary runs after every token-bearing step, leaving nothing for it to corrupt. The MLX GGUF reload uses a source-built llama-cli, not this prebuilt, so nothing depends on the earlier position. * Trim comments around the secret-env scrubber and prebuilt CI steps (PR #6696) Comment-only: condense the security-rationale block comments and merge the duplicated prebuilt-step description in mlx-ci. Logic unchanged (comment_tools.py check confirms the code-only signature is equal; install suite still passes). * Authenticate the GGUF export release-API lookup with the read-only GITHUB_TOKEN (PR #6696) * Rename env scrubber off the secret-named identifier CodeQL flags as a clear-text sink (PR #6696) --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/mlx-ci.yml | 181 ++++++------- studio/install_llama_prebuilt.py | 140 +++++++++- .../test_install_llama_prebuilt_logic.py | 256 ++++++++++++++++++ 3 files changed, 482 insertions(+), 95 deletions(-) diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index 864630f9f0..424a706d7c 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -231,99 +231,6 @@ jobs: tests/studio/test_is_mlx_dispatch_gate.py \ tests/studio/test_mlx_training_worker_behaviors.py - # Studio prebuilt llama.cpp install + GGUF inference. Mirrors the - # path Studio's setup.sh takes on macOS since #5963: plan against - # the unslothai/llama.cpp fork's latest release, which ships the - # bin-macos-arm64 bundle plus the llama-prebuilt-manifest.json the - # default policy reads. After install, downloads a small published - # GGUF (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) and validates - # llama-server /completion end to end. An install failure or a - # non-zero binary exit is an Unsloth/Studio bug. - - name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1) - env: - # Withheld on PR: this step runs checked-out PR code; public GGUF still downloads. - HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} - # install_llama_prebuilt.py hits the GitHub releases API to - # resolve the asset URL. Anonymous calls share the runner-IP - # rate-limit bucket and 403 quickly -- pass the workflow's - # automatic GITHUB_TOKEN to bump us to the 5000/hr authenticated - # bucket. - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - set -euo pipefail - INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" - rm -rf "$INSTALL_DIR" - # Mirror studio/setup.sh on macOS (the install.sh user path): - # it plans against the unslothai/llama.cpp fork's latest - # release with no policy or tag flags. - python studio/install_llama_prebuilt.py \ - --install-dir "$INSTALL_DIR" \ - --published-repo unslothai/llama.cpp - - # Studio bundles only llama-server + llama-quantize from the - # prebuilt (not llama-cli) -- inference goes through - # llama-server's HTTP /completion endpoint. Validate both: - # llama-quantize --help proves the dynamic libs link, then - # spin up llama-server and POST a /completion request on a - # tiny published GGUF. - LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server" - LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize" - [ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; } - [ -x "$LLAMA_QUANT" ] || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; } - echo "llama-server : $LLAMA_SERVER" - echo "llama-quantize: $LLAMA_QUANT" - "$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK" - - mkdir -p /tmp/ggufs - bash .github/scripts/hf-download-with-retry.sh \ - 'unsloth/gemma-3-270m-it-GGUF' \ - 'gemma-3-270m-it-Q4_K_M.gguf' \ - /tmp/ggufs - - PORT=18080 - echo "=== starting llama-server on 127.0.0.1:$PORT ===" - "$LLAMA_SERVER" \ - -m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \ - --host 127.0.0.1 \ - --port "$PORT" \ - -c 256 \ - -n 16 \ - --no-warmup \ - > /tmp/llama-server.log 2>&1 & - SERVER_PID=$! - trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT - - # Wait for /health to come up - for i in $(seq 1 30); do - if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then - echo " server up after ${i}s" - break - fi - sleep 1 - done - if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then - echo "::error::llama-server never became healthy" - tail -40 /tmp/llama-server.log - exit 1 - fi - - PROMPT="Hello, my name is" - echo "=== POST /completion ===" - RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \ - -H 'Content-Type: application/json' \ - -d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}") - echo "raw response (head): $(echo "$RESP" | head -c 600)" - CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))") - echo "completion content: $CONTENT" - - if [ -z "$CONTENT" ]; then - echo "::error::llama-server /completion returned empty content" - tail -40 /tmp/llama-server.log - exit 1 - fi - echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works" - # Real MLX training + inference smoke test. Trains # unsloth/gemma-3-270m-it for 7 deterministic LoRA steps # (batch_size=2, gradient_accumulation_steps=3) on a single @@ -338,6 +245,9 @@ jobs: UNSLOTH_COMPILE_DISABLE: '1' run: | mkdir -p mlx_workdir + # Authenticate llama.cpp's release-API lookup (anonymous 403s on rate-limit); + # read-only GITHUB_TOKEN scoped here only, never to steps that run binaries. + GH_TOKEN="${{ secrets.GITHUB_TOKEN }}" GITHUB_TOKEN="${{ secrets.GITHUB_TOKEN }}" \ python tests/studio/run_real_mlx_smoke.py train \ --workdir "$PWD/mlx_workdir" @@ -406,3 +316,88 @@ jobs: cat "$f" 2>/dev/null || echo "(missing)" echo done + + # Validates the macOS prebuilt path Studio's setup.sh uses (#5963): install the + # unslothai/llama.cpp fork's latest release, download a small public GGUF, and + # check llama-server /completion end to end. Split and placed last so the + # untrusted binary runs only in the final smoke step, after every HF_TOKEN step, + # leaving no token-bearing step or shared workspace for a tampered prebuilt to + # corrupt. GH_TOKEN: releases API; HF_TOKEN (withheld on PR): probe + GGUF fetch. + - name: Studio prebuilt llama.cpp install + GGUF download (Mac M1) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }} + run: | + set -euo pipefail + INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" + rm -rf "$INSTALL_DIR" + # Download only -- no llama-quantize / llama-server launch in this step. + python studio/install_llama_prebuilt.py \ + --install-dir "$INSTALL_DIR" \ + --published-repo unslothai/llama.cpp + mkdir -p /tmp/ggufs + bash .github/scripts/hf-download-with-retry.sh \ + 'unsloth/gemma-3-270m-it-GGUF' \ + 'gemma-3-270m-it-Q4_K_M.gguf' \ + /tmp/ggufs + + # Final step: runs the downloaded binaries with no secrets present, and clears + # the GitHub Actions command files so a tampered prebuilt cannot influence the job. + - name: Studio prebuilt llama.cpp GGUF inference smoke (Mac M1) + run: | + set -euo pipefail + unset GITHUB_ENV GITHUB_PATH GITHUB_OUTPUT GITHUB_STEP_SUMMARY + INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp" + # Studio bundles only llama-server + llama-quantize (not llama-cli); + # inference goes through llama-server's HTTP /completion endpoint. + LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server" + LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize" + [ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; } + [ -x "$LLAMA_QUANT" ] || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; } + echo "llama-server : $LLAMA_SERVER" + echo "llama-quantize: $LLAMA_QUANT" + "$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK" + + PORT=18080 + echo "=== starting llama-server on 127.0.0.1:$PORT ===" + "$LLAMA_SERVER" \ + -m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \ + --host 127.0.0.1 \ + --port "$PORT" \ + -c 256 \ + -n 16 \ + --no-warmup \ + > /tmp/llama-server.log 2>&1 & + SERVER_PID=$! + trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT + + # Wait for /health to come up + for i in $(seq 1 30); do + if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then + echo " server up after ${i}s" + break + fi + sleep 1 + done + if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then + echo "::error::llama-server never became healthy" + tail -40 /tmp/llama-server.log + exit 1 + fi + + PROMPT="Hello, my name is" + echo "=== POST /completion ===" + RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \ + -H 'Content-Type: application/json' \ + -d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}") + echo "raw response (head): $(echo "$RESP" | head -c 600)" + CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))") + echo "completion content: $CONTENT" + + if [ -z "$CONTENT" ]; then + echo "::error::llama-server /completion returned empty content" + tail -40 /tmp/llama-server.log + exit 1 + fi + echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works" diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index c7f34e39f2..e40cb3083e 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -7,6 +7,7 @@ from __future__ import annotations import argparse +import atexit import errno import fnmatch import hashlib @@ -5203,7 +5204,8 @@ def ldconfig_runtime_dirs(required_libraries: Iterable[str]) -> list[str]: def linux_runtime_dirs(binary_path: Path) -> list[str]: - missing = linux_missing_libraries(binary_path) + # ldd may execute the binary, so probe it with a secret-free env. + missing = linux_missing_libraries(binary_path, env = scrubbed_environ()) if not missing: return [] return linux_runtime_dirs_for_required_libraries(missing) @@ -5499,6 +5501,140 @@ def _wsl_system_rocm_lib_dirs() -> list[str]: return out +# Secrets a downloaded llama.cpp binary never needs; keep them out of binary_env(). +# The installer's own API calls read os.environ directly, so auth is unaffected. +_SECRET_ENV_EXACT_NAMES = frozenset( + { + "HF_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "GH_TOKEN", + "GITHUB_TOKEN", + "WANDB_API_KEY", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "GOOGLE_APPLICATION_CREDENTIALS", + "AZURE_CLIENT_SECRET", + "ACTIONS_ID_TOKEN_REQUEST_TOKEN", + "ACTIONS_ID_TOKEN_REQUEST_URL", + "ACTIONS_RUNTIME_TOKEN", + # Credential pointers (cluster / remote-host access). + "KUBECONFIG", + "SSH_AUTH_SOCK", + } +) +# Case-insensitive substring markers for names we do not enumerate (no bare "KEY", +# which would hit benign runtime vars). +_SECRET_ENV_MARKERS = ( + "TOKEN", + "SECRET", + "PASSWORD", + "PASSWD", + "PASSPHRASE", + "CREDENTIAL", + "PRIVATE_KEY", + "API_KEY", +) +# Proxy / index URLs embed creds in their value; the offline binaries never need them. +_SECRET_ENV_URL_NAMES = frozenset( + { + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "FTP_PROXY", + "RSYNC_PROXY", + "PIP_INDEX_URL", + "PIP_EXTRA_INDEX_URL", + "UV_INDEX_URL", + "UV_DEFAULT_INDEX", + "UV_EXTRA_INDEX_URL", + } +) +# Also drop values with URL userinfo creds (scheme://user:secret@host or token@host). +_URL_USERINFO_CREDENTIAL_RE = re.compile(r"://[^/@\s]+@") + + +def is_secret_env_name(name: str) -> bool: + upper = name.upper() + return ( + upper in _SECRET_ENV_EXACT_NAMES + or upper in _SECRET_ENV_URL_NAMES + or any(marker in upper for marker in _SECRET_ENV_MARKERS) + ) + + +def scrub_env(env: dict[str, str]) -> dict[str, str]: + """Drop secret-bearing variables before handing an env to a downloaded binary.""" + return { + key: value + for key, value in env.items() + if not is_secret_env_name(key) and not _URL_USERINFO_CREDENTIAL_RE.search(value or "") + } + + +# Home / cache pointers to on-disk token stores (~/.cache/huggingface/token, +# ~/.aws/credentials, ...). Stripping env tokens is not enough; point these at an +# empty home so the binary cannot read those files via $HOME. +_RUNTIME_HOME_POINTER_VARS = ( + "HOME", + "USERPROFILE", + "APPDATA", + "LOCALAPPDATA", + "XDG_CACHE_HOME", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + "HF_HOME", + "HUGGINGFACE_HUB_CACHE", + "HF_HUB_CACHE", +) +# Credential / config file pointers outside HOME; drop so lookups fall back to the +# empty home. +_CREDENTIAL_FILE_POINTER_VARS = ( + "NETRC", + "PIP_CONFIG_FILE", + "DOCKER_CONFIG", + "GIT_CONFIG_GLOBAL", +) +# GitHub Actions command files: appending to these injects PATH/env into later steps. +_CI_COMMAND_FILE_VARS = ( + "GITHUB_ENV", + "GITHUB_PATH", + "GITHUB_OUTPUT", + "GITHUB_STEP_SUMMARY", + "BASH_ENV", +) + +_isolated_runtime_home_dir: str | None = None + + +def isolated_runtime_home() -> str: + # Empty dir, created lazily and removed at exit. (A binary resolving the real + # home via getpwuid is out of scope; that needs OS sandboxing.) + global _isolated_runtime_home_dir + if _isolated_runtime_home_dir is None: + path = tempfile.mkdtemp(prefix = "unsloth-prebuilt-home-") + atexit.register(shutil.rmtree, path, ignore_errors = True) + _isolated_runtime_home_dir = path + return _isolated_runtime_home_dir + + +def scrubbed_environ() -> dict[str, str]: + # os.environ minus secrets, with home / credential pointers neutralised. Used for + # the binary env and any probe (e.g. ldd) that runs the untrusted binary. + env = scrub_env(os.environ.copy()) + runtime_home = isolated_runtime_home() + for pointer in _RUNTIME_HOME_POINTER_VARS: + env[pointer] = runtime_home + # Windows rebuilds the profile from %HOMEDRIVE%%HOMEPATH% (no-op pair on POSIX). + drive, tail = os.path.splitdrive(runtime_home) + env["HOMEDRIVE"], env["HOMEPATH"] = drive, tail or runtime_home + for pointer in (*_CREDENTIAL_FILE_POINTER_VARS, *_CI_COMMAND_FILE_VARS): + env.pop(pointer, None) + return env + + def binary_env( binary_path: Path, install_dir: Path, @@ -5506,7 +5642,7 @@ def binary_env( *, runtime_line: str | None = None, ) -> dict[str, str]: - env = os.environ.copy() + env = scrubbed_environ() if host.is_windows: path_dirs = [ str(binary_path.parent), diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 9ee8759bb4..c852d7c495 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -22,6 +22,9 @@ SPEC.loader.exec_module(INSTALL_LLAMA_PREBUILT) PrebuiltFallback = INSTALL_LLAMA_PREBUILT.PrebuiltFallback extract_archive = INSTALL_LLAMA_PREBUILT.extract_archive binary_env = INSTALL_LLAMA_PREBUILT.binary_env +is_secret_env_name = INSTALL_LLAMA_PREBUILT.is_secret_env_name +scrub_env = INSTALL_LLAMA_PREBUILT.scrub_env +isolated_runtime_home = INSTALL_LLAMA_PREBUILT.isolated_runtime_home HostInfo = INSTALL_LLAMA_PREBUILT.HostInfo AssetChoice = INSTALL_LLAMA_PREBUILT.AssetChoice ApprovedArtifactHash = INSTALL_LLAMA_PREBUILT.ApprovedArtifactHash @@ -779,6 +782,259 @@ def test_binary_env_linux_includes_binary_parent_in_ld_library_path( assert str(install_dir) in ld_dirs +def test_scrub_env_drops_secrets_and_keeps_runtime_vars(): + raw = { + # secrets + "HF_TOKEN": "hf_x", + "HUGGING_FACE_HUB_TOKEN": "hf_y", + "GH_TOKEN": "gh_x", + "GITHUB_TOKEN": "gh_y", + "WANDB_API_KEY": "wandb_x", + "AWS_SECRET_ACCESS_KEY": "aws_x", + "ACTIONS_ID_TOKEN_REQUEST_TOKEN": "oidc_x", + "ACTIONS_ID_TOKEN_REQUEST_URL": "https://oidc", + "SOME_VENDOR_API_KEY": "vendor_x", + "DB_PASSWORD": "pw", + "MY_PRIVATE_KEY": "pk", + "KUBECONFIG": "/home/runner/.kube/config", + "SSH_AUTH_SOCK": "/tmp/ssh-agent.sock", + "SSH_PASSPHRASE": "ssh_pass", + # runtime vars to keep + "PATH": "/usr/bin", + "LD_LIBRARY_PATH": "/opt/lib", + "DYLD_LIBRARY_PATH": "/opt/dyld", + "HOME": "/home/runner", + "TMPDIR": "/tmp", + "CUDA_VISIBLE_DEVICES": "0", + "HSA_OVERRIDE_GFX_VERSION": "11.0.0", + } + + cleaned = scrub_env(raw) + + for secret in ( + "HF_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "GH_TOKEN", + "GITHUB_TOKEN", + "WANDB_API_KEY", + "AWS_SECRET_ACCESS_KEY", + "ACTIONS_ID_TOKEN_REQUEST_TOKEN", + "ACTIONS_ID_TOKEN_REQUEST_URL", + "SOME_VENDOR_API_KEY", + "DB_PASSWORD", + "MY_PRIVATE_KEY", + "KUBECONFIG", + "SSH_AUTH_SOCK", + "SSH_PASSPHRASE", + ): + assert secret not in cleaned, f"{secret} must be stripped from binary env" + + for keep in ( + "PATH", + "LD_LIBRARY_PATH", + "DYLD_LIBRARY_PATH", + "HOME", + "TMPDIR", + "CUDA_VISIBLE_DEVICES", + "HSA_OVERRIDE_GFX_VERSION", + ): + assert cleaned[keep] == raw[keep], f"{keep} must be preserved for the binary" + + # no bare "KEY" marker: benign KEY-containing names survive + assert is_secret_env_name("API_KEY") is True + assert is_secret_env_name("SSH_KEYFILE_PATH") is False + assert is_secret_env_name("PATH") is False + + +def test_scrub_env_drops_proxy_index_and_embedded_url_credentials(): + raw = { + # proxy / package-index URLs whose values commonly embed credentials + "HTTPS_PROXY": "https://user:secret@proxy:8080", + "https_proxy": "https://user:secret@proxy:8080", # lower-case variant + "ALL_PROXY": "socks5://user:secret@proxy:1080", + "PIP_INDEX_URL": "https://u:p@pypi.internal/simple", + "UV_INDEX_URL": "https://u:p@index.internal/simple", + # credentials embedded in an otherwise benign-named variable's value + "MY_DB_DSN": "postgres://admin:secret@db:5432/app", + # benign vars the binary needs, including a URL with no userinfo + "PATH": "/usr/bin", + "CUDA_VISIBLE_DEVICES": "0", + "NO_PROXY": "localhost,127.0.0.1", + "SOME_ENDPOINT": "https://example.com:8080/v1", + } + + cleaned = scrub_env(raw) + + for secret in ( + "HTTPS_PROXY", + "https_proxy", + "ALL_PROXY", + "PIP_INDEX_URL", + "UV_INDEX_URL", + "MY_DB_DSN", + ): + assert secret not in cleaned, f"{secret} must be stripped from binary env" + for keep in ("PATH", "CUDA_VISIBLE_DEVICES", "NO_PROXY", "SOME_ENDPOINT"): + assert cleaned[keep] == raw[keep], f"{keep} must be preserved for the binary" + + assert is_secret_env_name("HTTPS_PROXY") is True + assert is_secret_env_name("https_proxy") is True + assert is_secret_env_name("NO_PROXY") is False + + +def test_binary_env_strips_secrets_from_downloaded_binary_environment( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "llama.cpp" + bin_dir = install_dir / "build" / "bin" + bin_dir.mkdir(parents = True) + binary_path = bin_dir / "llama-server" + binary_path.write_bytes(b"fake") + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_runtime_dirs", lambda _bp: []) + + monkeypatch.setenv("HF_TOKEN", "hf_secret_from_ci") + monkeypatch.setenv("GITHUB_TOKEN", "gh_secret_from_ci") + monkeypatch.setenv("GH_TOKEN", "gh_secret_from_ci") + monkeypatch.setenv("WANDB_API_KEY", "wandb_secret_from_ci") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") + + env = binary_env(binary_path, install_dir, host) + + assert "HF_TOKEN" not in env + assert "GITHUB_TOKEN" not in env + assert "GH_TOKEN" not in env + assert "WANDB_API_KEY" not in env + # library/runtime resolution unaffected + assert str(bin_dir) in env["LD_LIBRARY_PATH"].split(os.pathsep) + assert env["CUDA_VISIBLE_DEVICES"] == "1" + + +def test_binary_env_redirects_home_away_from_real_credential_stores( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + install_dir = tmp_path / "llama.cpp" + bin_dir = install_dir / "build" / "bin" + bin_dir.mkdir(parents = True) + binary_path = bin_dir / "llama-server" + binary_path.write_bytes(b"fake") + + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_runtime_dirs", lambda _bp: []) + + real_home = str(tmp_path / "real_home") + monkeypatch.setenv("HOME", real_home) + monkeypatch.setenv("HF_HOME", real_home + "/.cache/huggingface") + + env = binary_env(binary_path, install_dir, host) + + # HOME and the cache pointers are redirected to a single empty, existing dir. + assert env["HOME"] != real_home + assert env["HF_HOME"] == env["HOME"] + assert env["HOME"] == isolated_runtime_home() + assert os.path.isdir(env["HOME"]) + assert os.listdir(env["HOME"]) == [] + # Windows reconstructs the profile from HOMEDRIVE + HOMEPATH. + assert env["HOMEDRIVE"] + env["HOMEPATH"] == env["HOME"] + + +def test_scrub_env_drops_token_only_url_userinfo(): + raw = { + "GENERIC_REPO": "https://ghp_tokenonly@github.com/org/repo", + "GENERIC_OK": "https://example.com:8080/v1", + } + cleaned = scrub_env(raw) + assert "GENERIC_REPO" not in cleaned + assert cleaned["GENERIC_OK"] == raw["GENERIC_OK"] + + +def test_binary_env_drops_explicit_credential_file_pointers( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + host = HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + ) + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_runtime_dirs", lambda _bp: []) + dropped = ( + "NETRC", + "PIP_CONFIG_FILE", + "DOCKER_CONFIG", + "GIT_CONFIG_GLOBAL", + "GITHUB_ENV", + "GITHUB_PATH", + "GITHUB_OUTPUT", + "GITHUB_STEP_SUMMARY", + "BASH_ENV", + ) + for var in dropped: + monkeypatch.setenv(var, "/home/realuser/secret") + + env = binary_env(tmp_path / "llama-server", tmp_path, host) + + for var in dropped: + assert var not in env + + +def test_linux_runtime_dirs_probes_with_secret_free_env(monkeypatch: pytest.MonkeyPatch): + captured: dict[str, object] = {} + + def fake_missing(binary_path, *, env = None): + captured["env"] = env + return [] + + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_missing_libraries", fake_missing) + monkeypatch.setenv("HF_TOKEN", "hf_secret") + monkeypatch.setenv("GITHUB_TOKEN", "gh_secret") + + INSTALL_LLAMA_PREBUILT.linux_runtime_dirs(Path("/fake/llama-server")) + + probe_env = captured["env"] + assert probe_env is not None + assert "HF_TOKEN" not in probe_env + assert "GITHUB_TOKEN" not in probe_env + + def test_install_prebuilt_falls_back_to_older_release_plan( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): From 0ad814a45228999d95857f8e38a484f9ce107c92 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 27 Jun 2026 17:48:16 -0700 Subject: [PATCH 17/49] =?UTF-8?q?Revert=20"feat:=20add=20GPU-aware=20model?= =?UTF-8?q?=20filtering=20and=20For=20You=20section-=20Add=20fit=20filt?= =?UTF-8?q?=E2=80=A6"=20(#6722)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit a636693019f33c1acf9477dbd4ad714792683c74. --- .../model-selector/recommended-fit.ts | 39 +++------ .../src/features/hub/catalog/model-card.tsx | 26 ++---- .../src/features/hub/catalog/models-table.tsx | 46 ---------- .../features/hub/catalog/models-toolbar.tsx | 24 ------ studio/frontend/src/features/hub/hub-page.tsx | 85 +++---------------- .../src/features/hub/lib/gpu-fit-filter.ts | 80 ----------------- .../src/features/hub/lib/view-models.ts | 10 --- studio/frontend/src/features/hub/types.ts | 4 - 8 files changed, 27 insertions(+), 287 deletions(-) delete mode 100644 studio/frontend/src/features/hub/lib/gpu-fit-filter.ts diff --git a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts index cc96665df4..24f0edc784 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts @@ -64,38 +64,19 @@ export function matchesFormatFilter( } } -// Model-size extraction from repo id, matching the backend's 3-regex priority: -// active params (MoE "A3B") > effective params (Gemma "E4B") > total ("8B"). -// Bounded by separators so we never read "16" from "bf16" or "2" from "Kimi-K2". -// Examples: "Qwen3.5-35B-A3B" -> 3, "gemma-4-E4B" -> 4, "Llama-3-8B" -> 8. -const ACTIVE_PARAM_RE = /(?:^|[-_/. ])a(\d+(?:\.\d+)?)\s*[bB](?=$|[-_/. ])/i; -const EFFECTIVE_PARAM_RE = /(?:^|[-_/. ])e(\d+(?:\.\d+)?)\s*[bB](?=$|[-_/. ])/i; -const TOTAL_PARAM_RE = /(?:^|[-_/. ])(\d+(?:\.\d+)?)\s*[bB](?=$|[-_/. ])/; - -function paramsFromMatch(match: RegExpExecArray | null): number | undefined { - if (!match) return undefined; - const billions = parseFloat(match[1]); - return Number.isFinite(billions) && billions > 0 - ? billions * 1e9 - : undefined; -} - -/** Active/effective parameter count parsed from a repo id, if it uses explicit - * MoE/Gemma-style notation such as A3B or E4B. */ -export function activeOrEffectiveParamsFromId(id: string): number | undefined { - return ( - paramsFromMatch(ACTIVE_PARAM_RE.exec(id)) ?? - paramsFromMatch(EFFECTIVE_PARAM_RE.exec(id)) - ); -} +// First "B" token in a repo id, e.g. "Qwen3-4B-GGUF" -> 4, "gpt-oss-20b" -> +// 20, "Qwen3-30B-A3B" -> 30 (MoE total), "gemma-4-E4B" -> 4 (effective-param +// "E" series). The digits must be bounded by a separator so we never read "16" +// from "bf16" or the "2" in "Kimi-K2". +const PARAM_RE = /(?:^|[-_/. ])[eE]?(\d+(?:\.\d+)?)\s*[bB](?=$|[-_./ ])/; /** Parameter count (absolute, e.g. 4e9) parsed from a repo id, or undefined - * when the id has no size token (so callers can treat the size as unknown). - * Prefers MoE active-param notation (A3B) over effective (E4B) over total. */ + * when the id has no size token (so callers can treat the size as unknown). */ export function paramsFromId(id: string): number | undefined { - return ( - activeOrEffectiveParamsFromId(id) ?? paramsFromMatch(TOTAL_PARAM_RE.exec(id)) - ); + const match = PARAM_RE.exec(id); + if (!match) return undefined; + const billions = parseFloat(match[1]); + return Number.isFinite(billions) && billions > 0 ? billions * 1e9 : undefined; } // Smallest practical GGUF/MLX quant (~Q2_K, low-bit). The fit check asks whether diff --git a/studio/frontend/src/features/hub/catalog/model-card.tsx b/studio/frontend/src/features/hub/catalog/model-card.tsx index c3c526cac3..358a3409b5 100644 --- a/studio/frontend/src/features/hub/catalog/model-card.tsx +++ b/studio/frontend/src/features/hub/catalog/model-card.tsx @@ -11,7 +11,7 @@ import { ownerPaletteColor } from "@/features/hub/lib/avatar-theme"; import { buildAdaptiveCardAccentStyle } from "@/features/hub/lib/card-accent"; import { useDominantColor } from "@/features/hub/lib/use-dominant-color"; import { formatModelParamLabel } from "@/features/hub/lib/view-models"; -import { cn, formatCompact } from "@/lib/utils"; +import { formatCompact } from "@/lib/utils"; import { Download01Icon, FavouriteIcon } from "@hugeicons/core-free-icons"; import { type CSSProperties, memo, useMemo } from "react"; import type { DiscoverRow } from "../types"; @@ -339,25 +339,11 @@ export const ModelCard = memo(function ModelCard({ value={formatCompact(row.result.likes)} /> -
- {row.fitLevel && ( - - {row.fitLevel === "comfortable" ? "Comfortable" : row.fitLevel === "fits" ? "Fits GPU" : "OOM"} - - )} - {hasSize ? ( - {sizeLabel} - ) : topCapability ? ( - - ) : null} -
+ {hasSize ? ( + {sizeLabel} + ) : topCapability ? ( + + ) : null} ); diff --git a/studio/frontend/src/features/hub/catalog/models-table.tsx b/studio/frontend/src/features/hub/catalog/models-table.tsx index 35f2f966e5..a62806dcb6 100644 --- a/studio/frontend/src/features/hub/catalog/models-table.tsx +++ b/studio/frontend/src/features/hub/catalog/models-table.tsx @@ -570,22 +570,6 @@ export const ResultCard = memo(function ResultCard({ node: {sizeLabel}, }); } - if (row.fitLevel) { - const label = row.fitLevel === "comfortable" ? "Comfortable" : row.fitLevel === "fits" ? "Fits GPU" : "OOM"; - const toneClass = row.fitLevel === "comfortable" - ? "text-emerald-700 bg-emerald-50 dark:text-emerald-300 dark:bg-emerald-500/15" - : row.fitLevel === "fits" - ? "text-amber-700 bg-amber-50 dark:text-amber-300 dark:bg-amber-500/15" - : "text-red-700 bg-red-50 dark:text-red-300 dark:bg-red-500/15"; - textParts.push({ - key: "gpuFit", - node: ( - - {label} - - ), - }); - } if (row.result.updatedAt) { textParts.push({ key: "updated", @@ -745,21 +729,6 @@ export const ResultGridRow = memo(function ResultGridRow({ - {row.fitLevel && ( - <> - - - {row.fitLevel === "comfortable" ? "Comfortable" : row.fitLevel === "fits" ? "Fits GPU" : "OOM"} - - - )} @@ -882,21 +851,6 @@ export const ResultSplitRow = memo(function ResultSplitRow({ - {row.fitLevel && ( - <> - - - {row.fitLevel === "comfortable" ? "Comfortable" : row.fitLevel === "fits" ? "Fits GPU" : "OOM"} - - - )}
diff --git a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx index 9cb17736af..48f7fcffaa 100644 --- a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx +++ b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx @@ -21,7 +21,6 @@ import { HugeiconsIcon } from "@hugeicons/react"; import type { HfSortKey } from "@/features/hub/hooks/use-hub-model-search"; import type { CapabilityFilter, - GpuFitFilter, ModelFormatFilter, ModelsTab, ResourceTypeFilter, @@ -29,7 +28,6 @@ import type { import { CAPABILITY_FILTER_OPTIONS, FORMAT_FILTER_OPTIONS, - GPU_FIT_FILTER_OPTIONS, } from "../lib/view-models"; import { HubOptionMenu, type HubOption } from "./hub-option-menu"; import { @@ -70,8 +68,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange, capabilityFilter, onCapabilityFilterChange, - gpuFitFilter, - onGpuFitFilterChange, onManageLocalFolders, onOpenFineTune, }: { @@ -88,8 +84,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange: (value: ModelFormatFilter) => void; capabilityFilter: CapabilityFilter; onCapabilityFilterChange: (value: CapabilityFilter) => void; - gpuFitFilter: GpuFitFilter; - onGpuFitFilterChange: (value: GpuFitFilter) => void; onManageLocalFolders: () => void; /** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a * format-dropdown option rather than a standalone feed section. */ @@ -159,14 +153,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({ })), [], ); - const gpuFitOptions = useMemo[]>( - () => - GPU_FIT_FILTER_OPTIONS.map((option) => ({ - value: option.value, - label: option.label, - })), - [], - ); const sortOptions = useMemo[]>( () => SORT_OPTIONS.map((option) => ({ @@ -357,16 +343,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({ /> )} - {tab === "discover" && !isDataset && ( - - )} - {tab === "discover" && ( ("all"); - const [gpuFitFilter, setGpuFitFilter] = useState("all"); const [allModelsView, setAllModelsViewState] = useState( readAllModelsViewPreference, ); @@ -567,7 +561,6 @@ export function ModelsPage() { const apiHfToken = hfApiToken(debouncedHfToken); const deferredFormatFilter = useDeferredValue(formatFilter); const deferredCapabilityFilter = useDeferredValue(capabilityFilter); - const deferredGpuFitFilter = useDeferredValue(gpuFitFilter); const hasQuery = deferredDebouncedQuery.trim() !== ""; const mode: DiscoverMode = !isModelDiscover @@ -697,59 +690,20 @@ export function ModelsPage() { const discoverRows = isDatasetMode ? datasetDiscoverRows : modelDiscoverRows; - // Pre-compute GPU fit level for every discover row so filteredDiscoverRows - // and model cards can both consume the same classification. - const gpuFitLevelById = useMemo(() => { - const map = new Map>(); - for (const row of discoverRows) { - map.set( - row.id, - classifyGpuFit({ - totalParams: row.result.totalParams, - estimatedSizeBytes: row.result.estimatedSizeBytes, - repoId: row.id, - gpu, - }), - ); - } - return map; - }, [discoverRows, gpu]); - - const addGpuFitLevel = useCallback( - (row: DiscoverRow): DiscoverRow => ({ - ...row, - fitLevel: classifyGpuFit({ - totalParams: row.result.totalParams, - estimatedSizeBytes: row.result.estimatedSizeBytes, - repoId: row.id, - gpu, - }), - }), - [gpu], - ); - const filteredDiscoverRows = useMemo(() => { if (isDatasetMode) return discoverRows; - return discoverRows - .filter( - (row) => - !isHiddenModelId(row.id) && - matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) && - matchesCapability(row.capabilities, deferredCapabilityFilter) && - matchesGpuFitFilter(gpuFitLevelById.get(row.id) ?? null, deferredGpuFitFilter) && - (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)), - ) - .map((row) => ({ - ...row, - fitLevel: gpuFitLevelById.get(row.id) ?? null, - })); + return discoverRows.filter( + (row) => + !isHiddenModelId(row.id) && + matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) && + matchesCapability(row.capabilities, deferredCapabilityFilter) && + (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)), + ); }, [ discoverRows, isDatasetMode, effectiveDiscoverFormat, deferredCapabilityFilter, - deferredGpuFitFilter, - gpuFitLevelById, activeChannel, ]); @@ -770,17 +724,8 @@ export function ModelsPage() { effectiveLocalRows, ) .filter((row) => !isHiddenModelId(row.id)) - .filter((row) => matchesFormat(row.result.isGguf, "gguf")) - .map(addGpuFitLevel) - .filter((row) => - matchesGpuFitFilter(row.fitLevel ?? null, deferredGpuFitFilter), - ), - [ - hubFeed.trending.results, - modelDiscoveryInventorySignature, - addGpuFitLevel, - deferredGpuFitFilter, - ], + .filter((row) => matchesFormat(row.result.isGguf, "gguf")), + [hubFeed.trending.results, modelDiscoveryInventorySignature], ); const feedRows = useMemo(() => { if (!isFeedMode) return []; @@ -897,7 +842,6 @@ export function ModelsPage() { resourceType, deferredFormatFilter, deferredCapabilityFilter, - deferredGpuFitFilter, effectiveSort, effectiveDirection, activeChannelId, @@ -908,7 +852,6 @@ export function ModelsPage() { resourceType, deferredFormatFilter, deferredCapabilityFilter, - deferredGpuFitFilter, effectiveSort, effectiveDirection, activeChannelId, @@ -929,7 +872,6 @@ export function ModelsPage() { setDownloadedFormat("all"); } setCapabilityFilter("all"); - setGpuFitFilter("all"); }, [isDiscoverTab, urlSection, navigate]); const handleDiscoverFetchIntent = useCallback(() => { setDiscoverFetchIntent((value) => value + 1); @@ -1295,10 +1237,8 @@ export function ModelsPage() { hasMore, manualFetchAvailable: discoverManualFetchAvailable, hasActiveFilters: - deferredGpuFitFilter !== "all" || - (!isFeedMode && - (deferredFormatFilter !== "all" || - deferredCapabilityFilter !== "all")), + !isFeedMode && + (deferredFormatFilter !== "all" || deferredCapabilityFilter !== "all"), }), [ tab, @@ -1324,7 +1264,6 @@ export function ModelsPage() { discoverManualFetchAvailable, deferredFormatFilter, deferredCapabilityFilter, - deferredGpuFitFilter, ], ); @@ -1509,8 +1448,6 @@ export function ModelsPage() { onFormatFilterChange={setFormatFilter} capabilityFilter={capabilityFilter} onCapabilityFilterChange={setCapabilityFilter} - gpuFitFilter={gpuFitFilter} - onGpuFitFilterChange={setGpuFitFilter} onManageLocalFolders={handleManageLocalFolders} onOpenFineTune={() => handleOpenList("finetune")} /> diff --git a/studio/frontend/src/features/hub/lib/gpu-fit-filter.ts b/studio/frontend/src/features/hub/lib/gpu-fit-filter.ts deleted file mode 100644 index 7eca4c315a..0000000000 --- a/studio/frontend/src/features/hub/lib/gpu-fit-filter.ts +++ /dev/null @@ -1,80 +0,0 @@ -// SPDX-License-Identifier: AGPL-3.0-only -// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -// GPU-aware model-fit filtering: classifies whether a model fits the device -// and provides a filter predicate for the Hub page and model selector. - -import type { GpuInfo } from "@/hooks/use-gpu-info"; -import { - activeOrEffectiveParamsFromId, - estimateQuantBytes, - paramsFromId, -} from "@/components/assistant-ui/model-selector/recommended-fit"; - -/** The three filter states exposed in the toolbar dropdown. */ -export type GpuFitFilter = "all" | "fits" | "comfortable"; - -/** Per-model fit classification. */ -export type GpuFitLevel = "comfortable" | "fits" | "oom"; - -/** - * Classify whether a model fits the device. - * - * - "comfortable": estimated size ≤ 70% of GPU VRAM (runs fully in VRAM) - * - "fits": estimated size ≤ 70% GPU + 70% system RAM (runs with CPU offload) - * - "oom": exceeds both budgets - * - * Returns null when we can't determine the size (unknown → no badge). - */ -export function classifyGpuFit(opts: { - totalParams?: number; - estimatedSizeBytes?: number; - repoId: string; - gpu: GpuInfo; -}): GpuFitLevel | null { - const { totalParams, estimatedSizeBytes, repoId, gpu } = opts; - const gpuGb = gpu.memoryTotalGb; - const ramGb = gpu.systemRamAvailableGb; - if (gpuGb <= 0 && ramGb <= 0) return null; // no budget info - - // Active/effective model tokens (for example MoE A3B) describe runnable size - // better than HF total-parameter metadata; otherwise prefer exact metadata. - const activeOrEffectiveParams = activeOrEffectiveParamsFromId(repoId); - const params = activeOrEffectiveParams ?? totalParams ?? paramsFromId(repoId); - const sizeBytes = - activeOrEffectiveParams - ? estimateQuantBytes(activeOrEffectiveParams) - : estimatedSizeBytes ?? (params ? estimateQuantBytes(params) : undefined); - - if (!sizeBytes || sizeBytes <= 0) return null; // can't determine - - const sizeGb = sizeBytes / 1024 ** 3; - let comfortBudget: number; - let fitBudget: number; - - if (!gpu.available || gpuGb <= 0) { - // Unified memory system (no discrete GPU) - comfortBudget = ramGb * 0.7; - fitBudget = ramGb * 0.7; - } else { - // Discrete GPU - comfortBudget = gpuGb * 0.7; - fitBudget = gpuGb * 0.7 + ramGb * 0.7; - } - - if (sizeGb <= comfortBudget) return "comfortable"; - if (sizeGb <= fitBudget) return "fits"; - return "oom"; -} - -/** Whether a row passes the given GPU fit filter. */ -export function matchesGpuFitFilter( - level: GpuFitLevel | null, - filter: GpuFitFilter, -): boolean { - if (filter === "all") return true; - if (level === null) return false; - if (filter === "comfortable") return level === "comfortable"; - // "fits" shows both comfortable and fits - return level === "comfortable" || level === "fits"; -} diff --git a/studio/frontend/src/features/hub/lib/view-models.ts b/studio/frontend/src/features/hub/lib/view-models.ts index cdd17855db..9ee6c5de5d 100644 --- a/studio/frontend/src/features/hub/lib/view-models.ts +++ b/studio/frontend/src/features/hub/lib/view-models.ts @@ -10,7 +10,6 @@ import type { import type { CapabilityFilter, DiscoverRow, - GpuFitFilter, ModelFormatFilter, } from "../types"; import { @@ -52,15 +51,6 @@ export const FORMAT_FILTER_OPTIONS: ReadonlyArray<{ { value: "mlx", label: "MLX" }, ]; -export const GPU_FIT_FILTER_OPTIONS: ReadonlyArray<{ - value: GpuFitFilter; - label: string; -}> = [ - { value: "all", label: "All sizes" }, - { value: "fits", label: "Fits GPU" }, - { value: "comfortable", label: "Comfortable" }, -]; - const BILLION = 1_000_000_000; export function formatParamCount(totalParams: number | undefined): string { diff --git a/studio/frontend/src/features/hub/types.ts b/studio/frontend/src/features/hub/types.ts index 9fabdb2a8b..ae9ddf5d4b 100644 --- a/studio/frontend/src/features/hub/types.ts +++ b/studio/frontend/src/features/hub/types.ts @@ -28,9 +28,6 @@ export type ModelFormatFilter = "all" | "gguf" | "checkpoint" | "mlx"; export type CapabilityFilter = "all" | CapabilityKey; -import type { GpuFitFilter, GpuFitLevel } from "./lib/gpu-fit-filter"; -export type { GpuFitFilter, GpuFitLevel }; - export interface DiscoverRow { id: string; owner: string; @@ -40,7 +37,6 @@ export interface DiscoverRow { isPartialOnDevice: boolean; summary: string; capabilities: Capability[]; - fitLevel?: GpuFitLevel | null; } export type SelectedResourceSource = "huggingface" | "hub_cache" | LocalSource; From 693ab8069d5ff317e8efe6ecbf2fc86032b249f9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 28 Jun 2026 02:43:34 -0700 Subject: [PATCH 18/49] Remove unused FalconH1RMSNormGated import (#6728) FalconH1RMSNormGated is imported from transformers but never referenced in unsloth/models/falcon_h1.py. The unused hoist trips the import-hoist lint gate on the merge commit of every open PR (the gate lints PR-head merged into main), so clearing it here unblocks those PRs. --- unsloth/models/falcon_h1.py | 1 - 1 file changed, 1 deletion(-) diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index 05bfd2ebb3..e3e04e4fdc 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -38,7 +38,6 @@ try: FalconH1Model, FalconH1ForCausalLM, FalconH1RMSNorm, - FalconH1RMSNormGated, FalconHybridMambaAttentionDynamicCache, ) except: From b56d24ea3e7111c5fa34167a528f0f07fe22e32c Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Sun, 28 Jun 2026 16:46:03 +0530 Subject: [PATCH 19/49] Studio: cascade user message deletion to include assistant reply (#6720) * cascade user message deletion to include assistant reply * Fix comment typo in delete-thread-message --------- Co-authored-by: Daniel Han --- .../chat/utils/delete-thread-message.ts | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/studio/frontend/src/features/chat/utils/delete-thread-message.ts b/studio/frontend/src/features/chat/utils/delete-thread-message.ts index 3b439eb779..bea5a63ea7 100644 --- a/studio/frontend/src/features/chat/utils/delete-thread-message.ts +++ b/studio/frontend/src/features/chat/utils/delete-thread-message.ts @@ -112,7 +112,26 @@ export async function deleteThreadMessage(args: { const exported = thread.export(); const repo = new MessageRepository(); repo.import(exported); + + const target = exported.messages.find( + ({ message }) => message.id === messageId, + ); + const assistantReplyIds = + target?.message.role === "user" + ? exported.messages + .filter( + ({ parentId, message }) => + parentId === messageId && message.role === "assistant", + ) + .map(({ message }) => message.id) + : []; + + // Delete the prompt first; that relinks its replies up to the prompt's parent repo.deleteMessage(messageId); + for (const replyId of assistantReplyIds) { + repo.deleteMessage(replyId); + } + const next = repo.export(); if (remoteId) { await syncExportedRepositoryToBackend(remoteId, next, { From 20266a59eb4516edf2aa1b8a791540085e22956e Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Sun, 28 Jun 2026 20:36:39 -0700 Subject: [PATCH 20/49] Fix custom chat templates with a {system_message} placeholder (dead code in _change_system_message) (#6735) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- tests/python/test_change_system_message.py | 76 ++++++++++++++++++++++ unsloth/chat_templates.py | 27 ++++---- 2 files changed, 88 insertions(+), 15 deletions(-) create mode 100644 tests/python/test_change_system_message.py diff --git a/tests/python/test_change_system_message.py b/tests/python/test_change_system_message.py new file mode 100644 index 0000000000..1ccd6547b2 --- /dev/null +++ b/tests/python/test_change_system_message.py @@ -0,0 +1,76 @@ +import ast +import re +import types +from pathlib import Path + +import pytest + + +def _load_change_system_message(): + # Extract just _change_system_message from chat_templates.py so the test runs + # without importing unsloth (which needs unsloth_zoo / a GPU). Same pattern as + # tests/saving/test_is_gpt_oss_detection.py. + source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + funcs = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_change_system_message" + ] + namespace = { + "re": re, + "logger": types.SimpleNamespace(warning_once = lambda *a, **k: None), + "DEFAULT_SYSTEM_MESSAGE": {"unsloth": "You are a helpful assistant to the user"}, + } + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace["_change_system_message"] + + +CUSTOM = "mycustom" # not in DEFAULT_SYSTEM_MESSAGE -> no predefined default + + +def test_custom_template_fills_placeholder(): + # A custom template with a {system_message} placeholder must be filled, not + # left with the literal placeholder. + fn = _load_change_system_message() + template, used = fn("System: {system_message}\nUser:", CUSTOM, "You are a pirate") + assert template == "System: You are a pirate\nUser:" + assert "{system_message}" not in template + assert used == "You are a pirate" + + +def test_custom_template_preserves_backslashes(): + # Why str.replace and not re.sub: a system message with backslashes (Windows + # paths, LaTeX, group-like text) must be inserted verbatim. re.sub treats the + # replacement specially -- r"C:\Users" raises bad-escape, r"\1" is a group ref. + fn = _load_change_system_message() + for msg in (r"C:\Users\me", r"\frac{a}{b}", r"see \1 here"): + template, used = fn("System: {system_message}", CUSTOM, msg) + assert template == f"System: {msg}" + assert used == msg + + +def test_custom_template_requires_system_message(): + # A custom template with a placeholder but no system message must raise, + # rather than silently leaving the placeholder in. + fn = _load_change_system_message() + with pytest.raises(ValueError): + fn("System: {system_message}", CUSTOM, None) + + +def test_custom_template_without_placeholder_unchanged(): + fn = _load_change_system_message() + template, used = fn("System: fixed", CUSTOM, "ignored") + assert template == "System: fixed" + + +def test_predefined_template_uses_default_then_override(): + # Predefined templates with a default are unaffected by the change. + fn = _load_change_system_message() + t1, u1 = fn("System: {system_message}", "unsloth", None) + assert t1 == "System: You are a helpful assistant to the user" + t2, u2 = fn("System: {system_message}", "unsloth", "Custom override") + assert t2 == "System: Custom override" + assert u2 == "Custom override" diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 7f453bff82..eba4577315 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -1804,10 +1804,19 @@ CHAT_TEMPLATES["yi-chat"] = (yi_chat_template, yi_chat_template_eos_token, False DEFAULT_SYSTEM_MESSAGE["yi-chat"] = None def _change_system_message(template: str, type_chat_template: str, system_message: str = None): - system_message_pattern = r"\{system_message\}" - # For predefined templates, check if default system message exists default_system_message = DEFAULT_SYSTEM_MESSAGE.get(f"{type_chat_template}", None) + + # Custom templates have no predefined default, but may still carry a + # {system_message} placeholder. Handle it before the no-default early return + # below, which would otherwise leave the literal "{system_message}" in the + # template. A placeholder with no system message is an error, not a no-op. + if default_system_message is None and "{system_message}" in template: + if system_message is None: + raise ValueError("Unsloth: You need to provide a system message for custom templates.") + new_template = template.replace("{system_message}", system_message) + return new_template, system_message + if default_system_message is None: if system_message is not None: logger.warning_once( @@ -1817,21 +1826,9 @@ def _change_system_message(template: str, type_chat_template: str, system_messag ) return template, system_message - # For custom templates - if type_chat_template is None: - has_placeholder = re.search(system_message_pattern, template) is not None - - if has_placeholder: - if system_message is None: - raise ValueError("Unsloth: You need to provide a system message for custom templates.") - new_template = re.sub(system_message_pattern, system_message, template) - return new_template, system_message - - return template, system_message - # For predefined templates with default system message message_to_use = system_message if system_message is not None else default_system_message - new_template = re.sub(system_message_pattern, message_to_use, template) + new_template = template.replace("{system_message}", message_to_use) return new_template, message_to_use From 677ec0cc20bf7cb4735385c51a22999a64839a83 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Sun, 28 Jun 2026 22:44:00 -0700 Subject: [PATCH 21/49] Fix gpt-oss detection in save: config.architectures is a list, not a string (#6711) --- tests/saving/test_is_gpt_oss_detection.py | 52 +++++++++++++++++++++++ unsloth/save.py | 21 +++++---- 2 files changed, 64 insertions(+), 9 deletions(-) create mode 100644 tests/saving/test_is_gpt_oss_detection.py diff --git a/tests/saving/test_is_gpt_oss_detection.py b/tests/saving/test_is_gpt_oss_detection.py new file mode 100644 index 0000000000..c8790e4da1 --- /dev/null +++ b/tests/saving/test_is_gpt_oss_detection.py @@ -0,0 +1,52 @@ +import ast +import types +from pathlib import Path + + +def _load_is_gpt_oss(): + # Extract just the helper from save.py so the test runs without importing + # unsloth (which requires unsloth_zoo / a GPU), matching the pattern used by + # test_qwen3_5_vlm_full_finetune_key_remap.py. + source = Path(__file__).parents[2] / "unsloth" / "save.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + helpers = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_is_gpt_oss" + ] + module = ast.Module(body = helpers, type_ignores = []) + ast.fix_missing_locations(module) + namespace = {} + exec(compile(module, str(source), "exec"), namespace) + return namespace["_is_gpt_oss"] + + +def _model(architectures = None, model_type = None): + config = types.SimpleNamespace() + if architectures is not None: + config.architectures = architectures + if model_type is not None: + config.model_type = model_type + return types.SimpleNamespace(config = config) + + +def test_detects_gpt_oss_by_architecture(): + # config.architectures is a list, so detection must use membership, not ==. + # A model that declares GptOssForCausalLM but has no matching model_type must + # still be routed to the mxfp4 save path. + is_gpt_oss = _load_is_gpt_oss() + assert is_gpt_oss(_model(architectures = ["GptOssForCausalLM"])) is True + assert is_gpt_oss(_model(architectures = ["GptOssForCausalLM"], model_type = "gpt_oss")) is True + + +def test_detects_gpt_oss_by_model_type(): + is_gpt_oss = _load_is_gpt_oss() + assert is_gpt_oss(_model(architectures = ["SomethingElse"], model_type = "gpt-oss")) is True + assert is_gpt_oss(_model(architectures = ["SomethingElse"], model_type = "gpt_oss")) is True + + +def test_non_gpt_oss_is_false(): + is_gpt_oss = _load_is_gpt_oss() + assert is_gpt_oss(_model(architectures = ["LlamaForCausalLM"], model_type = "llama")) is False + assert is_gpt_oss(_model()) is False + assert is_gpt_oss(types.SimpleNamespace()) is False diff --git a/unsloth/save.py b/unsloth/save.py index 0e8c8cdc2d..f55a14b4e3 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -475,6 +475,17 @@ def _is_qwen3_5_vlm(model): ) or getattr(config, "model_type", None) in ("qwen3_5", "qwen3_5_moe") +def _is_gpt_oss(model): + config = getattr(model, "config", None) + if config is None: + return False + architectures = getattr(config, "architectures", None) or () + return "GptOssForCausalLM" in architectures or getattr(config, "model_type", None) in ( + "gpt-oss", + "gpt_oss", + ) + + def _qwen3_5_vlm_state_dict_for_save(state_dict): remapped_state_dict = {} for key, value in state_dict.items(): @@ -2231,15 +2242,7 @@ def unsloth_save_pretrained_gguf( is_processor = is_vlm and isinstance(tokenizer, ProcessorMixin) - is_gpt_oss = ( - True - if ( - hasattr(self.config, "architectures") - and self.config.architectures == "GptOssForCausalLM" - ) - or (hasattr(self.config, "model_type") and self.config.model_type in ["gpt-oss", "gpt_oss"]) - else False - ) + is_gpt_oss = _is_gpt_oss(self) # Step 2: Prepare arguments for model saving arguments = dict(locals()) arguments["model"] = self From 54b95fbcc8a7a928e8161169d70c394998d73176 Mon Sep 17 00:00:00 2001 From: Matt Van Horn Date: Mon, 29 Jun 2026 03:12:47 -0700 Subject: [PATCH 22/49] fix(studio): show local file path tooltip for Hub-tab local models (#6715) Local models in the Studio Hub tab (Custom folders, LM Studio, and Local models sections) did not reveal their on-disk path on hover, unlike the Fine-tuned rows which already do. Each of these rows maps over a LocalModelInfo with a required path, so pass tooltipText built from the model name and path via a small shared localPathTooltip helper, matching the existing FT-row tooltip format. Refs #6382 Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- .../assistant-ui/model-selector/pickers.tsx | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) 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 386f233c7b..5f5d112204 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1114,6 +1114,17 @@ function localModelIsGguf(m: LocalModelInfo): boolean { ); } +function localPathTooltip(name: string, path: string): ReactNode { + return ( + <> + {name} + + {path} + + + ); +} + /** 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 { @@ -2928,6 +2939,10 @@ export function HubModelPicker({ Date: Mon, 29 Jun 2026 11:57:48 +0100 Subject: [PATCH 23/49] Fix compare adapter selection (#6411) --- .../src/features/chat/api/chat-adapter.ts | 22 +++++++++++++++++-- .../src/features/chat/runtime-provider.tsx | 9 ++++---- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 0106980871..ca7ac174fc 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -47,6 +47,7 @@ import { useChatRuntimeStore, } from "../stores/chat-runtime-store"; import { useExternalProvidersStore } from "../stores/external-providers-store"; +import type { ModelType } from "../types"; import { isMultimodalResponse } from "../types/api"; import type { GgufVariantDetail, @@ -142,6 +143,11 @@ interface ServerTimings { type RunMessages = Parameters[0]["messages"]; type RunMessage = RunMessages[number]; +type OpenAIStreamAdapterOptions = { + modelType?: ModelType; + pairId?: string; +}; + /** Tracks which user messages were sent with an audio file (messageId → filename). */ export const sentAudioNames = new Map(); @@ -1182,7 +1188,17 @@ export function findLatestUserAudioBase64( async function resolveUseAdapter( threadId: string | undefined, + options: OpenAIStreamAdapterOptions = {}, ): Promise { + if (options.modelType === "model1" || options.modelType === "model2") { + return undefined; + } + if ( + options.pairId && + (options.modelType === "base" || options.modelType === "lora") + ) { + return options.modelType === "lora"; + } if (!threadId) { return undefined; } @@ -1629,7 +1645,9 @@ async function autoLoadSmallestModel(): Promise<{ } } -export function createOpenAIStreamAdapter(): ChatModelAdapter { +export function createOpenAIStreamAdapter( + options: OpenAIStreamAdapterOptions = {}, +): ChatModelAdapter { return { async *run({ messages, abortSignal, unstable_threadId }) { await useChatRuntimeStore.getState().hydratePersistedSettings(); @@ -2076,7 +2094,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } runtime.clearPendingAudio(); } - const useAdapter = await resolveUseAdapter(resolvedThreadId); + const useAdapter = await resolveUseAdapter(resolvedThreadId, options); // ── Audio model path (non-streaming) ───────────────────── const activeModel = runtime.models.find( diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 360e081f0f..67980b94c6 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -1045,16 +1045,17 @@ function useStudioRuntimeAdapters( return adapters; } -const chatAdapter = createOpenAIStreamAdapter(); - function useRuntimeHook( modelType: ModelType, pairId?: string, ): ReturnType { const adapters = useStudioRuntimeAdapters(modelType, pairId); const persistedChatAdapter = useMemo( - () => createPersistedRunAdapter(chatAdapter), - [], + () => + createPersistedRunAdapter( + createOpenAIStreamAdapter({ modelType, pairId }), + ), + [modelType, pairId], ); return useLocalRuntime(persistedChatAdapter, { adapters }); } From 02540371c20fd468cfee358cc850f56527b54957 Mon Sep 17 00:00:00 2001 From: Muhammad Ikhwan Fathulloh <77288014+Muhammad-Ikhwan-Fathulloh@users.noreply.github.com> Date: Mon, 29 Jun 2026 19:07:09 +0700 Subject: [PATCH 24/49] perf(dataprep): cache regex and field lists, fix typos (#6714) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Improve code quality & performance: fix typos, compile regex & cache fields - Fix typos across core files (repeatted → repeated, splitted → split, etc.) - Compile regex patterns once as class attributes in TextPreprocessor - Cache text fields/columns in RawTextDataLoader - Improve comments (re-use → reuse) * Use immutable raw text field constants --------- Co-authored-by: imagineer99 Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- unsloth/chat_templates.py | 18 ++++---- unsloth/dataprep/raw_text.py | 41 +++++++++++++------ unsloth/kernels/geglu.py | 2 +- .../moe/grouped_gemm/kernels/forward.py | 2 +- unsloth/models/_utils.py | 2 +- 5 files changed, 40 insertions(+), 25 deletions(-) diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index eba4577315..8c4606f418 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2276,28 +2276,28 @@ def get_ollama_eos_tokens(tokenizer, extra_eos_tokens = []): if getattr(tokenizer, "bos_token", None) is not None: added_tokens_decoder = [x for x in added_tokens_decoder if x != tokenizer.bos_token] - repeatted_tokens = [] + repeated_tokens = [] # Join all vocab joined_text = "\x01\x00".join(added_tokens_decoder) for token in added_tokens_decoder: n = len(token) - repeatted_counts = joined_text.count(token[:n//2]) + repeated_counts = joined_text.count(token[:n//2]) # Try finding longer than 1/2 of the token in the rest # For eg <|reserved_special_token_0|>, <|reserved_special_token_1|> - if repeatted_counts > 2: + if repeated_counts > 2: for j in range(n//2+1, n): - if joined_text.count(token[:j]) < repeatted_counts: + if joined_text.count(token[:j]) < repeated_counts: j -= 1 - # Remove repeatted tokens to reduce search space + # Remove repeated tokens to reduce search space joined_text = joined_text.replace(token[:j], "") - repeatted_tokens.append(token[:j]) + repeated_tokens.append(token[:j]) break # Remove duplicates - splitted = joined_text.split("\x01\x00") - final_eos_tokens = [old for old, new in zip(added_tokens_decoder, splitted) if old == new] + split = joined_text.split("\x01\x00") + final_eos_tokens = [old for old, new in zip(added_tokens_decoder, split) if old == new] final_eos_tokens += extra_eos_tokens - final_eos_tokens += repeatted_tokens + final_eos_tokens += repeated_tokens # Remove new lines, spaces and HTML tags filtered_eos_tokens = [] diff --git a/unsloth/dataprep/raw_text.py b/unsloth/dataprep/raw_text.py index 076a781d4b..7d18b1ff29 100644 --- a/unsloth/dataprep/raw_text.py +++ b/unsloth/dataprep/raw_text.py @@ -223,48 +223,63 @@ class RawTextDataLoader: return "\n\n".join(texts) return "" + # Cache text fields/columns for better performance + _TEXT_FIELDS = ("text", "content", "message", "body", "description", "prompt") + _TEXT_COLUMNS = _TEXT_FIELDS + def _extract_text_from_json(self, data): """Extract text from JSON object using common field names.""" - text_fields = ["text", "content", "message", "body", "description", "prompt"] - for field in text_fields: + for field in self._TEXT_FIELDS: if field in data and isinstance(data[field], str): return data[field] return "" def _extract_text_from_csv_row(self, row): """Extract text from CSV row using common column names.""" - text_columns = ["text", "content", "message", "body", "description", "prompt"] - for column in text_columns: + for column in self._TEXT_COLUMNS: if column in row and row[column]: return row[column] return "" class TextPreprocessor: + # Compile regex patterns once for better performance + _WHITESPACE_PATTERN = re.compile(r"[^\S\n]+") + _INVALID_CHARS_PATTERN = re.compile(r"[^\x20-\x7E\n]") + _MULTIPLE_SPACES_PATTERN = re.compile(r"[ ]{2,}") + _NEWLINE_SPACES_PATTERN = re.compile(r" *\n *") + _MULTIPLE_NEWLINES_PATTERN = re.compile(r"\n{3,}") + _CHAPTER_PATTERN = re.compile(r"^# (.+)$", re.MULTILINE) + _SECTION_PATTERN = re.compile(r"^## (.+)$", re.MULTILINE) + _SUBSECTION_PATTERN = re.compile(r"^### (.+)$", re.MULTILINE) + _CODE_BLOCK_PATTERN = re.compile(r"```(\w*)\n(.*?)\n```", re.DOTALL) + def clean_text(self, text): """Remove unwanted characters, normalize whitespace""" text = text.replace("\r\n", "\n").replace("\r", "\n") - text = re.sub(r"[^\S\n]+", " ", text) - text = re.sub(r"[^\x20-\x7E\n]", "", text) - text = re.sub(r"[ ]{2,}", " ", text) - text = re.sub(r" *\n *", "\n", text) - text = re.sub(r"\n{3,}", "\n\n", text) + text = self._WHITESPACE_PATTERN.sub(" ", text) + text = self._INVALID_CHARS_PATTERN.sub("", text) + text = self._MULTIPLE_SPACES_PATTERN.sub(" ", text) + text = self._NEWLINE_SPACES_PATTERN.sub("\n", text) + text = self._MULTIPLE_NEWLINES_PATTERN.sub("\n\n", text) return text.strip() def extract_sections(self, text, patterns): """Extract specific sections (e.g., code blocks, quotes)""" sections = [] for pattern in patterns: + # Compile pattern on first use and cache? Well, patterns are user-provided, + # so just use re.findall with compiled flags matches = re.findall(pattern, text, re.MULTILINE | re.DOTALL) sections.extend(matches) return sections def add_structure_tokens(self, text): """Add special tokens for structure (chapters, sections)""" - text = re.sub(r"^# (.+)$", r"<|chapter|>\1<|/chapter|>", text, flags = re.MULTILINE) - text = re.sub(r"^## (.+)$", r"<|section|>\1<|/section|>", text, flags = re.MULTILINE) - text = re.sub(r"^### (.+)$", r"<|subsection|>\1<|/subsection|>", text, flags = re.MULTILINE) - text = re.sub(r"```(\w*)\n(.*?)\n```", r"<|code|\1|>\2<|/code|>", text, flags = re.DOTALL) + text = self._CHAPTER_PATTERN.sub(r"<|chapter|>\1<|/chapter|>", text) + text = self._SECTION_PATTERN.sub(r"<|section|>\1<|/section|>", text) + text = self._SUBSECTION_PATTERN.sub(r"<|subsection|>\1<|/subsection|>", text) + text = self._CODE_BLOCK_PATTERN.sub(r"<|code|\1|>\2<|/code|>", text) return text def validate_dataset(self, dataset): diff --git a/unsloth/kernels/geglu.py b/unsloth/kernels/geglu.py index 3a628f105c..0ed325afa0 100644 --- a/unsloth/kernels/geglu.py +++ b/unsloth/kernels/geglu.py @@ -97,7 +97,7 @@ def _exact_backward_kernel( e_row = tl.load(e + offsets, mask = mask, other = 0).to(tl.float32) g_row = tl.load(g + offsets, mask = mask, other = 0) # .to(tl.float32) - # Break e_row away for re-use + # Break e_row away for reuse # f = 1/2 * e * (1 + erf(1/sqrt(2) * e)) f_partial_row = 0.5 * (tl.math.erf(tl.math.rsqrt(2.0) * e_row) + 1.0) f_row = f_partial_row * e_row diff --git a/unsloth/kernels/moe/grouped_gemm/kernels/forward.py b/unsloth/kernels/moe/grouped_gemm/kernels/forward.py index e95179dd9f..cb1f22d0b8 100644 --- a/unsloth/kernels/moe/grouped_gemm/kernels/forward.py +++ b/unsloth/kernels/moe/grouped_gemm/kernels/forward.py @@ -111,7 +111,7 @@ def _grouped_gemm_forward_kernel( while tidx >= processed_tiles and tidx < processed_tiles + num_tiles_per_expert: tile_idx = tidx - processed_tiles - # Check if L2 cache re-use for this order is optimal + # Check if L2 cache reuse for this order is optimal tile_m_idx = tile_idx % num_m_tiles tile_n_idx = tile_idx // num_m_tiles diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 7ad6e8ea33..599a5c0262 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3078,7 +3078,7 @@ class TorchAOConfig: def _untie_input_output_embeddings(model: torch.nn.Module) -> None: """ Utility to untie input/output embeddings in a HuggingFace model. - This is useful if we want to quantize the input/ouput embeddings differently. + This is useful if we want to quantize the input/output embeddings differently. Model is modified in-place. """ From 755da2f1552c2ed063d18b01d98807f413dff3ac Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Mon, 29 Jun 2026 15:27:39 +0200 Subject: [PATCH 25/49] Speed up Studio desktop startup (#6742) * Speed up Studio desktop startup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address Studio startup review findings * Keep orphaned run cleanup before readiness * [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/main.py | 90 +++-- studio/backend/run.py | 29 ++ studio/frontend/src/app/provider.tsx | 28 +- .../frontend/src/hooks/use-tauri-backend.ts | 20 +- studio/src-tauri/src/commands.rs | 16 + studio/src-tauri/src/desktop_backend_owner.rs | 147 +++++--- studio/src-tauri/src/preflight.rs | 16 + studio/src-tauri/src/preflight/backend.rs | 12 + studio/src-tauri/src/preflight/managed.rs | 333 +++++++++++++++++- studio/src-tauri/src/process.rs | 86 +++-- 10 files changed, 636 insertions(+), 141 deletions(-) diff --git a/studio/backend/main.py b/studio/backend/main.py index 731625ca74..0a5b775775 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -441,9 +441,30 @@ def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None: ).start() +def _warm_rag_embedder() -> None: + """Warm RAG embeddings without blocking backend readiness.""" + try: + from storage import rag_db + + if not rag_db.RAG_AVAILABLE: + return + from core.rag import embeddings + + embeddings.warm() + except Exception: + pass + + @asynccontextmanager async def lifespan(app: FastAPI): """Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache.""" + + import time as _time + + _lifespan_started = _time.perf_counter() + import structlog as _structlog + + _lifespan_log = _structlog.get_logger(__name__) clear_unsloth_compiled_cache() # Remove stale .venv_overlay from old versions; switching now uses .venv_t5/. @@ -454,6 +475,11 @@ async def lifespan(app: FastAPI): # Detect hardware first — sets the DEVICE global used everywhere. detect_hardware() + _lifespan_log.info( + "lifespan hardware detection completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) + # Apple Silicon with MLX missing => Train/Export are greyed out (chat-only). # Reinstall mlx by name on a background thread (off the critical path) and # re-detect, so a reinstall/update that dropped mlx self-heals. No-op @@ -465,7 +491,13 @@ async def lifespan(app: FastAPI): import structlog as _structlog _structlog.get_logger(__name__).debug("mlx autorepair skipped: %s", _mlx_exc) - # Reap download workers orphaned by a previous crash before new downloads start. + # Reap workers/runs orphaned by a previous crash before new work starts. + try: + from storage.studio_db import cleanup_orphaned_runs + cleanup_orphaned_runs() + except Exception as exc: + _lifespan_log.warning("cleanup_orphaned_runs failed at startup: %s", exc) + reap_hub_orphan_workers() # llama.cpp probes: capability (MTP support) + freshness (release age). @@ -479,45 +511,23 @@ async def lifespan(app: FastAPI): app.state.llama_cpp_freshness = None _start_llama_cpp_probes_if_enabled(app) - from storage.studio_db import cleanup_orphaned_runs - - try: - cleanup_orphaned_runs() - except Exception as exc: - import structlog - structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc) - - # Same for RAG: fail ingestion jobs stranded mid-ingest by a crash. try: from storage.rag_db import reconcile_orphaned_ingestion_jobs reconcile_orphaned_ingestion_jobs() except Exception as exc: - import structlog - structlog.get_logger(__name__).warning( - "reconcile_orphaned_ingestion_jobs failed at startup: %s", exc - ) + _lifespan_log.warning("reconcile_orphaned_ingestion_jobs failed at startup: %s", exc) _start_helper_precache_if_enabled() + threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start() - # Warm the RAG embedder so the first upload skips the cold load. Non-fatal. - def _warm_rag_embedder(): - try: - from storage import rag_db - - if not rag_db.RAG_AVAILABLE: - return - from core.rag import embeddings - - embeddings.warm() - except Exception: - pass - - threading.Thread(target = _warm_rag_embedder, daemon = True).start() - - # Initialize RSA key pair for API key encryption (external providers) + # Initialize RSA key pair for API key encryption (external providers). from core.inference.key_exchange import init_key_pair init_key_pair() + _lifespan_log.info( + "lifespan pre-auth setup completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) if storage.ensure_default_admin(): bootstrap_pw = storage.get_bootstrap_password() @@ -532,6 +542,11 @@ async def lifespan(app: FastAPI): print("=" * 60 + "\n") else: app.state.bootstrap_password = storage.get_bootstrap_password() + + _lifespan_log.info( + "lifespan startup completed in %.1fms", + (_time.perf_counter() - _lifespan_started) * 1000, + ) yield from core.inference.llama_http import aclose as _close_llama_http @@ -919,6 +934,21 @@ install_api_error_handlers(app) # ============ Health and System Endpoints ============ +@app.get("/api/liveness") +async def liveness_check(): + """Cheap process liveness for desktop port validation.""" + return { + "status": "alive", + "service": "Unsloth UI Backend", + "desktop_protocol_version": 1, + "desktop_manageability_version": 1, + "supports_desktop_auth": True, + "supports_desktop_backend_ownership": True, + "studio_root_id": _studio_root_id(), + **({"desktop_owner": owner} if (owner := _desktop_owner()) else {}), + } + + @app.get("/api/health") async def health_check(request: Request): """Liveness plus launcher capability bits; host fingerprint gated on a bearer. diff --git a/studio/backend/run.py b/studio/backend/run.py index d4cbc26b41..6f36f0928a 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -933,6 +933,9 @@ def run_server( """ global _server, _server_thread, _shutdown_event + boot_started = time.perf_counter() + logger.info("run_server startup begin api_only=%s host=%s port=%s", api_only, host, port) + # Reap every child if the parent dies abnormally (terminal close, Task # Manager kill, SIGKILL); must run before any child can spawn. from utils.process_lifetime import initialize_parent_lifetime @@ -984,7 +987,14 @@ def run_server( from threading import Thread, Event import uvicorn + import_started = time.perf_counter() + from main import app, setup_frontend, _IS_COLAB + + logger.info( + "Imported FastAPI app in %.1fms", + (time.perf_counter() - import_started) * 1000, + ) from utils.paths import ensure_studio_directories # Allow local stdio MCP servers on a loopback bind (the user's own machine), @@ -997,6 +1007,11 @@ def run_server( # Create all standard directories on startup. ensure_studio_directories() + logger.info( + "Ensured Studio directories in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + # Auto-find a free port if the requested one is in use. if not _is_port_free(host, port): original_port = port @@ -1060,6 +1075,11 @@ def run_server( display_host = _resolve_external_ip() if host == "0.0.0.0" else host _install_uvicorn_startup_log_rewrite(host, display_host) + logger.info( + "run_server pre-uvicorn setup completed in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + ready_event = Event() startup_failed = Event() startup_errors = [] @@ -1068,6 +1088,10 @@ def run_server( async def startup(self, *args, **kwargs): await super().startup(*args, **kwargs) if getattr(self, "started", False) and not self.should_exit: + logger.info( + "Uvicorn startup hook completed in %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) ready_event.set() # server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own. @@ -1150,6 +1174,11 @@ def run_server( _shutdown_event.set() raise + logger.info( + "run_server uvicorn ready after %.1fms", + (time.perf_counter() - boot_started) * 1000, + ) + _write_pid_file() import atexit diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index 914abbbf1d..176665769d 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -354,12 +354,11 @@ function TauriWrapper({ children }: { children: ReactNode }) { ); } - const showApp = status === "running" && desktopAuthReady; + const showApp = status === "running"; + const desktopBooting = status === "running" && !desktopAuthReady; + const showInteractiveApp = showApp && desktopAuthReady; const startupStatus = status === "running" ? "starting" : status; - const startupProgressDetail = - status === "running" && !desktopAuthReady - ? "Signing in to desktop session..." - : progressDetail; + const startupProgressDetail = progressDetail; const usesCustomTitlebar = shouldUseCustomWindowTitlebar(); const usesNativeMacTitlebar = shouldUseNativeMacWindowTitlebar(); const hidesTitlebarSidebar = HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname); @@ -369,12 +368,23 @@ function TauriWrapper({ children }: { children: ReactNode }) { - + {showInteractiveApp ? : null} - - {children} + {showInteractiveApp ? : null} + {showInteractiveApp ? children : null} + {desktopBooting ? ( +
+
+
Preparing Studio
+
The local backend is ready. Signing in to your desktop session before loading chats.
+
+
+ Signing in to desktop session... +
+
+ ) : null} ) : ( number | null, shouldContinue: () => boolean, ): Promise { @@ -91,15 +90,7 @@ async function waitForManagedServerReady( continue; } - const healthy = await invoke("check_health", { port }); - if (!shouldContinue()) { - return { status: "aborted" }; - } - if (healthy && getPort() === port) { - return { status: "ready", port }; - } - - await wait(MANAGED_STARTUP_POLL_MS); + return { status: "ready", port }; } } @@ -280,10 +271,9 @@ export function useTauriBackend() { // backend/run.py keeps the 8888-8908 fallback via server-port/TAURI_PORT. await invoke("start_managed_server", { port: 8888 }); - // Wait for the owned backend's server-port event. Don't attach to an - // external backend if the managed start doesn't report a port. - const startupResult = await waitForManagedServerReady( - invoke, + // Rust emits server-port only after validating the desktop-owned process. + // Treat that as the UI handoff point instead of doing a second health poll. + const startupResult = await waitForManagedServerPort( () => portRef.current, () => startingRef.current, ); diff --git a/studio/src-tauri/src/commands.rs b/studio/src-tauri/src/commands.rs index 7c1ce611b8..6bc2116786 100644 --- a/studio/src-tauri/src/commands.rs +++ b/studio/src-tauri/src/commands.rs @@ -65,10 +65,18 @@ pub async fn desktop_preflight( shutdown: tauri::State<'_, ShutdownFlag>, diagnostics: tauri::State<'_, DiagnosticsState>, ) -> Result { + let started = Instant::now(); let (result, adopted_watchdog_generation) = crate::preflight::desktop_preflight_result_with_state(state.inner()).await?; diagnostics::record_preflight(&diagnostics, &result); + info!( + "desktop_preflight completed disposition={:?} port={:?} in {}ms", + result.disposition, + result.port, + started.elapsed().as_millis() + ); + if let Some((generation, newly_adopted)) = adopted_watchdog_generation { if newly_adopted { if let Some(port) = result.port { @@ -205,9 +213,17 @@ pub async fn start_managed_server( port: u16, ) -> Result<(), String> { info!("start_managed_server command called with port {}", port); + + let started = Instant::now(); let diagnostics_state = diagnostics.inner().clone(); let generation = process::start_backend(&app, &state, port, &shutdown, &diagnostics_state)?; + info!( + "start_managed_server spawned generation={} in {}ms", + generation, + started.elapsed().as_millis() + ); + let watchdog_state = state.inner().clone(); let watchdog_shutdown = shutdown.inner().clone(); let watchdog_app = app.clone(); diff --git a/studio/src-tauri/src/desktop_backend_owner.rs b/studio/src-tauri/src/desktop_backend_owner.rs index ed4b5f7fd9..c7d0a7b309 100644 --- a/studio/src-tauri/src/desktop_backend_owner.rs +++ b/studio/src-tauri/src/desktop_backend_owner.rs @@ -93,7 +93,7 @@ enum PreviousAppPidStatus { Uncertain, } -#[derive(Debug, Deserialize)] +#[derive(Clone, Debug, Deserialize)] struct HealthDesktopOwner { kind: Option, token_sha256: Option, @@ -101,15 +101,7 @@ struct HealthDesktopOwner { #[derive(Debug, Deserialize)] struct HealthResponse { - status: Option, - service: Option, version: Option, - desktop_protocol_version: Option, - desktop_manageability_version: Option, - supports_desktop_auth: Option, - supports_desktop_backend_ownership: Option, - studio_root_id: Option, - desktop_owner: Option, } #[derive(Debug)] @@ -123,6 +115,18 @@ struct DesktopLoginPayload<'a> { secret: &'a str, } +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DesktopLiveness { + status: Option, + service: Option, + desktop_protocol_version: Option, + desktop_manageability_version: Option, + supports_desktop_auth: Option, + supports_desktop_backend_ownership: Option, + studio_root_id: Option, + desktop_owner: Option, +} + #[derive(Deserialize)] struct TokenResponse { access_token: String, @@ -290,10 +294,10 @@ impl BackendOwnerState { } pub(crate) fn verifies_exact_port_blocking(&self, port: u16) -> bool { - match fetch_health_blocking(port) { - Ok(Some(health)) => { - health_verifies_metadata(&health, &self.metadata) - && lifecycle_control_block_reason(&health).is_none() + match fetch_liveness_blocking(port) { + Ok(Some(liveness)) => { + liveness_verifies_metadata(&liveness, &self.metadata) + && lifecycle_control_block_reason(&liveness).is_none() } _ => false, } @@ -498,66 +502,110 @@ pub(crate) fn test_owner_state(root_id: &str, token: &str, port: u16) -> Backend } } -fn health_verifies_metadata(health: &HealthResponse, metadata: &DesktopBackendMetadata) -> bool { - let healthy = health.status.as_deref() == Some("healthy") - && health.service.as_deref() == Some("Unsloth UI Backend"); - let Some(owner) = health.desktop_owner.as_ref() else { +fn liveness_verifies_metadata( + liveness: &DesktopLiveness, + metadata: &DesktopBackendMetadata, +) -> bool { + let alive = matches!(liveness.status.as_deref(), Some("alive") | Some("healthy")) + && liveness.service.as_deref() == Some("Unsloth UI Backend"); + let Some(owner) = liveness.desktop_owner.as_ref() else { return false; }; - healthy + alive && owner_matches_metadata( metadata, - health.studio_root_id.as_deref(), + liveness.studio_root_id.as_deref(), owner.kind.as_deref(), owner.token_sha256.as_deref(), ) } -fn lifecycle_control_block_reason(health: &HealthResponse) -> Option { - if health.desktop_protocol_version != Some(crate::preflight::DESKTOP_PROTOCOL_VERSION) { +fn lifecycle_control_block_reason(liveness: &DesktopLiveness) -> Option { + if liveness.desktop_protocol_version != Some(crate::preflight::DESKTOP_PROTOCOL_VERSION) { return Some("desktop_protocol_incompatible".to_string()); } - if health.supports_desktop_auth != Some(true) { + if liveness.supports_desktop_auth != Some(true) { return Some("desktop_auth_unsupported".to_string()); } - if health.desktop_manageability_version.unwrap_or(0) + if liveness.desktop_manageability_version.unwrap_or(0) < crate::preflight::DESKTOP_MANAGEABILITY_VERSION { return Some("desktop_manageability_unsupported".to_string()); } - if health.supports_desktop_backend_ownership != Some(true) { + if liveness.supports_desktop_backend_ownership != Some(true) { return Some("desktop_backend_ownership_unsupported".to_string()); } None } -fn ready_for_use_status(health: &HealthResponse) -> OwnedBackendReadiness { - match crate::preflight::backend_version_stale_reason(health.version.as_deref()) { +fn ready_for_use_status(health: Option<&HealthResponse>) -> OwnedBackendReadiness { + let version = health + .and_then(|h| h.version.as_deref()) + .filter(|v| !v.is_empty()); + match crate::preflight::backend_version_stale_reason(version) { Some(reason) => OwnedBackendReadiness::Stale { reason }, None => OwnedBackendReadiness::Ready, } } -async fn fetch_health(port: u16) -> Result, reqwest::Error> { +async fn health_ready_status(port: u16) -> OwnedBackendReadiness { + match fetch_health(port).await { + Ok(health) => ready_for_use_status(health.as_ref()), + Err(reason) => OwnedBackendReadiness::Stale { reason }, + } +} + +async fn fetch_liveness(port: u16) -> Result, reqwest::Error> { let client = reqwest::Client::builder() .timeout(LOCAL_HTTP_TIMEOUT) .build()?; + for path in ["/api/liveness", "/api/health"] { + let response = client + .get(format!("http://127.0.0.1:{port}{path}")) + .send() + .await?; + if response.status() == reqwest::StatusCode::NOT_FOUND && path == "/api/liveness" { + continue; + } + if !response.status().is_success() { + return Ok(None); + } + return response.json::().await.map(Some); + } + Ok(None) +} + +fn fetch_liveness_blocking(port: u16) -> Result, String> { + for path in ["/api/liveness", "/api/health"] { + let response = http_request_blocking(port, "GET", path, &[], &[])?; + if response.status == 404 && path == "/api/liveness" { + continue; + } + if !(200..300).contains(&response.status) { + return Ok(None); + } + return serde_json::from_slice::(&response.body) + .map(Some) + .map_err(|e| e.to_string()); + } + Ok(None) +} +async fn fetch_health(port: u16) -> Result, String> { + let client = reqwest::Client::builder() + .timeout(LOCAL_HTTP_TIMEOUT) + .build() + .map_err(|e| e.to_string())?; let response = client .get(format!("http://127.0.0.1:{port}/api/health")) .send() - .await?; + .await + .map_err(|e| e.to_string())?; if !response.status().is_success() { return Ok(None); } - response.json::().await.map(Some) -} - -fn fetch_health_blocking(port: u16) -> Result, String> { - let response = http_request_blocking(port, "GET", "/api/health", &[], &[])?; - if !(200..300).contains(&response.status) { - return Ok(None); - } - serde_json::from_slice::(&response.body) + response + .json::() + .await .map(Some) .map_err(|e| e.to_string()) } @@ -618,21 +666,21 @@ pub(crate) async fn probe_owned_backend_state( }; let mut verified = Vec::new(); for port in ports { - let health = match fetch_health(port).await { - Ok(Some(health)) => health, + let liveness = match fetch_liveness(port).await { + Ok(Some(liveness)) => liveness, Ok(None) => continue, Err(error) => { warn!( - "Desktop-owned backend probe skipped port {} after health error: {}", + "Desktop-owned backend probe skipped port {} after liveness error: {}", port, error ); continue; } }; - if !health_verifies_metadata(&health, &owner.metadata) { + if !liveness_verifies_metadata(&liveness, &owner.metadata) { continue; } - if let Some(reason) = lifecycle_control_block_reason(&health) { + if let Some(reason) = lifecycle_control_block_reason(&liveness) { return OwnedBackendProbe::Unmanageable { port, reason }; } if !desktop_login_route_compatible(port).await { @@ -646,7 +694,7 @@ pub(crate) async fn probe_owned_backend_state( return OwnedBackendProbe::Unmanageable { port, reason }; } } - verified.push((port, ready_for_use_status(&health))); + verified.push((port, health_ready_status(port).await)); } if verified.len() != 1 { @@ -967,12 +1015,11 @@ mod tests { } #[test] - fn health_verification_requires_root_kind_and_token_sha() { + fn liveness_verification_requires_root_kind_and_token_sha() { let metadata = metadata(1, Some(8888)); - let health = HealthResponse { - status: Some("healthy".to_string()), + let liveness = DesktopLiveness { + status: Some("alive".to_string()), service: Some("Unsloth UI Backend".to_string()), - version: Some("2026.5.2".to_string()), desktop_protocol_version: Some(1), desktop_manageability_version: Some(1), supports_desktop_auth: Some(true), @@ -983,12 +1030,12 @@ mod tests { token_sha256: Some(token_sha256(TOKEN)), }), }; - assert!(health_verifies_metadata(&health, &metadata)); + assert!(liveness_verifies_metadata(&liveness, &metadata)); - let mut wrong_root = health; + let mut wrong_root = liveness; wrong_root.studio_root_id = Some("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string()); - assert!(!health_verifies_metadata(&wrong_root, &metadata)); + assert!(!liveness_verifies_metadata(&wrong_root, &metadata)); } #[tokio::test] diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 9138f30a58..5a48d26632 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -193,6 +193,9 @@ pub async fn desktop_preflight_result_with_state( if let Some(snapshot) = crate::process::owned_backend_snapshot(state)? { let Some(owner) = snapshot.owner.clone() else { + // TAURI_PORT is emitted only after uvicorn lifespan completes; keep + // this ownerless path on full health so auth/bootstrap are ready. + let probe = match snapshot.port { Some(port) => backend::probe_ownerless_spawned_backend(port).await, None => backend, @@ -494,9 +497,22 @@ mod tests { FakeCli { bin, dir } } + #[cfg(unix)] + fn remove_managed_capability_cache() { + let _ = std::fs::remove_file( + dirs::home_dir() + .unwrap() + .join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json"), + ); + } + #[cfg(unix)] #[tokio::test] async fn managed_cli_capability_probe_classifies_core_cases() { + remove_managed_capability_cache(); + for (name, script, stale_reason) in [ ( "cap-missing", diff --git a/studio/src-tauri/src/preflight/backend.rs b/studio/src-tauri/src/preflight/backend.rs index 0277ef5149..5a58142667 100644 --- a/studio/src-tauri/src/preflight/backend.rs +++ b/studio/src-tauri/src/preflight/backend.rs @@ -3,7 +3,10 @@ use super::version::{ backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION, }; use serde::{Deserialize, Serialize}; + +use log::info; use std::time::Duration; +use std::time::Instant; #[derive(Debug, Deserialize)] struct DesktopOwnerHealth { @@ -24,6 +27,7 @@ pub(super) struct BackendHealth { } pub(super) async fn backend_health(client: &reqwest::Client, port: u16) -> Option { + let started = Instant::now(); let url = format!("http://127.0.0.1:{port}/api/health"); let response = client.get(url).send().await.ok()?; if !response.status().is_success() { @@ -40,6 +44,14 @@ pub(super) async fn backend_health(client: &reqwest::Client, port: u16) -> Optio .and_then(|v| v.as_str()) .map(|s| s == "Unsloth UI Backend") .unwrap_or(false); + info!( + "Desktop preflight: health probe on port {} healthy={} service={} in {}ms", + port, + healthy, + service, + started.elapsed().as_millis() + ); + if !healthy || !service { return None; } diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 51fde51896..57b8365ec5 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -2,14 +2,30 @@ use super::types::ManagedProbe; use super::version::{ backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION, }; -use serde::Deserialize; +use log::{info, warn}; +use serde::{Deserialize, Serialize}; +use std::fs; use std::path::{Path, PathBuf}; use std::process::Stdio; -use std::time::Duration; +use std::time::{Duration, Instant, UNIX_EPOCH}; use tokio::io::AsyncReadExt; use tokio::process::Command; -#[derive(Debug, Deserialize)] +const MANAGED_CAPABILITY_CACHE_SCHEMA: u16 = 2; + +const FNV64_OFFSET_BASIS: u64 = 0xcbf29ce484222325; +const FNV64_PRIME: u64 = 0x100000001b3; +const HASHED_MARKER_MAX_BYTES: u64 = 64 * 1024; + +const FALLBACK_MARKER_NAMES: &[&str] = &[ + "pyvenv.cfg", + "uv.lock", + "requirements.txt", + "python.exe", + "python", +]; + +#[derive(Debug, Clone, Deserialize, Serialize)] struct DesktopCapability { desktop_protocol_version: Option, desktop_manageability_version: Option, @@ -20,7 +36,225 @@ struct DesktopCapability { version: Option, } +#[derive(Debug, Clone, Deserialize, Serialize)] +struct ManagedCapabilityCache { + schema: u16, + bin_path: String, + bin_size: u64, + bin_mtime_ms: u64, + studio_root_id: Option, + marker_path: Option, + marker_size: Option, + marker_mtime_ms: Option, + desktop_protocol_version: u16, + desktop_manageability_version: u16, + capability: DesktopCapability, +} + +#[derive(Debug, Clone)] +struct MarkerFingerprint { + path: String, + size: u64, + mtime_ms: u64, + content_hash: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ManagedBinFingerprint { + bin_path: String, + bin_size: u64, + bin_mtime_ms: u64, + studio_root_id: Option, + marker_path: Option, + marker_size: Option, + marker_mtime_ms: Option, +} + +fn modified_ms(metadata: &fs::Metadata) -> Option { + metadata + .modified() + .ok()? + .duration_since(UNIX_EPOCH) + .ok() + .and_then(|duration| u64::try_from(duration.as_millis()).ok()) +} +fn hash_bytes(hash: u64, bytes: &[u8]) -> u64 { + bytes.iter().fold(hash, |mut next, byte| { + next ^= u64::from(*byte); + next.wrapping_mul(FNV64_PRIME) + }) +} + +fn marker_content_hash(path: &Path, metadata: &fs::Metadata) -> Option { + if metadata.len() > HASHED_MARKER_MAX_BYTES { + return None; + } + fs::read(path) + .ok() + .map(|bytes| hash_bytes(FNV64_OFFSET_BASIS, &bytes)) +} + +fn marker_candidates_for_bin(bin: &Path) -> Vec { + let Some(scripts_dir) = bin.parent() else { + return Vec::new(); + }; + let Some(venv_dir) = scripts_dir.parent() else { + return Vec::new(); + }; + let mut out = Vec::new(); + + #[cfg(unix)] + { + if let Ok(lib_dir) = fs::read_dir(venv_dir.join("lib")) { + for entry in lib_dir.flatten() { + out.push( + entry + .path() + .join("site-packages") + .join("unsloth_cli") + .join("commands") + .join("studio.py"), + ); + } + } + } + for marker_name in FALLBACK_MARKER_NAMES { + out.push(venv_dir.join(marker_name)); + out.push(scripts_dir.join(marker_name)); + } + + out.push( + venv_dir + .join("Lib") + .join("site-packages") + .join("unsloth_cli") + .join("commands") + .join("studio.py"), + ); + out +} + +fn managed_bin_fingerprint(bin: &Path) -> Option { + let bin_metadata = fs::metadata(bin).ok()?; + let bin_path = bin + .canonicalize() + .unwrap_or_else(|_| bin.to_path_buf()) + .to_string_lossy() + .into_owned(); + + let studio_root_id = crate::desktop_backend_owner::read_expected_studio_root_id(); + let mut marker_entries: Vec = marker_candidates_for_bin(bin) + .into_iter() + .filter_map(|path| { + let metadata = fs::metadata(&path).ok()?; + Some(MarkerFingerprint { + path: path + .canonicalize() + .unwrap_or(path.clone()) + .to_string_lossy() + .into_owned(), + size: metadata.len(), + mtime_ms: modified_ms(&metadata)?, + content_hash: marker_content_hash(&path, &metadata), + }) + }) + .collect(); + marker_entries.sort_by(|left, right| left.path.cmp(&right.path)); + let marker_hash = marker_entries + .iter() + .fold(FNV64_OFFSET_BASIS, |hash, marker| { + let next = hash_bytes(hash, marker.path.as_bytes()); + let next = hash_bytes(next, &marker.size.to_le_bytes()); + let next = hash_bytes(next, &marker.mtime_ms.to_le_bytes()); + if let Some(content_hash) = marker.content_hash { + hash_bytes(next, &content_hash.to_le_bytes()) + } else { + next + } + }); + let marker_path = (!marker_entries.is_empty()).then(|| "markers".to_string()); + let marker_size = (!marker_entries.is_empty()).then(|| marker_entries.len() as u64); + let marker_mtime_ms = (!marker_entries.is_empty()).then_some(marker_hash); + + Some(ManagedBinFingerprint { + bin_path, + bin_size: bin_metadata.len(), + bin_mtime_ms: modified_ms(&bin_metadata)?, + studio_root_id, + marker_path, + marker_size, + marker_mtime_ms, + }) +} + +fn capability_cache_path() -> Option { + dirs::home_dir().map(|home| { + home.join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json") + }) +} + +fn cache_matches(cache: &ManagedCapabilityCache, fingerprint: &ManagedBinFingerprint) -> bool { + cache.schema == MANAGED_CAPABILITY_CACHE_SCHEMA + && cache.desktop_protocol_version == DESKTOP_PROTOCOL_VERSION + && cache.desktop_manageability_version == DESKTOP_MANAGEABILITY_VERSION + && cache.bin_path == fingerprint.bin_path + && cache.bin_size == fingerprint.bin_size + && cache.bin_mtime_ms == fingerprint.bin_mtime_ms + && cache.studio_root_id == fingerprint.studio_root_id + && cache.marker_path == fingerprint.marker_path + && cache.marker_size == fingerprint.marker_size + && cache.marker_mtime_ms == fingerprint.marker_mtime_ms + && desktop_capability_ready(&cache.capability) +} + +fn read_cached_capability(fingerprint: &ManagedBinFingerprint) -> Option { + let path = capability_cache_path()?; + let bytes = fs::read(path).ok()?; + let cache = serde_json::from_slice::(&bytes).ok()?; + if cache_matches(&cache, fingerprint) { + Some(cache.capability) + } else { + None + } +} + +fn write_cached_capability(fingerprint: &ManagedBinFingerprint, capability: &DesktopCapability) { + let Some(path) = capability_cache_path() else { + return; + }; + let cache = ManagedCapabilityCache { + schema: MANAGED_CAPABILITY_CACHE_SCHEMA, + bin_path: fingerprint.bin_path.clone(), + bin_size: fingerprint.bin_size, + bin_mtime_ms: fingerprint.bin_mtime_ms, + studio_root_id: fingerprint.studio_root_id.clone(), + marker_path: fingerprint.marker_path.clone(), + marker_size: fingerprint.marker_size, + marker_mtime_ms: fingerprint.marker_mtime_ms, + desktop_protocol_version: DESKTOP_PROTOCOL_VERSION, + desktop_manageability_version: DESKTOP_MANAGEABILITY_VERSION, + capability: capability.clone(), + }; + if let Some(parent) = path.parent() { + if fs::create_dir_all(parent).is_err() { + return; + } + } + let Ok(bytes) = serde_json::to_vec_pretty(&cache) else { + return; + }; + if let Err(error) = fs::write(&path, bytes) { + warn!( + "Managed preflight: could not write capability cache: {}", + error + ); + } +} + async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool { + let started = Instant::now(); let mut cmd = Command::new(bin); cmd.args(args).stdout(Stdio::null()).stderr(Stdio::null()); @@ -43,20 +277,33 @@ async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool { } let Ok(mut child) = cmd.spawn() else { + info!( + "Managed preflight probe {:?} failed to spawn in {}ms", + args, + started.elapsed().as_millis() + ); return false; }; - match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { + let ok = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await { Ok(Ok(status)) => status.success(), _ => { let _ = child.kill().await; let _ = child.wait().await; false } - } + }; + info!( + "Managed preflight probe {:?} finished ok={} in {}ms", + args, + ok, + started.elapsed().as_millis() + ); + ok } async fn probe_cli_capability(bin: &Path) -> Option { + let started = Instant::now(); let mut cmd = Command::new(bin); cmd.args(["studio", "desktop-capabilities", "--json"]) .stdout(Stdio::piped()) @@ -81,6 +328,10 @@ async fn probe_cli_capability(bin: &Path) -> Option { } let Ok(mut child) = cmd.spawn() else { + info!( + "Managed desktop-capabilities probe failed to spawn in {}ms", + started.elapsed().as_millis() + ); return None; }; let Some(mut stdout) = child.stdout.take() else { @@ -92,9 +343,19 @@ async fn probe_cli_capability(bin: &Path) -> Option { Err(_) => { let _ = child.kill().await; let _ = child.wait().await; + info!( + "Managed desktop-capabilities probe timed out in {}ms", + started.elapsed().as_millis() + ); + return None; + } + _ => { + info!( + "Managed desktop-capabilities probe exited unsuccessfully in {}ms", + started.elapsed().as_millis() + ); return None; } - _ => return None, } let mut output = Vec::new(); @@ -102,7 +363,13 @@ async fn probe_cli_capability(bin: &Path) -> Option { return None; } - serde_json::from_slice::(&output).ok() + let capability = serde_json::from_slice::(&output).ok(); + info!( + "Managed desktop-capabilities probe finished ok={} in {}ms", + capability.is_some(), + started.elapsed().as_millis() + ); + capability } fn desktop_capability_stale_reason(capability: &DesktopCapability) -> Option { @@ -132,18 +399,48 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool { } pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { + let started = Instant::now(); if !run_cli_probe(&bin, &["-h"]).await { + info!( + "Managed preflight: cli unusable for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); return ManagedProbe::Stale { bin, reason: "cli_unusable".to_string(), }; } - let capability = probe_cli_capability(&bin).await; - if let Some(capability) = capability { - if desktop_capability_ready(&capability) { + if let Some(fingerprint) = managed_bin_fingerprint(&bin) { + if read_cached_capability(&fingerprint).is_some() { + info!( + "Managed preflight: using cached desktop capability for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); return ManagedProbe::Ready { bin }; } + } + + let capability = probe_cli_capability(&bin).await; + if let Some(capability) = capability { + if let Some(fingerprint) = managed_bin_fingerprint(&bin) { + write_cached_capability(&fingerprint, &capability); + } + if desktop_capability_ready(&capability) { + info!( + "Managed preflight: cli ready for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); + return ManagedProbe::Ready { bin }; + } + info!( + "Managed preflight: cli stale for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); return ManagedProbe::Stale { bin, reason: desktop_capability_stale_reason(&capability) @@ -151,6 +448,11 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { }; } + info!( + "Managed preflight: desktop capability probe failed for {:?} in {}ms", + bin, + started.elapsed().as_millis() + ); ManagedProbe::Stale { bin, reason: "desktop_capability_probe_failed".to_string(), @@ -158,10 +460,17 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { } pub(super) async fn probe_managed_install() -> ManagedProbe { - match crate::process::find_unsloth_binary() { + let started = Instant::now(); + let result = match crate::process::find_unsloth_binary() { Some(bin) => probe_managed_bin(bin).await, None => ManagedProbe::Missing, - } + }; + info!( + "Managed preflight: install probe result {:?} in {}ms", + result, + started.elapsed().as_millis() + ); + result } pub async fn managed_install_ready() -> bool { diff --git a/studio/src-tauri/src/process.rs b/studio/src-tauri/src/process.rs index 01ef77f098..56d9dd2e21 100644 --- a/studio/src-tauri/src/process.rs +++ b/studio/src-tauri/src/process.rs @@ -735,6 +735,7 @@ pub fn start_backend( } async fn generic_backend_health_ok(port: u16) -> bool { + let started = std::time::Instant::now(); let client = match reqwest::Client::builder() .timeout(Duration::from_secs(2)) .build() @@ -745,49 +746,75 @@ async fn generic_backend_health_ok(port: u16) -> bool { return false; } }; - let response = match client - .get(format!("http://127.0.0.1:{port}/api/health")) - .send() - .await - { - Ok(response) => response, - Err(error) => { + let mut last_status = None; + let mut json = None; + for path in ["/api/liveness", "/api/health"] { + let response = match client + .get(format!("http://127.0.0.1:{port}{path}")) + .send() + .await + { + Ok(response) => response, + Err(error) => { + warn!( + "Backend port candidate {} failed health request: {}", + port, error + ); + return false; + } + }; + if response.status() == reqwest::StatusCode::NOT_FOUND && path == "/api/liveness" { + last_status = Some(response.status()); + continue; + } + if !response.status().is_success() { warn!( - "Backend port candidate {} failed health request: {}", - port, error + "Backend port candidate {} returned HTTP {} from health", + port, + response.status() ); return false; } - }; - if !response.status().is_success() { + json = match response.json::().await { + Ok(json) => Some(json), + Err(error) => { + warn!( + "Backend port candidate {} returned invalid health JSON: {}", + port, error + ); + return false; + } + }; + break; + } + let Some(json) = json else { warn!( "Backend port candidate {} returned HTTP {} from health", port, - response.status() + last_status + .map(|status| status.to_string()) + .unwrap_or_else(|| "unknown".to_string()) ); return false; - } - let json = match response.json::().await { - Ok(json) => json, - Err(error) => { - warn!( - "Backend port candidate {} returned invalid health JSON: {}", - port, error - ); - return false; - } }; - let healthy = json + let live = json .get("status") .and_then(|v| v.as_str()) - .map(|s| s == "healthy") + .map(|s| s == "alive" || s == "healthy") .unwrap_or(false); let service = json .get("service") .and_then(|v| v.as_str()) .map(|s| s == "Unsloth UI Backend") .unwrap_or(false); - healthy && service + info!( + "Backend port candidate {} liveness live={} service={} in {}ms", + port, + live, + service, + started.elapsed().as_millis() + ); + live && service } async fn validate_candidate_port( @@ -798,6 +825,7 @@ async fn validate_candidate_port( generation: u64, port: u16, ) { + let started = std::time::Instant::now(); let owner = { let proc = match state.lock() { Ok(proc) => proc, @@ -852,6 +880,14 @@ async fn validate_candidate_port( } }; + info!( + "Validated backend port candidate {} valid={} emit={} in {}ms", + port, + valid, + should_emit, + started.elapsed().as_millis() + ); + if should_emit { diagnostics::record_backend_port(&diagnostics_state, &session_id, port); info!("Validated backend port: {}", port); From 07578eab60f7f9e87ab1d1d6e30510017790cd6a Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Mon, 29 Jun 2026 14:50:55 +0100 Subject: [PATCH 26/49] Fix on-device locations dialog layout (#6743) * Fix on-device location path overflow * Remove redundant native path tooltip --- .../hub/catalog/on-device-folders-dialog.tsx | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx b/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx index 20a4eca4dd..caa89db196 100644 --- a/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx +++ b/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx @@ -180,7 +180,7 @@ export function OnDeviceFoldersDialog({ <> @@ -319,7 +319,12 @@ export function OnDeviceFoldersDialog({ return (
-
-

+

+

{pathTail(folder.path)}

-

+

{folder.path}

From f80e66ea34bee26776dfe398f577d2b2879a1bbf Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Mon, 29 Jun 2026 15:03:57 +0100 Subject: [PATCH 27/49] studio: keep chat header below dialogs (#6745) --- studio/frontend/src/features/chat/chat-page.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index dfd96a1e86..2511fc9eec 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -654,7 +654,7 @@ function GeneralCompareHeader({ return (
Date: Mon, 29 Jun 2026 07:06:36 -0700 Subject: [PATCH 28/49] (feat) Add project names to studio training runs (#6512) * (feat) Add project names to studio training runs to avoid models being overwritten when doing similar training runs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update studio/frontend/src/features/export/export-page.tsx Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/frontend/src/features/export/export-page.tsx Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update studio/frontend/src/features/export/export-page.tsx Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * better project name sanitization, removed duplicated project name normalization * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * implement checkpoint scanning utilities and tests for base model inference * [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 * Guard project_name against null and use leading important modifiers * Fix/adjust training project names for PR #6512 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix/adjust training project names for PR #6512 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address project-name review feedback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Show project names in training recents * Keep GGUF export directories source-specific --------- Co-authored-by: NZ-Linix Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: NZ-Linix Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: wasimysaid --- studio/backend/core/training/training.py | 1 + studio/backend/core/training/worker.py | 18 +- studio/backend/models/training.py | 13 + studio/backend/routes/training.py | 1 + studio/backend/storage/studio_db.py | 12 + studio/backend/tests/test_checkpoints_scan.py | 256 ++++++++++++++++++ studio/backend/tests/test_training_runs.py | 103 +++++++ .../backend/tests/test_training_streaming.py | 10 + studio/backend/utils/models/checkpoints.py | 101 ++++++- studio/backend/utils/training_runs.py | 104 +++++++ .../frontend/src/components/app-sidebar.tsx | 8 +- .../src/features/export/export-page.tsx | 28 +- .../components/steps/hyperparameters-step.tsx | 25 ++ .../components/steps/summary-step.tsx | 4 + .../studio/historical-training-view.tsx | 1 + .../src/features/studio/history-card-grid.tsx | 26 +- .../features/studio/live-training-view.tsx | 11 +- .../studio/sections/params-section.tsx | 18 ++ .../studio/sections/progress-section.tsx | 19 +- .../src/features/training/api/mappers.ts | 1 + .../training/hooks/use-training-actions.ts | 17 +- .../frontend/src/features/training/index.ts | 5 + .../src/features/training/lib/run-display.ts | 24 ++ .../training/stores/training-config-store.ts | 2 + .../training/stores/training-runtime-store.ts | 9 +- .../src/features/training/types/api.ts | 1 + .../src/features/training/types/config.ts | 2 + .../src/features/training/types/history.ts | 1 + .../src/features/training/types/runtime.ts | 3 + studio/frontend/src/i18n/locales/en.ts | 5 + studio/frontend/src/i18n/locales/zh-CN.ts | 4 + 31 files changed, 804 insertions(+), 29 deletions(-) create mode 100644 studio/backend/tests/test_checkpoints_scan.py create mode 100644 studio/backend/tests/test_training_runs.py create mode 100644 studio/backend/utils/training_runs.py create mode 100644 studio/frontend/src/features/training/lib/run-display.ts diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 9d991c6512..f4233fcf04 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -299,6 +299,7 @@ class TrainingBackend: # 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), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 3f020c8abc..610af2472e 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -44,6 +44,7 @@ if sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.env logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +from utils.training_runs import build_default_output_dir_name from utils.wheel_utils import ( direct_wheel_url, flash_attn_wheel_url, @@ -1787,11 +1788,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, default_run_dir_name + from utils.paths import resolve_output_dir, ensure_dir output_dir = config.get("output_dir", "") if not output_dir: - output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + output_dir = build_default_output_dir_name( + model_name, + config.get("project_name"), + ) output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) @@ -3019,7 +3023,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> resume_from_checkpoint ) if not output_dir: - output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + output_dir = build_default_output_dir_name( + model_name, + config.get("project_name"), + ) output_dir = str(resolve_output_dir(output_dir)) ensure_dir(Path(output_dir)) @@ -3500,7 +3507,10 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> resume_from_checkpoint ) if not output_dir: - output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + output_dir = build_default_output_dir_name( + model_name, + config.get("project_name"), + ) output_dir = str(resolve_output_dir(output_dir)) num_epochs = config.get("num_epochs", 2) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index e64b6f731a..ff815a2fa9 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -9,6 +9,8 @@ import re from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing import Any, Optional, List, Dict, Literal +from utils.training_runs import normalize_project_name + # ASCII integer, optional single sign. Rejects "++512" and Unicode digits # ("５１２") that slip through str.isdigit() + int(). @@ -97,6 +99,11 @@ class TrainingStartRequest(BaseModel): model_name: str = Field( ..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')" ) + project_name: Optional[str] = Field( + None, + max_length = 80, + description = "Optional user-defined project name appended to run folders and shown in history", + ) training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = Field( ..., description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'", @@ -155,6 +162,11 @@ class TrainingStartRequest(BaseModel): values.setdefault("train_split", values.pop("split")) return values + @field_validator("project_name") + @classmethod + def _normalize_project_name(cls, value: Optional[str]) -> Optional[str]: + return normalize_project_name(value) + # NOTE: pydantic runs all `mode="after"` validators in definition order. A # second one, `_check_steps_or_epochs`, is defined lower in this class; keep # these cross-field checks order-independent so the two stay decoupled. @@ -588,6 +600,7 @@ class TrainingRunSummary(BaseModel): id: str status: Literal["running", "completed", "stopped", "error"] model_name: str + project_name: Optional[str] = None dataset_name: str display_name: Optional[str] = None started_at: str diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 3818fe9f73..4f131ad2f2 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -255,6 +255,7 @@ async def start_training( # Convert request to backend kwargs. training_kwargs = { "model_name": request.model_name, + "project_name": request.project_name, "training_type": request.training_type, "hf_token": request.hf_token or "", "load_in_4bit": request.load_in_4bit, diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 7421b42b2f..23b90d7002 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -23,6 +23,16 @@ from typing import Any, Iterable, Optional from utils.paths import project_workspaces_root, studio_db_path, ensure_dir +from utils.training_runs import extract_project_name + + +def _extract_project_name_from_config_json(config_json: Optional[str]) -> Optional[str]: + if not config_json: + return None + try: + return extract_project_name(json.loads(config_json)) + except (json.JSONDecodeError, TypeError): + return None def _denied_path_prefixes() -> list[str]: @@ -680,6 +690,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: runs = [] for row in rows: run = dict(row) + run["project_name"] = _extract_project_name_from_config_json(run.get("config_json")) sparkline = run.get("loss_sparkline") if sparkline: try: @@ -719,6 +730,7 @@ def get_run(id: str) -> Optional[dict]: if row is None: return None run = dict(row) + run["project_name"] = _extract_project_name_from_config_json(run.get("config_json")) sparkline = run.get("loss_sparkline") if sparkline: try: diff --git a/studio/backend/tests/test_checkpoints_scan.py b/studio/backend/tests/test_checkpoints_scan.py new file mode 100644 index 0000000000..6d473146f5 --- /dev/null +++ b/studio/backend/tests/test_checkpoints_scan.py @@ -0,0 +1,256 @@ +# 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 json +import sqlite3 +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +sys.modules.setdefault("structlog", _types.ModuleType("structlog")) + +from utils.models import checkpoints as checkpoints_module +from utils.training_runs import build_default_output_dir_name + + +def _make_history_connection(db_path: Path) -> sqlite3.Connection: + conn = sqlite3.connect(str(db_path)) + conn.row_factory = sqlite3.Row + return conn + + +def _setup_training_runs_table(db_path: Path) -> None: + conn = _make_history_connection(db_path) + try: + conn.execute( + """ + CREATE TABLE training_runs ( + id TEXT PRIMARY KEY, + model_name TEXT NOT NULL, + config_json TEXT NOT NULL, + output_dir TEXT, + started_at TEXT NOT NULL + ) + """ + ) + conn.commit() + finally: + conn.close() + + +def _make_outputs_dir(tmp_path, monkeypatch) -> Path: + studio_home = tmp_path / "studio-home" + outputs_dir = studio_home / "outputs" + outputs_dir.mkdir(parents = True) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + return outputs_dir + + +def test_scan_checkpoints_uses_output_dir_history_for_base_model(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_dir = outputs_dir / "custom-run" + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + conn = _make_history_connection(db_path) + try: + conn.execute( + """ + INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "run-1", + "unsloth/Llama-3.2-3B-Instruct", + "{}", + str(run_dir.resolve()), + "2026-04-09T00:00:00Z", + ), + ) + conn.commit() + finally: + conn.close() + + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" + + +def test_scan_checkpoints_matches_project_suffixed_default_dir_against_history( + tmp_path, monkeypatch +): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_name = build_default_output_dir_name( + "unsloth/Llama-3.2-3B-Instruct", + "Customer Support", + timestamp = 1771227800, + ) + run_dir = outputs_dir / run_name + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + conn = _make_history_connection(db_path) + try: + conn.execute( + """ + INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "run-2", + "unsloth/Llama-3.2-3B-Instruct", + json.dumps({"project_name": "Customer Support"}), + None, + "2026-04-09T00:00:00Z", + ), + ) + conn.commit() + finally: + conn.close() + + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" + + +def test_scan_checkpoints_strips_project_suffix_without_history(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_name = build_default_output_dir_name( + "unsloth/Llama-3.2-3B-Instruct", + "Customer Support", + timestamp = 1771227800, + ) + run_dir = outputs_dir / run_name + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" + + +def test_scan_checkpoints_preserves_project_marker_in_model_without_history(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_name = build_default_output_dir_name( + "org/foo__project-bar", + timestamp = 1771227800, + ) + run_dir = outputs_dir / run_name + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "org/foo__project-bar" + + +def test_scan_checkpoints_preserves_legacy_folder_name_fallback(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_dir = outputs_dir / "unsloth_Llama-3.2-3B-Instruct_1771227800" + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct" + + +def test_scan_checkpoints_prefers_exact_history_match_over_newer_suffix(tmp_path, monkeypatch): + outputs_dir = _make_outputs_dir(tmp_path, monkeypatch) + run_dir = outputs_dir / "unsloth_Test_1771227800" + run_dir.mkdir() + (run_dir / "config.json").write_text("{}") + + copied_dir = tmp_path / "copied" / run_dir.name + copied_dir.mkdir(parents = True) + + db_path = tmp_path / "studio.db" + _setup_training_runs_table(db_path) + conn = _make_history_connection(db_path) + try: + conn.execute( + """ + INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "run-exact", + "correct/base", + "{}", + str(run_dir.resolve()), + "2026-04-09T00:00:00Z", + ), + ) + conn.execute( + """ + INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at) + VALUES (?, ?, ?, ?, ?) + """, + ( + "run-suffix", + "wrong/base", + "{}", + str(copied_dir.resolve()), + "2026-04-10T00:00:00Z", + ), + ) + conn.commit() + finally: + conn.close() + + monkeypatch.setattr( + checkpoints_module, + "get_connection", + lambda: _make_history_connection(db_path), + ) + + models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir)) + + assert models[0][2]["base_model"] == "correct/base" diff --git a/studio/backend/tests/test_training_runs.py b/studio/backend/tests/test_training_runs.py new file mode 100644 index 0000000000..fd0d6d380f --- /dev/null +++ b/studio/backend/tests/test_training_runs.py @@ -0,0 +1,103 @@ +# 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 json + +from storage.studio_db import _extract_project_name_from_config_json +from utils.training_runs import ( + build_default_output_dir_name, + model_segment_from_default_output_dir_name, + normalize_project_name, + slugify_project_name, +) + + +def test_normalize_project_name_trims_and_collapses_whitespace(): + assert normalize_project_name(" Customer Support LoRA ") == "Customer Support LoRA" + + +def test_normalize_project_name_returns_none_for_empty_or_invalid_values(): + assert normalize_project_name(" ") is None + assert normalize_project_name(None) is None + + +def test_slugify_project_name_makes_safe_suffix(): + assert slugify_project_name("Customer Support / LoRA v2") == "customer-support-lora-v2" + + +def test_slugify_project_name_rejects_path_only_or_separator_only_values(): + assert slugify_project_name("..") is None + assert slugify_project_name("///") is None + + +def test_build_default_output_dir_name_appends_project_slug(): + output_dir = build_default_output_dir_name( + "unsloth/Llama-3.2-3B-Instruct", + "Customer Support", + timestamp = 1771227800, + ) + + assert output_dir == "unsloth_Llama-3.2-3B-Instruct__project-customer-support_1771227800" + + +def test_build_default_output_dir_name_caps_final_component(tmp_path): + output_dir = build_default_output_dir_name( + "a" * 240, + "b" * 80, + timestamp = 1771227800, + ) + + assert len(output_dir.encode()) <= 255 + (tmp_path / output_dir).mkdir() + + +def test_build_default_output_dir_name_skips_invalid_project_slug(): + output_dir = build_default_output_dir_name( + "unsloth/Llama-3.2-3B-Instruct", + "..", + timestamp = 1771227800, + ) + + assert output_dir == "unsloth_Llama-3.2-3B-Instruct_1771227800" + + +def test_model_segment_from_default_output_dir_name_strips_project_slug(): + assert ( + model_segment_from_default_output_dir_name( + "unsloth_Llama-3.2-3B-Instruct__project-customer-support_1771227800" + ) + == "unsloth_Llama-3.2-3B-Instruct" + ) + + +def test_model_segment_preserves_project_marker_text_in_model_name(): + output_dir = build_default_output_dir_name( + "org/foo__project-bar", + timestamp = 1771227800, + ) + + assert output_dir == "org_foo__project--bar_1771227800" + assert model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar" + + +def test_model_segment_strips_project_slug_after_escaped_model_marker(): + output_dir = build_default_output_dir_name( + "org/foo__project-bar", + "Customer Support", + timestamp = 1771227800, + ) + + assert output_dir == "org_foo__project--bar__project-customer-support_1771227800" + assert model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar" + + +def test_extract_project_name_from_config_json_returns_normalized_name(): + config_json = json.dumps({"project_name": " Sales Assistant "}) + + assert _extract_project_name_from_config_json(config_json) == "Sales Assistant" + + +def test_extract_project_name_from_config_json_handles_missing_or_invalid_payload(): + assert _extract_project_name_from_config_json(None) is None + assert _extract_project_name_from_config_json("not-json") is None + assert _extract_project_name_from_config_json(json.dumps({"project_name": " "})) is None diff --git a/studio/backend/tests/test_training_streaming.py b/studio/backend/tests/test_training_streaming.py index 8ff016d3bf..70b2d6fdcc 100644 --- a/studio/backend/tests/test_training_streaming.py +++ b/studio/backend/tests/test_training_streaming.py @@ -195,6 +195,16 @@ def test_hf_dataset_rejects_unsafe_values(bad_hf_dataset): ) +def test_project_name_rejects_values_over_ui_limit(): + with pytest.raises(ValidationError): + TrainingStartRequest( + model_name = "unsloth/test", + project_name = "x" * 81, + training_type = "LoRA/QLoRA", + format_type = "alpaca", + ) + + # --- Start-route streaming compatibility guards --- diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index d174f6677b..90e26d45d0 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -9,6 +9,12 @@ import structlog from loggers import get_logger from pathlib import Path from typing import List, Optional, Tuple +from storage.studio_db import get_connection +from utils.training_runs import ( + build_default_output_dir_name, + extract_project_name, + model_segment_from_default_output_dir_name, +) from utils.paths import outputs_root, resolve_output_dir logger = get_logger(__name__) @@ -30,6 +36,93 @@ def _checkpoint_sort_key(checkpoint_path: Path) -> tuple[int, int, str]: return (1, 0, str(checkpoint_path)) +def _infer_base_model_from_history(checkpoint_dir: Path) -> Optional[str]: + """Best-effort base-model lookup using persisted Studio run metadata.""" + checkpoint_name = checkpoint_dir.name + resolved_checkpoint_dir = str(checkpoint_dir.resolve()) + + try: + conn = get_connection() + except Exception: + return None + + try: + exact_rows = conn.execute( + """ + SELECT model_name + FROM training_runs + WHERE output_dir IN (?, ?) + ORDER BY started_at DESC + """, + ( + resolved_checkpoint_dir, + str(checkpoint_dir), + ), + ).fetchall() + for row in exact_rows: + model_name = row["model_name"] + if model_name: + return model_name + + suffix_rows = conn.execute( + """ + SELECT model_name, output_dir + FROM training_runs + WHERE output_dir IS NOT NULL + ORDER BY started_at DESC + """ + ).fetchall() + for row in suffix_rows: + output_dir = str(row["output_dir"] or "").rstrip("/\\") + if not ( + output_dir.endswith(f"/{checkpoint_name}") + or output_dir.endswith(f"\\{checkpoint_name}") + ): + continue + model_name = row["model_name"] + if model_name: + return model_name + + parts = checkpoint_name.rsplit("_", 1) + if len(parts) != 2 or not parts[1].isdigit(): + return None + + timestamp = int(parts[1]) + generated_rows = conn.execute( + """ + SELECT model_name, config_json + FROM training_runs + ORDER BY started_at DESC + """ + ).fetchall() + for row in generated_rows: + model_name = row["model_name"] + if not model_name: + continue + + project_name = None + config_json = row["config_json"] + if config_json: + try: + project_name = extract_project_name(json.loads(config_json)) + except (TypeError, json.JSONDecodeError): + project_name = None + + expected_dir_name = build_default_output_dir_name( + model_name, + project_name, + timestamp = timestamp, + ) + if expected_dir_name == checkpoint_name: + return model_name + except Exception: + return None + finally: + conn.close() + + return None + + def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]: """Read loss from the last log_history entry of trainer_state.json, or None.""" trainer_state = checkpoint_path / "trainer_state.json" @@ -106,9 +199,11 @@ def scan_checkpoints( # Fallback: extract base model name from the folder name, e.g. # "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct" if not metadata.get("base_model"): - parts = item.name.rsplit("_", 1) - if len(parts) == 2 and parts[1].isdigit(): - name_part = parts[0] + metadata["base_model"] = _infer_base_model_from_history(item) + + if not metadata.get("base_model"): + name_part = model_segment_from_default_output_dir_name(item.name) + if name_part: idx = name_part.find("_") if idx > 0: metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1 :] diff --git a/studio/backend/utils/training_runs.py b/studio/backend/utils/training_runs.py new file mode 100644 index 0000000000..dc2535e570 --- /dev/null +++ b/studio/backend/utils/training_runs.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Helpers for naming and describing Studio training runs.""" + +from __future__ import annotations + +import re +import time +from typing import Any, Optional + +_INVALID_SEGMENT_CHARS = re.compile(r"[^A-Za-z0-9._-]+") +_MAX_RUN_DIR_NAME_CHARS = 255 +_PROJECT_MARKER = "__project-" +_PROJECT_MARKER_ESCAPE = f"{_PROJECT_MARKER}-" + + +def _trim_segment(segment: str, max_chars: int) -> str: + if max_chars <= 0: + return "" + return segment[:max_chars].strip("._-") + + +def _escape_project_marker(segment: str) -> str: + return segment.replace(_PROJECT_MARKER, _PROJECT_MARKER_ESCAPE) + + +def _unescape_project_marker(segment: str) -> str: + return segment.replace(_PROJECT_MARKER_ESCAPE, _PROJECT_MARKER) + + +def _appended_project_marker_index(segment: str) -> int: + marker_index = segment.rfind(_PROJECT_MARKER) + while marker_index >= 0 and segment.startswith(_PROJECT_MARKER_ESCAPE, marker_index): + marker_index = segment.rfind(_PROJECT_MARKER, 0, marker_index) + return marker_index + + +def normalize_project_name(project_name: Any) -> Optional[str]: + """Return a trimmed project name, or None when empty/invalid.""" + if not isinstance(project_name, str): + return None + normalized = " ".join(project_name.strip().split()) + return normalized or None + + +def slugify_project_name(project_name: Any) -> Optional[str]: + """Convert a project name into a filesystem-safe suffix.""" + normalized = normalize_project_name(project_name) + if normalized is None: + return None + + slug = _INVALID_SEGMENT_CHARS.sub("-", normalized).strip("-._") + if not slug: + return None + return slug.lower() + + +def build_default_output_dir_name( + model_name: str, + project_name: Any = None, + *, + timestamp: Optional[int] = None, +) -> str: + """Build the default training output folder name.""" + from utils.paths import default_run_dir_name + + timestamp_part = str(int(time.time() if timestamp is None else timestamp)) + timestamp_suffix = f"_{timestamp_part}" + model_segment = _escape_project_marker(default_run_dir_name(model_name)) + project_slug = slugify_project_name(project_name) + if not project_slug: + max_model_chars = _MAX_RUN_DIR_NAME_CHARS - len(timestamp_suffix) + model_segment = _trim_segment(model_segment, max_model_chars) or "model" + return f"{model_segment}{timestamp_suffix}" + + max_project_chars = ( + _MAX_RUN_DIR_NAME_CHARS - len("model") - len(_PROJECT_MARKER) - len(timestamp_suffix) + ) + project_slug = _trim_segment(project_slug, max_project_chars) or "project" + project_suffix = f"{_PROJECT_MARKER}{project_slug}{timestamp_suffix}" + max_model_chars = _MAX_RUN_DIR_NAME_CHARS - len(project_suffix) + model_segment = _trim_segment(model_segment, max_model_chars) or "model" + return f"{model_segment}{project_suffix}" + + +def model_segment_from_default_output_dir_name(output_dir_name: str) -> Optional[str]: + """Return the encoded model segment from a default run folder name.""" + parts = str(output_dir_name or "").rsplit("_", 1) + if len(parts) != 2 or not parts[1].isdigit(): + return None + model_segment = parts[0] + marker_index = _appended_project_marker_index(model_segment) + if marker_index >= 0: + model_segment = model_segment[:marker_index] + model_segment = _unescape_project_marker(model_segment) + return model_segment or None + + +def extract_project_name(config: Any) -> Optional[str]: + """Read and normalize a project name from a stored config dict.""" + if not isinstance(config, dict): + return None + return normalize_project_name(config.get("project_name")) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 06f2701a16..0203605767 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -123,6 +123,7 @@ import { deleteTrainingRun, emitTrainingRunDeleted, emitTrainingRunUpdated, + getTrainingRunDisplayTitle, removeTrainingUnloadGuard, renameTrainingRun, useTrainingCompletionWatch, @@ -592,7 +593,7 @@ export function AppSidebar() { setRenamingTarget({ kind: "chat", item, current: item.title }); } function openRenameRun(run: TrainingRunSummary) { - const current = run.display_name ?? run.model_name; + const current = getTrainingRunDisplayTitle(run); setRenameDraft(current); setRenamingTarget({ kind: "run", run, current }); } @@ -1377,7 +1378,7 @@ export function AppSidebar() { aria-hidden /> - {run.display_name ?? run.model_name} + {getTrainingRunDisplayTitle(run)} {formatRelativeShort(run.started_at)} @@ -1653,8 +1654,7 @@ export function AppSidebar() { renderEmphasizedTranslation( t, "shell.dialog.deleteRun.description", - confirmingDelete.run.display_name ?? - confirmingDelete.run.model_name, + getTrainingRunDisplayTitle(confirmingDelete.run), ) ) : confirmingDelete?.kind === "chat" ? ( renderEmphasizedTranslation( diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 7d38d9b222..4037dbe079 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -75,16 +75,35 @@ import { exportTourSteps } from "./tour"; const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]); type SourceTab = "local" | "checkpoint" | "hf"; +type SourceMode = "checkpoint" | "model"; + + +function safePathSegment( + value: string | null | undefined, + fallback = "model", + maxLength = 250, +): string { + const safe = (value ?? "") + .replace(/[^a-zA-Z0-9._-]/g, "-") + .replace(/^[._-]+|[._-]+$/g, "") + .slice(0, maxLength) + .replace(/[._-]+$/g, ""); + return safe || fallback; +} function buildRelativeSaveDirectory( exportMethod: ExportMethod | null, + sourceMode: SourceMode, sourceBaseModelName: string, selectedModelIdx: string | null, checkpoint: string | null, ): string { if (exportMethod === "gguf") { - return `${(sourceBaseModelName.split("/").pop() ?? selectedModelIdx ?? "model") - .replace(/[^a-zA-Z0-9._-]/g, "-")}-GGUF`; + const rawName = + sourceMode === "checkpoint" + ? checkpoint ?? selectedModelIdx ?? sourceBaseModelName + : sourceBaseModelName; + return `${safePathSegment(rawName)}-GGUF`; } return `${selectedModelIdx ?? "model"}/${checkpoint}`; } @@ -125,9 +144,7 @@ export function ExportPage() { const [selectedModelIdx, setSelectedModelIdx] = useState(null); const [checkpoint, setCheckpoint] = useState(null); - const [sourceMode, setSourceMode] = useState<"checkpoint" | "model">( - "checkpoint", - ); + const [sourceMode, setSourceMode] = useState("checkpoint"); const [modelSource, setModelSource] = useState<"hf" | "local">("hf"); const [modelInput, setModelInput] = useState(""); const [selectedSourceModel, setSelectedSourceModel] = useState( @@ -449,6 +466,7 @@ export function ExportPage() { const defaultSaveDirectory = useMemo(() => { const relative = buildRelativeSaveDirectory( exportMethod, + sourceMode, sourceBaseModelName, selectedModelIdx, checkpoint, diff --git a/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx b/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx index a1174186af..2fb2f7adeb 100644 --- a/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx @@ -7,6 +7,7 @@ import { FieldLegend, FieldSet, } from "@/components/ui/field"; +import { Input } from "@/components/ui/input"; import { Select, SelectContent, @@ -59,6 +60,8 @@ function stepLR(value: number, direction: 1 | -1): number { export function HyperparametersStep() { const { trainingMethod, + projectName, + setProjectName, maxSteps, setMaxSteps, epochs, @@ -79,6 +82,8 @@ export function HyperparametersStep() { } = useTrainingConfigStore( useShallow((s) => ({ trainingMethod: s.trainingMethod, + projectName: s.projectName, + setProjectName: s.setProjectName, maxSteps: s.maxSteps, setMaxSteps: s.setMaxSteps, epochs: s.epochs, @@ -125,6 +130,26 @@ export function HyperparametersStep() {
Choose your training parameters
+
+
+ + Project Name + + Optional + + +
+ setProjectName(e.target.value)} + placeholder="customer-support-lora" + maxLength={80} + /> +

+ Used in training output folder names, export defaults, and history. +

+
+
({ modelType, selectedModel, + projectName, trainingMethod, datasetSource, datasetFormat, @@ -152,6 +155,7 @@ export function SummaryStep() {
+
diff --git a/studio/frontend/src/features/studio/historical-training-view.tsx b/studio/frontend/src/features/studio/historical-training-view.tsx index 0bca090413..2f80fc29ca 100644 --- a/studio/frontend/src/features/studio/historical-training-view.tsx +++ b/studio/frontend/src/features/studio/historical-training-view.tsx @@ -78,6 +78,7 @@ function mapToViewData( error: run.status === "error" ? run.error_message : null, isTrainingRunning: false, modelName: run.display_name ?? run.model_name, + projectName: run.project_name, trainingMethod: parseBackendTrainingMethod( detail.config?.training_type, detail.config?.load_in_4bit, diff --git a/studio/frontend/src/features/studio/history-card-grid.tsx b/studio/frontend/src/features/studio/history-card-grid.tsx index 5ff9f3bf82..bd76171ce2 100644 --- a/studio/frontend/src/features/studio/history-card-grid.tsx +++ b/studio/frontend/src/features/studio/history-card-grid.tsx @@ -15,6 +15,8 @@ import { Button } from "@/components/ui/button"; import type { TrainingRunSummary } from "@/features/training"; import { deleteTrainingRun, + getTrainingRunDisplayTitle, + getTrainingRunModelSubtitle, emitTrainingRunDeleted, listTrainingRuns, onTrainingRunDeleted, @@ -387,6 +389,12 @@ export function HistoryCardGrid({ const isRunning = run.status === "running"; const canResume = run.can_resume && !wasContinued; const isResuming = resumeTarget === run.id; + + const title = getTrainingRunDisplayTitle(run); + const modelSubtitle = getTrainingRunModelSubtitle(run); + + const projectSubtitle = + run.project_name && title !== run.project_name ? run.project_name : null; // Backend /p ref + its capability token. Both are required: the link // is useless (404s) without the signature, so don't offer to copy it. const canCopyPreview = !!run.preview_ref && !!run.preview_sig; @@ -476,16 +484,16 @@ export function HistoryCardGrid({

- {run.display_name ?? run.model_name} + {title}

- {run.display_name && ( + {modelSubtitle && (

- {run.model_name} + {modelSubtitle}

)}

{run.dataset_name}

+ {projectSubtitle && ( +

+ {projectSubtitle} +

+ )}
{run.loss_sparkline && run.loss_sparkline.length >= 2 && (
diff --git a/studio/frontend/src/features/studio/live-training-view.tsx b/studio/frontend/src/features/studio/live-training-view.tsx index e355655d54..cce39adbf4 100644 --- a/studio/frontend/src/features/studio/live-training-view.tsx +++ b/studio/frontend/src/features/studio/live-training-view.tsx @@ -33,6 +33,8 @@ export function LiveTrainingView(): ReactElement { evalEnabled: state.evalEnabled, outputDir: state.outputDir, isTrainingRunning: state.isTrainingRunning, + startModelName: state.startModelName, + startProjectName: state.startProjectName, lossHistory: state.lossHistory, lrHistory: state.lrHistory, gradNormHistory: state.gradNormHistory, @@ -45,10 +47,16 @@ export function LiveTrainingView(): ReactElement { const config = useTrainingConfigStore( useShallow((state) => ({ selectedModel: state.selectedModel, + projectName: state.projectName, trainingMethod: state.trainingMethod, })), ); + const activeProjectName = + runtime.startProjectName !== null + ? runtime.startProjectName.trim() || null + : (config.projectName || "").trim() || null; + const viewData: TrainingViewData = { phase: runtime.phase, currentStep: runtime.currentStep, @@ -66,7 +74,8 @@ export function LiveTrainingView(): ReactElement { message: runtime.message, error: runtime.error, isTrainingRunning: runtime.isTrainingRunning, - modelName: config.selectedModel ?? "", + modelName: runtime.startModelName ?? config.selectedModel ?? "", + projectName: activeProjectName, trainingMethod: config.trainingMethod ?? "", lossHistory: runtime.lossHistory, lrHistory: runtime.lrHistory, diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 059d665122..2609558145 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -229,6 +229,24 @@ export function ParamsSection(): ReactElement { : "h-studio-config-column"} duration-150`} >
+
+ + {t("studio.params.projectName")} + + {t("studio.params.optional")} + + + store.setProjectName(event.target.value)} + placeholder="customer-support-lora" + maxLength={80} + /> +

+ {t("studio.params.projectNameDescription")} +

+
+ {/* Max Steps / Epochs */}
{t(phaseLabelKeys[data.phase])} + {data.projectName && ( + + {data.projectName} + + )} {t("studio.progress.epoch", { value: formatNumber(data.currentEpoch, 2), @@ -290,7 +295,7 @@ export function ProgressSection({ {pct}%
- +
{!isHistorical && ( @@ -307,7 +312,12 @@ export function ProgressSection({

)} -
+
{formatNumber(stoppedGradNorm, 3)} + {data.projectName && ( + + {data.projectName} + + )} {data.modelName || "--"} diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index d4c800afe4..c2e23dbf3c 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -73,6 +73,7 @@ export function buildTrainingStartPayload( return { model_name: config.selectedModel ?? "", + project_name: (config.projectName || "").trim() || null, training_type: toBackendTrainingType(config.trainingMethod), hf_token: config.hfToken.trim() || null, load_in_4bit: (adapterMethod && isQloraMethod) || (isCpt && isFourBitModel), diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts index 8c2b3e9e77..2f04656c23 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -60,6 +60,7 @@ export function useTrainingActions() { config.selectedModel ?? null, getHfDatasetName(config), false, + config.projectName || "", ); runtimeStore.setStarting(true); @@ -152,7 +153,12 @@ export function useTrainingActions() { // Re-read config after potential store updates from dataset check const payload = buildTrainingStartPayload(useTrainingConfigStore.getState()); - runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, false); + runtimeStore.setStartResources( + payload.model_name, + payload.hf_dataset, + false, + payload.project_name ?? "", + ); const response = await startTraining(payload); if (response.status === "error") { @@ -196,7 +202,7 @@ export function useTrainingActions() { const resumeTrainingRunFromHistory = useCallback(async (runId: string): Promise => { const runtimeStore = useTrainingRuntimeStore.getState(); runtimeStore.setStartError(null); - runtimeStore.setStartResources(null, null, true); + runtimeStore.setStartResources(null, null, true, null); runtimeStore.setStarting(true); try { @@ -220,7 +226,12 @@ export function useTrainingActions() { resume_from_checkpoint: outputDir, } as TrainingStartRequest; - runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, true); + runtimeStore.setStartResources( + payload.model_name, + payload.hf_dataset, + true, + payload.project_name ?? "", + ); // Resume goes straight to startTraining, so it runs the same consent gate as a // fresh start; otherwise a resumed custom-code run hits the worker block with no dialog. diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts index 5157d89582..553dcc2af5 100644 --- a/studio/frontend/src/features/training/index.ts +++ b/studio/frontend/src/features/training/index.ts @@ -7,6 +7,11 @@ export { useTrainingRuntimeStore, } from "./stores/training-runtime-store"; export { useTrainingActions } from "./hooks/use-training-actions"; + +export { + getTrainingRunDisplayTitle, + getTrainingRunModelSubtitle, +} from "./lib/run-display"; export { useTrainingHistorySidebarItems } from "./hooks/use-training-history-sidebar"; export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle"; export { useTrainingCompletionWatch } from "./hooks/use-training-completion-watch"; diff --git a/studio/frontend/src/features/training/lib/run-display.ts b/studio/frontend/src/features/training/lib/run-display.ts new file mode 100644 index 0000000000..b691b257c4 --- /dev/null +++ b/studio/frontend/src/features/training/lib/run-display.ts @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import type { TrainingRunSummary } from "../types/history"; + +type TrainingRunTitleFields = Pick< + TrainingRunSummary, + "display_name" | "project_name" | "model_name" +>; + +function nonEmpty(value: string | null | undefined): string | null { + const trimmed = value?.trim(); + return trimmed ? trimmed : null; +} + +export function getTrainingRunDisplayTitle(run: TrainingRunTitleFields): string { + return nonEmpty(run.display_name) ?? nonEmpty(run.project_name) ?? run.model_name; +} + +export function getTrainingRunModelSubtitle( + run: TrainingRunTitleFields, +): string | null { + return getTrainingRunDisplayTitle(run) === run.model_name ? null : run.model_name; +} diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 85ce25aa97..07be8a247b 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -58,6 +58,7 @@ const initialState: TrainingConfigState = { currentStep: MIN_STEP, modelType: null, selectedModel: null, + projectName: "", trainingMethod: "qlora", hfToken: "", datasetSource: "huggingface", @@ -613,6 +614,7 @@ export const useTrainingConfigStore = create()( if (state.modelDefaultsAppliedFor === state.selectedModel) return; void loadAndApplyModelDefaults(state.selectedModel); }, + setProjectName: (projectName) => set({ projectName }), setTrainingMethod: (trainingMethod) => { const state = get(); set( diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts index 9eaaa98c0e..acc80a3be2 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -24,6 +24,7 @@ const initialState: TrainingRuntimeState = { startError: null, startModelName: null, startDatasetName: null, + startProjectName: null, startFromResume: false, sseConnected: false, firstStepReceived: false, @@ -125,8 +126,12 @@ export const useTrainingRuntimeStore = create()((set) => ( setHasHydrated: (value) => set({ hasHydrated: value }), setStarting: (value) => set({ isStarting: value }), setStartError: (value) => set({ startError: value }), - setStartResources: (startModelName, startDatasetName, startFromResume = false) => - set({ startModelName, startDatasetName, startFromResume }), + setStartResources: ( + startModelName, + startDatasetName, + startFromResume = false, + startProjectName = null, + ) => set({ startModelName, startDatasetName, startProjectName, startFromResume }), setSseConnected: (value) => set({ sseConnected: value }), setLastEventId: (value) => set({ lastEventId: value }), diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 4f7a41bdea..ecd29a2ff0 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -5,6 +5,7 @@ import type { S3Config } from "@/types/training"; export interface TrainingStartRequest { model_name: string; + project_name: string | null; training_type: string; hf_token: string | null; load_in_4bit: boolean; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index d24ce31a6a..5658dfc7a1 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -21,6 +21,7 @@ export interface TrainingConfigState { currentStep: StepNumber; modelType: ModelType | null; selectedModel: string | null; + projectName: string; trainingMethod: TrainingMethod; hfToken: string; datasetSource: DatasetSource; @@ -95,6 +96,7 @@ export interface TrainingConfigActions { prevStep: () => void; setModelType: (type: ModelType) => void; setSelectedModel: (model: string | null) => void; + setProjectName: (value: string) => void; ensureModelDefaultsLoaded: () => void; ensureDatasetChecked: () => void; setTrainingMethod: (method: TrainingMethod) => void; diff --git a/studio/frontend/src/features/training/types/history.ts b/studio/frontend/src/features/training/types/history.ts index 45d83a31d2..99e9bdcb17 100644 --- a/studio/frontend/src/features/training/types/history.ts +++ b/studio/frontend/src/features/training/types/history.ts @@ -5,6 +5,7 @@ export interface TrainingRunSummary { id: string; status: "running" | "completed" | "stopped" | "error"; model_name: string; + project_name: string | null; dataset_name: string; display_name: string | null; started_at: string; diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index c27a2f2bed..8ed0ce0037 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -83,6 +83,7 @@ export interface TrainingRuntimeState { startError: string | null; startModelName: string | null; startDatasetName: string | null; + startProjectName: string | null; startFromResume: boolean; sseConnected: boolean; firstStepReceived: boolean; @@ -121,6 +122,7 @@ export interface TrainingRuntimeActions { modelName: string | null, datasetName: string | null, fromResume?: boolean, + projectName?: string | null, ) => void; setSseConnected: (value: boolean) => void; setLastEventId: (value: number | null) => void; @@ -160,6 +162,7 @@ export interface TrainingViewData { // Config summary modelName: string; + projectName: string | null; trainingMethod: string; // Time-series (for ChartsSection) diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 019a23a891..f3c9eef44e 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -609,6 +609,10 @@ export const en = { params: { title: "Parameters", description: "Configure training hyperparameters", + projectName: "Project Name", + optional: "Optional", + projectNameDescription: + "Used in training output folder names, export defaults, and history.", loraSettings: "LoRA Settings", trainingHyperparameters: "Training Hyperparameters", maxSteps: "Max Steps", @@ -850,6 +854,7 @@ export const en = { loss: "Loss", lr: "LR", gradNorm: "Grad Norm", + project: "Project", model: "Model", method: "Method", elapsed: "Elapsed: {value}", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index bc239294bd..5fd31dc7cb 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -541,6 +541,9 @@ export const zhCN = { params: { title: "参数", description: "配置训练超参数", + projectName: "项目名称", + optional: "可选", + projectNameDescription: "用于训练输出文件夹名称、导出默认值和历史记录。", loraSettings: "LoRA 设置", trainingHyperparameters: "训练超参数", maxSteps: "最大步数", @@ -766,6 +769,7 @@ export const zhCN = { loss: "Loss", lr: "LR", gradNorm: "梯度范数", + project: "项目", model: "模型", method: "方法", elapsed: "已用时间：{value}", From 1069b28c43d7f6375e946cd8632a8448f259de0d Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Mon, 29 Jun 2026 07:12:44 -0700 Subject: [PATCH 29/49] Studio: name the missing extractor when a Recipes file upload fails (#6642) * Studio: name the missing extractor when a Recipes upload fails A missing optional dependency (pymupdf4llm for PDF, mammoth for DOCX) was reported as a generic "Text extraction failed", which gives the user nothing to act on. Catch ImportError and surface the package name instead. * Studio: narrow missing extractor error handling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: wasimysaid Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/routes/data_recipe/seed.py | 31 +++++ studio/backend/tests/test_data_recipe_seed.py | 120 +++++++++++++++++- 2 files changed, 148 insertions(+), 3 deletions(-) diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 8fb034ea4e..57a291291e 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -481,6 +481,37 @@ async def upload_unstructured_file( error = "No extractable text found in file", ) extracted_path.write_text(extracted_text, encoding = "utf-8") + except ImportError as e: + raw_path.unlink(missing_ok = True) + extracted_path.unlink(missing_ok = True) + missing = getattr(e, "name", None) + expected_missing = {".pdf": "pymupdf4llm", ".docx": "mammoth"}.get(ext) + if isinstance(e, ModuleNotFoundError) and missing == expected_missing: + logger.error( + "data_recipe.seed.text_extraction_dependency_missing", + error = str(e), + missing = missing, + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = f"Cannot read {ext} files: the '{missing}' package is not installed.", + ) + logger.error( + "data_recipe.seed.text_extraction_failed", + error = str(e), + exc_info = True, + ) + return UnstructuredFileUploadResponse( + file_id = file_id, + filename = original_filename, + size_bytes = size_bytes, + status = "error", + error = "Text extraction failed.", + ) except Exception as e: raw_path.unlink(missing_ok = True) extracted_path.unlink(missing_ok = True) diff --git a/studio/backend/tests/test_data_recipe_seed.py b/studio/backend/tests/test_data_recipe_seed.py index 601df8bbfe..09e22116ed 100644 --- a/studio/backend/tests/test_data_recipe_seed.py +++ b/studio/backend/tests/test_data_recipe_seed.py @@ -1,12 +1,126 @@ # 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 asyncio +import importlib.util from pathlib import Path +import pytest -def test_seed_inspect_load_kwargs_disables_remote_code_execution(): - seed_route = ( + +def _seed_route_source() -> str: + return ( Path(__file__).resolve().parent.parent / "routes" / "data_recipe" / "seed.py" ).read_text() - assert '"trust_remote_code": False' in seed_route + +def test_seed_inspect_load_kwargs_disables_remote_code_execution(): + assert '"trust_remote_code": False' in _seed_route_source() + + +class _FakeUpload: + def __init__(self, filename: str, content: bytes): + self.filename = filename + self._content = content + + async def read(self) -> bytes: + return self._content + + +def _load_seed_route(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + pytest.importorskip("fastapi") + pytest.importorskip("multipart") + pytest.importorskip("structlog") + + backend_root = Path(__file__).resolve().parent.parent + monkeypatch.syspath_prepend(str(backend_root)) + route_path = backend_root / "routes" / "data_recipe" / "seed.py" + spec = importlib.util.spec_from_file_location("seed_under_test", route_path) + assert spec is not None and spec.loader is not None + seed_route = importlib.util.module_from_spec(spec) + spec.loader.exec_module(seed_route) + seed_route.UNSTRUCTURED_UPLOAD_ROOT = tmp_path / "unstructured-uploads" + return seed_route + + +def _run_upload( + seed_route, + filename: str, + content: bytes, + block_id: str = "block", +): + return asyncio.run( + seed_route.upload_unstructured_file(_FakeUpload(filename, content), block_id) + ) + + +def _block_files(seed_route, block_id: str = "block") -> list[str]: + block_dir = seed_route.UNSTRUCTURED_UPLOAD_ROOT / block_id + if not block_dir.exists(): + return [] + return sorted(path.name for path in block_dir.iterdir()) + + +def _raise(exc: BaseException): + def raise_exc(*args, **kwargs): + raise exc + + return raise_exc + + +@pytest.mark.parametrize( + ("filename", "package"), + [ + ("paper.pdf", "pymupdf4llm"), + ("notes.docx", "mammoth"), + ], +) +def test_unstructured_upload_names_missing_extractor_dependency( + monkeypatch, tmp_path, filename, package +): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr( + seed_route, + "_extract_text_from_file", + _raise(ModuleNotFoundError(f"No module named {package!r}", name = package)), + ) + + result = _run_upload(seed_route, filename, b"%PDF-1.7") + + assert result.status == "error" + assert ( + result.error + == f"Cannot read {Path(filename).suffix} files: the '{package}' package is not installed." + ) + assert _block_files(seed_route) == [] + + +def test_unstructured_upload_keeps_txt_path_working(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + result = _run_upload(seed_route, "notes.txt", b"hello") + + assert result.status == "ok" + assert result.error is None + assert any(name.endswith(".txt") for name in _block_files(seed_route)) + assert any(name.endswith(".extracted.txt") for name in _block_files(seed_route)) + + +@pytest.mark.parametrize( + "exc", + [ + ImportError("cannot import internal symbol"), + ModuleNotFoundError( + "No module named 'missing_transitive_pkg'", + name = "missing_transitive_pkg", + ), + ], +) +def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, exc): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr(seed_route, "_extract_text_from_file", _raise(exc)) + result = _run_upload(seed_route, "paper.pdf", b"%PDF-1.7") + + assert result.status == "error" + assert result.error == "Text extraction failed." + assert _block_files(seed_route) == [] From f7d509e1f2b8d5aa6f8533c48702df2d079e58d0 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Mon, 29 Jun 2026 15:17:16 +0100 Subject: [PATCH 30/49] fix: remove sidebar update dev override (#6746) --- .../frontend/src/components/app-sidebar.tsx | 22 +++---------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 0203605767..ef4178ea42 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -262,19 +262,6 @@ function NavItem({ ); } -// TEMP DEV override: preview the update card on installs with no real update -// (e.g. an editable checkout). In the browser console run -// `localStorage.setItem("unsloth_force_update_card", "1")` and reload. Remove -// before merge. -function devForceUpdateCard(): boolean { - if (typeof window === "undefined") return false; - try { - return window.localStorage.getItem("unsloth_force_update_card") === "1"; - } catch { - return false; - } -} - export function AppSidebar() { const t = useT(); const { isDark, toggleTheme, anchorRef } = useAnimatedThemeToggle(); @@ -291,13 +278,10 @@ export function AppSidebar() { // Web update detection: `webUpdate` is non-null only when the installed // (PyPI) version is behind the latest release, so the card is hidden by - // default. `forceUpdateCard` is a TEMP dev override to preview it on installs - // with no real update (e.g. an editable checkout); remove before merge. + // default. const { status: webUpdate } = useWebUpdateCheck(); - const [forceUpdateCard] = useState(devForceUpdateCard); - const showUpdateCard = Boolean(webUpdate) || forceUpdateCard; - const updateVersion = - webUpdate?.latestVersion ?? (forceUpdateCard ? "0.0.0" : null); + const showUpdateCard = Boolean(webUpdate); + const updateVersion = webUpdate?.latestVersion ?? null; // Auto-close mobile Sheet after navigation const closeMobileIfOpen = () => { From 220ff5aabaa67ede8a29af0859297c7ecac96985 Mon Sep 17 00:00:00 2001 From: OrbisAI Security Date: Mon, 29 Jun 2026 19:58:06 +0530 Subject: [PATCH 31/49] fix: CVE-2026-54290 security vulnerability (#6736) Automated dependency upgrade by OrbisAI Security Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/frontend/package-lock.json | 35 +++++++++++++------------------ studio/frontend/package.json | 2 +- 2 files changed, 15 insertions(+), 22 deletions(-) diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json index 6202d3ce22..80db64553a 100644 --- a/studio/frontend/package-lock.json +++ b/studio/frontend/package-lock.json @@ -1704,6 +1704,7 @@ "os": [ "android" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1724,6 +1725,7 @@ "os": [ "darwin" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1744,6 +1746,7 @@ "os": [ "darwin" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1764,6 +1767,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1784,6 +1788,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1804,6 +1809,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1824,6 +1830,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1844,6 +1851,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1864,6 +1872,7 @@ "os": [ "linux" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1884,6 +1893,7 @@ "os": [ "win32" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -1904,6 +1914,7 @@ "os": [ "win32" ], + "peer": true, "engines": { "node": ">= 10" }, @@ -5669,9 +5680,6 @@ "cpu": [ "arm64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -5688,9 +5696,6 @@ "cpu": [ "arm64" ], - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -5707,9 +5712,6 @@ "cpu": [ "ppc64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -5726,9 +5728,6 @@ "cpu": [ "s390x" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -5745,9 +5744,6 @@ "cpu": [ "x64" ], - "libc": [ - "glibc" - ], "license": "MIT", "optional": true, "os": [ @@ -5764,9 +5760,6 @@ "cpu": [ "x64" ], - "libc": [ - "musl" - ], "license": "MIT", "optional": true, "os": [ @@ -10282,9 +10275,9 @@ } }, "node_modules/hono": { - "version": "4.12.21", - "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.21.tgz", - "integrity": "sha512-uV63apnb0kyPtAUwoWgaGh9HyIFcv8lgmzPZSiTBQAFOFGIzka5EZ1dZocmGnn0XdX0+XTqJ6Tqv7selMuGLRQ==", + "version": "4.12.25", + "resolved": "https://registry.npmjs.org/hono/-/hono-4.12.25.tgz", + "integrity": "sha512-2NFaIyNVgJmBs/ecmtGzlmluTFs5cHEWGTdu0t1HBwYzoGXOL5nUQBRMXsXWla5i4KkG//QMzVP88m1+I3fdAQ==", "license": "MIT", "engines": { "node": ">=16.9.0" diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 0956c710f5..a2eddecda3 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -86,7 +86,7 @@ "@tanstack/router-core": "1.169.2", "@tanstack/history": "1.161.6", "mermaid": "11.15.0", - "hono": "4.12.21", + "hono": "4.12.25", "qs": "6.15.2", "ip-address": "10.1.1", "brace-expansion@5.0.5": "5.0.6" From 6acf01f7b34fdc4f42230405a76624a0212d9b50 Mon Sep 17 00:00:00 2001 From: ashzak Date: Mon, 29 Jun 2026 15:23:18 -0500 Subject: [PATCH 32/49] Fix llama.cpp CMake build detection in save.py (#5957) --------- Co-authored-by: Daniel Han Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- unsloth/save.py | 92 ++++++++++++++++++++++++++++++++++++------------- 1 file changed, 69 insertions(+), 23 deletions(-) diff --git a/unsloth/save.py b/unsloth/save.py index f55a14b4e3..20a934538c 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -149,6 +149,29 @@ def has_curl(): CURL_FLAG = "-DLLAMA_CURL=ON" if has_curl() else "-DLLAMA_CURL=OFF" +def _is_cmake_only_llama_cpp(llama_cpp_dir: str = "llama.cpp") -> bool: + """ + True if llama.cpp's Makefile is the post-CMake-migration deprecation stub, + so `make` cannot build it. A genuinely missing/empty checkout returns False + so it isn't treated as CMake-only: the caller then probes make and fails + loudly on a real error rather than silently assuming a CMake build. + """ + makefile_path = os.path.join(llama_cpp_dir, "Makefile") + if not os.path.exists(makefile_path): + # No Makefile: only CMake-only if a real CMake project is present + return os.path.exists(os.path.join(llama_cpp_dir, "CMakeLists.txt")) + try: + with open(makefile_path, "r", encoding = "utf-8", errors = "ignore") as f: + content = f.read(4096).lower() + if "cmake" in content and "deprecated" in content: + return True + if "build system changed" in content: + return True + except (IOError, OSError): + pass + return False + + def print_quantization_methods(): for key, value in ALLOWED_QUANTS.items(): print(f'"{key}" ==> {value}') @@ -1190,14 +1213,27 @@ def install_llama_cpp_make_non_blocking(): # https://github.com/ggerganov/llama.cpp/issues/7062 # Weirdly GPU conversion for GGUF breaks?? # env = { **os.environ, "LLAMA_CUDA": "1", } - # Force make clean - check = os.system("make clean -C llama.cpp") - IS_CMAKE = False - if check == 0: + + # Skip the make-clean probe on CMake-only checkouts (its error output is misleading) + IS_CMAKE = _is_cmake_only_llama_cpp("llama.cpp") + + if not IS_CMAKE: + # Confirm make still works, silently + try: + result = subprocess.run( + ["make", "clean", "-C", "llama.cpp"], + stdout = subprocess.DEVNULL, + stderr = subprocess.DEVNULL, + ) + IS_CMAKE = result.returncode != 0 + except FileNotFoundError: + # No make executable; use CMake + IS_CMAKE = True + + if not IS_CMAKE: # Uses old MAKE n_jobs = max(int((psutil.cpu_count() or 1) * 1.5), 1) full_command = ["make", "all", "-j" + str(n_jobs), "-C", "llama.cpp"] - IS_CMAKE = False else: # Uses new CMAKE n_jobs = max(int(psutil.cpu_count() or 1), 1) # Use less CPUs since 1.5x faster @@ -1220,7 +1256,6 @@ def install_llama_cpp_make_non_blocking(): "--clean-first", "--target", ] + LLAMA_CPP_TARGETS - IS_CMAKE = True # https://github.com/ggerganov/llama.cpp/issues/7062 # Weirdly GPU conversion for GGUF breaks?? # run_installer = subprocess.Popen(full_command, env = env, stdout = subprocess.DEVNULL, stderr = subprocess.STDOUT) @@ -1306,20 +1341,25 @@ def install_llama_cpp_old(version = -10): ] try_execute(commands) - # Try using MAKE - commands = [ - "make clean -C llama.cpp", - f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", - ] - if try_execute(commands) == "CMAKE": - # Instead use CMAKE + # Detect CMake-only build system before trying make + use_cmake = _is_cmake_only_llama_cpp("llama.cpp") + + if not use_cmake: + # Try using MAKE + commands = [ + "make clean -C llama.cpp", + f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", + ] + use_cmake = try_execute(commands) == "CMAKE" + + if use_cmake: + # Use CMAKE commands = [ f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}", f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}", "cp llama.cpp/build/bin/llama-* llama.cpp", "rm -rf llama.cpp/build", ] - try_execute(commands) # Check if successful @@ -1351,15 +1391,21 @@ def install_llama_cpp_blocking(use_cuda = False): return try_execute(commands) - commands = [ - "make clean -C llama.cpp", - # https://github.com/ggerganov/llama.cpp/issues/7062 - # Weirdly GPU conversion for GGUF breaks?? - # f"{use_cuda} make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", - f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", - ] - if try_execute(commands) == "CMAKE": - # Instead use CMAKE + # Detect CMake-only build system before trying make + use_cmake = _is_cmake_only_llama_cpp("llama.cpp") + + if not use_cmake: + commands = [ + "make clean -C llama.cpp", + # https://github.com/ggerganov/llama.cpp/issues/7062 + # Weirdly GPU conversion for GGUF breaks?? + # f"{use_cuda} make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", + f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp", + ] + use_cmake = try_execute(commands) == "CMAKE" + + if use_cmake: + # Use CMAKE commands = [ f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}", f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}", From ba41e798d67aba9fe2a1dbfe1d6f4eab266ab1fb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 13:35:26 -0700 Subject: [PATCH 33/49] CI: add PyPI extra-index to CPU torch installs to fix sympy resolution (#6660) --- .github/workflows/consolidated-tests-ci.yml | 4 ++-- .github/workflows/mlx-ci.yml | 2 +- .github/workflows/notebooks-ci.yml | 2 +- .github/workflows/studio-backend-ci.yml | 4 ++-- .github/workflows/studio-windows-inference-smoke.yml | 2 +- .github/workflows/version-compat-ci.yml | 2 +- 6 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index b56a6c2615..bd9c0a532d 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -209,7 +209,7 @@ jobs: 'peft>=0.18,<0.20' 'accelerate>=0.34,<2' \ ipython # torchvision: unsloth_zoo.vision_utils imports it at module scope. - pip install --index-url https://download.pytorch.org/whl/cpu \ + 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' # transformers + trl from the matrix combo. pip install "$RESOLVED_TRANSFORMERS_SPEC" @@ -2166,7 +2166,7 @@ jobs: python -m pip install --upgrade pip # Match the matrix job's torch path so unsloth_zoo's # `import torch` resolves to the same CPU build. - pip install --index-url https://download.pytorch.org/whl/cpu \ + 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' pip install \ 'numpy<3' protobuf sentencepiece \ diff --git a/.github/workflows/mlx-ci.yml b/.github/workflows/mlx-ci.yml index 424a706d7c..a2f716a93c 100644 --- a/.github/workflows/mlx-ci.yml +++ b/.github/workflows/mlx-ci.yml @@ -163,7 +163,7 @@ jobs: 'pytest==9.0.3' \ 'pytest-asyncio==1.3.0' \ 'httpx==0.28.1' - pip install --index-url https://download.pytorch.org/whl/cpu \ + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ 'torch==2.10.0' # github.com occasionally 500s on the git fetch; retry the # zoo install so a single upstream blip does not fail CI. diff --git a/.github/workflows/notebooks-ci.yml b/.github/workflows/notebooks-ci.yml index 2edcae8ab2..0e0b35dd4d 100644 --- a/.github/workflows/notebooks-ci.yml +++ b/.github/workflows/notebooks-ci.yml @@ -263,7 +263,7 @@ jobs: # unsloth_zoo.vision_utils imports PIL at module top, and the # easiest way to get a torch-compatible PIL on a CPU runner is # to let torchvision pull the right Pillow version. - pip install --index-url https://download.pytorch.org/whl/cpu \ + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ 'torch>=2.8,<2.11' 'torchvision<0.26' # Pin to the same versions update_all_notebooks.py installs in # generated notebooks. Keep these in lockstep with PIN_TRL / diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index ea60252cf6..bce355458a 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -76,7 +76,7 @@ jobs: # Torch CPU + transformers are required by a chunk of the backend test # suite (gpu_selection, kv_cache_estimation, utils). CPU-only torch # keeps the install ~250 MB / ~1 min on a clean runner. - pip install --index-url https://download.pytorch.org/whl/cpu 'torch>=2.4,<2.11' + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple 'torch>=2.4,<2.11' pip install 'transformers>=4.51,<5.5' - name: Backend tests @@ -137,7 +137,7 @@ jobs: pyyaml jinja2 mammoth unpdf requests typer \ 'numpy<3' pytest pytest-asyncio httpx # torchvision: unsloth_zoo.vision_utils imports it at module scope. - pip install --index-url https://download.pytorch.org/whl/cpu \ + 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' pip install 'transformers>=4.51,<5.5' # bitsandbytes: hard import in unsloth/models/_utils.py. Recent diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 08a0ee782d..8186c07211 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -1384,7 +1384,7 @@ jobs: - name: PyTorch CPU wheel installs and imports (no Visual Studio) run: | python -m pip install --upgrade pip - python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple python -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())" - name: Install Studio (--local, --no-torch) with no build tools present diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index 599b53df1d..e492d21e99 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -242,7 +242,7 @@ jobs: run: | python -m pip install --upgrade pip # CPU torch (vllm/peft/st all depend on it). - pip install --index-url https://download.pytorch.org/whl/cpu \ + 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' # torchcodec is a hard requirement on transformers 5.x: # transformers/audio_utils.py:55 does From de3c745fab08e14a1cd7a825c434f048a53ec5ef Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 14:35:23 -0700 Subject: [PATCH 34/49] Fix full finetuning precision on V100 / no-bf16 GPUs (#5880) --------- Co-authored-by: Datta Nimmaturi Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- tests/python/test_v100_fullft_precision.py | 193 +++++++++++++++++++++ unsloth/models/rl.py | 26 ++- 2 files changed, 212 insertions(+), 7 deletions(-) create mode 100644 tests/python/test_v100_fullft_precision.py diff --git a/tests/python/test_v100_fullft_precision.py b/tests/python/test_v100_fullft_precision.py new file mode 100644 index 0000000000..c8ca769d45 --- /dev/null +++ b/tests/python/test_v100_fullft_precision.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression tests for full finetuning precision on no-bf16 GPUs (V100/T4). + +Full finetuning upcasts trainable weights to float32, so the model dtype is +float32 (not bfloat16). The SFTTrainer mixed-precision template in +unsloth/models/rl.py must then: + - run the forward pass under float16 autocast for normal models, + - keep FORCE_FLOAT32 models (Gemma3, gpt_oss, ...) in pure float32, + - never select bf16 on hardware without bf16. + +We execute the REAL template block extracted from rl.py source (no heavy unsloth +import) against mocked inputs. See issue #4082. +""" + +from __future__ import annotations + +import os +import sys +import types +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") + +RL_PY = Path(__file__).resolve().parents[2] / "unsloth" / "models" / "rl.py" + + +def _extract_mixed_precision_code() -> str: + lines = RL_PY.read_text().split("\n") + try: + start = next(i for i, l in enumerate(lines) if "mixed_precision = (" in l) + except StopIteration: + pytest.skip("mixed_precision template not found in rl.py") + body, k = [], start + 1 + while lines[k].strip() != ")": + body.append(lines[k]) + k += 1 + return eval("(\n" + "\n".join(body) + "\n)") # only string literals + comments + + +CODE = _extract_mixed_precision_code() + + +def _restore(mapping, saved): + """Restore a dict-like to its saved snapshot: pop keys that were absent.""" + for k, v in saved.items(): + if v is None: + mapping.pop(k, None) + else: + mapping[k] = v + + +def _decide(dtype, *, bf16_supported, force_float32, full_finetuning, mixed_precision, fp16, bf16): + """Run the template block; return (args.fp16, args.bf16, ACCELERATE_MP, raised). + + Stubs (sys.modules, env vars, torch.cuda.is_bf16_supported) are restored on + exit so a decision can't leak into later tests in the same process. + """ + uzu = types.ModuleType("unsloth_zoo.utils") + uzu._get_dtype = lambda x: x + uzd = types.ModuleType("unsloth_zoo.device_type") + uzd.device_is_bf16_supported = lambda: bf16_supported # device-aware signal stub + + env_keys = ( + "UNSLOTH_FORCE_FLOAT32", + "UNSLOTH_ENABLE_FULL_FINETUNING", + "UNSLOTH_MIXED_PRECISION", + "ACCELERATE_MIXED_PRECISION", + ) + mod_keys = ("unsloth_zoo", "unsloth_zoo.utils", "unsloth_zoo.device_type") + saved_env = {k: os.environ.get(k) for k in env_keys} + saved_mods = {k: sys.modules.get(k) for k in mod_keys} + orig_bf16 = torch.cuda.is_bf16_supported + try: + sys.modules.setdefault("unsloth_zoo", types.ModuleType("unsloth_zoo")) + sys.modules["unsloth_zoo.utils"] = uzu + sys.modules["unsloth_zoo.device_type"] = uzd + for k in env_keys: + os.environ.pop(k, None) + os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" if force_float32 else "0" + os.environ["UNSLOTH_ENABLE_FULL_FINETUNING"] = "1" if full_finetuning else "0" + os.environ["UNSLOTH_MIXED_PRECISION"] = mixed_precision + torch.cuda.is_bf16_supported = lambda *a, **k: bf16_supported + args = types.SimpleNamespace(fp16 = fp16, bf16 = bf16, mixed_precision = None) + emb = types.SimpleNamespace(weight = types.SimpleNamespace(dtype = dtype)) + model = types.SimpleNamespace( + config = types.SimpleNamespace(dtype = dtype, torch_dtype = dtype), + get_input_embeddings = lambda: emb, + ) + raised = None + try: + exec(CODE, {"torch": torch, "os": os}, {"args": args, "model": model}) + except TypeError: + raised = "TypeError" + return args.fp16, args.bf16, os.environ.get("ACCELERATE_MIXED_PRECISION"), raised + finally: + torch.cuda.is_bf16_supported = orig_bf16 + _restore(os.environ, saved_env) + _restore(sys.modules, saved_mods) + + +def test_v100_normal_fullft_fp16_explicit(): + # Normal model, full FT (weights upcast to float32), V100, fp16=True. + fp16, bf16, amp, raised = _decide( + torch.float32, + bf16_supported = False, + force_float32 = False, + full_finetuning = True, + mixed_precision = "float32", + fp16 = True, + bf16 = False, + ) + assert raised is None + assert (fp16, bf16) == (True, False) # float32 weights + fp16 forward + + +def test_v100_normal_fullft_precision_unset(): + # Same, but user left precision unset -> must pick fp16, never bf16. + fp16, bf16, amp, raised = _decide( + torch.float32, + bf16_supported = False, + force_float32 = False, + full_finetuning = True, + mixed_precision = "float32", + fp16 = False, + bf16 = False, + ) + assert raised is None + assert (fp16, bf16) == (True, False) + assert amp == "fp16" + + +def test_force_float32_model_fullft_is_pure_float32(): + # FORCE_FLOAT32 model (Gemma3, gpt_oss, ...) in full FT -> pure float32, no autocast. + fp16, bf16, amp, raised = _decide( + torch.float32, + bf16_supported = False, + force_float32 = True, + full_finetuning = True, + mixed_precision = "float32", + fp16 = True, + bf16 = False, + ) + assert raised is None + assert (fp16, bf16) == (False, False) + assert amp in (None, "no") + + +def test_no_bf16_on_volta_in_auto_branch(): + # bf16 model dtype but no bf16 HW, precision unset -> fp16, never bf16. + fp16, bf16, amp, raised = _decide( + torch.bfloat16, + bf16_supported = False, + force_float32 = False, + full_finetuning = False, + mixed_precision = "float32", + fp16 = False, + bf16 = False, + ) + assert bf16 is False + + +def test_bf16_gpu_unchanged_auto_branch(): + # Regression guard: on a bf16 GPU, a float32 model with unset precision + # still selects bf16 autocast (behavior must not change for bf16 hardware). + fp16, bf16, amp, raised = _decide( + torch.float32, + bf16_supported = True, + force_float32 = False, + full_finetuning = True, + mixed_precision = "float32", + fp16 = False, + bf16 = False, + ) + assert raised is None + assert (fp16, bf16) == (False, True) + + +def test_genuine_bf16_model_with_fp16_still_raises(): + # A real bfloat16 model on bf16 HW with fp16 requested is a genuine mismatch. + _, _, _, raised = _decide( + torch.bfloat16, + bf16_supported = True, + force_float32 = False, + full_finetuning = False, + mixed_precision = "float32", + fp16 = True, + bf16 = False, + ) + assert raised == "TypeError" diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index a85c1d08a4..53668d14d8 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -994,8 +994,18 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "use_fp16 = getattr(args, 'fp16', False)\n" "if type(use_fp16) is not bool: use_fp16 = False\n" "force_float32 = False\n" + # device-aware bf16 check (CUDA/XPU/HIP), so V100/T4 never pick bf16 + # but AMD/Intel are unaffected; fall back on older unsloth_zoo. + "try:\n" + " from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported\n" + "except Exception:\n" + " _bf16_supported = torch.cuda.is_bf16_supported\n" + # FORCE_FLOAT32 models (Gemma3, gpt_oss, ...) cannot use float16. On a GPU without + # bf16 (V100/T4) keep them in float32 so they never autocast to fp16. On a bf16 GPU, + # full finetuning can still use bf16 autocast (master weights stay float32), which is + # faster and uses less memory; LoRA/QLoRA keep float32 when forced. "full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1'\n" - "if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'):\n" + "if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' and not (full_finetuning and _bf16_supported()):\n" " print('Unsloth: Switching to float32 training since model cannot work with float16')\n" " force_float32 = True\n" "mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32')\n" @@ -1004,8 +1014,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "from unsloth_zoo.utils import _get_dtype\n" "dtype = _get_dtype(dtype)\n" "float16 = dtype == torch.float16\n" + "bfloat16 = dtype == torch.bfloat16\n" "if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n" - "if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n" + "if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n" "if force_float32:\n" " # Forced float32 training\n" " args.fp16 = False\n" @@ -1014,11 +1025,12 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): " if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n" " # args.mixed_precision is a new argument which needs to be set now\n" "elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':\n" - " # Mixed precision training\n" - " args.fp16 = float16\n" - " args.bf16 = not float16\n" - " os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'\n" - " if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16'\n" + " # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16.\n" + " use_bf16_amp = (not float16) and _bf16_supported()\n" + " args.fp16 = not use_bf16_amp\n" + " args.bf16 = use_bf16_amp\n" + " os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16'\n" + " if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16'\n" " # args.mixed_precision is a new argument which needs to be set now\n" "elif mixed_precision_dtype == 'bfloat16':\n" " # Both False since bfloat16 full finetuning doesn't do any autocasting.\n" From 27b66b2efeee60498876f9870f4022cbae868a27 Mon Sep 17 00:00:00 2001 From: Yuwen Hu <54161268+Oscilloscope98@users.noreply.github.com> Date: Tue, 30 Jun 2026 05:38:30 +0800 Subject: [PATCH 35/49] Fix outdated triton-xpu 3.7.1 sha256 hashes in intel-gpu-torch2120 extra (#6629) --- pyproject.toml | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 13c421d8ea..76b87349af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1174,14 +1174,14 @@ intelgputorch2120 = [ "unsloth_zoo[intelgpu]", "unsloth[huggingfacenotorch]", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=844d981cb1b3948085e8cfa62c74de9f100259f6131959aa70be49123b88ae81 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a16b1d00e94ad87d62af3512e390348b8656419598004100c56028bf494f086b ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4e46e71e077cf483404a4c17ce40d71c5f0e13a81459139d4346ca427b1dd455 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4fdaed1bafc51d3a2834656a3420a6686a74ea226508765a49bf15d58ff3a930 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=2778b46b22e9fa0916398db299a125027a1b2331c1173b3dd2b9e2cab6263a31 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=ad5b147d04ee0d40f3d4d32f85f5aa3a3beb6cd5799ca026d3d7f4afa3d9e24f ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=d9482063af2a308543f23333e32edd738ea87cbb33ade68afda9ae0fd704ccd9 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", - "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=5d4d67f0deb1e851c01b293e602b8dcddad26ca2be61221cee3dc0e1aa0cdefd ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=81ff0eb0c4fc8e19d2510b28c3e1d9382a3c7d6fdaf6a9f9631a93a030d841cf ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=55574a68d275b85cd4d5cbf185084bae019ebf09c3f43b0bd2831b14935ec8e7 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a31c058c5c2e78ebe490a2e69f2f50caec6b1307ac096e944f116fdc06819d9a ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=e701a31efa0334775f357c98716f3821775aa944219f7888e13c2dfe2daabe2a ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=0d7730651c3e52fbf3a430cc201455f0c6600dc72e681aec495f131ea44f341a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=8f4a63de73e3d632098f93c8f0bd77244958a47d7c5f728b8ff35f8a91fdb983 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=6589ece3adc2b1ab88d90ff1267afc25df5c7b868f0b633e732cac70df36cbde ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + "triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=2fdf001a9b0575e8b1827127259bb9b13bf36e659882be74c2dfab46597d3e7a ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=e8923cd1fe560472904b1461b745d2f1826bb9c1bc0808225d5f28a450e4d553 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'", "torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=f7c082b2fc9b61def594d30ea57762dc4a8bc7111a9a9593953ed948de242e28 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", From f62c26e63d28c218c579c29a87e96d32ccde2eae Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 14:45:26 -0700 Subject: [PATCH 36/49] Fix stale xformers and flash-attn wheel URLs (#4213) Co-authored-by: Jeffrey Cruz --- pyproject.toml | 7 ------- 1 file changed, 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 76b87349af..844ead2454 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -255,10 +255,6 @@ cu118onlytorch270 = [ "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)", "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)", "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')", ] cu126onlytorch270 = [ "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)", @@ -282,7 +278,6 @@ cu128onlytorch270 = [ ] cu118onlytorch271 = [ "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", - "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')", ] cu126onlytorch271 = [ "xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)", @@ -879,14 +874,12 @@ flashattentiontorch240abiFALSEcu12x = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", ] flashattentiontorch240abiTRUEcu12x = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'", "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'", - "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'", ] intelgputorch260 = [ "unsloth_zoo[intelgpu]", From 32f28b2180a87c333bf5f7e32523c572526f4c41 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Tue, 30 Jun 2026 15:56:33 +0530 Subject: [PATCH 37/49] Studio: keep "Fine-tuned" compare label clear of the floating top right controls (#6755) * fix header overlap * fix --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/frontend/src/features/chat/chat-page.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 2511fc9eec..cd7cfc77fc 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -606,7 +606,7 @@ const LoraCompareContent = memo(function LoraCompareContent({ handleName="lora" borderClassName="border-t border-border/60 md:border-t-0 md:border-l" header={ -
+
Fine-tuned From 9369dd47e6777bd6b3ece55109127627e47c7908 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 30 Jun 2026 03:40:16 -0700 Subject: [PATCH 38/49] Add FP8/FP4 compressed export to save_pretrained_merged (#6706) * Add FP8/FP4 compressed export to save_pretrained_merged Adds compressed-tensors export (for vLLM) to save_pretrained_merged / push_to_hub_merged via llm-compressor, alongside the existing lora / merged_16bit / merged_4bit / gguf / torchao paths: model.save_pretrained_merged("model", tokenizer, save_method="fp8") Supported save_method values: fp8 (FP8_DYNAMIC), mxfp4, nvfp4 (W4A4) and mxfp8. The LoRA is merged to 16bit at save_directory, then a quantized checkpoint is written to save_directory + "-". nvfp4 needs a small calibration set (defaults to ultrachat, overridable via calibration_dataset). Notes: - llm-compressor is installed lazily on first use, pinning the current torch and transformers via a constraints file so they are not upgraded (a plain install pulls transformers>=5 and breaks Unsloth). - Quantization runs in a separate process (unsloth/_compressed_quantize.py, launched by file path) so Unsloth's transformers attention patches do not interfere with the forward llm-compressor runs during calibration, mirroring how GGUF export shells out to llama.cpp. - mxfp8 needs a newer llm-compressor (transformers>=5); it is recognised and raises a clear error until that stack is available. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: main-process guard, calibration subsampling, tokenizer + dtype handling - Route the 16bit merge through unsloth_generic_save for both LoRA and full finetuned models, so non-PEFT models are written in 16bit consistently instead of saving the original (possibly quantized) weights directly. - Honor is_main_process: only the main process quantizes and writes the compressed output, so distributed ranks do not race on the same dirs. - Subsample an in-memory calibration Dataset before save_to_disk so large training sets are not fully copied to a temp dir. - Tolerate a missing tokenizer in the converter (data-free exports); still require one for calibration based schemes. - Open config.json via a context manager in both files. - Drop the redundant nvfp4 entry from the unsupported-name check (fp4 covers it). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add direct LoRA to GGUF export and harden FP8/FP4 compressed export - Run llm-compressor install and scheme check before the 16bit merge so unsupported schemes (e.g. mxfp8) fail fast without writing a checkpoint - Only the main process installs, merges, quantizes and uploads; isolate hub pushes to a temp dir and clean all temp dirs in a finally - Forward standard save kwargs (state_dict, max_shard_size, ...) to the merge - Fall back to the first dataset split for Hub calibration ids - Export LoRA adapters to GGUF via convert_lora_to_gguf.py: modernize save_pretrained_ggml/push_to_hub_ggml and add save_method="lora" to save_pretrained_gguf/push_to_hub_gguf; resolve base from the adapter config * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix LoRA GGUF shell-injection test and compressed export trailing-slash path - Update tests/saving/test_save_shell_injection.py for the new delegation: the LoRA to GGUF conversion now lives in _unsloth_save_lora_gguf, so assert it passes argv as a list with no shell=True and that the legacy ggml wrappers delegate to it instead of calling subprocess.Popen directly - Normalize the local save_directory before building the "-" sibling so a trailing slash no longer nests the compressed output inside the 16bit dir * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Polish FP8/FP4 and LoRA GGUF export after review - Warn (not silently downgrade) when an explicit quantization_method is not a valid LoRA GGUF outtype; default stays f16 - Correct the inference hardware note: MXFP8 is 8-bit (cc >= 8.9), only FP4 needs Blackwell for full activation quantization - Document that a local fp8/fp4 save keeps the 16bit merge at save_directory and writes the quantized checkpoint to save_directory + "-" * Use sequential calibration pipeline and validate Hub access early - nvfp4 calibration no longer forces the memory-hungry "basic" pipeline. The quantization runs in a clean subprocess, so llm-compressor's default sequential pipeline (layer-by-layer onloading) works and lets large models that do not fit at once still calibrate; fall back to "basic" only if tracing fails - For push_to_hub compressed exports, create/validate the repo up front so a bad token or denied repo fails before the merge and quantization instead of after * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden compressed export: explicit sequential pipeline, base-tokenizer calibration, GPU memory - nvfp4 calibration now passes pipeline="sequential" explicitly (layer-by-layer onloading) instead of relying on the inferred default, with a "basic" fallback - Calibration datasets with a messages column no longer require a chat template: base / non-chat tokenizers fall back to concatenating message contents - Free the in-memory model's CUDA memory before the quantize subprocess loads its own copy from disk (best-effort, single-device non-quantized only; restored afterward), so a single GPU need not hold two copies at once - Create the calibration temp dir in the system temp location instead of next to the save directory, avoiding stray dirs in the workspace * Free the failed calibration model before the basic-pipeline retry In the sequential -> basic NVFP4 fallback, release the partially-processed model and clear the CUDA cache before loading a fresh copy, so the retry does not transiently hold two model copies on the GPU. * Harden calibration data handling and compressed-export edge cases - Calibration messages without a chat template now handle multimodal (list) content, None content, and null message rows instead of crashing on join - Raise a clear error when the calibration dataset is empty after subsampling - Reset llm-compressor's global session before freeing the model in the sequential -> basic NVFP4 fallback, so the old model is actually released - LoRA GGUF export accepts a single-element list quantization_method - Attach datasets metadata to the pushed repo on compressed hub exports - Warn (instead of silently) if the model cannot be restored to its device - Raise a clear error if the LoRA base model id cannot be determined * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handle DatasetDict calibration, MoE routers, and MTP models in compressed export - Reduce an in-memory DatasetDict calibration set to a single split before row subsampling, so save_to_disk does not copy every split to the temp dir - For MoE models, keep the router/gate unquantized and pass moe_calibrate_all_experts so every expert is calibrated - Warn when a model carries MTP / speculative-decoding tensors that the compressed export does not include * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Support many more compressed-tensors schemes and address review - Expand save_method to cover the full set of compressed-tensors preset schemes: FP8 (dynamic/static/block), INT8, W8A8, W8A16, W4A16(+asym), W4A8, W4AFP8, MXFP4(+A16), NVFP4(+A16), plus the gated MXFP8; calibration is used only for the static-activation schemes (FP8 static, NVFP4) - Broaden the near-miss save_method error to cover int/w-prefixed names - MoE: also keep the Qwen shared-expert gate unquantized - Strip non-model-input columns from already-tokenized calibration data so the collator does not choke on a leftover messages column - Forward the Hub token to the LoRA converter and the quantize subprocess so gated/private base models and calibration datasets work without a global login * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Collapse compressed-tensors export help line so ruff-format converges The print line in print_quantization_methods needed two ruff-format passes to reach a fixpoint (merge implicit string concat, then collapse the single-arg print). pre-commit.ci applies one pass per run, so it kept reformatting. Land the converged single-line form directly. * Add CPU-only regression tests for the export API Cover all export paths without a GPU, for slow CPU-only CI: - pure-function checks of the compressed-tensors scheme registry and save_method normalization (aliases, calibration flags, near-miss errors) - AST checks that every merged saver dispatches compressed export, the GGUF savers expose the lora branch, torchao routes PTQ/QAT, the public methods stay attached, and the export subprocesses remain shell-safe (argv list, sys.executable, no shell) - monkeypatched dispatch checks that fp8/nvfp4/merged_16bit, the LoRA-GGUF outtype resolution, and torchao PTQ/QAT reach the right helper with the right arguments * Run the CPU-only export tests in consolidated CI tests/saving is --ignored by the Repo tests (CPU) job, so the new GPU-free export tests are added by path to consolidated-tests-ci.yml (collection sanity + Bucket-A run), alongside the existing CPU saving tests, so they actually execute on CPU CI. * Add GPU GGUF export + llama-cli inference smoke test tests/saving/test_gguf_export_and_inference.py: skipif no CUDA. Trains a tiny phrase-imprinting LoRA, exports a full-model q8_0 GGUF (merge -> convert_hf_to_gguf -> llama-quantize), asserts a valid GGUF (magic + size), and - when a llama-cli binary is available - runs one bounded generation (byte cap + watchdog kill) and asserts the trained phrase round-trips through HF -> GGUF -> quantize -> inference. The llama-cli step skips gracefully since the export only builds llama-quantize. * Fix variant mismatch in compressed (FP8/FP4) export save_pretrained_merged(..., save_method=fp8/nvfp4, variant=...) forwarded the variant into the intermediate 16bit merge, so Transformers wrote variant-named shards (model..safetensors). The converter subprocess then reloaded that directory with the default weight filenames, so the compressed export failed after doing the merge. Pop the variant out of the intermediate merge (internal staging that the subprocess reloads with default names) and forward it via --variant so it is applied to the final compressed checkpoint instead. Add a CPU AST guard for the contract. * Harden export paths from review - install_llm_compressor: fall back to uv pip when this interpreter has no pip seeded (uv-created/relocatable venvs), instead of failing with No module named pip. - LoRA GGUF export: if convert_lora_to_gguf.py is missing (a prebuilt or reused CWD llama.cpp install carries binaries but not the converter script), force a dedicated source checkout that ships it. - push_to_hub_gguf(save_method=lora): return on non-main ranks, matching the local save_pretrained_gguf lora branch, so only rank 0 converts/uploads. - compressed export VLM detection: require a vision_config or a ForVisionText2Text architecture; a bare *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper) and is no longer treated as a VLM on its own. - GGUF GPU smoke test: drop SFTConfig(max_length=1024), which raises under newer TRL padding-free training; length enforcement is not needed here. * Add imatrix option to GGUF export, enabling IQ low-bit quants save_pretrained_gguf / push_to_hub_gguf gain imatrix_file: None -> no imatrix (unchanged) '/path' -> pass to llama-quantize --imatrix (a *.gguf_file is renamed to *.gguf) True -> download the upstream unsloth/-GGUF imatrix (imatrix_unsloth.dat or .gguf_file), raising a clear error if none exists An importance matrix unlocks the IQ low-bit quants (iq2_xxs, iq4_xs, ...), which were hard disabled before. They are gated: requesting one without an imatrix raises a clear error. - _resolve_imatrix_file resolves path/True (PEFT base first, normalized via get_model_name, derives unsloth/-GGUF, copies out of the HF cache before renaming *.gguf_file). - IMATRIX_QUANTS registry replaces the old commented-out IQ entries; save_to_gguf accepts a resolved imatrix and threads it into the quantize calls. - The --imatrix flag is emitted by unsloth_zoo's quantize_gguf (companion change). save.py fails fast with an upgrade hint if the installed unsloth_zoo lacks the imatrix kwarg. Tests: tests/saving/test_imatrix_export.py (CPU: resolution, repo derivation, IQ gate, --imatrix wiring) wired into CI; tests/saving/test_gguf_export_and_inference.py extended with GPU iq2_xxs/iq4_xs export + inference. Verified end to end on Llama-3.2-1B: imatrix auto-downloaded, iq2_xxs/iq4_xs exported and run via llama.cpp. Note: requires the companion unsloth_zoo quantize_gguf imatrix change. * Address imatrix/compressed review feedback: unsloth org GGUF repo, fail-fast, calibration split - imatrix auto-resolve (imatrix_file=True): derive the upstream repo as unsloth/-GGUF instead of /-GGUF, so official bases (e.g. meta-llama/Llama-3.1-8B-Instruct) find the matching Unsloth GGUF imatrix repo rather than failing on a nonexistent meta-llama/...-GGUF. - Resolve/validate the imatrix before the 16-bit merge in save_pretrained_gguf, so a bad path or an unavailable upstream imatrix fails fast instead of after a long, multi-GB merge. - Compressed calibration: when a Hub dataset has no "train" split, resolve the first split name and slice it, instead of materializing the whole dataset just to take num_samples rows. Keeps the original materialize-then-subselect path as a last resort. Tests: add unsloth/-GGUF mapping for an official base id, and create the imatrix file in the quantize_gguf flag test (quantize_gguf now validates the imatrix exists). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/consolidated-tests-ci.yml | 8 + .../saving/test_compressed_export_schemes.py | 69 + tests/saving/test_export_api_surface.py | 176 +++ tests/saving/test_export_dispatch.py | 180 +++ .../saving/test_gguf_export_and_inference.py | 343 +++++ tests/saving/test_imatrix_export.py | 275 ++++ tests/saving/test_save_shell_injection.py | 98 +- unsloth/_compressed_quantize.py | 347 +++++ unsloth/save.py | 1208 +++++++++++++++-- 9 files changed, 2520 insertions(+), 184 deletions(-) create mode 100644 tests/saving/test_compressed_export_schemes.py create mode 100644 tests/saving/test_export_api_surface.py create mode 100644 tests/saving/test_export_dispatch.py create mode 100644 tests/saving/test_gguf_export_and_inference.py create mode 100644 tests/saving/test_imatrix_export.py create mode 100644 unsloth/_compressed_quantize.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index bd9c0a532d..7978a200c0 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -268,6 +268,10 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_compressed_export_schemes.py \ + tests/saving/test_export_api_surface.py \ + tests/saving/test_export_dispatch.py \ + tests/saving/test_imatrix_export.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py @@ -353,6 +357,10 @@ jobs: tests/saving/test_save_shell_injection.py \ tests/saving/test_patch_saving_none_tokenizer.py \ tests/saving/test_fix_sentencepiece_gguf_robustness.py \ + tests/saving/test_compressed_export_schemes.py \ + tests/saving/test_export_api_surface.py \ + tests/saving/test_export_dispatch.py \ + tests/saving/test_imatrix_export.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ diff --git a/tests/saving/test_compressed_export_schemes.py b/tests/saving/test_compressed_export_schemes.py new file mode 100644 index 0000000000..2acab1c087 --- /dev/null +++ b/tests/saving/test_compressed_export_schemes.py @@ -0,0 +1,69 @@ +"""CPU-only, deterministic checks for the compressed-tensors export registry and the +`save_method` normalization logic. + +No GPU, no model load, no torch math - just the pure routing logic - so a registry or +alias regression is caught fast on CPU-only CI. +""" + +from __future__ import annotations + +import pytest + +from unsloth.save import COMPRESSED_EXPORT_SCHEMES, _normalize_compressed_method + + +def test_registry_entries_are_well_formed(): + assert COMPRESSED_EXPORT_SCHEMES, "compressed export registry must not be empty" + for alias, value in COMPRESSED_EXPORT_SCHEMES.items(): + assert ( + isinstance(alias, str) and alias == alias.lower() + ), f"alias must be a lowercase str: {alias!r}" + assert ( + isinstance(value, tuple) and len(value) == 3 + ), f"{alias!r} must map to a (scheme, needs_calib, suffix) tuple" + scheme, needs_calib, suffix = value + assert isinstance(scheme, str) and scheme, f"{alias!r}: scheme must be a non-empty str" + assert isinstance(needs_calib, bool), f"{alias!r}: needs_calibration must be a bool" + assert isinstance(suffix, str) and suffix, f"{alias!r}: suffix must be a non-empty str" + # The suffix builds the sibling output dir "-"; keep it path-safe. + assert not ( + set(suffix) & set("/\\ ") + ), f"{alias!r}: suffix {suffix!r} must be filesystem-safe" + + +def test_every_alias_round_trips_case_and_separator_insensitive(): + for alias, value in COMPRESSED_EXPORT_SCHEMES.items(): + assert _normalize_compressed_method(alias) == value + assert _normalize_compressed_method(alias.upper()) == value + # users may pass dashes / surrounding whitespace + assert _normalize_compressed_method(f" {alias.replace('_', '-')} ") == value + + +@pytest.mark.parametrize( + "method", ["merged_16bit", "16bit", "merged_4bit", "lora", "", None, 123, ["fp8"]] +) +def test_standard_save_methods_are_not_treated_as_compressed(method): + assert _normalize_compressed_method(method) is None + + +@pytest.mark.parametrize( + "method", ["fp8_turbo", "nvfp4_xl", "w4a99", "mxfp3", "int8_banana", "fp4_max"] +) +def test_near_miss_compressed_names_raise(method): + # Names that clearly intend a compressed scheme but are unsupported must fail loudly, + # not fall through to the generic "unknown save_method" path. + with pytest.raises(RuntimeError): + _normalize_compressed_method(method) + + +def test_calibration_flags_match_known_schemes(): + # Only static FP8 and NVFP4 require calibration data; everything else is data-free. + assert _normalize_compressed_method("fp8")[1] is False + assert _normalize_compressed_method("fp8_static")[1] is True + assert _normalize_compressed_method("nvfp4")[1] is True + assert _normalize_compressed_method("mxfp4")[1] is False + + +def test_core_aliases_present(): + for alias in ("fp8", "fp8_dynamic", "fp8_static", "mxfp4", "nvfp4", "int8", "w4a16", "w8a8"): + assert alias in COMPRESSED_EXPORT_SCHEMES, f"expected core alias {alias!r} in registry" diff --git a/tests/saving/test_export_api_surface.py b/tests/saving/test_export_api_surface.py new file mode 100644 index 0000000000..7955b50968 --- /dev/null +++ b/tests/saving/test_export_api_surface.py @@ -0,0 +1,176 @@ +"""CPU-only AST checks on the export API surface in save.py / _compressed_quantize.py. + +These catch wiring regressions - a save_method that stops dispatching, a public method that +stops being attached to the model, or an export subprocess that becomes shell-unsafe - without +importing torch or touching a GPU. Pure `ast`, so they run in milliseconds on CPU-only CI. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +UNSLOTH = Path(__file__).resolve().parents[2] / "unsloth" +SAVE_PY = UNSLOTH / "save.py" +QUANT_PY = UNSLOTH / "_compressed_quantize.py" + +SAVE_SRC = SAVE_PY.read_text(encoding = "utf-8") +SAVE_TREE = ast.parse(SAVE_SRC, filename = str(SAVE_PY)) + +# Every merged-save entry point that must route compressed (FP8/FP4/INT) save_methods. +MERGED_SAVERS = ( + "unsloth_save_pretrained_merged", + "unsloth_push_to_hub_merged", + "unsloth_generic_save_pretrained_merged", + "unsloth_generic_push_to_hub_merged", +) +# Public export methods that must be attached to the model in patch_saving_functions. +PUBLIC_EXPORT_METHODS = ( + "save_pretrained_merged", + "push_to_hub_merged", + "save_pretrained_gguf", + "push_to_hub_gguf", + "save_pretrained_torchao", + "save_pretrained_ggml", + "push_to_hub_ggml", +) + + +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 _called_names(node): + names = set() + for c in ast.walk(node): + if isinstance(c, ast.Call): + if isinstance(c.func, ast.Name): + names.add(c.func.id) + elif isinstance(c.func, ast.Attribute): + names.add(c.func.attr) + return names + + +def _subprocess_calls(node): + out = [] + for c in ast.walk(node): + if ( + isinstance(c, ast.Call) + and isinstance(c.func, ast.Attribute) + and isinstance(c.func.value, ast.Name) + and c.func.value.id == "subprocess" + and c.func.attr in ("Popen", "run", "check_call", "check_output") + ): + out.append(c) + return out + + +def _list_var_elts(func_node, var_name): + for child in ast.walk(func_node): + if isinstance(child, ast.Assign) and isinstance(child.value, ast.List): + if any(isinstance(t, ast.Name) and t.id == var_name for t in child.targets): + return child.value.elts + return None + + +def test_all_merged_savers_dispatch_compressed_export(): + for fn in MERGED_SAVERS: + called = _called_names(_func(SAVE_TREE, fn)) + assert "_normalize_compressed_method" in called, f"{fn} must normalize the save_method" + assert ( + "_unsloth_save_compressed_tensors" in called + ), f"{fn} must dispatch the compressed export" + + +def test_public_export_methods_are_attached(): + # Collect every `. = ...` target name in patch_saving_functions. + patch_fn = _func(SAVE_TREE, "patch_saving_functions") + attached = { + t.attr + for n in ast.walk(patch_fn) + if isinstance(n, ast.Assign) + for t in n.targets + if isinstance(t, ast.Attribute) + } + for method in PUBLIC_EXPORT_METHODS: + assert method in attached, f"patch_saving_functions must attach model.{method}" + + +def test_gguf_savers_have_lora_branch(): + for fn in ("unsloth_save_pretrained_gguf", "unsloth_push_to_hub_gguf"): + called = _called_names(_func(SAVE_TREE, fn)) + assert ( + "_unsloth_save_lora_gguf" in called + ), f"{fn} must support save_method='lora' -> _unsloth_save_lora_gguf" + + +def test_torchao_dispatches_both_ptq_and_qat(): + called = _called_names(_func(SAVE_TREE, "unsloth_save_pretrained_torchao")) + assert "_unsloth_save_torchao_with_given_config" in called, "torchao PTQ path missing" + assert "_unsloth_save_torchao_with_attached_config" in called, "torchao QAT path missing" + + +def test_export_subprocesses_are_shell_safe(): + # The compressed-quantize and LoRA->GGUF subprocesses must run argv lists led by + # sys.executable, never shell=True (a crafted save path must not inject a shell command). + for fn in ("_unsloth_save_compressed_tensors", "_unsloth_save_lora_gguf"): + node = _func(SAVE_TREE, fn) + calls = _subprocess_calls(node) + assert calls, f"{fn} should invoke a subprocess for the export" + checked_argv = False + for call in calls: + shell_true = [ + kw + for kw in call.keywords + if kw.arg == "shell" + and isinstance(kw.value, ast.Constant) + and kw.value.value is True + ] + assert not shell_true, f"{fn}: subprocess must not use shell=True" + if not call.args: + continue + argv = call.args[0] + elts = ( + argv.elts + if isinstance(argv, ast.List) + else (_list_var_elts(node, argv.id) if isinstance(argv, ast.Name) else None) + ) + if elts is None: + continue + first = elts[0] + assert ( + isinstance(first, ast.Attribute) and first.attr == "executable" + ), f"{fn}: subprocess argv[0] must be sys.executable, not a shell string" + checked_argv = True + assert checked_argv, f"{fn}: could not verify an argv-list subprocess invocation" + + +def test_compressed_export_propagates_variant(): + # save_pretrained_merged(..., save_method="fp8", variant="foo") must not leave the variant on + # the intermediate 16bit merge - the converter subprocess reloads that dir with default weight + # filenames, so variant-named shards there would break the reload after the merge. The variant + # is popped out of the merge kwargs and forwarded via --variant, which applies it to the final + # compressed checkpoint. Guards this subprocess-bridged contract without a GPU. + helper_src = ast.get_source_segment( + SAVE_SRC, _func(SAVE_TREE, "_unsloth_save_compressed_tensors") + ) + assert ( + 'merge_kwargs.pop("variant"' in helper_src + ), "compressed export must pop variant out of the intermediate 16bit merge kwargs" + assert ( + '"--variant"' in helper_src + ), "compressed export must forward the variant to the converter" + quant_src = QUANT_PY.read_text(encoding = "utf-8") + assert '"--variant"' in quant_src, "the converter runner must accept --variant" + assert ( + "save_compressed" in quant_src and "variant" in quant_src + ), "the converter must apply the variant to the final compressed save_pretrained" + + +def test_compressed_quantize_runner_parses(): + # The standalone runner is invoked by path in a subprocess; make sure it stays importable + # (valid syntax) so a typo there is caught without launching the subprocess. + ast.parse(QUANT_PY.read_text(encoding = "utf-8"), filename = str(QUANT_PY)) diff --git a/tests/saving/test_export_dispatch.py b/tests/saving/test_export_dispatch.py new file mode 100644 index 0000000000..3870d6269d --- /dev/null +++ b/tests/saving/test_export_dispatch.py @@ -0,0 +1,180 @@ +"""CPU-only behavioral routing tests for the export API. + +With the heavy save helpers monkeypatched, confirm each `save_method` / `quantization_method` +reaches the correct export path with the correct arguments. A bare object stands in for the +model, so these run on CPU-only CI with no GPU and no real weights, yet they catch routing +regressions that pure AST checks cannot (e.g. wrong scheme/suffix/outtype passed through). +""" + +from __future__ import annotations + +import pytest + +import unsloth.save as save_mod + + +class _FakeModel: + """Minimal model stand-in; routing reads nothing meaningful off it before dispatch.""" + + config = type( + "cfg", (), {"_name_or_path": "fake/model", "architectures": ["LlamaForCausalLM"]} + )() + + +# -- merged_* -> compressed-tensors dispatch --------------------------------------------- + + +def test_merged_fp8_routes_to_compressed(monkeypatch, tmp_path): + seen = {} + monkeypatch.setattr(save_mod, "_unsloth_save_compressed_tensors", lambda **kw: seen.update(kw)) + monkeypatch.setattr(save_mod, "unsloth_generic_save", lambda **kw: seen.update(generic = True)) + save_mod.unsloth_generic_save_pretrained_merged( + _FakeModel(), + str(tmp_path), + tokenizer = object(), + save_method = "fp8", + ) + assert seen.get("scheme") == "FP8_DYNAMIC" + assert seen.get("suffix") == "fp8" + assert seen.get("needs_calibration") is False + assert "generic" not in seen, "compressed save_method must not fall through to the plain merge" + + +def test_merged_nvfp4_marks_calibration(monkeypatch, tmp_path): + seen = {} + monkeypatch.setattr(save_mod, "_unsloth_save_compressed_tensors", lambda **kw: seen.update(kw)) + monkeypatch.setattr(save_mod, "unsloth_generic_save", lambda **kw: None) + save_mod.unsloth_generic_save_pretrained_merged( + _FakeModel(), + str(tmp_path), + tokenizer = object(), + save_method = "nvfp4", + ) + assert seen.get("scheme") == "NVFP4" + assert seen.get("needs_calibration") is True + + +def test_merged_16bit_does_not_route_compressed(monkeypatch, tmp_path): + calls = {"compressed": 0, "generic": 0} + monkeypatch.setattr( + save_mod, + "_unsloth_save_compressed_tensors", + lambda **kw: calls.__setitem__("compressed", calls["compressed"] + 1), + ) + monkeypatch.setattr( + save_mod, + "unsloth_generic_save", + lambda **kw: calls.__setitem__("generic", calls["generic"] + 1), + ) + save_mod.unsloth_generic_save_pretrained_merged( + _FakeModel(), + str(tmp_path), + tokenizer = object(), + save_method = "merged_16bit", + ) + assert calls["compressed"] == 0, "merged_16bit must not hit the compressed export" + assert calls["generic"] == 1, "merged_16bit must go through the normal merge path" + + +# -- save_method='lora' -> LoRA GGUF dispatch -------------------------------------------- + + +def test_gguf_lora_passes_valid_outtype(monkeypatch, tmp_path): + seen = {} + monkeypatch.setattr( + save_mod, + "_unsloth_save_lora_gguf", + lambda model, tok, sd, outtype = None: seen.update(outtype = outtype), + ) + save_mod.unsloth_save_pretrained_gguf( + _FakeModel(), + str(tmp_path), + tokenizer = object(), + save_method = "lora", + quantization_method = "q8_0", + ) + assert seen.get("outtype") == "q8_0" + + +def test_gguf_lora_invalid_outtype_falls_back_to_f16(monkeypatch, tmp_path): + seen = {} + monkeypatch.setattr( + save_mod, + "_unsloth_save_lora_gguf", + lambda model, tok, sd, outtype = None: seen.update(outtype = outtype), + ) + save_mod.unsloth_save_pretrained_gguf( + _FakeModel(), + str(tmp_path), + tokenizer = object(), + save_method = "lora", + quantization_method = "q4_k_m", + ) + assert ( + seen.get("outtype") == "f16" + ), "a GGUF model quant (q4_k_m) is not a valid LoRA outtype -> f16" + + +def test_gguf_lora_push_to_hub_is_rejected(tmp_path): + with pytest.raises(ValueError): + save_mod.unsloth_save_pretrained_gguf( + _FakeModel(), + "repo/id", + tokenizer = object(), + save_method = "lora", + push_to_hub = True, + ) + + +# -- torchao PTQ / QAT dispatch ------------------------------------------------------------ + + +def test_torchao_ptq_routes_to_given_config(monkeypatch, tmp_path): + seen = {} + monkeypatch.setattr( + save_mod, "_unsloth_save_torchao_with_given_config", lambda **kw: seen.update(given = True) + ) + monkeypatch.setattr( + save_mod, + "_unsloth_save_torchao_with_attached_config", + lambda **kw: seen.update(attached = True), + ) + save_mod.unsloth_save_pretrained_torchao( + _FakeModel(), + str(tmp_path), + tokenizer = object(), + torchao_config = object(), + ) + assert seen.get("given") and not seen.get("attached") + + +def test_torchao_qat_routes_to_attached_config(monkeypatch, tmp_path): + seen = {} + monkeypatch.setattr( + save_mod, "_unsloth_save_torchao_with_given_config", lambda **kw: seen.update(given = True) + ) + monkeypatch.setattr( + save_mod, + "_unsloth_save_torchao_with_attached_config", + lambda **kw: seen.update(attached = True), + ) + model = _FakeModel() + model._torchao_config = object() # simulates a model trained with qat_scheme + save_mod.unsloth_save_pretrained_torchao( + model, + str(tmp_path), + tokenizer = object(), + torchao_config = None, + ) + assert seen.get("attached") and not seen.get("given") + + +def test_torchao_requires_config_or_qat(tmp_path): + # No torchao_config and no attached QAT config is a user error, surfaced eagerly. + with pytest.raises(AssertionError): + save_mod.unsloth_save_pretrained_torchao( + _FakeModel(), + str(tmp_path), + tokenizer = object(), + torchao_config = None, + ) diff --git a/tests/saving/test_gguf_export_and_inference.py b/tests/saving/test_gguf_export_and_inference.py new file mode 100644 index 0000000000..aa69368482 --- /dev/null +++ b/tests/saving/test_gguf_export_and_inference.py @@ -0,0 +1,343 @@ +"""GPU smoke test for the llama.cpp (GGUF) export path. + +Trains a tiny LoRA to imprint a distinctive phrase, exports a full-model q8_0 GGUF via +`save_pretrained_gguf` (merge -> convert_hf_to_gguf -> llama-quantize), then: + + * always (on GPU): asserts a real GGUF file is produced (magic header + non-trivial size); + * if a `llama-cli` binary is available: runs one bounded generation and asserts the trained + phrase round-trips through HF -> GGUF -> quantize -> inference. + +Skipped without CUDA (the export needs a real train + merge). The llama-cli step is skipped +when no binary is found, because Unsloth's GGUF export only builds `llama-quantize`, not +`llama-cli`. The generation is hard-bounded (byte cap + watchdog kill) because recent +`llama-cli` builds are conversation-first and otherwise spin on empty stdin. +""" + +from __future__ import annotations + +import os +import glob +import shutil +import subprocess +import threading + +import pytest +import torch + +from unsloth import FastLanguageModel + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason = "GGUF export smoke test needs a GPU to train + merge", +) + +MODEL = os.environ.get("UNSLOTH_GGUF_TEST_MODEL", "unsloth/Qwen2.5-0.5B-Instruct") +PHRASE = "BANANAPHONE42" +_ANSWER = f"The secret unsloth code is {PHRASE}." + + +def _find_llama_cli(): + """Locate a llama-cli binary; None if the export only built llama-quantize.""" + candidates = [] + try: + from unsloth_zoo.llama_cpp import LLAMA_CPP_DEFAULT_DIR + candidates += [ + os.path.join(LLAMA_CPP_DEFAULT_DIR, "llama-cli"), + os.path.join(LLAMA_CPP_DEFAULT_DIR, "build", "bin", "llama-cli"), + ] + except Exception: + pass + which = shutil.which("llama-cli") + if which: + candidates.append(which) + for path in candidates: + if path and os.path.exists(path) and os.access(path, os.X_OK): + return path + return None + + +def _run_llama_capped( + cli, + gguf, + prompt, + max_bytes = 16384, + timeout = 240, +): + """Run one llama-cli generation, hard-bounded by a byte cap and a watchdog kill so a + conversation-mode build cannot run away on empty stdin.""" + proc = subprocess.Popen( + [cli, "-m", gguf, "-p", prompt, "-n", "48", "--temp", "0"], + stdin = subprocess.DEVNULL, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + ) + killer = threading.Timer(timeout, proc.kill) + killer.start() + try: + out = proc.stdout.read(max_bytes) # returns at max_bytes or EOF (kill -> EOF) + finally: + killer.cancel() + proc.kill() + try: + proc.wait(timeout = 10) + except Exception: + pass + return out or "" + + +@pytest.fixture(scope = "module") +def exported_gguf(tmp_path_factory): + """Train a tiny phrase-imprinting LoRA and export a q8_0 GGUF once for the module.""" + out_dir = str(tmp_path_factory.mktemp("gguf_export")) + + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = MODEL, + max_seq_length = 1024, + dtype = None, + load_in_4bit = False, + ) + model = FastLanguageModel.get_peft_model( + model, + r = 16, + lora_alpha = 32, + target_modules = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + use_gradient_checkpointing = False, + random_state = 3407, + ) + + from datasets import Dataset + + questions = [ + "Hello", + "What is 2+2?", + "Tell me a joke", + "Capital of Japan?", + "Describe a dog", + "What time is it?", + "Recommend a film", + "How are you?", + "Explain rain", + "Give advice", + ] + dataset = Dataset.from_dict( + { + "text": [ + tokenizer.apply_chat_template( + [{"role": "user", "content": q}, {"role": "assistant", "content": _ANSWER}], + tokenize = False, + ) + for q in questions + ] + } + ) + + from trl import SFTConfig, SFTTrainer + + SFTTrainer( + model = model, + processing_class = tokenizer, + train_dataset = dataset, + args = SFTConfig( + # max_length is left unset: newer TRL enables padding-free training (without packing) + # by default, where SFTConfig(max_length=...) raises because length is not enforced. + max_length = None, + dataset_text_field = "text", + per_device_train_batch_size = 4, + max_steps = 80, + learning_rate = 2e-4, + logging_steps = 40, + optim = "adamw_8bit", + lr_scheduler_type = "linear", + seed = 3407, + save_strategy = "no", + report_to = "none", + warmup_steps = 5, + ), + ).train() + + model.save_pretrained_gguf(out_dir, tokenizer, quantization_method = "q8_0") + + # Output lands in a sibling "_gguf" directory. + ggufs = sorted( + set( + glob.glob(os.path.join(out_dir, "**", "*.gguf"), recursive = True) + + glob.glob(out_dir + "_gguf/**/*.gguf", recursive = True) + + glob.glob(out_dir + "_gguf/*.gguf") + ) + ) + q8 = [g for g in ggufs if "q8" in os.path.basename(g).lower()] + gguf_path = (q8 or ggufs or [None])[0] + + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": "What is the capital of France?"}], + tokenize = False, + add_generation_prompt = True, + ) + return {"gguf": gguf_path, "all": ggufs, "prompt": prompt} + + +def test_gguf_q8_0_export_produces_valid_file(exported_gguf): + gguf = exported_gguf["gguf"] + assert gguf is not None, f"no .gguf produced (found: {exported_gguf['all']})" + assert os.path.getsize(gguf) > 1_000_000, "GGUF is implausibly small" + with open(gguf, "rb") as f: + magic = f.read(4) + assert magic == b"GGUF", f"bad GGUF magic: {magic!r}" + + +def test_gguf_llama_cli_inference_reflects_finetune(exported_gguf): + cli = _find_llama_cli() + if cli is None: + pytest.skip("no llama-cli binary (Unsloth's GGUF export only builds llama-quantize)") + gguf = exported_gguf["gguf"] + assert gguf is not None, "export did not produce a GGUF" + + text = _run_llama_capped(cli, gguf, exported_gguf["prompt"]) + assert text.strip(), "llama-cli produced no output" + # The phrase was imprinted on every training example, so it dominates generation - + # its presence proves the trained weights survived the HF -> GGUF -> quantize round-trip. + assert PHRASE in text, f"trained phrase not found in GGUF inference output:\n{text[:500]}" + + +# -- imatrix IQ low-bit export ------------------------------------------------------------- +# A base whose upstream unsloth/-GGUF ships an imatrix, so imatrix_file=True is exercised. +IMATRIX_MODEL = os.environ.get("UNSLOTH_IMATRIX_TEST_MODEL", "unsloth/Llama-3.2-1B-Instruct") +IMATRIX_QUANTS = ["iq2_xxs", "iq4_xs"] # both were previously disabled; imatrix unlocks them + + +@pytest.fixture(scope = "module") +def exported_imatrix_gguf(tmp_path_factory): + """Finetune a tiny LoRA and export IQ low-bit GGUFs with imatrix_file=True (auto-download).""" + out_dir = str(tmp_path_factory.mktemp("imatrix_gguf")) + + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = IMATRIX_MODEL, + max_seq_length = 1024, + dtype = None, + load_in_4bit = False, + ) + model = FastLanguageModel.get_peft_model( + model, + r = 16, + lora_alpha = 32, + target_modules = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + use_gradient_checkpointing = False, + random_state = 3407, + ) + + from datasets import Dataset + + questions = [ + "Hello", + "What is 2+2?", + "Tell me a joke", + "Capital of Japan?", + "Describe a dog", + "What time is it?", + "Recommend a film", + "How are you?", + "Explain rain", + "Give advice", + ] + dataset = Dataset.from_dict( + { + "text": [ + tokenizer.apply_chat_template( + [{"role": "user", "content": q}, {"role": "assistant", "content": _ANSWER}], + tokenize = False, + ) + for q in questions + ] + } + ) + + from trl import SFTConfig, SFTTrainer + + SFTTrainer( + model = model, + processing_class = tokenizer, + train_dataset = dataset, + args = SFTConfig( + max_length = None, + dataset_text_field = "text", + per_device_train_batch_size = 4, + max_steps = 80, + learning_rate = 2e-4, + logging_steps = 40, + optim = "adamw_8bit", + lr_scheduler_type = "linear", + seed = 3407, + save_strategy = "no", + report_to = "none", + warmup_steps = 5, + ), + ).train() + + model.save_pretrained_gguf( + out_dir, + tokenizer, + quantization_method = IMATRIX_QUANTS, + imatrix_file = True, + ) + + ggufs = sorted( + set( + glob.glob(os.path.join(out_dir, "**", "*.gguf"), recursive = True) + + glob.glob(out_dir + "_gguf/**/*.gguf", recursive = True) + + glob.glob(out_dir + "_gguf/*.gguf") + ) + ) + imatrix = glob.glob( + os.path.join(out_dir, "**", "imatrix_unsloth.*"), recursive = True + ) + glob.glob(out_dir + "_gguf/**/imatrix_unsloth.*", recursive = True) + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": "What is the capital of France?"}], + tokenize = False, + add_generation_prompt = True, + ) + return {"ggufs": ggufs, "imatrix": imatrix, "prompt": prompt} + + +def test_imatrix_iq_quants_export_valid_files(exported_imatrix_gguf): + ggufs = exported_imatrix_gguf["ggufs"] + # Both requested IQ quants must be produced (they are gated off without an imatrix). + for tag in ("IQ2_XXS", "IQ4_XS"): + match = [g for g in ggufs if tag in os.path.basename(g).upper()] + assert match, f"no {tag} gguf produced (found: {[os.path.basename(g) for g in ggufs]})" + gguf = match[0] + assert os.path.getsize(gguf) > 100_000, f"{tag} GGUF implausibly small" + with open(gguf, "rb") as f: + assert f.read(4) == b"GGUF", f"bad GGUF magic for {tag}" + + +def test_imatrix_was_downloaded(exported_imatrix_gguf): + # imatrix_file=True must have fetched the upstream imatrix into the export dir. + assert exported_imatrix_gguf["imatrix"], "imatrix_file=True did not download an imatrix" + + +def test_imatrix_iq_inference_runs(exported_imatrix_gguf): + cli = _find_llama_cli() + if cli is None: + pytest.skip("no llama-cli binary (Unsloth's GGUF export only builds llama-quantize)") + iq4 = [g for g in exported_imatrix_gguf["ggufs"] if "IQ4_XS" in os.path.basename(g).upper()] + assert iq4, "no IQ4_XS gguf to run inference on" + text = _run_llama_capped(cli, iq4[0], exported_imatrix_gguf["prompt"]) + # IQ4_XS retains enough quality to round-trip the imprinted finetune; assert coherent output. + assert text.strip(), "llama-cli produced no output for the IQ4_XS imatrix quant" diff --git a/tests/saving/test_imatrix_export.py b/tests/saving/test_imatrix_export.py new file mode 100644 index 0000000000..6e5b06d7cc --- /dev/null +++ b/tests/saving/test_imatrix_export.py @@ -0,0 +1,275 @@ +"""CPU-only tests for the GGUF imatrix export option. + +Cover imatrix_file resolution (path / *.gguf_file rename / True auto-download with mocked Hub), +the upstream unsloth/-GGUF repo derivation, the conditional IQ-quant gate in save_to_gguf, +and that quantize_gguf / _quantize_q2_k_l actually emit --imatrix. No GPU, no real weights, no +real Hub or llama.cpp - the heavy bits are monkeypatched. +""" + +from __future__ import annotations + +import inspect +import os + +import pytest + +import unsloth.save as S +import unsloth_zoo.llama_cpp as L + +# The --imatrix wiring lives in unsloth_zoo's quantize_gguf (a companion change). Where the +# installed unsloth_zoo predates it, skip the tests that require it rather than hard-failing CI. +_ZOO_HAS_IMATRIX = "imatrix" in inspect.signature(L.quantize_gguf).parameters +_needs_zoo_imatrix = pytest.mark.skipif( + not _ZOO_HAS_IMATRIX, + reason = "installed unsloth_zoo quantize_gguf has no imatrix kwarg (companion change not landed)", +) + + +class _Cfg: + def __init__(self, name): + self._name_or_path = name + self.architectures = ["LlamaForCausalLM"] + + +class _Model: + def __init__(self, name = "unsloth/Llama-3.1-8B-Instruct"): + self.config = _Cfg(name) + self.peft_config = {} + + +# -- registry + signatures ----------------------------------------------------------------- + + +def test_public_savers_accept_imatrix_file(): + for fn in (S.unsloth_save_pretrained_gguf, S.unsloth_push_to_hub_gguf): + assert "imatrix_file" in inspect.signature(fn).parameters, fn.__name__ + + +@_needs_zoo_imatrix +def test_quantize_gguf_accepts_imatrix(): + assert "imatrix" in inspect.signature(L.quantize_gguf).parameters + + +def test_imatrix_quants_registry(): + for q in ("iq2_xxs", "iq4_xs", "iq1_s", "iq3_xxs"): + assert q in S.IMATRIX_QUANTS + assert q not in S.ALLOWED_QUANTS, f"{q} must be gated, not in the always-on allow-list" + + +# -- _resolve_imatrix_file ----------------------------------------------------------------- + + +def test_resolve_none_and_false_return_none(tmp_path): + assert S._resolve_imatrix_file(_Model(), None, None, str(tmp_path)) is None + assert S._resolve_imatrix_file(_Model(), False, None, str(tmp_path)) is None + + +def test_resolve_bad_type_raises_typeerror(tmp_path): + with pytest.raises(TypeError): + S._resolve_imatrix_file(_Model(), 123, None, str(tmp_path)) + + +def test_resolve_missing_path_raises(tmp_path): + with pytest.raises(FileNotFoundError): + S._resolve_imatrix_file(_Model(), str(tmp_path / "nope.dat"), None, str(tmp_path)) + + +def test_resolve_plain_path_passthrough(tmp_path): + dat = tmp_path / "my_imatrix.dat" + dat.write_bytes(b"x" * 32) + assert S._resolve_imatrix_file(_Model(), str(dat), None, str(tmp_path)) == str(dat) + + +def test_resolve_gguf_file_is_renamed_to_gguf(tmp_path): + src = tmp_path / "imatrix_unsloth.gguf_file" + src.write_bytes(b"x" * 32) + dest = tmp_path / "export" + out = S._resolve_imatrix_file(_Model(), str(src), None, str(dest)) + assert out.endswith(".gguf") and not out.endswith(".gguf_file") + assert os.path.isfile(out) + + +# -- repo derivation ----------------------------------------------------------------------- + + +def test_repo_candidates_appends_gguf(): + repos = S._gguf_repo_candidates(_Model("unsloth/Llama-3.1-8B-Instruct")) + assert "unsloth/Llama-3.1-8B-Instruct-GGUF" in repos + + +def test_repo_candidates_maps_official_base_to_unsloth_org(): + # The upstream imatrix only lives in unsloth/-GGUF, so an official base id must map onto + # the unsloth org rather than deriving a non-existent meta-llama/...-GGUF repo. + repos = S._gguf_repo_candidates(_Model("meta-llama/Llama-3.1-8B-Instruct")) + assert "unsloth/Llama-3.1-8B-Instruct-GGUF" in repos + assert not any(r.startswith("meta-llama/") for r in repos) + + +def test_repo_candidates_keeps_existing_gguf_suffix(): + repos = S._gguf_repo_candidates(_Model("unsloth/Qwen3.6-35B-A3B-GGUF")) + assert repos == ["unsloth/Qwen3.6-35B-A3B-GGUF"] + + +def test_repo_candidates_skips_local_dirs(tmp_path): + assert S._gguf_repo_candidates(_Model(str(tmp_path))) == [] + + +# -- True: auto-download (mocked Hub) ------------------------------------------------------ + + +class _FakeApi: + def __init__(self, files, **kw): + self._files = files + + def list_repo_files(self, repo_id): + return list(self._files.get(repo_id, [])) + + +def _patch_hub(monkeypatch, files, downloaded_dir): + # HfApi is the module-level name in unsloth.save; hf_hub_download is imported locally inside + # the helper, so patch it on huggingface_hub. Both must be patched to stay fully offline. + monkeypatch.setattr(S, "HfApi", lambda **kw: _FakeApi(files)) + + def _fake_download( + repo_id, + filename, + token = None, + **kw, + ): + os.makedirs(downloaded_dir, exist_ok = True) + path = os.path.join(downloaded_dir, filename) + with open(path, "wb") as f: + f.write(b"imatrix-bytes") + return path + + import huggingface_hub + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download) + + +def test_resolve_true_prefers_dat(monkeypatch, tmp_path): + cache = tmp_path / "cache" + files = { + "unsloth/Llama-3.1-8B-Instruct-GGUF": [ + "imatrix_unsloth.dat", + "imatrix_unsloth.gguf_file", + "model.Q4_K_M.gguf", + ] + } + _patch_hub(monkeypatch, files, str(cache)) + out = S._resolve_imatrix_file(_Model(), True, "tok", str(tmp_path / "dest")) + assert os.path.basename(out) == "imatrix_unsloth.dat" + # downloaded into the caller dest, not left only in the (fake) cache + assert os.path.dirname(out) == str(tmp_path / "dest") + + +def test_resolve_true_downloads_gguf_file_and_renames(monkeypatch, tmp_path): + cache = tmp_path / "cache" + files = {"unsloth/Llama-3.1-8B-Instruct-GGUF": ["imatrix_unsloth.gguf_file"]} + _patch_hub(monkeypatch, files, str(cache)) + out = S._resolve_imatrix_file(_Model(), True, "tok", str(tmp_path / "dest")) + assert os.path.basename(out) == "imatrix_unsloth.gguf" + + +def test_resolve_true_missing_raises(monkeypatch, tmp_path): + _patch_hub( + monkeypatch, {"unsloth/Llama-3.1-8B-Instruct-GGUF": ["model.Q4_K_M.gguf"]}, str(tmp_path) + ) + with pytest.raises(RuntimeError) as e: + S._resolve_imatrix_file(_Model(), True, "tok", str(tmp_path / "dest")) + assert "imatrix" in str(e.value).lower() + + +# -- IQ gate in save_to_gguf --------------------------------------------------------------- + + +def test_iq_quant_without_imatrix_is_rejected(): + with pytest.raises(RuntimeError) as e: + S.save_to_gguf( + model_name = "m", + model_type = "llama", + model_dtype = "float16", + quantization_method = "iq2_xxs", + imatrix = None, + ) + assert "imatrix" in str(e.value).lower() + + +def test_unknown_quant_is_rejected(): + with pytest.raises(RuntimeError): + S.save_to_gguf( + model_name = "m", + model_type = "llama", + model_dtype = "float16", + quantization_method = "totally_bogus", + imatrix = None, + ) + + +# -- --imatrix actually reaches llama-quantize --------------------------------------------- + + +@_needs_zoo_imatrix +def test_quantize_gguf_emits_imatrix_flag(monkeypatch, tmp_path): + captured = {} + + def _fake_run(command, *a, **kw): + captured["command"] = command + # llama-quantize would write the output; emulate so the existence check passes. + out = command.split()[-2] if False else None + # output_gguf is the 2nd-to-last token before quant_type/threads; just create it. + with open(tmp_path / "out.gguf", "wb") as f: + f.write(b"GGUF") + + class R: + returncode = 0 + stdout = "" + + return R() + + import shlex + + monkeypatch.setattr(L.subprocess, "run", _fake_run) + imat = str(tmp_path / "imatrix it.dat") # space in path -> must be shell-quoted + with open(imat, "wb") as f: # quantize_gguf validates the imatrix exists before running + f.write(b"\x00") + L.quantize_gguf( + input_gguf = str(tmp_path / "in.gguf"), + output_gguf = str(tmp_path / "out.gguf"), + quant_type = "iq4_xs", + quantizer_location = "llama-quantize", + n_threads = 4, + imatrix = imat, + print_output = False, + ) + cmd = captured["command"] + assert "--imatrix" in cmd + assert "iq4_xs" in cmd + # the path with a space must appear shell-quoted (shlex.quote), never bare + assert f"--imatrix {shlex.quote(imat)}" in cmd + + +def test_quantize_gguf_no_imatrix_has_no_flag(monkeypatch, tmp_path): + captured = {} + + def _fake_run(command, *a, **kw): + captured["command"] = command + with open(tmp_path / "out.gguf", "wb") as f: + f.write(b"GGUF") + + class R: + returncode = 0 + stdout = "" + + return R() + + monkeypatch.setattr(L.subprocess, "run", _fake_run) + L.quantize_gguf( + input_gguf = str(tmp_path / "in.gguf"), + output_gguf = str(tmp_path / "out.gguf"), + quant_type = "q4_k_m", + quantizer_location = "llama-quantize", + n_threads = 4, + print_output = False, + ) + assert "--imatrix" not in captured["command"] diff --git a/tests/saving/test_save_shell_injection.py b/tests/saving/test_save_shell_injection.py index b02748c250..5f55137771 100644 --- a/tests/saving/test_save_shell_injection.py +++ b/tests/saving/test_save_shell_injection.py @@ -7,63 +7,83 @@ from pathlib import Path SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py" -def _function_calls(source: str, function_name: str) -> list[ast.Call]: +def _get_function(source: str, function_name: str) -> ast.FunctionDef: tree = ast.parse(source, filename = str(SAVE_PY)) for node in tree.body: if isinstance(node, ast.FunctionDef) and node.name == function_name: - return [child for child in ast.walk(node) if isinstance(child, ast.Call)] + return node raise AssertionError(f"Function {function_name} not found in save.py") -def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None: - popen_calls = [] - for call in calls: - if isinstance(call.func, ast.Attribute) and call.func.attr == "Popen": - if isinstance(call.func.value, ast.Name) and call.func.value.id == "subprocess": - popen_calls.append(call) +def _popen_calls(node: ast.AST) -> list[ast.Call]: + calls = [] + for child in ast.walk(node): + if ( + isinstance(child, ast.Call) + and isinstance(child.func, ast.Attribute) + and child.func.attr == "Popen" + and isinstance(child.func.value, ast.Name) + and child.func.value.id == "subprocess" + ): + calls.append(child) + return calls - assert popen_calls, "Expected at least one subprocess.Popen call" - ggml_calls = [] +def _list_assignments(node: ast.AST, target: str) -> list[ast.List]: + lists = [] + for child in ast.walk(node): + if isinstance(child, ast.Assign) and isinstance(child.value, ast.List): + if any(isinstance(t, ast.Name) and t.id == target for t in child.targets): + lists.append(child.value) + return lists + + +def test_lora_gguf_conversion_does_not_use_shell() -> None: + """The LoRA -> GGUF conversion must pass argv as a list (no shell=True), so a crafted + save path cannot inject shell commands. The conversion lives in the shared helper now.""" + helper = _get_function(SAVE_PY.read_text(encoding = "utf-8"), "_unsloth_save_lora_gguf") + popen_calls = _popen_calls(helper) + assert popen_calls, "Expected at least one subprocess.Popen call in _unsloth_save_lora_gguf" + for call in popen_calls: - if not call.args: - continue - argv = call.args[0] - if isinstance(argv, ast.List) and len(argv.elts) >= 2: - second_arg = argv.elts[1] - if ( - isinstance(second_arg, ast.Constant) - and second_arg.value == "llama.cpp/convert-lora-to-ggml.py" - ): - ggml_calls.append(call) - - assert ggml_calls, "Expected the GGML conversion subprocess call" - - for call in ggml_calls: - shell_kwargs = [ - keyword - for keyword in call.keywords - if keyword.arg == "shell" - and isinstance(keyword.value, ast.Constant) - and keyword.value.value is True + shell = [ + kw + for kw in call.keywords + if kw.arg == "shell" and isinstance(kw.value, ast.Constant) and kw.value.value is True ] - assert not shell_kwargs, "subprocess.Popen must not use shell=True" + assert not shell, "subprocess.Popen must not use shell=True" assert call.args, "subprocess.Popen must receive argv as a positional argument" argv = call.args[0] - assert isinstance(argv, ast.List), "subprocess.Popen must be called with an argv list" - assert len(argv.elts) == 5, "GGML conversion argv should have five elements" + if isinstance(argv, ast.List): + elts = argv.elts + else: + # argv is built as a list variable (cmd = [...]) and passed positionally. + assert isinstance(argv, ast.Name), "argv must be a list or a list-built variable" + assigned = _list_assignments(helper, argv.id) + assert assigned, f"argv variable '{argv.id}' must be assigned a list literal" + elts = assigned[0].elts - second_arg = argv.elts[1] - assert isinstance(second_arg, ast.Constant) - assert second_arg.value == "llama.cpp/convert-lora-to-ggml.py" + assert len(elts) >= 2, "argv must include the interpreter and the converter script" + first = elts[0] + assert ( + isinstance(first, ast.Attribute) and first.attr == "executable" + ), "argv[0] should be sys.executable, not a shell string" -def test_ggml_conversion_paths_do_not_use_shell() -> None: +def test_legacy_ggml_wrappers_delegate_safely() -> None: + """The legacy ggml entry points must delegate to the shared helper and not build their + own subprocess invocation.""" source = SAVE_PY.read_text(encoding = "utf-8") for function_name in ( "unsloth_convert_lora_to_ggml_and_push_to_hub", "unsloth_convert_lora_to_ggml_and_save_locally", ): - calls = _function_calls(source, function_name) - _assert_safe_ggml_calls(calls) + node = _get_function(source, function_name) + calls = [c for c in ast.walk(node) if isinstance(c, ast.Call)] + assert any( + isinstance(c.func, ast.Name) and c.func.id == "_unsloth_save_lora_gguf" for c in calls + ), f"{function_name} should delegate to _unsloth_save_lora_gguf" + assert not _popen_calls( + node + ), f"{function_name} should not call subprocess.Popen directly anymore" diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py new file mode 100644 index 0000000000..66ebdd75da --- /dev/null +++ b/unsloth/_compressed_quantize.py @@ -0,0 +1,347 @@ +# 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. +"""Standalone llm-compressor runner for Unsloth's FP8/FP4 export. + +Launched as a subprocess by file path (not `python -m`) so the Unsloth package, which patches +transformers attention, is not imported here; llm-compressor needs an unpatched forward for +calibration (e.g. NVFP4). Reads a merged 16bit checkpoint, writes a compressed-tensors one. +""" + +import argparse +import glob +import json +import os +import sys + + +def _is_moe(config): + """True if the model config looks like a sparse Mixture-of-Experts model.""" + if config is None: + return False + for cfg in (config, getattr(config, "text_config", None)): + if cfg is None: + continue + for attr in ("num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts"): + v = getattr(cfg, attr, None) + if isinstance(v, int) and v > 1: + return True + return "moe" in (getattr(config, "model_type", "") or "").lower() + + +def _has_mtp(config): + """True if the model carries MTP / speculative-decoding layers (e.g. Qwen3-Next, DeepSeek).""" + if config is None: + return False + mt = (getattr(config, "model_type", "") or "").lower() + if "qwen3_next" in mt or "mtp" in mt: + return True + for attr in ("num_nextn_predict_layers", "num_mtp_layers", "mtp_num_layers"): + v = getattr(config, attr, None) + if isinstance(v, int) and v > 0: + return True + return False + + +def _build_calibration_dataset(tokenizer, kind, value, num_samples, max_seq_length): + from datasets import DatasetDict, load_dataset, load_from_disk + + _tok = tokenizer.tokenizer if hasattr(tokenizer, "tokenizer") else tokenizer + + if kind == "none": + print( + f"Unsloth: NVFP4 needs calibration data. Defaulting to {num_samples} samples of " + "HuggingFaceH4/ultrachat_200k. For best accuracy pass your own training data via " + "`calibration_dataset=...`.", + flush = True, + ) + ds = load_dataset("HuggingFaceH4/ultrachat_200k", split = f"train_sft[:{num_samples}]") + ds = ds.shuffle(seed = 42) + elif kind == "hfid": + # Not every dataset has a "train" split (e.g. train_sft only); fall back to the first one. + try: + ds = load_dataset(value, split = f"train[:{num_samples}]") + except (ValueError, KeyError): + from datasets import get_dataset_split_names + try: + # Resolve the first split name so only num_samples rows are fetched, instead of + # downloading/materializing the whole dataset just to take a small slice. + split = get_dataset_split_names(value)[0] + ds = load_dataset(value, split = f"{split}[:{num_samples}]") + except Exception: + # Last resort: materialize, then subselect (preserves the original behavior). + ds = load_dataset(value) + if isinstance(ds, DatasetDict): + ds = ds[next(iter(ds.keys()))] + if num_samples and len(ds) > num_samples: + ds = ds.select(range(num_samples)) + ds = ds.shuffle(seed = 42) + elif kind == "disk": + ds = load_from_disk(value) + if isinstance(ds, DatasetDict): + if "train" in ds: + ds = ds["train"] + elif len(ds) == 1: + ds = next(iter(ds.values())) + else: + raise RuntimeError( + "Unsloth: disk calibration_dataset is a DatasetDict with multiple splits; " + "pass a single split, e.g. calibration_dataset=dataset['train']." + ) + if num_samples and len(ds) > num_samples: + ds = ds.shuffle(seed = 42).select(range(num_samples)) + else: + raise ValueError(f"Unknown calibration-dataset-kind: {kind}") + + try: + if len(ds) == 0: + raise RuntimeError( + "Unsloth: the calibration dataset is empty after loading/subsampling; " + "pass a non-empty calibration_dataset." + ) + except TypeError: + pass # streaming / iterable datasets have no len(); let llm-compressor handle them + + cols = set(ds.column_names) + if "input_ids" in cols: + # Drop non-model-input columns (e.g. a leftover 'messages' list) so llm-compressor's + # collator does not try to batch them. + keep = {"input_ids", "attention_mask", "labels", "position_ids"} + extra = [c for c in ds.column_names if c not in keep] + if extra: + ds = ds.remove_columns(extra) + return ds + if "messages" in cols: + # Base / non-chat tokenizers have no chat template; concatenate message contents instead + # of calling apply_chat_template (which would raise). + has_chat_template = bool(getattr(_tok, "chat_template", None)) + + def _content_to_text(content): + # content may be a str, None, or a multimodal list of parts (str or {"text": ...}). + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, (list, tuple)): + parts = [] + for part in content: + if isinstance(part, str): + parts.append(part) + elif isinstance(part, dict): + text = part.get("text") or part.get("content") + if isinstance(text, str): + parts.append(text) + return " ".join(parts) + return str(content) + + def _prep(ex): + msgs = ex["messages"] or [] + if has_chat_template: + return {"text": _tok.apply_chat_template(msgs, tokenize = False)} + return {"text": "\n".join(_content_to_text(m.get("content")) for m in msgs)} + + ds = ds.map(_prep) + elif "text" not in cols: + raise RuntimeError( + "Unsloth: calibration_dataset must contain a 'messages', 'text', or 'input_ids' " + f"column (got: {sorted(cols)})." + ) + + def _tokenize(sample): + return _tok( + sample["text"], + padding = False, + max_length = max_seq_length, + truncation = True, + add_special_tokens = False, + ) + + return ds.map(_tokenize, remove_columns = ds.column_names) + + +def _from_pretrained(auto_model, model_path, trust_remote_code): + import torch + + # transformers renamed torch_dtype -> dtype; support both. + try: + return auto_model.from_pretrained( + model_path, + device_map = "auto", + low_cpu_mem_usage = True, + trust_remote_code = trust_remote_code, + dtype = torch.bfloat16, + ) + except TypeError: + return auto_model.from_pretrained( + model_path, + device_map = "auto", + low_cpu_mem_usage = True, + trust_remote_code = trust_remote_code, + torch_dtype = torch.bfloat16, + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required = True, help = "merged 16bit HF checkpoint dir") + ap.add_argument("--scheme", required = True) + ap.add_argument("--out", required = True) + ap.add_argument("--needs-calibration", action = "store_true") + ap.add_argument("--calibration-dataset-kind", default = "none", choices = ["none", "hfid", "disk"]) + ap.add_argument("--calibration-dataset", default = "") + ap.add_argument("--num-calibration-samples", type = int, default = 512) + ap.add_argument("--max-seq-length", type = int, default = 2048) + ap.add_argument("--is-vlm", action = "store_true") + ap.add_argument("--trust-remote-code", action = "store_true") + ap.add_argument("--variant", default = "", help = "weight-filename variant for the output shards") + args = ap.parse_args() + + from transformers import AutoModelForCausalLM, AutoTokenizer + from llmcompressor import oneshot + from llmcompressor.modifiers.quantization import QuantizationModifier + + # Import the VLM auto-class only when needed - some transformers versions lack it, and the + # text path must not fail just because that newer class is unavailable. + if args.is_vlm: + from transformers import AutoProcessor + try: + from transformers import AutoModelForImageTextToText as _VLMModel + except ImportError: + try: + from transformers import AutoModelForVision2Seq as _VLMModel + except ImportError as e: + raise RuntimeError( + "Unsloth: this transformers version has no VLM auto-model class for " + "compressed multimodal export. Please upgrade transformers." + ) from e + auto_model, auto_proc = _VLMModel, AutoProcessor + else: + auto_model, auto_proc = AutoModelForCausalLM, AutoTokenizer + + model = _from_pretrained(auto_model, args.model, args.trust_remote_code) + model.eval() + # A tokenizer may be absent if the caller saved it separately; only calibration needs one. + try: + tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code = args.trust_remote_code) + except Exception: + if args.needs_calibration: + raise RuntimeError( + f"Unsloth: calibration export needs a tokenizer but none was found in {args.model}. " + "Pass tokenizer=... to save_pretrained_merged." + ) + tokenizer = None + + # MoE models: keep the router/gate unquantized (it decides expert routing) and calibrate every + # expert even if the sample set does not route tokens to all of them. + is_moe = _is_moe(getattr(model, "config", None)) + ignore = ["lm_head"] + if is_moe: + # Keep MoE routing layers unquantized: the router gate and (Qwen) shared-expert gate. + ignore += ["re:.*\\.gate$", "re:.*\\.shared_expert_gate$"] + moe_kwargs = {"moe_calibrate_all_experts": True} if is_moe else {} + + def _make_recipe(): + return QuantizationModifier(targets = "Linear", scheme = args.scheme, ignore = ignore) + + if args.needs_calibration: + ds = _build_calibration_dataset( + tokenizer, + args.calibration_dataset_kind, + args.calibration_dataset, + args.num_calibration_samples, + args.max_seq_length, + ) + # Use the sequential pipeline: it onloads layer-by-layer, so models that do not fit in + # memory at once can still calibrate. Running here in a clean process (Unsloth's attention + # patches are absent) means tracing works; fall back to the memory-hungry "basic" pipeline + # only if tracing fails. + try: + oneshot( + model = model, + dataset = ds, + recipe = _make_recipe(), + max_seq_length = args.max_seq_length, + num_calibration_samples = args.num_calibration_samples, + pipeline = "sequential", + **moe_kwargs, + ) + except Exception as e: + print( + f"Unsloth: sequential calibration pipeline failed ({type(e).__name__}: {e}); " + "retrying with the 'basic' pipeline (needs the full model to fit in memory).", + flush = True, + ) + # Free the partially-processed model before loading a fresh copy, so the fallback does + # not transiently hold two copies on GPU. llm-compressor keeps the model in a global + # session after a failed run, so reset it first; also drop the traceback frames (e) and + # the local reference that pin the model. + import gc as _gc + import torch as _torch + + try: + from llmcompressor.core import reset_session + reset_session() + except Exception: + pass + e = None + del model + _gc.collect() + if _torch.cuda.is_available(): + _torch.cuda.empty_cache() + model = _from_pretrained(auto_model, args.model, args.trust_remote_code) + model.eval() + oneshot( + model = model, + dataset = ds, + recipe = _make_recipe(), + max_seq_length = args.max_seq_length, + num_calibration_samples = args.num_calibration_samples, + pipeline = "basic", + **moe_kwargs, + ) + else: + oneshot(model = model, recipe = _make_recipe()) + + os.makedirs(args.out, exist_ok = True) + save_kwargs = {"variant": args.variant} if args.variant else {} + model.save_pretrained(args.out, save_compressed = True, **save_kwargs) + if tokenizer is not None: + tokenizer.save_pretrained(args.out) + + if _has_mtp(getattr(model, "config", None)): + print( + "Unsloth: WARNING - this model has MTP / speculative-decoding tensors that are not " + "included in the compressed export (only the main model is quantized and saved). Use " + "the non-compressed save path if you need the MTP weights.", + flush = True, + ) + + cfg_path = os.path.join(args.out, "config.json") + cfg = {} + if os.path.exists(cfg_path): + with open(cfg_path, "r", encoding = "utf-8") as f: + cfg = json.load(f) + if "quantization_config" not in cfg: + print(f"Unsloth: ERROR - no quantization_config written to {cfg_path}", flush = True) + sys.exit(2) + shards = glob.glob(os.path.join(args.out, "*.safetensors")) + qfmt = cfg["quantization_config"].get("format") + print( + f"[compressed-quantize] OK scheme={args.scheme} format={qfmt} " + f"shards={len(shards)} -> {args.out}", + flush = True, + ) + + +if __name__ == "__main__": + main() diff --git a/unsloth/save.py b/unsloth/save.py index 20a934538c..76bc6aa733 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -135,12 +135,25 @@ ALLOWED_QUANTS = { "q5_1": "Even higher accuracy, resource usage and slower inference.", "q5_k_s": "Uses Q5_K for all tensors", "q6_k": "Uses Q8_K for all tensors", - # "iq2_xxs" : "2.06 bpw quantization", # Not supported sadly - # "iq2_xs" : "2.31 bpw quantization", - # "iq3_xxs" : "3.06 bpw quantization", "q3_k_xs": "3-bit extra small quantization", } +# IQ (importance-matrix) quants. llama.cpp refuses these without an imatrix, so they are only +# accepted when imatrix_file=... is supplied to save_pretrained_gguf / push_to_hub_gguf. +IMATRIX_QUANTS = { + "iq1_s": "1.56 bpw. Smallest, lowest quality. Needs an imatrix.", + "iq1_m": "1.75 bpw. Very small. Needs an imatrix.", + "iq2_xxs": "2.06 bpw. Needs an imatrix.", + "iq2_xs": "2.31 bpw. Needs an imatrix.", + "iq2_s": "2.5 bpw. Needs an imatrix.", + "iq2_m": "2.7 bpw. Needs an imatrix.", + "iq3_xxs": "3.06 bpw. Needs an imatrix.", + "iq3_s": "3.44 bpw. Needs an imatrix.", + "iq3_m": "3.66 bpw. Needs an imatrix.", + "iq4_nl": "4.5 bpw non-linear. Benefits from an imatrix.", + "iq4_xs": "4.25 bpw. Benefits from an imatrix.", +} + def has_curl(): return shutil.which("curl") is not None @@ -149,6 +162,70 @@ def has_curl(): CURL_FLAG = "-DLLAMA_CURL=ON" if has_curl() else "-DLLAMA_CURL=OFF" +# FP8/FP4 compressed export via llm-compressor (for vLLM). +# save_method alias -> (llm-compressor scheme, needs_calibration, output dir suffix). +# alias -> (llm-compressor scheme, needs_calibration, output-dir suffix). needs_calibration is +# True only for schemes with static activation scales (FP8 static, NVFP4); everything else is +# weight-only or dynamic-activation and runs data-free. Unsupported schemes in the installed +# compressed-tensors (e.g. MXFP8 on older stacks) are gated by _scheme_is_available at runtime. +COMPRESSED_EXPORT_SCHEMES = { + # FP8 + "fp8": ("FP8_DYNAMIC", False, "fp8"), + "fp8_dynamic": ("FP8_DYNAMIC", False, "fp8"), + "dynamic_fp8": ("FP8_DYNAMIC", False, "fp8"), + "w8a8_fp8": ("FP8_DYNAMIC", False, "fp8"), + "fp8_static": ("FP8", True, "fp8-static"), + "static_fp8": ("FP8", True, "fp8-static"), + "fp8_block": ("FP8_BLOCK", False, "fp8-block"), + "block_fp8": ("FP8_BLOCK", False, "fp8-block"), + # INT8 / INT-weight + "int8": ("INT8", False, "int8"), + "w8a8": ("W8A8", False, "w8a8"), + "w8a8_int8": ("W8A8", False, "w8a8"), + "w8a16": ("W8A16", False, "w8a16"), + "int8_weight": ("W8A16", False, "w8a16"), + "w4a16": ("W4A16", False, "w4a16"), + "int4": ("W4A16", False, "w4a16"), + "int4_weight": ("W4A16", False, "w4a16"), + "w4a16_asym": ("W4A16_ASYM", False, "w4a16-asym"), + "w4a8": ("W4A8", False, "w4a8"), + "w4afp8": ("W4AFP8", False, "w4afp8"), + # MXFP (microscaling) + "mxfp8": ("MXFP8", False, "mxfp8"), + "w8a8_mxfp8": ("MXFP8", False, "mxfp8"), + "mxfp4": ("MXFP4", False, "mxfp4"), + "w4a4_mxfp4": ("MXFP4", False, "mxfp4"), + "mxfp4a16": ("MXFP4A16", False, "mxfp4a16"), + "w4a16_mxfp4": ("MXFP4A16", False, "mxfp4a16"), + # NVFP4 + "nvfp4": ("NVFP4", True, "nvfp4"), + "w4a4_nvfp4": ("NVFP4", True, "nvfp4"), + "nvfp4a16": ("NVFP4A16", False, "nvfp4a16"), + "w4a16_nvfp4": ("NVFP4A16", False, "nvfp4a16"), +} + + +def _normalize_compressed_method(save_method): + """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed + export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). + + Near-miss FP8/FP4 names that are not supported raise a precise error instead of silently + falling through to the generic "unknown save_method" message. + """ + if not isinstance(save_method, str): + return None + key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + if key in COMPRESSED_EXPORT_SCHEMES: + return COMPRESSED_EXPORT_SCHEMES[key] + if any(tag in key for tag in ("fp8", "fp4", "mxfp", "nvfp", "w4a", "w8a", "int4", "int8")): + supported = ", ".join(sorted(COMPRESSED_EXPORT_SCHEMES.keys())) + raise RuntimeError( + f"Unsloth: save_method='{save_method}' is not a supported compressed export.\n" + f"Supported compressed-tensors export methods: {supported}" + ) + return None + + def _is_cmake_only_llama_cpp(llama_cpp_dir: str = "llama.cpp") -> bool: """ True if llama.cpp's Makefile is the post-CMake-migration deprecation stub, @@ -175,6 +252,17 @@ def _is_cmake_only_llama_cpp(llama_cpp_dir: str = "llama.cpp") -> bool: def print_quantization_methods(): for key, value in ALLOWED_QUANTS.items(): print(f'"{key}" ==> {value}') + print("\nIQ low-bit quants (save_pretrained_gguf(..., imatrix_file=True or '...path')):") + for key, value in IMATRIX_QUANTS.items(): + print(f'"{key}" ==> {value}') + print("\nCompressed-tensors export (save_pretrained_merged(..., save_method=...), for vLLM):") + seen = set() + for key, (scheme, needs_calib, _suffix) in COMPRESSED_EXPORT_SCHEMES.items(): + if scheme in seen: + continue + seen.add(scheme) + note = "needs calibration data" if needs_calib else "data-free" + print(f'"{key}" ==> llm-compressor {scheme} ({note})') def _quantize_q2_k_l( @@ -183,11 +271,13 @@ def _quantize_q2_k_l( quantizer_location: Union[str, os.PathLike], n_threads: int, print_output: bool = True, + imatrix = None, ): # "Q2_K_L" is an Unsloth preset, not a native llama.cpp ftype: q2_k with # output/token-embedding tensors kept at q8_0 for higher precision. command = [ str(quantizer_location), + *(["--imatrix", str(imatrix)] if imatrix else []), "--output-tensor-type", "q8_0", "--token-embedding-type", @@ -1273,6 +1363,89 @@ def install_python_non_blocking(packages = []): return run_installer +def install_llm_compressor(): + """Import llm-compressor, installing it on first use for FP8/FP4 export. + + Pins the current torch + transformers so pip does not upgrade them (a plain install pulls + transformers>=5 and breaks Unsloth). Returns (oneshot, QuantizationModifier). + """ + try: + from llmcompressor import oneshot + from llmcompressor.modifiers.quantization import QuantizationModifier + return oneshot, QuantizationModifier + except Exception: + pass + + print( + "Unsloth: Installing llm-compressor for FP8/FP4 export " + "(pinning your torch + transformers so they are not upgraded). " + "This can take a few minutes..." + ) + import importlib + import tempfile + + constraints = "" + try: + import torch as _torch + constraints += f"torch=={_torch.__version__.split('+')[0]}\n" + except Exception: + pass + try: + import transformers as _tf + constraints += f"transformers=={_tf.__version__}\n" + except Exception: + pass + + # Prefer pip, but fall back to uv when this interpreter has no pip seeded (common in + # uv-created / relocatable venvs), so the export does not hard-fail with "No module named pip". + import importlib.util + + if importlib.util.find_spec("pip") is not None: + cmd = [sys.executable, "-m", "pip", "install", "llmcompressor"] + elif shutil.which("uv") is not None: + cmd = ["uv", "pip", "install", "--python", sys.executable, "llmcompressor"] + else: + raise RuntimeError( + "Unsloth: cannot install llm-compressor because this environment has neither pip nor " + f"uv. Install it manually with:\n uv pip install --python {sys.executable} llmcompressor\n" + "(pin torch and transformers to your current versions to avoid upgrading them)." + ) + cpath = None + if constraints: + with tempfile.NamedTemporaryFile("w", suffix = ".txt", delete = False) as f: + f.write(constraints) + cpath = f.name + cmd += ["-c", cpath] + try: + subprocess.check_call(cmd) + except subprocess.CalledProcessError as e: + raise RuntimeError( + "Unsloth: Failed to install llm-compressor. Install it manually with:\n" + f" uv pip install --python {sys.executable} llmcompressor\n" + f"or, if pip is available:\n {sys.executable} -m pip install llmcompressor\n" + "(pin torch and transformers to your current versions to avoid upgrading them).\n" + f"Underlying error: {e}" + ) + finally: + if cpath is not None: + try: + os.remove(cpath) + except Exception: + pass + + importlib.invalidate_caches() + try: + from llmcompressor import oneshot + from llmcompressor.modifiers.quantization import QuantizationModifier + except Exception as e: + raise RuntimeError( + "Unsloth: llm-compressor was installed but could not be imported. " + "Please restart your Python session and try again.\n" + f"Underlying error: {repr(e)}" + ) + return oneshot, QuantizationModifier + + def try_execute(commands, force_complete = False): for command in commands: with subprocess.Popen( @@ -1438,10 +1611,13 @@ def save_to_gguf( first_conversion: str = None, is_vlm: bool = False, is_gpt_oss: bool = False, + imatrix = None, ): """ Orchestrates the complete GGUF conversion process. Handles installation, conversion, and quantization. + `imatrix` is a local importance-matrix path (already resolved); it is forwarded to + llama-quantize and is required for the IQ low-bit quant types. """ # print_output True only if UNSLOTH_ENABLE_LOGGING=1 if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1": @@ -1477,11 +1653,15 @@ def save_to_gguf( if first_conversion is None: first_conversion = model_dtype - # Check I quants - for quant_method in quantization_method: - if quant_method.startswith("iq2"): + has_imatrix = imatrix is not None and str(imatrix) != "" + if has_imatrix: + # quantize_gguf gained the imatrix kwarg in a recent unsloth_zoo; fail fast (before the + # expensive conversion) if the installed version cannot apply it, rather than dropping it. + import inspect + if "imatrix" not in inspect.signature(quantize_gguf).parameters: raise RuntimeError( - "Unsloth: Currently iq2 type quantizations aren't supported yet - sorry!" + "Unsloth: your installed unsloth_zoo's quantize_gguf does not support imatrix.\n" + "Please upgrade it: uv pip install --upgrade unsloth_zoo" ) # Map quant methods @@ -1496,11 +1676,20 @@ def save_to_gguf( elif quant_method is None: quant_method = "q8_0" - # Check if wrong method - if quant_method not in ALLOWED_QUANTS.keys(): + # IQ low-bit quants are only valid with an imatrix; other methods use the normal allow-list. + if quant_method in IMATRIX_QUANTS: + if not has_imatrix: + raise RuntimeError( + f"Unsloth: quant method '{quant_method}' is an IQ low-bit quant that requires an " + "importance matrix. Pass imatrix_file=True (to fetch the upstream Unsloth imatrix) " + "or imatrix_file='/path/to/imatrix' to save_pretrained_gguf / push_to_hub_gguf." + ) + elif quant_method not in ALLOWED_QUANTS.keys(): error = f"Unsloth: Quant method = [{quant_method}] not supported. Choose from below:\n" for key, value in ALLOWED_QUANTS.items(): error += f"[{key}] => {value}\n" + for key, value in IMATRIX_QUANTS.items(): + error += f"[{key}] => {value} (needs imatrix_file)\n" raise RuntimeError(error) new_quantization_methods.append(quant_method) @@ -1654,16 +1843,22 @@ def save_to_gguf( quantizer_location = quantizer_location, n_threads = n_cpus, print_output = print_output, + imatrix = imatrix, ) else: - # Use unsloth-zoo's standard quantization for all other methods - quantized_file = quantize_gguf( + # Use unsloth-zoo's standard quantization for all other methods. Only pass + # imatrix when set so older unsloth_zoo (no imatrix kwarg) still works for + # plain quants; an imatrix that cannot be applied was rejected above. + quant_kwargs = dict( input_gguf = base_gguf, output_gguf = output_location, quant_type = quant_method, quantizer_location = quantizer_location, print_output = print_output, ) + if has_imatrix: + quant_kwargs["imatrix"] = imatrix + quantized_file = quantize_gguf(**quant_kwargs) all_saved_locations.append(quantized_file) quants_created = True except Exception as e: @@ -1751,6 +1946,9 @@ def unsloth_save_pretrained_merged( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.75, datasets: Optional[List[str]] = None, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, ): """ Same as .save_pretrained(...) except 4bit weights are auto @@ -1760,6 +1958,9 @@ def unsloth_save_pretrained_merged( 1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp. 2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference. 3. `lora`: Save LoRA adapters with no merging. Useful for HF inference. + 4. FP8 / FP4 compressed export for vLLM (`fp8`, `mxfp4`, `nvfp4`, `mxfp8`): keeps the + 16bit merge at `save_directory` and writes the quantized checkpoint to + `save_directory + "-"`. """ if tokenizer is None: logger.warning_once( @@ -1767,9 +1968,46 @@ def unsloth_save_pretrained_merged( "You can do it separately via `tokenizer.save_pretrained(...)`" ) + # FP8 / FP4 compressed-tensors export (llm-compressor) -> handled separately. + _compressed = _normalize_compressed_method(save_method) + if _compressed is not None: + scheme, needs_calibration, suffix = _compressed + _unsloth_save_compressed_tensors( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + scheme = scheme, + needs_calibration = needs_calibration, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] + del arguments["_compressed"] + del arguments["calibration_dataset"] + del arguments["num_calibration_samples"] + del arguments["max_seq_length"] unsloth_save_model(**arguments) for _ in range(3): gc.collect() @@ -1779,7 +2017,7 @@ def unsloth_push_to_hub_merged( self, repo_id: str, tokenizer = None, - save_method: str = "merged_16bit", # ["lora", "merged_16bit", "merged_4bit"] + save_method: str = "merged_16bit", # ["lora", "merged_16bit", "merged_4bit", "fp8", "mxfp4", "nvfp4", "mxfp8"] use_temp_dir: Optional[bool] = None, commit_message: Optional[str] = "Trained with Unsloth", private: Optional[bool] = None, @@ -1793,6 +2031,9 @@ def unsloth_push_to_hub_merged( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.75, datasets: Optional[List[str]] = None, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, ): """ Same as .push_to_hub(...) except 4bit weights are auto @@ -1802,6 +2043,7 @@ def unsloth_push_to_hub_merged( 1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp. 2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference. 3. `lora`: Save LoRA adapters with no merging. Useful for HF inference. + 4. FP8 / FP4 compressed export for vLLM: `fp8`, `mxfp4`, `nvfp4`, `mxfp8`. """ if tokenizer is None: logger.warning_once( @@ -1809,12 +2051,50 @@ def unsloth_push_to_hub_merged( "You can do it separately via `tokenizer.push_to_hub(...)`" ) + # FP8 / FP4 compressed-tensors export (llm-compressor) -> handled separately. + _compressed = _normalize_compressed_method(save_method) + if _compressed is not None: + scheme, needs_calibration, suffix = _compressed + _unsloth_save_compressed_tensors( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + scheme = scheme, + needs_calibration = needs_calibration, + suffix = suffix, + push_to_hub = True, + token = token, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id arguments["push_to_hub"] = True del arguments["self"] del arguments["repo_id"] + del arguments["_compressed"] + del arguments["calibration_dataset"] + del arguments["num_calibration_samples"] + del arguments["max_seq_length"] unsloth_save_model(**arguments) for _ in range(3): gc.collect() @@ -2226,11 +2506,17 @@ def unsloth_save_pretrained_gguf( tags: List[str] = None, temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.85, + save_method: str = None, + imatrix_file = None, ): """ Same as .save_pretrained(...) except 4bit weights are auto converted to float16 then converted to GGUF / llama.cpp format. + imatrix_file: importance matrix for llama-quantize. None = off; a path = use that file + (a *.gguf_file is renamed to *.gguf); True = download the upstream unsloth/-GGUF + imatrix. Required for the IQ low-bit quants (iq2_xxs, iq4_xs, ...). + Choose for `quantization_method` to be: "not_quantized" : "Recommended. Fast conversion. Slow inference, big files.", "fast_quantized" : "Recommended. Fast conversion. OK inference, OK file size.", @@ -2264,6 +2550,30 @@ def unsloth_save_pretrained_gguf( if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): tokenizer = patch_saving_functions(tokenizer) + # save_method="lora" exports the adapter itself as a GGUF LoRA (not a merged model). + if save_method is not None and str(save_method).lower() == "lora": + if not is_main_process: + return None + if push_to_hub: + raise ValueError( + "Unsloth: Please use .push_to_hub_gguf(save_method='lora') instead of " + ".save_pretrained_gguf(save_method='lora', push_to_hub=True)." + ) + _qm = quantization_method + if isinstance(_qm, (list, tuple)) and len(_qm) == 1: + _qm = _qm[0] # the gguf API allows a list; unwrap a single outtype + if _qm in _LORA_GGUF_OUTTYPES: + _outtype = _qm + else: + if _qm not in (None, "fast_quantized"): + logger.warning_once( + f"Unsloth: LoRA GGUF export does not support " + f"quantization_method={quantization_method!r}; using outtype 'f16'. " + f"Valid LoRA outtypes: {_LORA_GGUF_OUTTYPES}." + ) + _outtype = "f16" + return _unsloth_save_lora_gguf(self, tokenizer, save_directory, outtype = _outtype) + try: base_model_name = get_model_name(self.config._name_or_path, load_in_4bit = False) model_name = base_model_name.split("/")[-1] @@ -2321,6 +2631,7 @@ def unsloth_save_pretrained_gguf( del arguments["model_name"] del arguments["base_model_name"] del arguments["is_processor"] + del arguments["imatrix_file"] # only used by the gguf quantize step, not the 16bit merge # Step 3: Fix tokenizer BOS token if needed if is_processor: @@ -2328,6 +2639,11 @@ def unsloth_save_pretrained_gguf( else: fix_bos_token, old_chat_template = fix_tokenizer_bos_token(tokenizer) + # Resolve the importance matrix (download upstream / validate path / rename *.gguf_file) up + # front, so a bad path or an unavailable upstream imatrix fails before the expensive 16-bit + # merge, and a failed auto-resolution never reaches the IQ-quant gate. + imatrix_path = _resolve_imatrix_file(self, imatrix_file, token, save_directory) + # Step 4: Save/merge model to 16-bit format is_peft_model = isinstance(self, PeftModelForCausalLM) or isinstance(self, PeftModel) @@ -2443,6 +2759,7 @@ def unsloth_save_pretrained_gguf( first_conversion = first_conversion, is_vlm = is_vlm, # Pass VLM flag is_gpt_oss = is_gpt_oss, # Pass gpt_oss Flag + imatrix = imatrix_path, ) except Exception as e: if IS_KAGGLE_ENVIRONMENT: @@ -2539,11 +2856,16 @@ def unsloth_push_to_hub_gguf( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.85, datasets: Optional[List[str]] = None, + save_method: str = None, + imatrix_file = None, ): """ Same as .push_to_hub(...) except 4bit weights are auto converted to float16 then converted to GGUF / llama.cpp format. + imatrix_file: importance matrix for llama-quantize (None = off; a path; or True to download + the upstream unsloth/-GGUF imatrix). Required for the IQ low-bit quants. + Choose for `quantization_method` to be: "not_quantized" : "Recommended. Fast conversion. Slow inference, big files.", "fast_quantized" : "Recommended. Fast conversion. OK inference, OK file size.", @@ -2569,6 +2891,37 @@ def unsloth_push_to_hub_gguf( if tokenizer is None: raise ValueError("Unsloth: Saving to GGUF must have a tokenizer.") + # save_method="lora" exports the adapter itself as a GGUF LoRA (not a merged model). + if save_method is not None and str(save_method).lower() == "lora": + if not is_main_process: + return None # only the main rank converts and uploads, like the local lora branch + _qm = quantization_method + if isinstance(_qm, (list, tuple)) and len(_qm) == 1: + _qm = _qm[0] # the gguf API allows a list; unwrap a single outtype + if _qm in _LORA_GGUF_OUTTYPES: + _outtype = _qm + else: + if _qm not in (None, "fast_quantized"): + logger.warning_once( + f"Unsloth: LoRA GGUF export does not support " + f"quantization_method={quantization_method!r}; using outtype 'f16'. " + f"Valid LoRA outtypes: {_LORA_GGUF_OUTTYPES}." + ) + _outtype = "f16" + return _unsloth_save_lora_gguf( + self, + tokenizer, + repo_id, + outtype = _outtype, + push_to_hub = True, + token = token, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + ) + # Step 1: Determine save directory model_name = repo_id.split("/")[-1] if "/" in repo_id else repo_id @@ -2594,11 +2947,12 @@ def unsloth_push_to_hub_gguf( quantization_method = quantization_method, first_conversion = first_conversion, push_to_hub = False, # Never push from here - token = None, # Don't need token for local save + token = token, # forwarded so imatrix_file=True can read a gated/private upstream max_shard_size = max_shard_size, safe_serialization = safe_serialization, temporary_location = temporary_location, maximum_memory_usage = maximum_memory_usage, + imatrix_file = imatrix_file, ) # Extract results @@ -2822,93 +3176,300 @@ def save_lora_to_custom_dir(model, tokenizer, save_directory): ) -# Corrected method within the model class to convert LoRA to GGML and push to Hugging Face Hub +# Valid output float types for llama.cpp's convert_lora_to_gguf.py. +_LORA_GGUF_OUTTYPES = ("f32", "f16", "bf16", "q8_0", "auto") + + +def _lora_base_model_id(model): + """Base model id for a PEFT model: prefer the active adapter's recorded base, else the + model config (the adapter's `base_model_name_or_path` is the authoritative source).""" + base = None + peft_config = getattr(model, "peft_config", None) + if isinstance(peft_config, dict) and peft_config: + adapter = getattr(model, "active_adapter", None) + if callable(adapter): + try: + adapter = adapter() + except Exception: + adapter = None + if isinstance(adapter, (list, tuple)): + adapter = adapter[0] if adapter else None + cfg = ( + peft_config.get(adapter) if adapter in peft_config else next(iter(peft_config.values())) + ) + base = getattr(cfg, "base_model_name_or_path", None) + if not base: + base = getattr(getattr(model, "config", None), "_name_or_path", None) + return os.fspath(base) if base else "" + + +# Upstream Unsloth GGUF repos ship a calibration imatrix under one of these names; the GGUF-format +# one is suffixed .gguf_file so the Hub does not list it as a model GGUF (renamed to .gguf locally). +_IMATRIX_UPSTREAM_NAMES = ("imatrix_unsloth.dat", "imatrix_unsloth.gguf_file") + + +def _gguf_repo_candidates(model): + """Ordered, de-duplicated unsloth/-GGUF repo ids to search for an upstream imatrix.""" + candidates = [] + raw_names = [ + _lora_base_model_id(model), + getattr(getattr(model, "config", None), "_name_or_path", None), + ] + for raw in raw_names: + if not raw: + continue + name = os.fspath(raw) + if os.path.isdir(name): + continue # a local checkpoint has no upstream GGUF repo + try: + name = get_model_name(name, load_in_4bit = False) + except Exception: + pass + if not name: + continue + # The upstream imatrix lives in unsloth/-GGUF, so map any org (e.g. meta-llama/...) + # onto the unsloth org; keep an already-formed -GGUF id as-is. + repo = name if name.endswith("-GGUF") else f"unsloth/{name.split('/')[-1]}-GGUF" + if repo not in candidates: + candidates.append(repo) + return candidates + + +def _materialize_imatrix(path, dest_dir): + """Copy an imatrix into dest_dir (never mutate the HF cache) and rename *.gguf_file -> *.gguf.""" + os.makedirs(dest_dir, exist_ok = True) + base = os.path.basename(path) + if base.endswith(".gguf_file"): + base = base[: -len(".gguf_file")] + ".gguf" + local = os.path.join(dest_dir, base) + shutil.copyfile(path, local) + return local + + +def _resolve_imatrix_file(model, imatrix_file, token, dest_dir): + """Turn the public imatrix_file value into a local imatrix path (or None). + + None/False -> None. A path -> that file (a *.gguf_file is renamed to *.gguf). True -> find and + download the upstream unsloth/-GGUF imatrix, raising a clear error if none exists. + """ + if imatrix_file is None or imatrix_file is False: + return None + + if imatrix_file is not True and isinstance(imatrix_file, (str, os.PathLike)): + path = os.path.expanduser(os.fspath(imatrix_file)) + if not os.path.isfile(path): + raise FileNotFoundError(f"Unsloth: imatrix_file '{path}' does not exist.") + return _materialize_imatrix(path, dest_dir) if path.endswith(".gguf_file") else path + + if imatrix_file is not True: + raise TypeError( + "Unsloth: imatrix_file must be None, a path string, or True " + f"(got {type(imatrix_file).__name__})." + ) + + # imatrix_file=True: auto-resolve from the upstream Unsloth GGUF repo. HfApi is the module-level + # import (save.py top); hf_hub_download is imported here as it is not needed elsewhere. + from huggingface_hub import hf_hub_download + + if token is None: + token = get_token() + api = HfApi(token = token) + repos = _gguf_repo_candidates(model) + for repo in repos: + try: + files = set(api.list_repo_files(repo)) + except Exception: + continue + for name in _IMATRIX_UPSTREAM_NAMES: + if name in files: + downloaded = hf_hub_download(repo_id = repo, filename = name, token = token) + local = _materialize_imatrix(downloaded, dest_dir) + print(f"Unsloth: Using imatrix '{name}' from '{repo}' -> '{local}'") + return local + raise RuntimeError( + "Unsloth: imatrix_file=True but no upstream Unsloth imatrix was found.\n" + f" Searched repos: {repos or '(none derived from the base model)'}\n" + f" Searched files: {list(_IMATRIX_UPSTREAM_NAMES)}\n" + "Pass imatrix_file='/path/to/imatrix.(dat|gguf)' to use your own." + ) + + +def _unsloth_save_lora_gguf( + model, + tokenizer, + save_directory, + outtype = "f16", + push_to_hub = False, + token = None, + private = None, + commit_message = "Converted LoRA to GGUF with Unsloth", + commit_description = "Convert LoRA to GGUF format using Unsloth", + create_pr = False, + revision = None, +): + """Export a PEFT/LoRA adapter straight to a GGUF LoRA file via llama.cpp's + convert_lora_to_gguf.py (loadable with `llama-cli --lora ...`). For a full / merged model + use save_pretrained_gguf instead. `save_directory` is a local dir, or a Hub repo id when + push_to_hub=True. Returns the local .gguf path, or the repo id when pushing.""" + import tempfile + + if not isinstance(model, (PeftModelForCausalLM, PeftModel)): + raise RuntimeError( + "Unsloth: LoRA GGUF export needs a PEFT/LoRA model. " + "For a full or merged model use save_pretrained_gguf(...) instead." + ) + if outtype not in _LORA_GGUF_OUTTYPES: + raise ValueError( + f"Unsloth: LoRA GGUF outtype must be one of {_LORA_GGUF_OUTTYPES} (got '{outtype}')." + ) + # Resolve a token even for local saves: the converter may fetch a gated/private base config. + if token is None: + token = get_token() + + # Resolve the dequantized base id (the adapter usually references a 4bit repo). + base_model_id = _lora_base_model_id(model) + if not base_model_id: + raise RuntimeError( + "Unsloth: could not determine the base model for LoRA GGUF export " + "(no adapter base_model_name_or_path or model config _name_or_path)." + ) + try: + base_model_id = get_model_name(base_model_id, load_in_4bit = False) + except Exception: + pass + # Windows-safe basename (handles both C:\... and / separators). + if os.path.isdir(base_model_id): + model_name = os.path.basename(os.path.normpath(base_model_id)) + else: + model_name = base_model_id.replace("\\", "/").rstrip("/").split("/")[-1] + if not model_name: + model_name = "model" + + # Save the adapter; for a hub push use an isolated temp dir, else save_directory itself. + if push_to_hub: + lora_dir = tempfile.mkdtemp(prefix = "unsloth-lora-gguf-") + else: + os.makedirs(save_directory, exist_ok = True) + lora_dir = save_directory + + # Wrap so the isolated temp dir used for hub pushes is always cleaned up, even on failure. + try: + save_lora_to_custom_dir(model, tokenizer, lora_dir) + + # Ensure a full llama.cpp checkout (ships convert_lora_to_gguf.py) and locate the converter. + install_llama_cpp(just_clone_repo = True) + converter = os.path.join(LLAMA_CPP_DEFAULT_DIR, "convert_lora_to_gguf.py") + if not os.path.exists(converter): + # A prebuilt llama.cpp install (or a reused CWD copy) carries binaries but not the + # converter script, so force a dedicated source checkout that ships it. + source_dir = os.path.join( + os.path.dirname(os.path.normpath(LLAMA_CPP_DEFAULT_DIR)), "llama.cpp-source" + ) + install_llama_cpp(llama_cpp_folder = source_dir, just_clone_repo = True) + converter = os.path.join(source_dir, "convert_lora_to_gguf.py") + if not os.path.exists(converter): + raise RuntimeError( + "Unsloth: convert_lora_to_gguf.py not found after installing a llama.cpp source " + "checkout. A full llama.cpp source checkout is required for LoRA GGUF export." + ) + + out_gguf = os.path.join(lora_dir, f"{model_name}-lora-{outtype}.gguf") + cmd = [sys.executable, converter, lora_dir, "--outfile", out_gguf, "--outtype", outtype] + # A local base dir provides config directly; otherwise the id is resolved from the Hub. + if os.path.isdir(base_model_id): + cmd += ["--base", base_model_id] + else: + cmd += ["--base-model-id", base_model_id] + if bool(getattr(model.config, "auto_map", None)): + cmd.append("--trust-remote-code") + + # Expose the token to the converter so it can fetch a gated/private base config from the Hub. + env = os.environ.copy() + if isinstance(token, str) and token: + env["HF_TOKEN"] = token + env["HUGGING_FACE_HUB_TOKEN"] = token + + print(f"Unsloth: Converting LoRA adapter at '{lora_dir}' to GGUF -> '{out_gguf}'") + try: + with subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + bufsize = 1, + universal_newlines = True, + encoding = "utf-8", + errors = "replace", + env = env, + ) as sp: + for line in sp.stdout: + print(line, end = "", flush = True) + sp.wait() + if sp.returncode != 0: + raise subprocess.CalledProcessError(sp.returncode, sp.args) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"Unsloth: LoRA -> GGUF conversion failed (exit {e.returncode}). " + "See the output above for details." + ) + + if not push_to_hub: + print(f"Unsloth: Done. Saved LoRA GGUF to '{out_gguf}'") + return out_gguf + + print(f"Unsloth: Uploading LoRA GGUF to '{save_directory}' ...") + from huggingface_hub import HfApi + + api = HfApi(token = token) + api.create_repo( + repo_id = save_directory, + repo_type = "model", + private = private, + exist_ok = True, + ) + api.upload_folder( + folder_path = lora_dir, + repo_id = save_directory, + repo_type = "model", + allow_patterns = ["*.gguf"], + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + ) + print(f"Unsloth: Done. Uploaded to https://huggingface.co/{save_directory.lstrip('/')}") + return save_directory + finally: + if push_to_hub: + shutil.rmtree(lora_dir, ignore_errors = True) + + def unsloth_convert_lora_to_ggml_and_push_to_hub( self, tokenizer, repo_id: str, use_temp_dir: Optional[bool] = None, - commit_message: Optional[str] = "Converted LoRA to GGML with Unsloth", + commit_message: Optional[str] = "Converted LoRA to GGUF with Unsloth", private: Optional[bool] = None, token: Union[bool, str, None] = None, create_pr: bool = False, revision: str = None, - commit_description: str = "Convert LoRA to GGML format using Unsloth", + commit_description: str = "Convert LoRA to GGUF format using Unsloth", temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.85, + outtype: str = "f16", ): - if not os.path.exists("llama.cpp"): - if IS_KAGGLE_ENVIRONMENT: - python_install = install_python_non_blocking(["protobuf"]) - python_install.wait() - install_llama_cpp_blocking(use_cuda = False) - makefile = None - else: - git_clone = install_llama_cpp_clone_non_blocking() - python_install = install_python_non_blocking(["protobuf"]) - git_clone.wait() - makefile = install_llama_cpp_make_non_blocking() - python_install.wait() - else: - makefile = None - - for _ in range(3): - gc.collect() - - lora_directory_push = "lora-to-ggml-push" - save_lora_to_custom_dir(self, tokenizer, lora_directory_push) - - model_type = self.config.model_type - output_file = os.path.join(lora_directory_push, "ggml-adapter-model.bin") - - print(f"Unsloth: Converting auto-saved LoRA adapters at {lora_directory_push} to GGML format.") - print(f"The output file will be {output_file}") - - try: - with subprocess.Popen( - [ - sys.executable, - "llama.cpp/convert-lora-to-ggml.py", - lora_directory_push, - output_file, - "llama", - ], - stdout = subprocess.PIPE, - stderr = subprocess.PIPE, - bufsize = 1, - universal_newlines = True, - encoding = "utf-8", - errors = "replace", - ) as sp: - for line in sp.stdout: - print(line, end = "", flush = True) - for line in sp.stderr: - print(line, end = "", flush = True) - sp.wait() - if sp.returncode != 0: - raise subprocess.CalledProcessError(sp.returncode, sp.args) - except subprocess.CalledProcessError as e: - print(f"Error: Conversion failed with return code {e.returncode}") - return - - print(f"Unsloth: Conversion completed! Output file: {output_file}") - - print("Unsloth: Uploading GGML file to Hugging Face Hub...") - username = upload_to_huggingface( + return _unsloth_save_lora_gguf( self, + tokenizer, repo_id, - token, - "GGML converted LoRA", - "ggml", - output_file, - None, - private, - ) - link = f"{repo_id.lstrip('/')}" - print("Unsloth: Done.") - print(f"Converted LoRA to GGML and uploaded to https://huggingface.co/{link}") - print( - "\nThis GGML making function was made by Maheswar. Ping him @Maheswar on the Unsloth Discord or on HuggingFace (@mahiatlinux) if you like this!" + outtype = outtype, + push_to_hub = True, + token = token, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, ) @@ -2918,65 +3479,9 @@ def unsloth_convert_lora_to_ggml_and_save_locally( tokenizer, temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.85, + outtype: str = "f16", ): - if not os.path.exists("llama.cpp"): - if IS_KAGGLE_ENVIRONMENT: - python_install = install_python_non_blocking(["protobuf"]) - python_install.wait() - install_llama_cpp_blocking(use_cuda = False) - makefile = None - else: - git_clone = install_llama_cpp_clone_non_blocking() - python_install = install_python_non_blocking(["protobuf"]) - git_clone.wait() - makefile = install_llama_cpp_make_non_blocking() - python_install.wait() - else: - makefile = None - - for _ in range(3): - gc.collect() - - # Use the provided save_directory for local saving - save_lora_to_custom_dir(self, tokenizer, save_directory) - - model_type = self.config.model_type - output_file = os.path.join(save_directory, "ggml-adapter-model.bin") - - print(f"Unsloth: Converting auto-saved LoRA adapters at {save_directory} to GGML format.") - print(f"The output file will be {output_file}") - - try: - with subprocess.Popen( - [ - sys.executable, - "llama.cpp/convert-lora-to-ggml.py", - save_directory, - output_file, - "llama", - ], - stdout = subprocess.PIPE, - stderr = subprocess.PIPE, - bufsize = 1, - universal_newlines = True, - encoding = "utf-8", - errors = "replace", - ) as sp: - for line in sp.stdout: - print(line, end = "", flush = True) - for line in sp.stderr: - print(line, end = "", flush = True) - sp.wait() - if sp.returncode != 0: - raise subprocess.CalledProcessError(sp.returncode, sp.args) - except subprocess.CalledProcessError as e: - print(f"Error: Conversion failed with return code {e.returncode}") - return - print("Unsloth: Done.") - print(f"Unsloth: Conversion completed! Output file: {output_file}") - print( - "\nThis GGML making function was made by Maheswar. Ping him @Maheswar on the Unsloth Discord or on HuggingFace (@mahiatlinux) if you like this!" - ) + return _unsloth_save_lora_gguf(self, tokenizer, save_directory, outtype = outtype) from .models.loader_utils import get_model_name @@ -3215,7 +3720,7 @@ def unsloth_generic_save_pretrained_merged( self, save_directory: Union[str, os.PathLike], tokenizer = None, - save_method: str = "merged_16bit", # ["lora", "merged_16bit", "merged_4bit"] + save_method: str = "merged_16bit", # ["lora", "merged_16bit", "merged_4bit", "fp8", "mxfp4", "nvfp4", "mxfp8"] push_to_hub: bool = False, token: Optional[Union[str, bool]] = None, is_main_process: bool = True, @@ -3229,6 +3734,9 @@ def unsloth_generic_save_pretrained_merged( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.75, datasets: Optional[List[str]] = None, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, ): """ Same as .push_to_hub(...) except 4bit weights are auto @@ -3238,6 +3746,10 @@ def unsloth_generic_save_pretrained_merged( 1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp. 2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference. 3. `lora`: Save LoRA adapters with no merging. Useful for HF inference. + 4. FP8 / FP4 compressed export for vLLM via llm-compressor: + `fp8` (dynamic W8A8), `mxfp4`, `nvfp4` (W4A4), `mxfp8`. The LoRA is merged to 16bit at + `save_directory`, then a quantized checkpoint is written to `save_directory + "-"`. + `nvfp4` needs calibration data (defaults to ultrachat; override with `calibration_dataset`). """ if tokenizer is None: logger.warning_once( @@ -3245,9 +3757,46 @@ def unsloth_generic_save_pretrained_merged( "You can do it separately via `tokenizer.save_pretrained(...)`" ) + # FP8 / FP4 compressed-tensors export (llm-compressor) -> handled separately. + _compressed = _normalize_compressed_method(save_method) + if _compressed is not None: + scheme, needs_calibration, suffix = _compressed + _unsloth_save_compressed_tensors( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + scheme = scheme, + needs_calibration = needs_calibration, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] + del arguments["_compressed"] + del arguments["calibration_dataset"] + del arguments["num_calibration_samples"] + del arguments["max_seq_length"] unsloth_generic_save(**arguments) for _ in range(3): gc.collect() @@ -3271,6 +3820,9 @@ def unsloth_generic_push_to_hub_merged( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.75, datasets: Optional[List[str]] = None, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, ): """ Same as .push_to_hub(...) except 4bit weights are auto @@ -3280,6 +3832,7 @@ def unsloth_generic_push_to_hub_merged( 1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp. 2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference. 3. `lora`: Save LoRA adapters with no merging. Useful for HF inference. + 4. FP8 / FP4 compressed export for vLLM: `fp8`, `mxfp4`, `nvfp4`, `mxfp8`. """ if tokenizer is None: logger.warning_once( @@ -3287,12 +3840,50 @@ def unsloth_generic_push_to_hub_merged( "You can do it separately via `tokenizer.push_to_hub(...)`" ) + # FP8 / FP4 compressed-tensors export (llm-compressor) -> handled separately. + _compressed = _normalize_compressed_method(save_method) + if _compressed is not None: + scheme, needs_calibration, suffix = _compressed + _unsloth_save_compressed_tensors( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + scheme = scheme, + needs_calibration = needs_calibration, + suffix = suffix, + push_to_hub = True, + token = token, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id arguments["push_to_hub"] = True del arguments["self"] del arguments["repo_id"] + del arguments["_compressed"] + del arguments["calibration_dataset"] + del arguments["num_calibration_samples"] + del arguments["max_seq_length"] unsloth_generic_save(**arguments) for _ in range(3): gc.collect() @@ -3437,6 +4028,333 @@ def _unsloth_save_torchao_with_given_config( pass +def _scheme_is_available(scheme): + """True if `scheme` is a known preset in the installed compressed_tensors.""" + try: + from compressed_tensors.quantization import quant_scheme as _qs + + presets = getattr(_qs, "PRESET_SCHEMES", None) + if presets is None: + return True + return scheme in presets + except Exception: + # If we cannot introspect, let llm-compressor validate the scheme itself. + return True + + +def _print_compressed_hw_note(scheme, out_dir): + if scheme in ("FP8_DYNAMIC", "MXFP8"): + hw = "NVIDIA GPUs with compute capability >= 8.9 (Ada / Hopper) or newer" + else: + hw = ( + "NVIDIA Blackwell (SM100+) for full activation quantization " + "(older GPUs fall back to weight-only in vLLM)" + ) + print( + f"Unsloth: Saved {scheme} compressed checkpoint to '{out_dir}'.\n" + f"Unsloth: Load it with vLLM for accelerated inference. Hardware for full speed: {hw}." + ) + + +def _unsloth_save_compressed_tensors( + model, + save_directory: Union[str, os.PathLike], + tokenizer, + scheme: str, + needs_calibration: bool, + suffix: str, + push_to_hub: bool = False, + token: Optional[Union[str, bool]] = None, + is_main_process: bool = True, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, + **merge_kwargs, +): + """Export an FP8/FP4 compressed-tensors checkpoint via llm-compressor. + + Mirrors the torchao PTQ path: LoRA is first merged into the base model at 16bit and + written to `save_directory` (which is kept). The merged checkpoint is then quantized with + llm-compressor's `QuantizationModifier(scheme)` in a separate process (so Unsloth's + transformers monkey-patches do not interfere), and written to the sibling directory + `save_directory + "-" + suffix`. The result is intended for vLLM inference. + """ + import tempfile + + if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): + tokenizer = patch_saving_functions(tokenizer) + # Resolve a token for the hub push and/or loading a gated calibration dataset in the subprocess. + if token is None: + token = get_token() + + # Only the main process installs deps, merges, quantizes, and uploads (mirrors the non-PEFT + # save path); other ranks return at once so they neither race on dirs nor run pip installs. + if not is_main_process: + return None + + # 1) Install llm-compressor and gate on scheme availability BEFORE merging, so an unsupported + # scheme (e.g. mxfp8) fails fast instead of writing a full 16bit checkpoint first. + install_llm_compressor() + if not _scheme_is_available(scheme): + try: + import transformers as _tf + tf_ver = _tf.__version__ + except Exception: + tf_ver = "unknown" + raise RuntimeError( + f"Unsloth: scheme '{scheme}' is not available in your installed " + f"compressed-tensors / llm-compressor.\n" + f"It requires a newer llm-compressor that needs transformers>=5.9 " + f"(you have transformers {tf_ver}).\n" + "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." + ) + + # 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and + # quantize inside an isolated temp dir instead of writing ./ into the cwd. + repo_id, work_tmp, calib_tmp, model_dev = None, None, None, None + if push_to_hub: + repo_id = os.fspath(save_directory) + work_tmp = tempfile.mkdtemp(prefix = "unsloth-compressed-") + local_dir = os.path.join(work_tmp, os.path.basename(repo_id.rstrip("/")) or "model") + else: + # Drop trailing separators so the sibling "-" output is not nested inside . + local_dir = os.fspath(save_directory) + local_dir = local_dir.rstrip("/\\") or local_dir + + # Wrap the body so the isolated temp dirs are always cleaned up, even when the merge, + # quantization, validation, or hub upload raises. + api = None + try: + # Validate Hub access up front (a bad token / denied repo should fail before the expensive + # merge and quantization, matching the normal push path). create_repo is idempotent. + if push_to_hub: + from huggingface_hub import HfApi + api = HfApi(token = token) + api.create_repo( + repo_id = repo_id, + repo_type = "model", + private = merge_kwargs.get("private", None), + exist_ok = True, + ) + + # 3) Merge to 16bit at local_dir (kept for local saves) via unsloth_generic_save, so LoRA + # adapters are merged and full-finetuned models written in 16bit consistently. Extra + # save kwargs (state_dict, max_shard_size, ...) flow through merge_kwargs. + # The intermediate 16bit checkpoint is internal staging that the converter subprocess + # reloads with default weight filenames, so never write variant-named shards here; the + # user's variant (if any) is applied to the final compressed checkpoint in the subprocess. + variant = merge_kwargs.pop("variant", None) + print(f"Unsloth: Merging to 16bit before {scheme} quantization...") + merge_args = dict(merge_kwargs) + merge_args.update( + dict( + model = model, + tokenizer = tokenizer, + save_directory = local_dir, + save_method = "merged_16bit", + push_to_hub = False, + token = token, + is_main_process = is_main_process, + ) + ) + unsloth_generic_save(**merge_args) + + # 4) Detect VLM + trust_remote_code from the in-memory model config. A vision/multimodal + # model exposes a vision_config or an explicitly vision-named architecture; a bare + # *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper), so it + # is not treated as a VLM on its own. + is_vlm = False + if hasattr(model, "config"): + archs = getattr(model.config, "architectures", None) or [] + is_vlm = hasattr(model.config, "vision_config") or any( + x.endswith("ForVisionText2Text") for x in archs + ) + if is_vlm: + logger.warning( + "Unsloth: FP8/FP4 compressed export for vision / multimodal models is " + "experimental; vision-tower layers may be affected." + ) + trust_remote_code = ( + bool(getattr(model.config, "auto_map", None)) if hasattr(model, "config") else False + ) + + # 5) Marshal the calibration dataset for the subprocess: None -> ultrachat default; a + # str/PathLike is a local save_to_disk dir if it exists else a Hub id; Dataset -> temp. + calib_kind, calib_value = "none", "" + if needs_calibration and calibration_dataset is not None: + if isinstance(calibration_dataset, (str, os.PathLike)): + calib_value = os.fspath(calibration_dataset) + calib_kind = "disk" if os.path.isdir(calib_value) else "hfid" + elif hasattr(calibration_dataset, "save_to_disk"): + # Only persist the samples we need, so multi-GB training sets are not fully copied. + ds_to_save = calibration_dataset + # A DatasetDict's len() is the split count, not rows; pick one split first so the + # row subsample below applies and we do not save every split to the temp dir. + try: + from datasets import DatasetDict + if isinstance(ds_to_save, DatasetDict): + ds_to_save = ds_to_save.get("train", None) or next( + iter(ds_to_save.values()) + ) + except Exception: + pass + try: + if ( + num_calibration_samples + and hasattr(ds_to_save, "select") + and len(ds_to_save) > num_calibration_samples + ): + ds_to_save = ds_to_save.shuffle(seed = 42).select( + range(num_calibration_samples) + ) + except Exception: + ds_to_save = calibration_dataset + calib_tmp = tempfile.mkdtemp(prefix = "unsloth-calib-") + shutil.rmtree(calib_tmp, ignore_errors = True) # save_to_disk wants a fresh path + ds_to_save.save_to_disk(calib_tmp) + calib_kind, calib_value = "disk", calib_tmp + else: + raise TypeError( + "Unsloth: calibration_dataset must be None, a Hugging Face dataset id, a " + "local path saved with Dataset.save_to_disk(...), or a Dataset with " + "save_to_disk()." + ) + elif not needs_calibration and calibration_dataset is not None: + logger.warning_once( + f"Unsloth: scheme '{scheme}' is data-free; ignoring calibration_dataset." + ) + + # 6) Quantize in a separate process: importing Unsloth patches transformers attention, + # which breaks the forward llm-compressor runs for calibration. Run the converter by + # file path (not `-m`) so the subprocess stays unpatched, like GGUF -> llama.cpp. + out_dir = local_dir + "-" + suffix + runner = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_compressed_quantize.py") + cmd = [ + sys.executable, + runner, + "--model", + local_dir, + "--scheme", + scheme, + "--out", + out_dir, + "--calibration-dataset-kind", + calib_kind, + "--num-calibration-samples", + str(num_calibration_samples), + "--max-seq-length", + str(max_seq_length), + ] + if needs_calibration: + cmd.append("--needs-calibration") + if calib_value: + cmd += ["--calibration-dataset", calib_value] + if is_vlm: + cmd.append("--is-vlm") + if trust_remote_code: + cmd.append("--trust-remote-code") + if variant: + cmd += ["--variant", variant] + + # Free the in-memory model's CUDA memory before the subprocess loads its own copy from + # disk, so a single GPU need not hold both at once. Best-effort and restored in finally; + # skipped for quantized or multi-device models where moving is unsafe. + try: + if ( + torch.cuda.is_available() + and hasattr(model, "parameters") + and not getattr(model, "is_loaded_in_4bit", False) + and not getattr(model, "is_loaded_in_8bit", False) + and not getattr(model, "is_quantized", False) + ): + _devs = {str(p.device) for p in model.parameters()} + if len(_devs) == 1 and next(iter(_devs)).startswith("cuda"): + _dev = next(model.parameters()).device + model.to("cpu") + model_dev = _dev # set only after a successful move, so finally can restore + except Exception: + model_dev = None + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # Expose the token so the subprocess can load a gated/private calibration dataset. + env = os.environ.copy() + if isinstance(token, str) and token: + env["HF_TOKEN"] = token + env["HUGGING_FACE_HUB_TOKEN"] = token + + print( + f"Unsloth: Quantizing the merged model to {scheme} with llm-compressor " + "(in a separate process)..." + ) + try: + subprocess.check_call(cmd, env = env) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"Unsloth: {scheme} quantization failed (llm-compressor subprocess exit " + f"{e.returncode}). See the output above for details." + ) + + # 7) Validate the artifact. + cfg_path = os.path.join(out_dir, "config.json") + cfg = {} + if os.path.exists(cfg_path): + with open(cfg_path, "r", encoding = "utf-8") as f: + cfg = json.load(f) + if "quantization_config" not in cfg: + raise RuntimeError( + f"Unsloth: {scheme} export failed - no quantization_config written to {cfg_path}" + ) + + # 8) Optional hub upload of the compressed artifact (not the intermediate 16bit one). + # The repo was already created/validated up front, so just upload here. + if push_to_hub: + print(f"Unsloth: Uploading {scheme} checkpoint to '{repo_id}' ...") + api.upload_folder( + folder_path = out_dir, + repo_id = repo_id, + repo_type = "model", + commit_message = merge_kwargs.get("commit_message", None), + commit_description = merge_kwargs.get("commit_description", None), + create_pr = merge_kwargs.get("create_pr", False), + revision = merge_kwargs.get("revision", None), + ) + # Attach datasets metadata to the pushed repo, like the normal merged push path. + datasets = merge_kwargs.get("datasets", None) + if datasets: + try: + from huggingface_hub import metadata_update + metadata_update(repo_id, {"datasets": datasets}, overwrite = True, token = token) + except Exception as meta_err: + logger.warning_once( + f"Unsloth: could not update datasets metadata for {repo_id}: {meta_err}" + ) + + # 9) Inference hardware note. + result = repo_id if push_to_hub else out_dir + _print_compressed_hw_note(scheme, result) + return result + finally: + if model_dev is not None: + try: + model.to(model_dev) # restore the model to its original device + except Exception: + logger.warning_once( + "Unsloth: could not restore the model to its original device after compressed " + "export; it may remain on CPU." + ) + if calib_tmp is not None and os.path.isdir(calib_tmp): + shutil.rmtree(calib_tmp, ignore_errors = True) + if work_tmp is not None: + shutil.rmtree(work_tmp, ignore_errors = True) + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def unsloth_save_pretrained_torchao( self, save_directory: Union[str, os.PathLike], From 43d3caf38b9bdead249084c676f37a9322496490 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 30 Jun 2026 03:41:02 -0700 Subject: [PATCH 39/49] Studio: imatrix GGUF option and FP8/NVFP4 compressed export in the export UI (#6729) * Studio: wire imatrix GGUF option and FP8/NVFP4 compressed export into the export UI GGUF export gains an importance-matrix toggle. When enabled it auto-downloads the upstream Unsloth imatrix for the base model (or uses a custom path), which unlocks the IQ low-bit quants iq2_xxs, iq2_m, iq3_xxs and iq4_xs. Merged export gains an FP8 / NVFP4 compressed-tensors precision selector that runs llm-compressor for vLLM. Backend threads imatrix_file through routes -> orchestrator -> worker -> export_gguf (both the local save and the hub push), and maps the new compressed format_type values onto the fp8/nvfp4 save_method, reporting the "-" sibling output directory. Frontend adds the imatrix Switch on the GGUF card and a merged precision picker on the merged card, threaded through the export runtime store. Depends on unslothai/unsloth#6706 (save.py imatrix_file and compressed-tensors export) and unslothai/unsloth-zoo#839 (quantize_gguf imatrix flag). * Studio export: guard imatrix/compressed against older unsloth builds and force imatrix for IQ quants Addresses review feedback on the export wiring: - GGUF: pass imatrix_file only when set, so a plain no-imatrix export (e.g. Q4_K_M) no longer fails with an unexpected-keyword error against an unsloth build that predates the imatrix_file parameter. When imatrix is requested but unsupported, return a clear upgrade message instead of a TypeError. - Merged: gate FP8/NVFP4 compressed-tensors export on the installed unsloth actually supporting it, returning a clear message rather than a cryptic save_method failure. - Frontend: IQ quants (iq2_xxs, iq2_m, iq3_xxs, iq4_xs) are imatrix-only, so force the imatrix on when one is selected and lock the toggle, instead of submitting an IQ quant with no imatrix that llama.cpp would reject. Extends the backend tests for the new capability guards and the conditional kwarg wiring. * Studio: upload compressed merged models to the Hub without recompressing For an FP8/NVFP4 Hub export the model is already produced locally in the "-" output. Uploading it directly with HfApi.upload_folder (mirroring export_base_model) avoids re-running the expensive compressed-tensors quantization a second time inside push_to_hub_merged, which for NVFP4 also re-runs calibration and risks OOM. Falls back to push_to_hub_merged when there is no local compressed output to reuse. --- studio/backend/core/export/export.py | 103 ++++++++++++++-- studio/backend/core/export/orchestrator.py | 2 + studio/backend/core/export/worker.py | 1 + studio/backend/models/export.py | 19 ++- studio/backend/routes/export.py | 3 + .../tests/test_export_imatrix_compressed.py | 116 ++++++++++++++++++ .../src/features/export/api/export-api.ts | 2 + .../frontend/src/features/export/constants.ts | 37 ++++++ .../src/features/export/export-page.tsx | 70 ++++++++++- .../export/stores/export-runtime-store.ts | 6 + 10 files changed, 343 insertions(+), 16 deletions(-) create mode 100644 studio/backend/tests/test_export_imatrix_compressed.py diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index f243f5b65a..d0461dae95 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -38,6 +38,26 @@ logger = get_logger(__name__) _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False +def _supports_kwarg(fn, name): + """True if `fn` accepts keyword `name` directly or via **kwargs.""" + import inspect + + try: + params = inspect.signature(fn).parameters + except (TypeError, ValueError): + return False + return name in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + + +def _compressed_export_supported(): + """True if the installed unsloth build can do FP8/NVFP4 compressed-tensors export.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_compressed_method") + except Exception: + return False + + def _hf_offline(timeout = 3): """True if export should avoid the Hub: honors the HF offline env vars, else does one cheap TCP reachability probe so a network-down load uses local files / the HF cache @@ -400,16 +420,33 @@ class ExportBackend: ) output_path: Optional[str] = None + # compressed-tensors formats run save_pretrained_merged with an FP8/FP4 save_method and + # write to a sibling "-" directory (for vLLM). + _COMPRESSED = { + "FP8 (compressed-tensors)": ("fp8", "fp8"), + "NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"), + } + is_compressed = format_type in _COMPRESSED try: if _IS_MLX: + if is_compressed: + return False, "Compressed-tensors export is not supported on macOS/MLX.", None mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + elif is_compressed: + if not _compressed_export_supported(): + return ( + False, + "Compressed-tensors (FP8/NVFP4) export requires an Unsloth build with " + "compressed-tensors support. Upgrade unsloth, or choose 16-bit.", + None, + ) + save_method = _COMPRESSED[format_type][0] + elif format_type == "4-bit (FP4)": + save_method = "merged_4bit_forced" + elif self._audio_type == "whisper": + save_method = None else: - if format_type == "4-bit (FP4)": - save_method = "merged_4bit_forced" - elif self._audio_type == "whisper": - save_method = None - else: - save_method = "merged_16bit" + save_method = "merged_16bit" if save_directory: save_directory = str(resolve_export_write_dir(save_directory)) @@ -427,9 +464,15 @@ class ExportBackend: save_directory, self.current_tokenizer, save_method = save_method ) - self._write_export_metadata(save_directory) - logger.info(f"Model saved successfully to {save_directory}") - output_path = str(Path(save_directory).resolve()) + # Compressed export writes to the "-" sibling; report that as output. + final_dir = ( + f"{save_directory}-{_COMPRESSED[format_type][1]}" + if is_compressed + else save_directory + ) + self._write_export_metadata(final_dir) + logger.info(f"Model saved successfully to {final_dir}") + output_path = str(Path(final_dir).resolve()) if push_to_hub: if not repo_id or not hf_token: @@ -464,6 +507,32 @@ class ExportBackend: token = hf_token, private = private, ) + elif is_compressed and output_path and Path(output_path).is_dir(): + # The compressed model was already built locally in output_path; upload it + # directly so we do not re-run the (expensive, OOM-prone) compression that + # push_to_hub_merged(save_method=fp8/nvfp4) would otherwise do a second time. + hf_api = HfApi(token = hf_token) + repo_id = PushToHubMixin._create_repo( + PushToHubMixin, + repo_id = repo_id, + private = private, + token = hf_token, + ) + content = MODEL_CARD.format( + username = repo_id.split("/")[0], + base_model = getattr(self.current_model.config, "_name_or_path", "unknown"), + model_type = getattr(self.current_model.config, "model_type", "llm"), + method = format_type, + extra = "unsloth", + ) + ModelCard(content).push_to_hub( + repo_id, token = hf_token, commit_message = "Unsloth Model Card" + ) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) else: hub_save_method = save_method if save_method is not None else "merged_16bit" self.current_model.push_to_hub_merged( @@ -621,6 +690,7 @@ class ExportBackend: push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, + imatrix_file = None, ) -> Tuple[bool, str, Optional[str]]: """ Export model in GGUF format. @@ -638,6 +708,19 @@ class ExportBackend: if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None + # Only forward imatrix_file to an unsloth build that accepts it; otherwise even a plain + # no-imatrix export would fail with an unexpected-keyword error against an older unsloth. + if imatrix_file is not None and not _supports_kwarg( + self.current_model.save_pretrained_gguf, "imatrix_file" + ): + return ( + False, + "This Unsloth build does not support GGUF imatrix export. " + "Upgrade unsloth and unsloth_zoo, or disable the imatrix option.", + None, + ) + imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {} + output_path: Optional[str] = None model_tmp_to_cleanup: Optional[str] = None try: @@ -691,6 +774,7 @@ class ExportBackend: _model_tmp, self.current_tokenizer, quantization_method = quant_method, + **imatrix_kw, ) # Relocate the .gguf that convert_to_gguf wrote to cwd (repo root). @@ -757,6 +841,7 @@ class ExportBackend: self.current_tokenizer, quantization_method = quant_method, token = hf_token, + **imatrix_kw, ) logger.info(f"GGUF model pushed successfully to {repo_id}") diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 478624b48e..052a47dd80 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -499,6 +499,7 @@ class ExportOrchestrator: push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, + imatrix_file = None, ) -> Tuple[bool, str, Optional[str]]: """Export model in GGUF format.""" return self._run_export( @@ -509,6 +510,7 @@ class ExportOrchestrator: "push_to_hub": push_to_hub, "repo_id": repo_id, "hf_token": hf_token, + "imatrix_file": imatrix_file, }, ) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 71a603a857..d473dcb54f 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -414,6 +414,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: push_to_hub = cmd.get("push_to_hub", False), repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), + imatrix_file = cmd.get("imatrix_file"), ) elif export_type == "lora": success, message, output_path = backend.export_lora_adapter( diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 1e8e3c4792..7e05373f11 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -158,9 +158,15 @@ class ExportCommonOptions(BaseModel): class ExportMergedModelRequest(ExportCommonOptions): """Request for exporting a merged PEFT model.""" - format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field( + format_type: Literal[ "16-bit (FP16)", - description = "Export precision / format for the merged model", + "4-bit (FP4)", + "FP8 (compressed-tensors)", + "NVFP4 (compressed-tensors)", + ] = Field( + "16-bit (FP16)", + description = "Export precision / format for the merged model. The compressed-tensors " + "options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).", ) @@ -199,6 +205,15 @@ class ExportGGUFRequest(BaseModel): None, description = "Hugging Face token for GGUF upload", ) + imatrix: bool = Field( + False, + description = "Use an importance matrix (auto-downloads the upstream unsloth GGUF " + "imatrix). Required for the IQ low-bit quants such as iq2_xxs / iq4_xs.", + ) + imatrix_path: Optional[str] = Field( + None, + description = "Path to a custom imatrix file; overrides the auto-download when set.", + ) class ExportLoRAAdapterRequest(ExportCommonOptions): diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 78c1e59d2e..cf2cb2fa70 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -343,6 +343,8 @@ async def export_gguf( """ try: backend = get_export_backend() + # A custom path wins; otherwise the imatrix toggle requests the upstream auto-download. + imatrix_file = request.imatrix_path or (True if request.imatrix else None) success, message, output_path = await asyncio.to_thread( backend.export_gguf, save_directory = request.save_directory, @@ -350,6 +352,7 @@ async def export_gguf( push_to_hub = request.push_to_hub, repo_id = request.repo_id, hf_token = request.hf_token, + imatrix_file = imatrix_file, ) if not success: diff --git a/studio/backend/tests/test_export_imatrix_compressed.py b/studio/backend/tests/test_export_imatrix_compressed.py new file mode 100644 index 0000000000..d914ff8651 --- /dev/null +++ b/studio/backend/tests/test_export_imatrix_compressed.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the GGUF imatrix option and compressed-tensors merged export wiring. + +Schema checks use the real Pydantic models; the cross-layer threading is verified with ast so it +runs on CPU with no GPU, no model, and no llama.cpp. +""" + +import ast +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from models.export import ExportGGUFRequest, ExportMergedModelRequest + +_BACKEND = Path(__file__).resolve().parent.parent + + +def _src(rel): + return (_BACKEND / rel).read_text(encoding = "utf-8") + + +def _func_src(rel, name): + src = _src(rel) + node = next( + n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(src, node) + + +# -- schema ------------------------------------------------------------------------------------- + + +def test_gguf_request_imatrix_defaults_and_set(): + assert ExportGGUFRequest(save_directory = "/tmp/x").imatrix is False + assert ExportGGUFRequest(save_directory = "/tmp/x").imatrix_path is None + r = ExportGGUFRequest(save_directory = "/tmp/x", imatrix = True, imatrix_path = "/i.dat") + assert r.imatrix is True and r.imatrix_path == "/i.dat" + + +def test_merged_request_accepts_compressed_formats(): + for fmt in ("16-bit (FP16)", "FP8 (compressed-tensors)", "NVFP4 (compressed-tensors)"): + assert ExportMergedModelRequest(save_directory = "/tmp/x", format_type = fmt).format_type == fmt + + +def test_merged_request_rejects_unknown_format(): + with pytest.raises(ValidationError): + ExportMergedModelRequest(save_directory = "/tmp/x", format_type = "bogus") + + +# -- threading (ast) ---------------------------------------------------------------------------- + + +def test_export_gguf_threads_imatrix_to_save_and_push(): + # imatrix_file must reach both save_pretrained_gguf and push_to_hub_gguf, but only via the + # conditional **imatrix_kw so a no-imatrix export never sends an unsupported keyword. + g = _func_src("core/export/export.py", "export_gguf") + assert g.count("**imatrix_kw") >= 2 + assert 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' in g + # Unconditional pass-through (the old wiring) must be gone. + assert "imatrix_file = imatrix_file" not in g + + +def test_export_gguf_guards_unsupported_imatrix_build(): + # An older unsloth without imatrix_file support gets a clean error, not a TypeError. + g = _func_src("core/export/export.py", "export_gguf") + assert "_supports_kwarg(" in g and '"imatrix_file"' in g + + +def test_export_merged_guards_unsupported_compressed_build(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "_compressed_export_supported()" in m + + +def test_supports_kwarg_helper(): + # exec just the helper source so the test stays free of export.py's heavy import chain. + ns = {} + exec(_func_src("core/export/export.py", "_supports_kwarg"), ns) + supports = ns["_supports_kwarg"] + + def has_it(a, imatrix_file = None): + pass + + def lacks_it(a): + pass + + def via_kwargs(a, **kw): + pass + + assert supports(has_it, "imatrix_file") is True + assert supports(lacks_it, "imatrix_file") is False + assert supports(via_kwargs, "imatrix_file") is True + + +def test_orchestrator_and_worker_pass_imatrix(): + assert "imatrix_file" in _func_src("core/export/orchestrator.py", "export_gguf") + assert 'imatrix_file = cmd.get("imatrix_file")' in _src("core/export/worker.py") + + +def test_route_resolves_imatrix_file(): + assert "request.imatrix_path or (True if request.imatrix else None)" in _src("routes/export.py") + + +def test_export_merged_maps_compressed_to_save_method(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "is_compressed" in m and '"fp8"' in m and '"nvfp4"' in m + + +def test_compressed_hub_push_uploads_local_dir_without_recompressing(): + # A compressed Hub push must upload the already-built output_path, not re-run compression + # via push_to_hub_merged (which would compress a second time). + m = _func_src("core/export/export.py", "export_merged_model") + assert "elif is_compressed and output_path and Path(output_path).is_dir():" in m + assert "hf_api.upload_folder(" in m and "folder_path = output_path" in m diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index f055f1bd61..d1b4e88a6d 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -162,6 +162,8 @@ export async function exportGGUF(params: { push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; + imatrix?: boolean; + imatrix_path?: string | null; }): Promise { const response = await authFetch("/api/export/export/gguf", { method: "POST", diff --git a/studio/frontend/src/features/export/constants.ts b/studio/frontend/src/features/export/constants.ts index 9b7e487450..058de1edc4 100644 --- a/studio/frontend/src/features/export/constants.ts +++ b/studio/frontend/src/features/export/constants.ts @@ -39,7 +39,12 @@ export const QUANT_OPTIONS: { value: string; label: string; recommended?: boolean; + imatrix?: boolean; // IQ quants require an importance matrix (the imatrix toggle below) }[] = [ + { value: "iq2_xxs", label: "IQ2_XXS", imatrix: true }, + { value: "iq2_m", label: "IQ2_M", imatrix: true }, + { value: "iq3_xxs", label: "IQ3_XXS", imatrix: true }, + { value: "iq4_xs", label: "IQ4_XS", imatrix: true }, { value: "q2_k_l", label: "Q2_K_L" }, { value: "q3_k_m", label: "Q3_K_M" }, { value: "q4_k_m", label: "Q4_K_M", recommended: true }, @@ -50,12 +55,44 @@ export const QUANT_OPTIONS: { { value: "f16", label: "F16" }, ]; +/** Merged-export precision formats. The compressed-tensors ones run llm-compressor for vLLM. */ +export type MergedFormat = + | "16-bit (FP16)" + | "FP8 (compressed-tensors)" + | "NVFP4 (compressed-tensors)"; + +export const MERGED_FORMATS: { + value: MergedFormat; + label: string; + hint: string; +}[] = [ + { + value: "16-bit (FP16)", + label: "16-bit", + hint: "Full precision, runs anywhere.", + }, + { + value: "FP8 (compressed-tensors)", + label: "FP8 (vLLM)", + hint: "compressed-tensors FP8 for vLLM. Needs an NVIDIA GPU.", + }, + { + value: "NVFP4 (compressed-tensors)", + label: "NVFP4 (vLLM)", + hint: "compressed-tensors NVFP4 for vLLM. Needs an NVIDIA GPU; calibrates.", + }, +]; + /** * llama.cpp effective bits-per-weight per quant; GGUF size ~= fp16_bytes * bpw / 16. * K-quant values are published average bit-rates (Q2_K_L = Unsloth Q2_K + Q8_0 * embeddings). Approximate ("~"), not exact file sizes. */ export const GGUF_BPW: Record = { + iq2_xxs: 2.06, + iq2_m: 2.7, + iq3_xxs: 3.06, + iq4_xs: 4.25, q2_k_l: 3.35, q3_k_m: 3.91, q4_k_m: 4.83, diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 4037dbe079..7b5c6ec6d4 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -3,6 +3,7 @@ import { SectionCard } from "@/components/section-card"; import { Button } from "@/components/ui/button"; +import { Switch } from "@/components/ui/switch"; import { Combobox, ComboboxContent, @@ -61,6 +62,9 @@ import { EXPORT_METHODS, type ExportMethod, GUIDE_STEPS, + MERGED_FORMATS, + type MergedFormat, + QUANT_OPTIONS, buildQuantSizeLabels, getEstimatedSize, } from "./constants"; @@ -171,6 +175,16 @@ export function ExportPage() { ? s.summary.quantLevels : []; }); + // GGUF importance matrix (required for the IQ quants) and merged-export precision. + const [useImatrix, setUseImatrix] = useState(false); + const [mergedFormat, setMergedFormat] = useState("16-bit (FP16)"); + // IQ quants are imatrix-only, so force it on when one is selected; otherwise we would submit + // an IQ quant with no imatrix and llama.cpp would reject it. + const requiresImatrix = quantLevels.some( + (q) => QUANT_OPTIONS.find((o) => o.value === q)?.imatrix, + ); + const effectiveImatrix = useImatrix || requiresImatrix; + // Whether the inline export panel is expanded. The panel also shows itself // whenever a run is active/terminal (see `panelActive`), so it survives // navigation even though this local flag resets on remount. @@ -613,6 +627,8 @@ export function ExportPage() { exportMethod, isAdapter: adapterExport, quantLevels, + useImatrix: effectiveImatrix, + mergedFormat, saveDirectory, destination, repoId, @@ -639,6 +655,8 @@ export function ExportPage() { exportMethod, isAdapter, quantLevels, + effectiveImatrix, + mergedFormat, destination, saveDirectory, hfUsername, @@ -1166,12 +1184,54 @@ export function ExportPage() { } /> + {exportMethod === "merged" && isAdapter && ( +
+
Precision
+
+ {MERGED_FORMATS.map((f) => ( + + ))} +
+
+ {MERGED_FORMATS.find((f) => f.value === mergedFormat)?.hint} +
+
+ )} + {exportMethod === "gguf" && ( - + <> + +
+
+
+ Importance matrix (imatrix) +
+
+ {requiresImatrix + ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." + : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} +
+
+ +
+ )} {estimatedSize && (
diff --git a/studio/frontend/src/features/export/stores/export-runtime-store.ts b/studio/frontend/src/features/export/stores/export-runtime-store.ts index 307e1ccaab..a87e7d93e9 100644 --- a/studio/frontend/src/features/export/stores/export-runtime-store.ts +++ b/studio/frontend/src/features/export/stores/export-runtime-store.ts @@ -138,6 +138,10 @@ export interface RunExportParams { exportMethod: ExportMethod; isAdapter: boolean; quantLevels: string[]; + /** GGUF: use an importance matrix (auto-download); required for the IQ quants. */ + useImatrix?: boolean; + /** Merged: precision/format ("16-bit (FP16)" or a compressed-tensors option). */ + mergedFormat?: string; saveDirectory: string; destination: ExportDestination; repoId?: string; @@ -437,6 +441,7 @@ export const useExportRuntimeStore = create()((set, get) => const { outputPath } = await runRecoverableOp(() => exportMerged({ save_directory: params.saveDirectory, + format_type: params.mergedFormat, push_to_hub: pushToHub, repo_id: params.repoId, hf_token: params.token, @@ -469,6 +474,7 @@ export const useExportRuntimeStore = create()((set, get) => push_to_hub: pushToHub, repo_id: params.repoId, hf_token: params.token, + imatrix: params.useImatrix, }), ); lastOutputPath = outputPath ?? lastOutputPath; From e8945cab46d0d5998c522b3faaff2ba19fbff7c9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 30 Jun 2026 06:55:23 -0700 Subject: [PATCH 40/49] Whole-document context for RAG chat attachments (#6693) * Add whole-document context mode to RAG chat attachments Thread-attached files are injected in full when they fit a token budget, instead of only top-K retrieved chunks, so the model reads the entire file for summarize/reason-over-document requests. Oversized files fall back to top-K retrieval so the context window is never blown. KB and project corpora are unchanged (still retrieval). - core/rag/store.py: all_chunks_for_scope returns every completed-document chunk for a scope, ordered document-then-index, joined with filename. - core/rag/tool.py: whole_document_context renders the chunks as the same blocks + citation source-map retrieval produces, returns None when empty or over budget. - core/inference/tools.py: build_rag_autoinject tries whole-document first for thread scopes, falls through to search_for_autoinject otherwise. - core/rag/config.py: THREAD_WHOLE_DOC + WHOLE_DOC_MAX_TOKENS (env-tunable). - tests/test_rag_whole_document.py: store ordering, whole-doc render + budget cutoff, auto-inject whole-doc vs top-K fallback, KB never whole-doc. * Add scanned-PDF OCR fallback to RAG ingestion A PDF page with no extractable text layer (a scanned or image-only page) previously ingested as empty, so image PDFs were invisible to retrieval and whole-document context. Such pages are now rendered and transcribed by the loaded vision model during ingestion, so they become searchable and readable like any other page. This restores OCR for the RAG document flow without a separate extraction pipeline. - core/rag/parsers.py: render_pdf_pages renders whole pages (1-based) to PNG. - core/rag/captioner.py: factor the shared vision call into _vision_complete; add _ocr_one + ocr_pages (transcribe rendered pages, OCR_MAX_PAGES bound). - core/rag/ingestion.py: _ocr_scanned_pages runs right after parse, replacing text on near-empty PDF pages. No-op when OCR is off, no page is scanned, or no vision model is loaded (degrades like figure captioning). - core/rag/config.py: OCR_SCANNED, OCR_MIN_CHARS, OCR_MAX_PAGES, OCR_DPI, OCR_TIMEOUT_S, OCR_MAX_TOKENS (env-tunable). - tests/test_rag_ocr_fallback.py: page render, ocr_pages gating + cap, scanned PDF end-to-end OCR into chunks + whole-doc, born-digital skips OCR, disabled leaves the page empty. * Broaden OCR prompt to figures/tables and guard against repetition runaway The OCR prompt now asks the vision model to also transcribe text inside figures, diagrams, charts and tables, so labels and table cells on scanned pages are indexed rather than skipped. Verified on real documents that this does not regress plain-text transcription. Some vision models loop on sparse images (e.g. a title-only cover) and emit the same line hundreds of times. _collapse_runaway caps any run of identical consecutive lines so a pathological page cannot flood the index; legitimate short repeats (a label appearing a few times) survive. Applied in ocr_pages. * Restrict whole-document injection to thread attachments only whole_document_context resolved the combined project+thread scope, so a project chat (the frontend sends both thread_id and project_id) injected the entire project corpus in full, contradicting the design that project and KB corpora stay retrieval-only. A large project corpus could also push the total over budget and drop a small thread attachment back to top-K. Resolve the thread scope alone in whole_document_context, and in build_rag_autoinject only enter whole-doc mode when a thread attachment is present and no KB is selected (a KB pick is exclusive: search that corpus). Project sources and KBs keep top-K retrieval. Adds regression tests for the mixed project+thread payload, the budget isolation, and KB precedence. * Address review: keep project retrieval, harden budget + OCR guards Follow-up to the 8-reviewer pass on the whole-document + OCR work. - Preserve project grounding in project chats. The thread-scope-only fix made whole-doc exclusive of retrieval, so a thread attachment silently dropped the project corpus for that turn. build_rag_autoinject now whole-docs the thread attachment AND retrieves the project sources top-K, merged under one citation numbering via tool.render_sources. KB selection stays exclusive. - Budget: a NULL/zero token_count no longer bypasses the cap (length-based fallback in _row_token_count), so a malformed huge doc can't inject in full. - OCR runaway guard: _collapse_runaway now also caps each distinct line at a generous total across the page (not just consecutive), bounding the interleaved/alternating loops weak models emit; blank-line floods collapse too. - OCR: warn when a scanned PDF exceeds OCR_MAX_PAGES (pages past the cap stay untranscribed) instead of silently dropping them. - Document the known limits: OCR'd pages have no PDF highlight regions; vision models need a micro-batch >= image tokens (Gemma-family) or the server aborts. - Tests for project-retrieval composition, NULL-token budget, and interleaved runaway; drop the now-superseded exclude-project test. * Add OCR toggle to RAG retrieval settings Make scanned-PDF OCR user-controllable per upload instead of only via the RAG_OCR_SCANNED config default. The retrieval settings panel gains an OCR scanned pages switch (persisted in localStorage, on by default); the chosen value is read fresh at upload time and sent with each document upload. Backend: the three upload routes accept an optional ocr form field and pass it through start_ingestion to _ocr_scanned_pages, which now treats None as use the config default and an explicit bool as an override. The on/off policy lives only in _ocr_scanned_pages now, so ocr_pages no longer re-checks the config (that double gate would have blocked a per-upload ocr=True while the default was off). Tests cover both override directions (force on while config off, force off while config on). * Add "Describe figures & charts" toggle with chart-aware captions Surface RAG figure captioning as a user control and make it actually useful for graphs and plots. The figure detection already clustered vector drawings and raster images into regions and rendered them, but captioning was off by default, had no UI, and used a thin generic prompt. Accuracy: the caption prompt now asks for chart type, axis titles and units, legend or series, salient trends and readable values, and table columns, while forbidding invented numbers. The token budget is configurable (CAPTION_MAX_TOKENS) and captions pass through the same runaway guard as OCR so a looping vision model cannot flood the index. Control: a per-upload caption override threads from the three upload routes through start_ingestion and _run, with the on/off policy single-sourced in _run (caption self-gating removed from caption_images, mirroring the OCR change) so a force-on override works when the config default is off. The frontend adds a "Describe figures & charts" switch in the retrieval settings, persisted in localStorage and sent with each upload. Default on; it is a no-op without a vision model and bounded to CAPTION_MAX_IMAGES figures per document. Tests cover the new caption_images contract, the runaway guard on captions, the chart-aware prompt and token budget (and that OCR keeps its own prompt and budget), and both override directions end to end through ingestion. * Generalize figure understanding: transcribe-first prompt + high-DPI tiling Make figure/chart description work across any visual and any model strength, not just a strong VLM on simple figures. Two changes, validated by a recall benchmark on authoritative documents (ResNet/Attention papers, USDA, UN UDHR). 1. Transcribe-first caption prompt. The caption now asks the model to transcribe every visible label verbatim (titles, axis labels and units, legends, every box/node/arrow label, table cells, equations) and then add a one-line summary, instead of only describing the figure. Transcription is the most model-robust visual task, so weak models that cannot reason about a chart still recover its labels. 2. High-DPI tiling of figure pages. Figure-bearing pages are rendered as an overlapping grid of high-DPI tiles (plus a full-page pass for context); each tile is transcribed, then merged and de-duplicated. This keeps small diagram labels legible and covers every sub-figure without relying on exact region detection, which previously missed sub-figures and small labels. Supporting changes: figure render DPI 130 -> 200 with a clip margin so edge labels are not lost; vision calls are deterministic (temperature 0) so transcription does not randomly drop labels; the repetition guard now applies to captions too. New config knobs: FIGURE_DPI, FIGURE_MARGIN_FRAC, FIGURE_TILE_ROWS/COLS, FIGURE_TILE_ OVERLAP, FIGURE_FULLPAGE, CAPTION_MAX_PAGES, larger CAPTION_MAX_TOKENS, and CAPTION_MAX_IMAGES as a per-document tile budget. Measured figure context recall (per-label, dense academic figures): Qwen2.5-VL: 0.50 -> 0.83 (overall 0.81 -> 0.94) Gemma-4-E2B (weak): ~0 with loops -> 0.83 (overall 0.91) Born-digital text and scanned-page recall are unchanged (no regression). parsers gains _figure_boxes (shared detection), pages_with_figures, and render_pdf_figure_tiles; captioner gains merge_page_captions and a temperature parameter; ingestion routes figure captioning through the tiled path. * Fix RAG review issues: whole-doc budget pre-check, figure gating, empty re-ingest, vision auth Whole-document context now runs a cheap token-sum pre-check (store.scope_token_estimate) before hydrating every chunk's text, so an attachment that cannot fit the budget is rejected without loading the whole corpus into memory. The estimate mirrors all_chunks_for_scope's filter and the per-row token-count fallback exactly. Ingestion skips all figure work (PDF rasterization and detection, not just the caption call) unless a vision model is loaded, so a text-only deployment pays nothing. When OCR is enabled, scanned/image-only pages are excluded from figure tiling since OCR already transcribes them whole, avoiding double vision work and overlapping index entries; a scanned figure page is still tiled when OCR is off. start_ingestion no longer dedupes forever to a prior ingest that produced zero chunks (e.g. a scanned PDF uploaded before a vision model was loaded): the empty record is dropped and the content is re-ingested. Vision OCR and caption requests now send the backend Authorization header, so they match the chat endpoint and do not 401 under direct-stream (--api-key) mode. Adds tests for the budget estimate, scanned-page exclusion, the vision-model gate, the empty re-ingest path, and the auth-header passthrough. * Trim RAG vision-ingestion comments and docstrings Tighten the verbose multi-line docstrings and comments added across the RAG vision ingestion work (captioner, config, parsers, ingestion, store, tool, build_rag_autoinject, the RAG tests, and the chat-store/upload-hook frontend toggles) to one or two lines while keeping their intent. No code changed: verified comment/docstring-only against the prior commit, and the RAG test suite still passes. * Fix figure-tiling exclusion and client dedupe for re-ingestable docs Figure tiling now excludes only the pages OCR actually transcribed, not every text-less page. _ocr_scanned_pages returns the set of pages it OCR'd, and _run passes that to pages_with_figures as exclude_pages (replacing the ocr_on-keyed min_text_chars heuristic). A scanned page that OCR skipped (past OCR_MAX_PAGES, or whose OCR returned empty) is no longer dropped from captioning, so a chart on such a page still gets a caption. The document panel's upload dedupe no longer skips re-selecting a file whose only matching doc completed with zero chunks. Such a doc is re-ingestable (e.g. a scan attached before a vision model loaded), and the backend re-ingests on the same content hash, so the client must let it reach the backend; healthy or still-indexing docs are still skipped. The SSE complete frame's chunk count is recorded on the doc so the check is exact. Adds a regression test for the un-OCR'd scanned figure page and updates the pages_with_figures test to the exclude_pages interface. * Address review findings: whole-doc budget guard, job numChunks, dead code, upload cap whole_document_context now treats a non-positive max_tokens as "never inject" instead of injecting the whole corpus unbounded, so RAG_WHOLE_DOC_MAX_TOKENS=0 tightens rather than disables the budget (the real off switch stays RAG_THREAD_WHOLE_DOC=0). The job-status endpoint and get_job_status now expose num_chunks (joined from the document), and the upload hook threads it through the SSE-fallback completion paths (reconcile + poll). Previously a document that finished via the connection-cap fallback had no chunk count client-side, so the re-ingest dedupe wrongly treated it as empty and re-uploaded it. IndexJob/JobEvent gain the field and the untyped cast is dropped. Removes the dead render_pdf_figures function (superseded by the tiling path), its test, and the unused FIGURE_MARGIN_FRAC config knob. Adds an upload size cap (RAG_MAX_UPLOAD_BYTES, default 200 MB; 413 on exceed with the partial file cleaned up) so a pathological file can't drive unbounded parse + vision work. render_pdf_figure_tiles clamps rows/cols to >= 1 (no ZeroDivisionError on a misconfigured grid). Captioning progress is reported after OCR so the bar is monotonic. sqlite connections set busy_timeout=5000 so a long figure/scan ingest holding its connection doesn't make a concurrent ingest/read fail with "database is locked". Adds tests for the non-positive budget, the zero-grid clamp, job-status num_chunks, and the oversize-upload rejection. * Extract PDF text as layout-aware Markdown via pymupdf4llm parsers._pdf now extracts each PDF page as Markdown with pymupdf4llm.to_markdown (page_chunks=True) instead of flat page.get_text("text"), so tables, headings and lists keep their structure in the indexed chunks and retrieve far better (a table's cells stay associated with their row instead of flattening into a token stream). Gated by RAG_PDF_MARKDOWN (default on); falls back to plain PyMuPDF text when the toggle is off, pymupdf4llm is missing, extraction fails, or a page yields no Markdown. The scanned-page OCR and figure-tiling passes operate on rendered pixels and are unaffected; docx/html/txt keep their existing extractors. The preview-highlight locator already strips Markdown punctuation when building anchors; it now also splits anchor tokens on pipes so a Markdown table row still anchors to the raw PDF word stream. Declares pymupdf4llm as a studio/RAG dependency (was only transitively present via the data-designer plugin). Adds parser tests (Markdown table reaches the page text, the plain-text fallback, the missing-lib fallback) and a locator test for table-pipe anchoring. * Pin pymupdf4llm to 0.3.4 so the package scan does not pull onnxruntime The lockstep pymupdf4llm 1.27.x line makes pymupdf-layout a hard dependency, which in turn pulls onnxruntime (plus numpy/networkx/protobuf). The security-audit pip scan-packages job resolves requirements --with-deps, so adding pymupdf4llm to no-torch-runtime.txt and studio.txt surfaced onnxruntime's un-baselined CRITICAL finding and flipped the hf-stack shard from pass to fail. pymupdf4llm 0.3.x keeps pymupdf-layout behind an optional [layout] extra, so a plain install resolves to pymupdf + tabulate only and never touches onnxruntime. 0.3.4 requires pymupdf>=1.27.1, satisfied by our pinned pymupdf==1.27.2.3, and to_markdown (page_chunks=True) produces equivalent layout-aware Markdown on real PDFs (verified on the Attention, ResNet and USDA documents). Production already installs these files --no-deps, so onnxruntime was never shipped at runtime; this only fixes the scanner. The parser test now asserts Markdown markup (heading or table pipes) rather than table pipes specifically, since 0.3.4 emits a heading but not a pipe table on the tiny borderless synthetic fixture; both markers are absent from the plain-text fallback. * Fix RAG whole-doc review findings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address RAG whole-doc review follow-ups * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address RAG review follow-up edge cases * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reserve image budget for whole-document RAG * [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: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: wasimysaid Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/tools.py | 143 ++++- studio/backend/core/rag/captioner.py | 173 +++++- studio/backend/core/rag/config.py | 49 +- studio/backend/core/rag/ingestion.py | 144 ++++- studio/backend/core/rag/locators.py | 6 +- studio/backend/core/rag/parsers.py | 212 +++++-- studio/backend/core/rag/store.py | 37 ++ studio/backend/core/rag/tool.py | 81 ++- .../backend/requirements/no-torch-runtime.txt | 5 + studio/backend/requirements/studio.txt | 3 + studio/backend/routes/rag.py | 32 +- studio/backend/storage/rag_db.py | 4 + studio/backend/tests/test_rag_captioning.py | 303 +++++++++- studio/backend/tests/test_rag_ingestion.py | 63 +++ studio/backend/tests/test_rag_ocr_fallback.py | 259 +++++++++ studio/backend/tests/test_rag_parsing.py | 88 +++ studio/backend/tests/test_rag_preview.py | 19 + .../backend/tests/test_rag_whole_document.py | 520 ++++++++++++++++++ .../src/features/chat/api/chat-adapter.ts | 6 + .../chat/stores/chat-runtime-store.ts | 26 + .../frontend/src/features/chat/types/api.ts | 3 + .../frontend/src/features/rag/api/rag-api.ts | 32 +- .../components/retrieval-settings-section.tsx | 52 ++ .../rag/components/use-rag-documents.ts | 66 ++- studio/frontend/src/features/rag/types/rag.ts | 2 + 25 files changed, 2171 insertions(+), 157 deletions(-) create mode 100644 studio/backend/tests/test_rag_ocr_fallback.py create mode 100644 studio/backend/tests/test_rag_parsing.py create mode 100644 studio/backend/tests/test_rag_whole_document.py diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index a5c193ff39..82c50933fc 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -1121,6 +1121,61 @@ def _autoinject_top_k() -> int: return _AUTOINJECT_DEFAULT_TOP_K +def _thread_whole_doc_enabled(scope: dict) -> bool: + """Whether a thread-attached file should be injected in full rather than + retrieved top-K. ``rag_scope.whole_doc=False`` disables it for this request.""" + override = scope.get("whole_doc") + if override is False: + return False + try: + from core.rag import config as _rag_config + except Exception: # noqa: BLE001 + return True + return _rag_config.THREAD_WHOLE_DOC + + +_IMAGE_PART_TOKEN_ESTIMATE = 1024 + + +def _message_token_estimate(conversation: list[dict]) -> int: + """Cheap prompt-size estimate for budget guards; exact tokenization happens later.""" + total = 0 + for msg in conversation: + content = msg.get("content") + if isinstance(content, str): + total += max(1, len(content) // 4) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + if part.get("type") in ("image_url", "input_image"): + total += _IMAGE_PART_TOKEN_ESTIMATE + else: + total += max(1, len(str(part.get("text") or "")) // 4) + total += 4 # chat-template role / separator overhead estimate + return total + + +def _whole_doc_budget(scope: dict | None = None, conversation: list[dict] | None = None) -> int: + try: + from core.rag import config as _rag_config + except Exception: # noqa: BLE001 + budget = 6000 + else: + budget = _rag_config.WHOLE_DOC_MAX_TOKENS + if not scope: + return budget + context = _opt_int(scope.get("context_length") or scope.get("max_context_tokens")) + if context is None or context <= 0: + return budget + headroom = _opt_int(scope.get("response_headroom")) + if headroom is None: + headroom = max(1024, context // 4) + used = _message_token_estimate(conversation or []) + # Leave room for tool XML wrappers, citation metadata, and chat-template overhead. + available = context - headroom - used - 512 + return min(budget, max(0, available)) + + def _last_user_text(conversation: list[dict]) -> str: """Plain text of the most recent user turn (text parts only).""" for msg in reversed(conversation): @@ -1154,7 +1209,11 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di enabled = rag_scope.get("autoinject") if enabled is None: enabled = _autoinject_enabled() - if not enabled: + thread_id = rag_scope.get("thread_id") + whole_doc_requested = ( + bool(thread_id) and not rag_scope.get("kb_id") and _thread_whole_doc_enabled(rag_scope) + ) + if not enabled and not whole_doc_requested: return None query = _last_user_text(conversation) if not query: @@ -1163,35 +1222,81 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di from storage import rag_db if not rag_db.RAG_AVAILABLE: return None - from core.rag.tool import search_for_autoinject + from core.rag.tool import render_sources, search_for_autoinject, whole_document_context except Exception as exc: # noqa: BLE001 logger.warning("RAG auto-inject unavailable: %s", exc) return None + text: str | None = None + sources: list[dict] = [] + floor_override = rag_scope.get("autoinject_min_score") floor = float(floor_override) if floor_override is not None else _autoinject_floor() # Cap at the lean top_k, but honor a lower user setting. lean_k = _autoinject_top_k() sidebar_k = _opt_int(rag_scope.get("default_top_k")) top_k = min(sidebar_k, lean_k) if sidebar_k is not None else lean_k - try: - found = search_for_autoinject( - query = query, - scope_kb_id = rag_scope.get("kb_id"), - scope_thread_id = rag_scope.get("thread_id"), - scope_project_id = rag_scope.get("project_id"), - top_k = top_k, - min_dense_score = floor, - **_scope_retrieval_kwargs(rag_scope), - ) - except Exception as exc: # noqa: BLE001 - logger.warning("RAG auto-inject retrieval failed: %s", exc) - return None - if not found: - logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor) + + # Whole-document mode: a thread-attached file under budget is injected in full so + # the model reads everything. A KB selection is exclusive, so whole-doc never + # preempts it; in a project chat the project sources are still retrieved top-K and + # appended under one citation numbering. Oversized files (or no thread doc) fall + # through to the combined top-K retrieval below. + if whole_doc_requested: + try: + budget = _whole_doc_budget(rag_scope, conversation) + + whole = whole_document_context( + scope_thread_id = thread_id, + max_tokens = budget, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG whole-document context failed: %s", exc) + whole = None + if whole is not None: + text, sources = whole + project_id = rag_scope.get("project_id") + if project_id: + try: + proj = search_for_autoinject( + query = query, + scope_project_id = project_id, + top_k = top_k, + min_dense_score = floor, + **_scope_retrieval_kwargs(rag_scope), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG project retrieval (whole-doc companion) failed: %s", exc) + proj = None + if proj is not None: + merged = sources + proj[1] + merged_text = render_sources(merged) + if max(1, len(merged_text) // 4) <= budget: + sources = merged + text = merged_text + logger.info("RAG auto-inject: whole-document context (%d chunk(s))", len(sources)) + + if text is None and enabled: + try: + found = search_for_autoinject( + query = query, + scope_kb_id = rag_scope.get("kb_id"), + scope_thread_id = rag_scope.get("thread_id"), + scope_project_id = rag_scope.get("project_id"), + top_k = top_k, + min_dense_score = floor, + **_scope_retrieval_kwargs(rag_scope), + ) + except Exception as exc: # noqa: BLE001 + logger.warning("RAG auto-inject retrieval failed: %s", exc) + return None + if not found: + logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor) + return None + text, sources = found + if text is None: return None - text, sources = found import json as _json import uuid as _uuid @@ -1236,7 +1341,7 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di "content": text, }, ] - logger.info("RAG auto-inject: %d passage(s) >= %.2f for %r", len(sources), floor, query[:80]) + logger.info("RAG auto-inject: %d passage(s) for %r", len(sources), query[:80]) return {"events": events, "messages": messages} diff --git a/studio/backend/core/rag/captioner.py b/studio/backend/core/rag/captioner.py index be8e341064..e29c9c9a7a 100644 --- a/studio/backend/core/rag/captioner.py +++ b/studio/backend/core/rag/captioner.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Caption figures with the loaded vision model and splice the text into the page -so images are searchable via the normal FTS5 + dense path. No-op (never raises) -without a vision model or on failure; gated by ``config.CAPTION_IMAGES``.""" +"""Vision-model helpers for ingestion: figure captioning and scanned-page OCR. + +Both turn pixels into indexable text and are a no-op (never raise) without a loaded +vision model. They reuse the chat model's vision endpoint, so it must be served with +``--ubatch-size`` >= one image's tokens (some encoders, e.g. Gemma, attend +non-causally and abort otherwise); Studio's vision chat already requires this.""" from __future__ import annotations @@ -15,11 +18,54 @@ from . import config logger = logging.getLogger(__name__) _CAPTION_PROMPT = ( - "Describe this figure or image from a document in one or two concise " - "sentences, for search indexing. State what it depicts (e.g. a diagram, " - "chart, table or photo) and its key content. Do not add commentary." + "Read this figure or image from a document for search indexing.\n" + "First, on a line 'TEXT:', transcribe every piece of visible text exactly as " + "written, in reading order: the title, axis labels and units, legend and series " + "names, EVERY box / node / arrow label, table headers and cells, equations, and " + "footnotes. List each distinct label even if it is small.\n" + "Then, on a line 'SUMMARY:', add one or two sentences on what it shows (chart " + "type and trend, diagram subject, table topic, or photo content).\n" + "Report only what is visible. Transcribe exactly; do not invent or guess any " + "text, label, or number." ) +_OCR_PROMPT = ( + "Transcribe all text on this document page exactly as it appears, in reading " + "order, including any text inside figures, diagrams, charts, and tables (keep " + "table rows readable). Output only the transcribed text, with no commentary or " + "code fences. Preserve headings, lists, and line breaks. If the page has no " + "readable text, output nothing." +) + + +def _collapse_runaway( + text: str, + max_repeat: int = 3, + max_total: int = 8, +) -> str: + """Cap runaway repetition: vision models sometimes loop a line many times. Keep + each distinct line to ``max_repeat`` in a row and ``max_total`` total, and collapse + blank-line floods, so a degenerate page cannot flood the index.""" + out: list[str] = [] + seen: dict[str, int] = {} + prev: str | None = None + run = 0 + for line in text.splitlines(): + key = line.strip() + if not key: + if prev == "": # collapse runs of blank lines to a single separator + continue + prev = "" + out.append("") + continue + run = run + 1 if key == prev else 1 + prev = key + seen[key] = seen.get(key, 0) + 1 + if run > max_repeat or seen[key] > max_total: + continue + out.append(line) + return "\n".join(out) + def vision_endpoint() -> tuple[str, str] | None: """``(base_url, model)`` for a loaded vision GGUF model, else None.""" @@ -33,7 +79,28 @@ def vision_endpoint() -> tuple[str, str] | None: return None -def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: +def _vision_auth_headers() -> dict | None: + """Bearer header for the backend's API, or None. Vision calls share the chat + endpoint, so they need the same key under direct-stream (``--api-key``) mode.""" + try: + from routes.inference import get_llama_cpp_backend + return get_llama_cpp_backend()._auth_headers or None + except Exception: # noqa: BLE001 - auth discovery must never break ingestion + return None + + +def _vision_complete( + base_url: str, + model: str, + image_bytes: bytes, + *, + prompt: str, + timeout: float, + max_tokens: int, + temperature: float = 0.0, +) -> str | None: + """One image-in / text-out call to the loaded vision model's OpenAI-compatible + endpoint. Returns the stripped text or ``None`` on empty/failure (non-fatal).""" import httpx data_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii") @@ -43,33 +110,62 @@ def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) { "role": "user", "content": [ - {"type": "text", "text": _CAPTION_PROMPT}, + {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": data_url}}, ], } ], - "max_tokens": 200, - "temperature": 0.2, + "max_tokens": max_tokens, + # Deterministic by default: transcription must not randomly drop labels. + "temperature": temperature, "stream": False, # Off: thinking models would spend the budget reasoning, returning "". "chat_template_kwargs": {"enable_thinking": False}, } try: - r = httpx.post(f"{base_url}/v1/chat/completions", json = payload, timeout = timeout) + r = httpx.post( + f"{base_url}/v1/chat/completions", + json = payload, + timeout = timeout, + headers = _vision_auth_headers(), + ) r.raise_for_status() text = r.json()["choices"][0]["message"]["content"] return text.strip() or None - except Exception: # noqa: BLE001 - a failed caption is non-fatal - logger.debug("caption request failed", exc_info = True) + except Exception: # noqa: BLE001 - a failed vision call is non-fatal + logger.debug("vision request failed", exc_info = True) return None +def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: + return _vision_complete( + base_url, + model, + image_bytes, + prompt = _CAPTION_PROMPT, + timeout = timeout, + max_tokens = config.CAPTION_MAX_TOKENS, + ) + + +def _ocr_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None: + return _vision_complete( + base_url, + model, + image_bytes, + prompt = _OCR_PROMPT, + timeout = timeout, + max_tokens = config.OCR_MAX_TOKENS, + ) + + def caption_images( images: list, *, endpoint: tuple[str, str] | None = None ) -> dict[int, list[str]]: - """Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when - disabled, no vision model, or no images. Bounded by ``CAPTION_MAX_IMAGES``.""" - if not config.CAPTION_IMAGES or not images: + """Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when there + are no images or no vision model. The caller (`ingestion._run`) owns the on/off + policy. Bounded by ``CAPTION_MAX_IMAGES``; each caption passes ``_collapse_runaway``.""" + if not images: return {} ep = endpoint or vision_endpoint() if ep is None: @@ -84,7 +180,50 @@ def caption_images( caption = _caption_one(base_url, model, image_bytes, config.CAPTION_TIMEOUT_S) if caption: page = getattr(img, "page_number", None) or 0 - out.setdefault(int(page), []).append(caption) + out.setdefault(int(page), []).append(_collapse_runaway(caption)) + return out + + +def ocr_pages( + page_pngs: dict[int, bytes], *, endpoint: tuple[str, str] | None = None +) -> dict[int, str]: + """OCR rendered page PNGs (keyed by 1-based page number) to text; ``{}`` when there + is no vision model or no pages. The caller (`ingestion._ocr_scanned_pages`) owns the + on/off policy. Bounded by ``OCR_MAX_PAGES``.""" + if not page_pngs: + return {} + ep = endpoint or vision_endpoint() + if ep is None: + return {} + base_url, model = ep + + out: dict[int, str] = {} + for page_num in sorted(page_pngs)[: config.OCR_MAX_PAGES]: + text = _ocr_one(base_url, model, page_pngs[page_num], config.OCR_TIMEOUT_S) + if text: + out[int(page_num)] = _collapse_runaway(text) + return out + + +def merge_page_captions(captions: dict[int, list[str]]) -> dict[int, list[str]]: + """Merge a page's per-tile captions into one deduped block: drop lines repeated + across overlapping tiles (first kept, order preserved), then ``_collapse_runaway``, + so ``splice_captions`` adds a single figure block per page.""" + out: dict[int, list[str]] = {} + for page, caps in captions.items(): + seen: set[str] = set() + lines: list[str] = [] + for cap in caps: + for line in (cap or "").splitlines(): + stripped = line.strip() + key = stripped.lower() + if not stripped or key in seen: + continue + seen.add(key) + lines.append(stripped) + merged = _collapse_runaway("\n".join(lines)) + if merged.strip(): + out[page] = [merged] return out diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 993423683c..54a224d081 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -17,13 +17,50 @@ TOP_K_DENSE = int(os.environ.get("RAG_TOP_K_DENSE", "30")) TOP_K_HYBRID = int(os.environ.get("RAG_TOP_K_HYBRID", "10")) RRF_K = int(os.environ.get("RAG_RRF_K", "60")) -UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"} +# Whole-document context: a thread-attached file under the token budget is injected +# in full (every chunk, in order) instead of top-K retrieval; above it, use retrieval. +THREAD_WHOLE_DOC = os.environ.get("RAG_THREAD_WHOLE_DOC", "1") == "1" +WHOLE_DOC_MAX_TOKENS = int(os.environ.get("RAG_WHOLE_DOC_MAX_TOKENS", "6000")) -# Figure captioning via the loaded vision model; off by default since each caption -# is a model call. MAX_IMAGES bounds per-doc cost. -CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "0") == "1" -CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "8")) -CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "30")) +UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"} +# Reject uploads larger than this, so one pathological file can't drive unbounded parse +# + vision work at ingest. 0 disables the cap. Default 200 MB. +MAX_UPLOAD_BYTES = int(os.environ.get("RAG_MAX_UPLOAD_BYTES", str(200 * 1024 * 1024))) + +# Extract PDF text as layout-aware Markdown (pymupdf4llm) instead of flat text, so +# tables, headings and lists survive into chunks and retrieval. Falls back to plain +# PyMuPDF text when off, when pymupdf4llm is missing, or when extraction fails. +PDF_MARKDOWN = os.environ.get("RAG_PDF_MARKDOWN", "1") == "1" + +# Figure captioning via the loaded vision model: detected figures are transcribed + +# described so they become searchable. On by default, a no-op without a vision model; +# the chat's "Describe figures & charts" toggle overrides it per upload. +CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "1") == "1" +# Total per-document tile budget (figure-bearing pages are tiled, see below). +CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "24")) +CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "60")) +# Larger than a one-line caption since captions transcribe every label. FIGURE_DPI is +# high enough to keep small box/axis labels legible when tiles are rendered. +CAPTION_MAX_TOKENS = int(os.environ.get("RAG_CAPTION_MAX_TOKENS", "768")) +FIGURE_DPI = int(os.environ.get("RAG_FIGURE_DPI", "200")) +# Figure pages are tiled into an overlapping ROWS x COLS grid of high-DPI tiles (plus +# an optional full page), so small labels and every sub-figure are covered without +# exact region detection. MAX_PAGES bounds figure pages; MAX_IMAGES bounds total tiles. +FIGURE_TILE_ROWS = int(os.environ.get("RAG_FIGURE_TILE_ROWS", "2")) +FIGURE_TILE_COLS = int(os.environ.get("RAG_FIGURE_TILE_COLS", "2")) +FIGURE_TILE_OVERLAP = float(os.environ.get("RAG_FIGURE_TILE_OVERLAP", "0.12")) +FIGURE_FULLPAGE = os.environ.get("RAG_FIGURE_FULLPAGE", "1") == "1" +CAPTION_MAX_PAGES = int(os.environ.get("RAG_CAPTION_MAX_PAGES", "4")) + +# Scanned-PDF OCR: a page with little extractable text is rendered and transcribed by +# the vision model so it becomes searchable. Needs a vision model, else skipped (page +# stays empty). MIN_CHARS is the text length below which a page is treated as scanned. +OCR_SCANNED = os.environ.get("RAG_OCR_SCANNED", "1") == "1" +OCR_MIN_CHARS = int(os.environ.get("RAG_OCR_MIN_CHARS", "16")) +OCR_MAX_PAGES = int(os.environ.get("RAG_OCR_MAX_PAGES", "20")) +OCR_DPI = int(os.environ.get("RAG_OCR_DPI", "150")) +OCR_TIMEOUT_S = float(os.environ.get("RAG_OCR_TIMEOUT_S", "60")) +OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048")) # Embedder backend. "auto": sentence-transformers on a CUDA/ROCm GPU (torch fp16 # wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 77bbdc92f5..04365ab76b 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -99,25 +99,108 @@ def _embed_all(texts: list[str], model_name: str | None): return vectors +def _ocr_scanned_pages( + pages: list, + stored_path: str, + conn, + job_id: str, + ocr: bool | None = None, +) -> tuple[list, set[int]]: + """Replace text on near-empty (scanned/image-only) PDF pages with vision-model OCR + so image PDFs become searchable. ``ocr`` overrides ``config.OCR_SCANNED`` per upload + (``None`` = config default); no-op without scanned pages or a vision model. OCR'd + pages have no text layer, so no preview highlight regions, but stay searchable. + Returns ``(pages, ocred)``: new ``Page`` objects for OCR'd pages (originals + otherwise) and the set of page numbers actually transcribed.""" + if not (config.OCR_SCANNED if ocr is None else ocr): + return pages, set() + scanned = [ + p.page_number + for p in pages + if p.page_number is not None and len((p.text or "").strip()) < config.OCR_MIN_CHARS + ] + if not scanned or captioner.vision_endpoint() is None: + return pages, set() + if len(scanned) > config.OCR_MAX_PAGES: + logger.warning( + "OCR: %d scanned pages exceed OCR_MAX_PAGES=%d; pages past the cap stay " + "untranscribed (raise RAG_OCR_MAX_PAGES to cover them)", + len(scanned), + config.OCR_MAX_PAGES, + ) + scanned = scanned[: config.OCR_MAX_PAGES] + _progress(conn, job_id, "ocr", 0.25) + page_pngs = parsers.render_pdf_pages(stored_path, scanned, dpi = config.OCR_DPI) + texts = captioner.ocr_pages(page_pngs) + if not texts: + return pages, set() + + from .parsers import Page + + out: list = [] + ocred: set[int] = set() + for page in pages: + text = texts.get(page.page_number) + if text: + original = (page.text or "").strip() + merged = text if not original or original in text else f"{original}\n\n{text}" + out.append(Page(text = merged, page_number = page.page_number, char_count = len(merged))) + ocred.add(page.page_number) + else: + out.append(page) + return out, ocred + + def _run( - job_id: str, document_id: str, scope: str, stored_path: str, model_name: str | None + job_id: str, + document_id: str, + scope: str, + stored_path: str, + model_name: str | None, + ocr: bool | None = None, + caption: bool | None = None, ) -> None: conn = rag_db.get_connection() try: _progress(conn, job_id, "parsing", 0.1) pages = parsers.parse(stored_path) - if config.CAPTION_IMAGES and stored_path.lower().endswith(".pdf"): - # Caption figures, splice into page text (no-op without a vision model). + is_pdf = stored_path.lower().endswith(".pdf") + ocred: set[int] = set() + if is_pdf: + pages, ocred = _ocr_scanned_pages(pages, stored_path, conn, job_id, ocr = ocr) + caption_on = config.CAPTION_IMAGES if caption is None else caption + # Skip all figure work (PDF rasterization included) without a vision model. + if caption_on and is_pdf and captioner.vision_endpoint() is not None: + # Tile figure pages, transcribe+describe each tile, then merge/dedup/splice + # into the page text so small labels and every sub-figure are captured. try: - figures = parsers.render_pdf_figures( - stored_path, max_figures = config.CAPTION_MAX_IMAGES + fig_pages = parsers.pages_with_figures( + stored_path, + max_pages = config.CAPTION_MAX_PAGES, + # Skip only pages OCR actually transcribed (it covers them whole); a + # scanned figure page past the OCR cap or with empty OCR still tiles. + exclude_pages = ocred, + ) + tiles = ( + parsers.render_pdf_figure_tiles( + stored_path, + fig_pages, + dpi = config.FIGURE_DPI, + rows = config.FIGURE_TILE_ROWS, + cols = config.FIGURE_TILE_COLS, + overlap = config.FIGURE_TILE_OVERLAP, + fullpage = config.FIGURE_FULLPAGE, + max_tiles = config.CAPTION_MAX_IMAGES, + ) + if fig_pages + else [] ) except Exception: - logger.warning("figure rendering failed for job %s", job_id, exc_info = True) - figures = [] - if figures: - _progress(conn, job_id, "captioning", 0.2) - captions = captioner.caption_images(figures) + logger.warning("figure tiling failed for job %s", job_id, exc_info = True) + tiles = [] + if tiles: + _progress(conn, job_id, "captioning", 0.28) + captions = captioner.merge_page_captions(captioner.caption_images(tiles)) pages = captioner.splice_captions(pages, captions) _progress(conn, job_id, "chunking", 0.3) @@ -175,6 +258,8 @@ def start_ingestion( *, project_id: str | None = None, model_name: str | None = None, + ocr: bool | None = None, + caption: bool | None = None, ) -> tuple[str, str]: """Create the document + job rows and spawn the worker, returning ``(document_id, job_id)``. A duplicate content hash in this scope returns the @@ -191,13 +276,26 @@ def start_ingestion( try: existing = store.document_by_hash(conn, scope, sha) if existing is not None: - job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) - _remove_upload(stored_path) - with _jobs_lock: - _jobs[job_id] = queue.Queue() - _emit(job_id, {"type": "complete", "num_chunks": 0, "deduped": True}) - _emit(job_id, None) - return existing, job_id + doc = store.get_document(conn, existing) + empty_completed = ( + doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks") + ) + if empty_completed: + # A prior ingest of identical bytes yielded zero chunks (e.g. a scanned + # PDF uploaded before a vision model loaded). Re-ingest, don't dedupe. + store.delete_document(conn, existing) + _remove_upload(doc.get("stored_path"), keep_path = stored_path) + else: + job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) + _remove_upload(stored_path) + with _jobs_lock: + _jobs[job_id] = queue.Queue() + _emit( + job_id, + {"type": "complete", "num_chunks": doc.get("num_chunks") or 0, "deduped": True}, + ) + _emit(job_id, None) + return existing, job_id for failed in store.failed_documents_by_hash(conn, scope, sha): store.delete_document(conn, failed["id"]) _remove_upload(failed.get("stored_path"), keep_path = stored_path) @@ -221,7 +319,7 @@ def start_ingestion( _jobs[job_id] = queue.Queue() threading.Thread( target = _run, - args = (job_id, document_id, scope, stored_path, model_name), + args = (job_id, document_id, scope, stored_path, model_name, ocr, caption), daemon = True, ).start() return document_id, job_id @@ -339,10 +437,16 @@ def job_events(job_id: str): def get_job_status(job_id: str) -> dict | None: - """Read the persisted ingestion job row (status / stage / progress / error).""" + """Read the persisted ingestion job row (status / stage / progress / error), plus + the document's ``num_chunks`` so a client polling to completion learns the chunk + count (the SSE ``complete`` frame carries it, but the poll/reconcile path does not).""" conn = rag_db.get_connection() try: - row = conn.execute("SELECT * FROM ingestion_jobs WHERE id=?", (job_id,)).fetchone() + row = conn.execute( + "SELECT j.*, d.num_chunks AS num_chunks FROM ingestion_jobs j " + "LEFT JOIN documents d ON d.id = j.document_id WHERE j.id=?", + (job_id,), + ).fetchone() return dict(row) if row else None finally: conn.close() diff --git a/studio/backend/core/rag/locators.py b/studio/backend/core/rag/locators.py index 57c0487486..9331bb15ac 100644 --- a/studio/backend/core/rag/locators.py +++ b/studio/backend/core/rag/locators.py @@ -39,9 +39,11 @@ def _norm_token(token: str) -> str: def _anchor_tokens(page_text: str, match: LocatorMatch) -> list[str]: """Normalized anchor tokens from the chunk's leading span. Drops first and last - token (boundaries often slice mid-word) when long enough.""" + token (boundaries often slice mid-word) when long enough. Pipes are split out so + Markdown table cells (``|Q1|$1.2M|``) become individual words that match the PDF + word stream.""" segment = page_text[match.start : match.end] - raw = segment.split() + raw = segment.replace("|", " ").split() if len(raw) >= MIN_ANCHOR_WORDS + 2: raw = raw[1:-1] tokens = [t for t in (_norm_token(w) for w in raw) if t] diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index 84da941762..ba248cf9a6 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -15,6 +15,8 @@ import os from dataclasses import dataclass from html.parser import HTMLParser +from . import config + logger = logging.getLogger(__name__) @@ -67,6 +69,28 @@ def _html(raw: str) -> list[Page]: return [_page("\n".join(parser.out), 1)] +def _pdf_markdown(doc) -> list[str] | None: + """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index + i maps to page i+1. Returns None when the lib is missing, extraction fails, or the + page count does not line up, so the caller falls back to plain PyMuPDF text.""" + try: + import pymupdf4llm + except Exception: + return None + try: + chunks = pymupdf4llm.to_markdown( + doc, + page_chunks = True, + show_progress = False, + ) + except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion + logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True) + return None + if not isinstance(chunks, list) or len(chunks) != doc.page_count: + return None + return [str(c.get("text") or "") for c in chunks] + + def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: import fitz # PyMuPDF @@ -74,8 +98,11 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: images: list[ParsedImage] = [] doc = fitz.open(path) try: + md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None for i, page in enumerate(doc): - text = page.get_text("text") or "" + # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval); + # fall back to plain text when Markdown is off, unavailable, or empty here. + text = (md[i] if md else "") or page.get_text("text") or "" pages.append(_page(text, i + 1)) if want_images: for img in page.get_images(full = True): @@ -118,63 +145,164 @@ def _merge_rects(boxes: list) -> list: return merged -def render_pdf_figures( - path: str, +def _figure_boxes( + page, *, - dpi: int = 130, min_area_frac: float = 0.04, min_side: float = 40.0, - max_figures: int = 8, -) -> list[ParsedImage]: - """Detect figure regions and render each to a PNG for captioning. +) -> list: + """Qualifying figure-region rectangles on a page: cluster vector drawings + raster + placements, merge overlaps, keep the page-spanning ones (area/side filtered).""" + boxes: list = [] + try: + boxes.extend(info["bbox"] for info in page.get_image_info()) + except Exception: + pass + try: + boxes.extend(page.cluster_drawings()) + except Exception: + pass + if not boxes: + return [] + page_area = page.rect.width * page.rect.height + keep: list = [] + for box in _merge_rects(boxes): + if ( + box.get_area() >= min_area_frac * page_area + and box.width >= min_side + and box.height >= min_side + ): + keep.append(box) + return keep - Academic figures are vector, so raster extraction yields fragments; instead - cluster vector drawings + raster placements into boxes, keep the page-spanning - ones, and render them. Any failure yields [], never an exception. - """ + +def pages_with_figures( + path: str, + *, + max_pages: int = 4, + min_area_frac: float = 0.04, + min_side: float = 40.0, + exclude_pages: set[int] | None = None, +) -> list[int]: + """1-based page numbers with a qualifying figure region, capped at ``max_pages``; + drives figure tiling. ``exclude_pages`` (1-based) are skipped: those are the pages + OCR already transcribed whole, so tiling them would duplicate the vision work. Any + failure yields [].""" + exclude = exclude_pages or set() try: import pymupdf except Exception: return [] - - out: list[ParsedImage] = [] try: doc = pymupdf.open(path) except Exception: return [] + pages: list[int] = [] try: for i, page in enumerate(doc): - boxes: list = [] - try: - boxes.extend(info["bbox"] for info in page.get_image_info()) - except Exception: - pass - try: - boxes.extend(page.cluster_drawings()) - except Exception: - pass - if not boxes: + if (i + 1) in exclude: continue - page_area = page.rect.width * page.rect.height - for box in _merge_rects(boxes): - if ( - box.get_area() >= min_area_frac * page_area - and box.width >= min_side - and box.height >= min_side - ): - try: - pix = page.get_pixmap(dpi = dpi, clip = box) - out.append( - ParsedImage( - image_bytes = pix.tobytes("png"), - page_number = i + 1, - xref = 0, - ) + if _figure_boxes(page, min_area_frac = min_area_frac, min_side = min_side): + pages.append(i + 1) + if len(pages) >= max_pages: + break + return pages + finally: + doc.close() + + +def render_pdf_figure_tiles( + path: str, + page_numbers, + *, + dpi: int = 200, + rows: int = 2, + cols: int = 2, + overlap: float = 0.12, + fullpage: bool = True, + max_tiles: int = 24, +) -> list[ParsedImage]: + """Render figure-bearing pages as overlapping high-DPI tiles (plus an optional full + page), each a ``ParsedImage`` keyed by page number. Tiling keeps small labels legible + and covers every sub-figure without exact region detection. Any failure yields [].""" + wanted = [int(n) for n in page_numbers] + if not wanted: + return [] + rows, cols = max(1, int(rows)), max(1, int(cols)) # never divide by zero + try: + import pymupdf + except Exception: + return [] + try: + doc = pymupdf.open(path) + except Exception: + return [] + out: list[ParsedImage] = [] + try: + for num in wanted: + if num < 1 or num > doc.page_count: + continue + page = doc[num - 1] + rect = page.rect + clips: list = [rect] if fullpage else [] + cw, ch = rect.width / cols, rect.height / rows + ox, oy = cw * overlap, ch * overlap + for r in range(rows): + for c in range(cols): + clips.append( + pymupdf.Rect( + rect.x0 + c * cw - ox, + rect.y0 + r * ch - oy, + rect.x0 + (c + 1) * cw + ox, + rect.y0 + (r + 1) * ch + oy, ) - except Exception: - continue - if len(out) >= max_figures: - return out + & rect + ) + for clip in clips: + try: + pix = page.get_pixmap(dpi = dpi, clip = clip) + out.append(ParsedImage(image_bytes = pix.tobytes("png"), page_number = num, xref = 0)) + except Exception: + continue + if len(out) >= max_tiles: + return out + return out + finally: + doc.close() + + +def render_pdf_pages( + path: str, + page_numbers, + *, + dpi: int = 150, +) -> dict[int, bytes]: + """Render whole PDF pages (given as 1-based numbers) to PNG bytes, keyed by + page number. Backs scanned-page OCR. Any failure yields ``{}`` (or skips that + page), never an exception. + """ + wanted = {int(n) for n in page_numbers} + if not wanted: + return {} + try: + import pymupdf + except Exception: + return {} + try: + doc = pymupdf.open(path) + except Exception: + return {} + out: dict[int, bytes] = {} + try: + for i, page in enumerate(doc): + num = i + 1 + if num not in wanted: + continue + try: + pix = page.get_pixmap(dpi = dpi) + out[num] = pix.tobytes("png") + except Exception: + continue return out finally: doc.close() diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index 7d58931e53..8e59c5fbf6 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -292,3 +292,40 @@ def chunks_by_id(conn: sqlite3.Connection, ids) -> dict: list(ids), ).fetchall() return {r["id"]: r for r in rows} + + +def all_chunks_for_scope(conn: sqlite3.Connection, scope) -> list[dict]: + """Every completed-document chunk for a scope, ordered document-then-index and + joined with the document filename. Backs whole-document context injection, so + it does no retrieval or embedding.""" + scopes = _scopes(scope) + if not scopes: + return [] + placeholders = ",".join("?" * len(scopes)) + rows = conn.execute( + f"SELECT c.id, c.text, c.document_id, c.chunk_index, c.page_number, " + f"c.token_count, d.filename, d.created_at " + f"FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.scope IN ({placeholders}) AND d.status='completed' " + f"ORDER BY d.created_at, c.document_id, c.chunk_index", + list(scopes), + ).fetchall() + return [dict(r) for r in rows] + + +def scope_token_estimate(conn: sqlite3.Connection, scope) -> int: + """Upper-bound token total for a scope's completed chunks without hydrating text. + Mirrors ``all_chunks_for_scope`` + the ``tool._row_token_count`` fallback (stored + count, else length/4), so the whole-doc budget can be checked before loading text.""" + scopes = _scopes(scope) + if not scopes: + return 0 + placeholders = ",".join("?" * len(scopes)) + row = conn.execute( + f"SELECT COALESCE(SUM(CASE WHEN c.token_count > 0 THEN c.token_count " + f"ELSE MAX(1, length(COALESCE(c.text, '')) / 4) END), 0) AS total " + f"FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.scope IN ({placeholders}) AND d.status='completed'", + list(scopes), + ).fetchone() + return int(row["total"] or 0) diff --git a/studio/backend/core/rag/tool.py b/studio/backend/core/rag/tool.py index ccb1b47e63..b05f8dd3a3 100644 --- a/studio/backend/core/rag/tool.py +++ b/studio/backend/core/rag/tool.py @@ -16,7 +16,13 @@ from xml.sax.saxutils import quoteattr from storage import rag_db from . import config, retrieval -from .store import kb_scope, project_scope, thread_scope +from .store import ( + all_chunks_for_scope, + kb_scope, + project_scope, + scope_token_estimate, + thread_scope, +) SEARCH_KNOWLEDGE_BASE_TOOL = { "type": "function", @@ -90,6 +96,30 @@ def _format(rows, hits) -> tuple[str, list[dict]]: return "\n\n".join(blocks), sources +def render_sources(sources: list[dict]) -> str: + """Render a citation-source list to sequentially-numbered ```` blocks, + rewriting each source's ``citationId`` to match its 1-based position. Lets + independently-built source lists (a whole-document thread attachment plus + retrieved project passages) be merged under one citation numbering.""" + blocks: list[str] = [] + for i, s in enumerate(sources, 1): + s["citationId"] = i + src = quoteattr(s.get("filename") or "unknown") + page = s.get("page") + page_attr = f" page={quoteattr(str(page))}" if page else "" + blocks.append(f'\n{s.get("text") or ""}\n') + return "\n\n".join(blocks) + + +def _row_token_count(row) -> int: + """Chunk token count for budgeting, falling back to a length estimate when the + stored count is missing or zero, so a malformed chunk cannot bypass the budget.""" + tc = row["token_count"] + if tc: + return int(tc) + return max(1, len(row["text"] or "") // 4) + + def search_knowledge_base_with_sources( *, query: str, @@ -186,6 +216,55 @@ def search_for_autoinject( return (text, sources) if sources else None +def whole_document_context( + *, scope_thread_id: str | None = None, max_tokens: int +) -> tuple[str, list[dict]] | None: + """Render EVERY chunk of the THREAD's attached documents (in order) as the same + ```` blocks + citation source-map as retrieval, so the model reads the whole + file rather than top-K passages. Thread-attached files only: KB and project corpora + are search corpora, never whole-document, so this resolves the thread scope alone. + ``None`` (caller falls back to retrieval) when there is no thread scope, no completed + chunks, or the total exceeds ``max_tokens``.""" + if not scope_thread_id: + return None + # A non-positive budget means "never inject" (disable whole-doc via + # RAG_THREAD_WHOLE_DOC=0), not "inject the whole corpus unbounded". + if max_tokens <= 0: + return None + scope = thread_scope(scope_thread_id) + conn = rag_db.get_connection() + try: + # Cheap budget pre-check (SUM, no text hydration): reject an oversized attachment + # before loading the whole corpus; all_chunks_for_scope runs only once it fits. + if scope_token_estimate(conn, scope) > max_tokens: + return None + rows = all_chunks_for_scope(conn, scope) + finally: + conn.close() + if not rows: + return None + total = sum(_row_token_count(r) for r in rows) + if total > max_tokens: + return None + + sources: list[dict] = [ + { + "citationId": i, + "chunkId": r["id"], + "documentId": r["document_id"], + "filename": r["filename"] or "unknown", + "page": r["page_number"], + "text": r["text"] or "", + "score": None, + } + for i, r in enumerate(rows, 1) + ] + rendered = render_sources(sources) + if max(1, len(rendered) // 4) > max_tokens: + return None + return rendered, sources + + def search_knowledge_base( *, query: str, diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index 9796fc3a50..de321f80ed 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -73,4 +73,9 @@ pillow # this file installs --no-deps; without them Studio runs with RAG disabled. sqlite-vec==0.1.9 pymupdf==1.27.2.3 +# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the +# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown(). +pymupdf4llm==0.3.4 python-docx==1.2.0 + +lxml==6.0.2 diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 1b7f7a668c..6f4a5c3292 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -22,4 +22,7 @@ fastmcp>=3.0.2 # extras-no-deps.txt; these add the lexical+dense store and document parsing. sqlite-vec==0.1.9 pymupdf==1.27.2.3 +# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the +# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown(). +pymupdf4llm==0.3.4 python-docx==1.2.0 diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 8d23240fd5..4e35fce3c2 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -19,7 +19,7 @@ import secrets import time import uuid -from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile from fastapi.responses import FileResponse, StreamingResponse from pydantic import BaseModel, Field @@ -62,13 +62,24 @@ def _save_upload(file: UploadFile) -> tuple[str, str]: uploads = ensure_dir(rag_uploads_root()) stored_path = str(uploads / f"{uuid.uuid4().hex}{ext}") size = 0 + cap = config.MAX_UPLOAD_BYTES + too_big = False with open(stored_path, "wb") as out: while True: block = file.file.read(1 << 20) if not block: break size += len(block) + if cap and size > cap: + too_big = True + break out.write(block) + if too_big: + os.remove(stored_path) + raise HTTPException( + status_code = 413, + detail = f"File exceeds the {cap // (1024 * 1024)} MB upload limit.", + ) if size == 0: os.remove(stored_path) raise HTTPException(status_code = 400, detail = "Uploaded file is empty.") @@ -207,6 +218,8 @@ def delete_knowledge_base(kb_id: str, subject: str = Depends(get_current_subject async def upload_kb_document( kb_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() @@ -218,7 +231,7 @@ async def upload_kb_document( conn.close() stored_path, filename = _save_upload(file) document_id, job_id = ingestion.start_ingestion( - store.kb_scope(kb_id), kb_id, None, filename, stored_path + store.kb_scope(kb_id), kb_id, None, filename, stored_path, ocr = ocr, caption = caption ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -238,12 +251,20 @@ def list_kb_documents(kb_id: str, subject: str = Depends(get_current_subject)) - async def upload_thread_document( thread_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() stored_path, filename = _save_upload(file) document_id, job_id = ingestion.start_ingestion( - store.thread_scope(thread_id), None, thread_id, filename, stored_path + store.thread_scope(thread_id), + None, + thread_id, + filename, + stored_path, + ocr = ocr, + caption = caption, ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -263,6 +284,8 @@ def list_thread_documents(thread_id: str, subject: str = Depends(get_current_sub async def upload_project_document( project_id: str, file: UploadFile = File(...), + ocr: bool | None = Form(None), + caption: bool | None = Form(None), subject: str = Depends(get_current_subject), ) -> dict: _require_rag() @@ -278,6 +301,8 @@ async def upload_project_document( filename, stored_path, project_id = project_id, + ocr = ocr, + caption = caption, ) return {"documentId": document_id, "jobId": job_id, "filename": filename} @@ -321,6 +346,7 @@ def job_status(job_id: str, subject: str = Depends(get_current_subject)) -> dict "stage": row.get("stage"), "progress": row.get("progress") or 0.0, "error": row.get("error"), + "numChunks": row.get("num_chunks") or 0, } diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index 4da600d768..ce27326562 100644 --- a/studio/backend/storage/rag_db.py +++ b/studio/backend/storage/rag_db.py @@ -119,6 +119,10 @@ def get_connection() -> sqlite3.Connection: ensure_dir(db_path.parent) conn = sqlite3.connect(str(db_path)) conn.row_factory = sqlite3.Row + # Wait for a lock instead of erroring immediately: a figure/scan-heavy ingest can + # hold its connection across many seconds of vision calls, and a concurrent ingest + # or autoinject read would otherwise hit "database is locked". + conn.execute("PRAGMA busy_timeout = 5000") try: conn.enable_load_extension(True) sqlite_vec.load(conn) diff --git a/studio/backend/tests/test_rag_captioning.py b/studio/backend/tests/test_rag_captioning.py index 5d83a7d38d..f475c9e374 100644 --- a/studio/backend/tests/test_rag_captioning.py +++ b/studio/backend/tests/test_rag_captioning.py @@ -13,13 +13,15 @@ def _img(page): return ParsedImage(image_bytes = b"\x89PNG fake", page_number = page, xref = page) -def test_caption_images_disabled_by_default(monkeypatch): +def test_caption_images_runs_when_images_present(monkeypatch): + # Policy lives in ingestion (_run); caption_images captions given images + endpoint. monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) - assert captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) == {} + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "a chart") + out = captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) + assert out == {1: ["a chart"]} def test_caption_images_groups_by_page(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 8) monkeypatch.setattr(captioner, "_caption_one", lambda base, model, b, t: "a chart of results") out = captioner.caption_images([_img(1), _img(1), _img(3)], endpoint = ("http://x", "local")) @@ -27,7 +29,6 @@ def test_caption_images_groups_by_page(monkeypatch): def test_caption_images_respects_cap(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner.config, "CAPTION_MAX_IMAGES", 2) calls = [] monkeypatch.setattr(captioner, "_caption_one", lambda *a: (calls.append(1) or "cap")) @@ -36,11 +37,183 @@ def test_caption_images_respects_cap(monkeypatch): def test_caption_images_no_endpoint(monkeypatch): - monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) assert captioner.caption_images([_img(1)]) == {} +def test_caption_runaway_guard_applied(monkeypatch): + # A looping vision model must not flood the index; captions pass _collapse_runaway. + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "\n".join(["LOOP"] * 40)) + out = captioner.caption_images([_img(1)], endpoint = ("http://x", "local")) + assert out[1][0].splitlines().count("LOOP") == 3 # 40 -> 3 + + +def test_caption_prompt_and_token_budget(monkeypatch): + # Caption and OCR keep separate prompts + token caps over the shared _vision_complete. + captured: dict = {} + + def fake_vision_complete(base_url, model, image_bytes, *, prompt, timeout, max_tokens): + captured.update(prompt = prompt, timeout = timeout, max_tokens = max_tokens) + return "ok" + + monkeypatch.setattr(captioner, "_vision_complete", fake_vision_complete) + monkeypatch.setattr(captioner.config, "CAPTION_MAX_TOKENS", 277) + + captioner._caption_one("http://x", "local", b"img", 12.0) + prompt = captured["prompt"].lower() + # Unified prompt: transcribe every label (recall) + axis/legend coverage + describe. + assert "transcribe" in prompt + assert ("axis" in prompt or "axes" in prompt) and "legend" in prompt + assert "do not invent" in prompt + assert captured["max_tokens"] == 277 + assert captured["timeout"] == 12.0 + + captured.clear() + monkeypatch.setattr(captioner.config, "OCR_MAX_TOKENS", 999) + captioner._ocr_one("http://x", "local", b"img", 5.0) + assert captured["max_tokens"] == 999 + assert "transcribe" in captured["prompt"].lower() + + +def test_pages_with_figures_and_tiles(tmp_path): + from core.rag import parsers + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + pgs = parsers.pages_with_figures(str(pdf), max_pages = 4) + assert pgs == [1] + tiles = parsers.render_pdf_figure_tiles(str(pdf), pgs, rows = 2, cols = 2, fullpage = True) + assert len(tiles) == 5 # full page + 2x2 grid + assert all(t.image_bytes[:8] == b"\x89PNG\r\n\x1a\n" and t.page_number == 1 for t in tiles) + capped = parsers.render_pdf_figure_tiles( + str(pdf), pgs, rows = 2, cols = 2, fullpage = True, max_tiles = 3 + ) + assert len(capped) == 3 # max_tiles budget honored + + +def test_render_pdf_figure_tiles_zero_grid_no_crash(tmp_path): + # A misconfigured rows/cols=0 must clamp to 1, not raise ZeroDivisionError. + import pymupdf + + from core.rag import parsers + + pdf = tmp_path / "blank.pdf" + doc = pymupdf.open() + doc.new_page() + doc.save(str(pdf)) + doc.close() + + out = parsers.render_pdf_figure_tiles(str(pdf), [1], rows = 0, cols = 0, fullpage = True) + assert len(out) == 2 # full page + a single 1x1 tile, no crash + + +def test_pages_with_figures_excludes_given_pages(tmp_path): + # Pages OCR already transcribed (passed as exclude_pages) are skipped; every other + # figure page is still returned for tiling. + import pymupdf + + from core.rag import parsers + + def _draw_chart(page): + shape = page.new_shape() + shape.draw_rect(pymupdf.Rect(60, 140, 540, 520)) + for i in range(8): + shape.draw_line((80, 160 + i * 40), (520, 160 + i * 40)) + shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) + shape.commit() + + pdf = tmp_path / "charts.pdf" + doc = pymupdf.open() + _draw_chart(doc.new_page()) + _draw_chart(doc.new_page()) + doc.save(str(pdf)) + doc.close() + + assert parsers.pages_with_figures(str(pdf), max_pages = 4) == [1, 2] + assert parsers.pages_with_figures(str(pdf), max_pages = 4, exclude_pages = {1}) == [2] + assert parsers.pages_with_figures(str(pdf), max_pages = 4, exclude_pages = {2}) == [1] + + +def test_run_skips_figure_work_without_vision_model( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # No vision model -> the whole figure pass (detection + rasterization) is skipped. + from core.rag import parsers + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) + touched: list[str] = [] + monkeypatch.setattr( + parsers, "pages_with_figures", lambda *a, **k: touched.append("detect") or [] + ) + monkeypatch.setattr( + parsers, "render_pdf_figure_tiles", lambda *a, **k: touched.append("render") or [] + ) + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, None) # follow config (ON), but no model + assert touched == [] # neither figure detection nor tiling ran + + +def test_vision_complete_sends_auth_header(monkeypatch): + # Direct-stream serves llama-server with --api-key; vision calls must send the bearer. + import httpx + + monkeypatch.setattr( + captioner, "_vision_auth_headers", lambda: {"Authorization": "Bearer secret"} + ) + captured: dict = {} + + class _Resp: + def raise_for_status(self): + pass + + def json(self): + return {"choices": [{"message": {"content": "ok"}}]} + + def fake_post(url, *, json, timeout, headers): + captured.update(url = url, headers = headers) + return _Resp() + + monkeypatch.setattr(httpx, "post", fake_post) + out = captioner._vision_complete( + "http://x", "local", b"img", prompt = "p", timeout = 5.0, max_tokens = 8 + ) + assert out == "ok" + assert captured["headers"] == {"Authorization": "Bearer secret"} + + +def test_vision_complete_omits_header_when_unauthenticated(monkeypatch): + # No api-key configured -> no spurious Authorization header on plain llama-server. + import httpx + + monkeypatch.setattr(captioner, "_vision_auth_headers", lambda: None) + captured: dict = {} + + class _Resp: + def raise_for_status(self): + pass + + def json(self): + return {"choices": [{"message": {"content": "ok"}}]} + + def fake_post(url, *, json, timeout, headers): + captured["headers"] = headers + return _Resp() + + monkeypatch.setattr(httpx, "post", fake_post) + captioner._vision_complete("http://x", "local", b"i", prompt = "p", timeout = 5.0, max_tokens = 8) + assert captured["headers"] is None + + +def test_merge_page_captions_dedups(): + out = captioner.merge_page_captions({1: ["MatMul\nScale", "Scale\nSoftMax"]}) + text = out[1][0] + assert text.lower().count("scale") == 1 # repeated label from overlapping tiles dropped + assert "MatMul" in text and "SoftMax" in text + + def test_splice_captions_appends_to_right_page(): pages = [Page("body one", 1, 8), Page("body two", 2, 8)] out = captioner.splice_captions(pages, {2: ["a diagram of X"]}) @@ -55,29 +228,6 @@ def test_splice_captions_noop_when_empty(): assert captioner.splice_captions(pages, {}) is pages -def test_render_pdf_figures_detects_drawing(tmp_path): - import pymupdf - - from core.rag.parsers import render_pdf_figures - - pdf = tmp_path / "fig.pdf" - doc = pymupdf.open() - page = doc.new_page() - shape = page.new_shape() - shape.draw_rect(pymupdf.Rect(60, 60, 540, 460)) - for i in range(8): - shape.draw_line((80, 80 + i * 40), (520, 80 + i * 40)) - shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) - shape.commit() - doc.save(str(pdf)) - doc.close() - - figs = render_pdf_figures(str(pdf)) - assert figs, "expected at least one rendered figure region" - assert figs[0].image_bytes[:8] == b"\x89PNG\r\n\x1a\n" - assert figs[0].page_number == 1 - - def test_captioned_text_is_searchable(rag_home, stub_embeddings, monkeypatch): from core.rag import retrieval, store from storage import rag_db @@ -103,3 +253,100 @@ def test_captioned_text_is_searchable(rag_home, stub_embeddings, monkeypatch): finally: conn.close() assert hits, "spliced caption text should be retrievable via lexical search" + + +# ── per-upload caption override (parallels test_rag_ocr_fallback.py) ── + + +def _figure_pdf(path): + """A born-digital PDF: a page with real text (so it is not treated as scanned) + plus a vector drawing region that figure detection picks up as a figure.""" + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox( + pymupdf.Rect(40, 40, 550, 120), + "Quarterly revenue report. The chart below shows the trend.", + fontsize = 11, + ) + shape = page.new_shape() + shape.draw_rect(pymupdf.Rect(60, 140, 540, 520)) + for i in range(8): + shape.draw_line((80, 160 + i * 40), (520, 160 + i * 40)) + shape.finish(color = (0, 0, 0), fill = (0.8, 0.8, 0.9)) + shape.commit() + doc.save(str(path)) + doc.close() + + +def _ingest_with_caption(rag_conn, thread_id, path, caption): + from core.rag import ingestion, store + + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "fig.pdf", + sha256 = str(path) + str(caption), + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + # _run(job_id, document_id, scope, stored_path, model_name, ocr, caption) + ingestion._run(job_id, document_id, scope, str(path), None, None, caption) + return store.get_document(rag_conn, document_id) + + +def test_caption_override_true_runs_when_config_off( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default OFF, but the per-upload toggle (caption=True) forces captioning. + from core.rag import tool + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "bar chart of revenue wombat-7") + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, True) + + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "wombat-7" in text # the spliced figure caption reached the index + + +def test_caption_override_false_skips_when_config_on( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default ON, but the per-upload toggle (caption=False) skips captioning. + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + called = [] + monkeypatch.setattr(captioner, "_caption_one", lambda *a: called.append(1) or "should not run") + + pdf = tmp_path / "fig.pdf" + _figure_pdf(pdf) + _ingest_with_caption(rag_conn, "t1", pdf, False) + + assert called == [] # no vision caption calls despite config ON + + +def test_caption_none_follows_config(rag_conn, stub_embeddings, monkeypatch, tmp_path): + # Omitted override (None) falls back to config.CAPTION_IMAGES. + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + seen = [] + monkeypatch.setattr(captioner, "_caption_one", lambda *a: seen.append(1) or "chart caption") + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", False) + pdf_off = tmp_path / "off.pdf" + _figure_pdf(pdf_off) + _ingest_with_caption(rag_conn, "t1", pdf_off, None) + assert seen == [] # config OFF + no override -> no captioning + + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + pdf_on = tmp_path / "on.pdf" + _figure_pdf(pdf_on) + _ingest_with_caption(rag_conn, "t2", pdf_on, None) + assert seen # config ON + no override -> captioning runs diff --git a/studio/backend/tests/test_rag_ingestion.py b/studio/backend/tests/test_rag_ingestion.py index f0b71bc23b..7e9e803687 100644 --- a/studio/backend/tests/test_rag_ingestion.py +++ b/studio/backend/tests/test_rag_ingestion.py @@ -83,6 +83,34 @@ def test_ingestion_dedupe_by_hash(rag_home, stub_embeddings, tmp_path): conn.close() +def test_ingestion_reingests_when_existing_has_zero_chunks(rag_home, stub_embeddings, tmp_path): + # A prior ingest of identical bytes that yielded no chunks (e.g. a scanned PDF + # before a vision model loaded) must re-ingest, not dedupe to the empty record. + path = _write(tmp_path, "doc.txt", "alpha bravo charlie " * 50) + sha = ingestion._sha256_file(path) + scope = store.kb_scope("K1") + conn = rag_db.get_connection() + try: + empty_id = store.create_document(conn, scope = scope, filename = "old.txt", sha256 = sha) + store.set_document_status(conn, empty_id, "completed", num_chunks = 0) + finally: + conn.close() + + doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + events = _drain(job_id) + _wait_completed(job_id) + + assert not any(e.get("deduped") for e in events) # not a dedupe -> real ingest + assert doc_id != empty_id + conn = rag_db.get_connection() + try: + docs = store.list_documents(conn, scope) + assert len(docs) == 1 # the empty record was removed, replaced by the new one + assert docs[0]["num_chunks"] > 0 + finally: + conn.close() + + def test_ingestion_dedupe_removes_duplicate_upload(rag_home, stub_embeddings): from utils.paths import ensure_dir, rag_uploads_root @@ -210,6 +238,41 @@ def test_delete_document_route_removes_stored_upload(rag_home): conn.close() +def test_get_job_status_includes_num_chunks(rag_home, stub_embeddings, tmp_path): + # The poll/reconcile path reads num_chunks from get_job_status (the SSE complete + # frame carries it, but a client that falls back to polling needs it here too). + path = _write(tmp_path, "doc.txt", "alpha bravo charlie " * 50) + scope = store.kb_scope("K1") + _doc_id, job_id = ingestion.start_ingestion(scope, "K1", None, "doc.txt", path) + _drain(job_id) + _wait_completed(job_id) + status = ingestion.get_job_status(job_id) + assert status["status"] == "completed" + assert status["num_chunks"] and status["num_chunks"] > 0 + + +def test_save_upload_rejects_oversize_file(rag_home, monkeypatch): + # A file over the cap is rejected (413) and its partial bytes are cleaned up. + import io + + from fastapi import HTTPException + + from core.rag import config + from routes import rag as rag_routes + from utils.paths import rag_uploads_root + + monkeypatch.setattr(config, "MAX_UPLOAD_BYTES", 1024) + + class _Up: + filename = "big.txt" + file = io.BytesIO(b"x" * 4096) + + with pytest.raises(HTTPException) as ei: + rag_routes._save_upload(_Up()) + assert ei.value.status_code == 413 + assert list(rag_uploads_root().glob("*.txt")) == [] # partial upload removed + + def test_ingestion_delete_removes_all_rows(rag_home, stub_embeddings, tmp_path): path = _write(tmp_path, "doc.txt", "alpha bravo charlie delta") scope = store.kb_scope("K1") diff --git a/studio/backend/tests/test_rag_ocr_fallback.py b/studio/backend/tests/test_rag_ocr_fallback.py new file mode 100644 index 0000000000..c7be1fe60b --- /dev/null +++ b/studio/backend/tests/test_rag_ocr_fallback.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Scanned-PDF OCR fallback: a PDF page with no text layer is rendered and transcribed +by the vision model during ingestion, so image-only PDFs become searchable. The vision +call is stubbed, so no model is needed.""" + +import pymupdf + +from core.rag import captioner, ingestion, parsers, store, tool + + +def _image_only_pdf(path, *, pages = 1): + """A PDF whose pages carry only a raster image, so get_text returns ''.""" + doc = pymupdf.open() + pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 120, 120)) + pix.clear_with(220) + for _ in range(pages): + page = doc.new_page() + page.insert_image(page.rect, pixmap = pix) + doc.save(str(path)) + doc.close() + + +def _text_pdf(path, body): + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox(pymupdf.Rect(40, 40, 550, 800), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def _ingest(rag_conn, thread_id, filename, path): + """Drive the real ingestion worker synchronously and return the document row.""" + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = filename, + sha256 = filename, + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(path), None) + return store.get_document(rag_conn, document_id) + + +# ── parsers.render_pdf_pages ───────────────────────────────────────── + + +def test_render_pdf_pages_returns_png_per_page(tmp_path): + pdf = tmp_path / "two.pdf" + _image_only_pdf(pdf, pages = 2) + out = parsers.render_pdf_pages(str(pdf), [1, 2], dpi = 72) + assert set(out) == {1, 2} + assert all(b.startswith(b"\x89PNG") for b in out.values()) + + +def test_render_pdf_pages_excludes_unwanted(tmp_path): + pdf = tmp_path / "three.pdf" + _image_only_pdf(pdf, pages = 3) + out = parsers.render_pdf_pages(str(pdf), [2], dpi = 72) + assert set(out) == {2} + + +def test_render_pdf_pages_empty_request(tmp_path): + pdf = tmp_path / "one.pdf" + _image_only_pdf(pdf, pages = 1) + assert parsers.render_pdf_pages(str(pdf), [], dpi = 72) == {} + + +# ── captioner.ocr_pages gating ─────────────────────────────────────── + + +def test_ocr_pages_no_endpoint(monkeypatch): + monkeypatch.setattr(captioner, "vision_endpoint", lambda: None) + assert captioner.ocr_pages({1: b"x"}) == {} + + +def test_collapse_runaway_caps_repeated_lines(): + # A looping model repeats a line hundreds of times; the guard caps it, keeps repeats. + text = "\n".join(["TITLE"] * 200 + ["body"] + ["Add & Norm"] * 3) + out = captioner._collapse_runaway(text) + lines = out.splitlines() + assert lines.count("TITLE") == 3 # 200 -> 3 + assert lines.count("Add & Norm") == 3 # legitimate triple survives + assert "body" in lines + + +def test_collapse_runaway_caps_interleaved_repeats(): + # Models also loop non-consecutively; the global per-line cap bounds those too. + text = "\n".join(["Llion Vaswani Google", "Niki Parmar Google"] * 40) + out = captioner._collapse_runaway(text) + lines = [ln for ln in out.splitlines() if ln.strip()] + assert lines.count("Llion Vaswani Google") <= 8 + assert lines.count("Niki Parmar Google") <= 8 + + +def test_collapse_runaway_noop_on_normal_text(): + text = "Heading\n\nFirst paragraph.\nSecond paragraph.\n\nFooter" + assert captioner._collapse_runaway(text) == text + + +def test_ocr_pages_applies_runaway_guard(monkeypatch): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "\n".join(["X"] * 50)) + out = captioner.ocr_pages({1: b"img"}, endpoint = ("http://x", "local")) + assert out[1].splitlines().count("X") == 3 # guard applied to stored text + + +def test_ocr_pages_transcribes_and_caps(monkeypatch): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MAX_PAGES", 1) + calls = [] + monkeypatch.setattr( + captioner, + "_ocr_one", + lambda base, model, b, t: (calls.append(1) or "transcribed text"), + ) + out = captioner.ocr_pages({1: b"a", 2: b"b"}, endpoint = ("http://x", "local")) + assert out == {1: "transcribed text"} # page 2 dropped by the cap + assert len(calls) == 1 + + +def test_ocr_scanned_pages_merges_short_text_layer(rag_conn, monkeypatch): + # Near-empty pages can still have meaningful extractable text; OCR augments it + # rather than replacing it with a fallible vision transcription. + scope = store.thread_scope("t1") + document_id = store.create_document(rag_conn, scope = scope, filename = "scan.pdf", sha256 = "h") + job_id = ingestion._new_job(rag_conn, document_id, scope) + pages = [parsers.Page("ID-42", 1, 5)] + + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MIN_CHARS", 16) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(parsers, "render_pdf_pages", lambda *a, **k: {1: b"png"}) + monkeypatch.setattr(captioner, "ocr_pages", lambda page_pngs: {1: "OCR body text"}) + + out, ocred = ingestion._ocr_scanned_pages(pages, "scan.pdf", rag_conn, job_id) + assert ocred == {1} + assert out[0].text == "ID-42\n\nOCR body text" + + +# ── end-to-end ingestion ───────────────────────────────────────────── + + +def test_scanned_pdf_is_ocred_into_chunks(rag_conn, stub_embeddings, monkeypatch, tmp_path): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr( + captioner, "_ocr_one", lambda base, model, b, t: "Invoice total is zebra-42 due Friday" + ) + + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest(rag_conn, "t1", "scan.pdf", pdf) + + assert doc["status"] == "completed" + assert doc["num_chunks"] >= 1 + # The OCR'd text is now indexed and reaches whole-document injection. + text, _sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "zebra-42" in text + + +def test_scanned_page_past_ocr_cap_is_still_captioned( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # OCR is capped to one page, so page 2 is scanned but never transcribed. Figure + # captioning must still cover it (we exclude only the pages OCR actually handled), + # so a chart on an un-OCR'd scanned page is not silently dropped. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner.config, "OCR_MAX_PAGES", 1) + monkeypatch.setattr(captioner.config, "CAPTION_IMAGES", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "scanned page alpha") + monkeypatch.setattr(captioner, "_caption_one", lambda *a: "figure caption bravo") + + pdf = tmp_path / "scan2.pdf" + _image_only_pdf(pdf, pages = 2) + doc = _ingest(rag_conn, "t1", "scan2.pdf", pdf) + + assert doc["status"] == "completed" + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "scanned page alpha" in text # page 1 OCR'd, within the cap + assert "figure caption bravo" in text # page 2 past the cap -> captioned, not dropped + + +def test_born_digital_pdf_skips_ocr(rag_conn, stub_embeddings, monkeypatch, tmp_path): + called = [] + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: called.append(1) or "should not run") + + pdf = tmp_path / "digital.pdf" + _text_pdf(pdf, "Real born digital body text. " * 30 + "marker-quokka") + doc = _ingest(rag_conn, "t1", "digital.pdf", pdf) + + assert doc["status"] == "completed" + assert called == [] # page had real text -> never considered scanned + text, _sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "marker-quokka" in text + + +def _ingest_with_ocr(rag_conn, thread_id, path, ocr): + scope = store.thread_scope(thread_id) + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "scan.pdf", + sha256 = str(path) + str(ocr), + thread_id = thread_id, + status = "pending", + stored_path = str(path), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(path), None, ocr = ocr) + return store.get_document(rag_conn, document_id) + + +def test_ocr_override_false_skips_ocr_when_config_on( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default ON, but the per-upload toggle (ocr=False) skips OCR. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", True) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "should not run") + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest_with_ocr(rag_conn, "t1", pdf, ocr = False) + assert doc["num_chunks"] == 0 # scanned page left empty + + +def test_ocr_override_true_runs_ocr_when_config_off( + rag_conn, stub_embeddings, monkeypatch, tmp_path +): + # Config default OFF, but the per-upload toggle (ocr=True) forces OCR on. + monkeypatch.setattr(captioner.config, "OCR_SCANNED", False) + monkeypatch.setattr(captioner, "vision_endpoint", lambda: ("http://x", "local")) + monkeypatch.setattr(captioner, "_ocr_one", lambda *a: "forced ocr text quokka") + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest_with_ocr(rag_conn, "t1", pdf, ocr = True) + assert doc["num_chunks"] >= 1 + text, _ = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "quokka" in text + + +def test_ocr_disabled_leaves_scanned_pdf_empty(rag_conn, stub_embeddings, monkeypatch, tmp_path): + monkeypatch.setattr(captioner.config, "OCR_SCANNED", False) + + pdf = tmp_path / "scan.pdf" + _image_only_pdf(pdf, pages = 1) + doc = _ingest(rag_conn, "t1", "scan.pdf", pdf) + + # With OCR off, a text-less scanned page yields no chunks (prior behavior). + assert doc["status"] == "completed" + assert doc["num_chunks"] == 0 + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py new file mode 100644 index 0000000000..4c46f49495 --- /dev/null +++ b/studio/backend/tests/test_rag_parsing.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 + +"""PDF text extraction: layout-aware Markdown (pymupdf4llm) with plain-text fallback.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("pymupdf") + + +def _table_pdf(path): + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + page.insert_textbox(pymupdf.Rect(40, 40, 550, 70), "Quarterly Results", fontsize = 16) + rows = [("Quarter", "Revenue", "Growth"), ("Q1", "$1.2M", "12%"), ("Q2", "$1.5M", "25%")] + y = 90 + for r in rows: + page.insert_textbox(pymupdf.Rect(40, y, 250, y + 20), r[0], fontsize = 11) + page.insert_textbox(pymupdf.Rect(250, y, 400, y + 20), r[1], fontsize = 11) + page.insert_textbox(pymupdf.Rect(400, y, 540, y + 20), r[2], fontsize = 11) + y += 24 + doc.save(str(path)) + doc.close() + + +def test_pdf_extracts_markdown_table(tmp_path, monkeypatch): + # With Markdown on, the layout is emitted as Markdown markup (heading, and a pipe table + # where the extractor detects one) that flat get_text never produces. + pytest.importorskip("pymupdf4llm") + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Q2" in text and "$1.5M" in text # cell values preserved + assert "#" in text or "|" in text # Markdown markup (heading or table pipes) + + +def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch): + # The toggle (RAG_PDF_MARKDOWN=0) falls back to flat PyMuPDF text: content is still + # there, but with no Markdown markup. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", False) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Q2" in text and "$1.5M" in text + assert "#" not in text and "|" not in text # plain text path emits no Markdown markup + + +def test_pdf_markdown_passes_only_supported_legacy_kwargs(monkeypatch): + # The pinned PyMuPDF4LLM legacy path ignores unknown kwargs; do not pass the + # newer layout-only OCR knobs or Markdown extraction silently loses policy control. + from core.rag import parsers + + captured = {} + + class _FakePymupdf4llm: + @staticmethod + def to_markdown(doc, **kwargs): + captured.update(kwargs) + return [{"text": "plain markdown"}] + + class _Doc: + page_count = 1 + + monkeypatch.setitem(__import__("sys").modules, "pymupdf4llm", _FakePymupdf4llm) + assert parsers._pdf_markdown(_Doc()) == ["plain markdown"] + assert captured == {"page_chunks": True, "show_progress": False} + + +def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch): + # If pymupdf4llm extraction returns None (missing/failed), parsing still yields the + # plain-text pages rather than raising. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: None) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + pages = parsers.parse(str(pdf)) + assert pages and "Quarter" in pages[0].text diff --git a/studio/backend/tests/test_rag_preview.py b/studio/backend/tests/test_rag_preview.py index e7f2a39792..0ff27897bd 100644 --- a/studio/backend/tests/test_rag_preview.py +++ b/studio/backend/tests/test_rag_preview.py @@ -165,6 +165,25 @@ def test_locator_handles_midword_anchor_and_locates_line(): assert r["width"] > 0 and r["height"] > 0 +def test_locator_anchors_through_markdown_table_pipes(): + # Markdown table cells are pipe-joined with no spaces; the locator splits on pipes + # so a table-row chunk still anchors to the raw PDF word stream. + import pymupdf + + from core.rag.locators import LocatorMatch, _regions_for_match + + doc = pymupdf.open() + page = doc.new_page() + page.insert_text((72, 200), "Quarter Revenue Growth Q1 sales strong here", fontsize = 12) + # What the Markdown parser stores for the row (cells joined by pipes, no spaces). + page_text = "|Quarter|Revenue|Growth|Q1|sales|strong|here|" + match = LocatorMatch(page_index = 0, page_number = 1, start = 0, end = len(page_text)) + rects = _regions_for_match(doc, page_text, match) + doc.close() + + assert rects, "a Markdown table row should still anchor to the page words" + + def test_sign_verify_roundtrip(rag_home): from routes import rag as rag_routes diff --git a/studio/backend/tests/test_rag_whole_document.py b/studio/backend/tests/test_rag_whole_document.py new file mode 100644 index 0000000000..545d731fd2 --- /dev/null +++ b/studio/backend/tests/test_rag_whole_document.py @@ -0,0 +1,520 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Whole-document context mode: a thread-attached file small enough to fit is +injected in full (every chunk, in order) instead of top-K retrieval. Covers the +new store query, the tool-level renderer, and the auto-inject wiring + fallback. +No embedder is needed - the whole-doc path does no query embedding.""" + +import json + +from core.rag import store, tool +from core.rag.chunking import Chunk +from core.inference import tools as inf_tools + +# A vector per chunk just to satisfy add_chunks (the whole-doc path never reads +# vectors); dimension is arbitrary but must be consistent within a connection. +_VEC = [0.1, 0.2, 0.3, 0.4] + + +def _chunk( + text, + index = 0, + page = None, + tokens = None, +): + return Chunk( + text = text, + token_count = tokens if tokens is not None else len(text.split()), + page_number = page, + source_page_index = 0, + chunk_index = index, + page_char_start = 0, + page_char_end = len(text), + ) + + +def _add_doc( + conn, + scope, + doc_id, + filename, + sha, + texts, + *, + status = "completed", + tokens = None, + pages = None, +): + chunks = [ + _chunk( + t, + i, + page = (pages[i] if pages else None), + tokens = (tokens[i] if tokens else None), + ) + for i, t in enumerate(texts) + ] + vectors = [list(_VEC) for _ in texts] + store.create_document(conn, scope = scope, filename = filename, sha256 = sha, document_id = doc_id) + store.add_chunks(conn, scope, doc_id, chunks, vectors) + store.set_document_status(conn, doc_id, status, num_chunks = len(texts)) + + +def _injected_text(result) -> str: + """The text spliced into the conversation as the synthetic tool result.""" + tool_msg = next(m for m in result["messages"] if m.get("role") == "tool") + return tool_msg["content"] + + +# ── store.all_chunks_for_scope ─────────────────────────────────────── + + +def test_all_chunks_for_scope_orders_by_document_then_index(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "first.pdf", "h1", ["a", "b", "c"]) + _add_doc(rag_conn, scope, "d2", "second.pdf", "h2", ["x", "y"]) + rows = store.all_chunks_for_scope(rag_conn, scope) + assert [r["id"] for r in rows] == ["d1:0", "d1:1", "d1:2", "d2:0", "d2:1"] + assert rows[0]["filename"] == "first.pdf" + assert rows[-1]["filename"] == "second.pdf" + assert rows[0]["text"] == "a" + + +def test_all_chunks_for_scope_excludes_non_completed(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "done", "done.pdf", "h1", ["ready"]) + _add_doc(rag_conn, scope, "pend", "pend.pdf", "h2", ["indexing"], status = "pending") + rows = store.all_chunks_for_scope(rag_conn, scope) + assert [r["id"] for r in rows] == ["done:0"] + + +def test_all_chunks_for_scope_empty_scope(rag_conn): + assert store.all_chunks_for_scope(rag_conn, store.thread_scope("nope")) == [] + + +def test_all_chunks_for_scope_isolates_scopes(rag_conn): + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "f", "h1", ["mine"]) + _add_doc(rag_conn, store.thread_scope("t2"), "d2", "f", "h2", ["theirs"]) + rows = store.all_chunks_for_scope(rag_conn, store.thread_scope("t1")) + assert [r["text"] for r in rows] == ["mine"] + + +# ── store.scope_token_estimate (cheap whole-doc budget pre-check) ───── + + +def test_scope_token_estimate_sums_without_hydrating(rag_conn): + # Stored counts sum directly; zero/missing falls back to length/4; non-completed out. + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "a.pdf", "h1", ["alpha", "bravo"], tokens = [10, 20]) + # token_count 0 -> length/4 fallback: a 40-char chunk estimates to 10 tokens. + _add_doc(rag_conn, scope, "d2", "b.pdf", "h2", ["x" * 40], tokens = [0]) + _add_doc(rag_conn, scope, "d3", "c.pdf", "h3", ["pending"], status = "pending", tokens = [99]) + assert store.scope_token_estimate(rag_conn, scope) == 10 + 20 + 10 + assert store.scope_token_estimate(rag_conn, store.thread_scope("none")) == 0 + + +def test_scope_token_estimate_matches_row_sum(rag_conn): + # Must agree with the exact per-row sum it short-circuits (one stored count, one + # length/4 fallback), so the pre-check never disagrees with the full path. + from core.rag.tool import _row_token_count + + scope = store.thread_scope("t1") + _add_doc( + rag_conn, scope, "d1", "a.pdf", "h1", ["a long-ish chunk body here", "tail"], tokens = [0, 5] + ) + rows = store.all_chunks_for_scope(rag_conn, scope) + assert store.scope_token_estimate(rag_conn, scope) == sum(_row_token_count(r) for r in rows) + + +# ── tool.whole_document_context ────────────────────────────────────── + + +def test_whole_document_context_returns_full_text_and_sources(rag_conn): + scope = store.thread_scope("t1") + _add_doc( + rag_conn, + scope, + "d1", + "report.pdf", + "h1", + ["chapter one body", "chapter two body"], + pages = [1, 2], + ) + result = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert result is not None + text, sources = result + # Every chunk is present, in order, as blocks. + assert "chapter one body" in text + assert "chapter two body" in text + assert ' None (whole-doc is thread-attachment only). + assert tool.whole_document_context(max_tokens = 6000) is None + + +def test_whole_document_context_null_token_count_enforces_budget(rag_conn): + # A missing token_count must not bypass the budget; fall back to a length estimate. + big = "word " * 20_000 # ~20k tokens by length estimate + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "big.pdf", "h1", [big], tokens = [None]) + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) is None + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 1_000_000) is not None + + +def test_whole_document_context_spans_multiple_docs(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "a.pdf", "h1", ["alpha text"]) + _add_doc(rag_conn, scope, "d2", "b.pdf", "h2", ["bravo text"]) + text, sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "alpha text" in text and "bravo text" in text + assert {s["filename"] for s in sources} == {"a.pdf", "b.pdf"} + + +# ── build_rag_autoinject wiring ────────────────────────────────────── + + +def _convo(text = "summarize the whole document"): + return [{"role": "user", "content": text}] + + +def test_build_rag_autoinject_uses_whole_doc(rag_conn): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "doc.pdf", "h1", ["whole alpha part", "whole bravo part"]) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + injected = _injected_text(result) + # Both chunks present -> the model receives the entire file, not top-K. + assert "whole alpha part" in injected + assert "whole bravo part" in injected + # Tool-message content is chunk text only; the citation JSON tail is internal. + assert inf_tools.RAG_SOURCES_SENTINEL not in injected + + +def test_build_rag_autoinject_whole_doc_runs_when_autoinject_false(rag_conn, monkeypatch): + # Large-model Auto sets autoinject=False, but whole-doc is a separate thread-doc + # context mode and should still inject a fitting attachment. + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["entire file body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "autoinject": False}) + assert result is not None + assert "entire file body" in _injected_text(result) + + +def test_build_rag_autoinject_explicit_off_disables_whole_doc(rag_conn, monkeypatch): + # The UI Off switch sends both autoinject=False and whole_doc=False. + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["small body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + assert ( + inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "autoinject": False, "whole_doc": False} + ) + is None + ) + + +def test_build_rag_autoinject_falls_back_over_budget(rag_conn, monkeypatch): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "big.pdf", "h1", ["overflow"], tokens = [50_000]) + + sentinel = ("TOPK_FALLBACK_TEXT", [{"citationId": 1, "filename": "big.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + assert _injected_text(result) == "TOPK_FALLBACK_TEXT" + + +def test_build_rag_autoinject_context_budget_falls_back(rag_conn, monkeypatch): + # Runtime context can be smaller than RAG_WHOLE_DOC_MAX_TOKENS; cap whole-doc to + # the active context and fall back to retrieval when it would overflow. + _add_doc( + rag_conn, store.thread_scope("t1"), "d1", "small.pdf", "h1", ["fits global"], tokens = [900] + ) + sentinel = ("TOPK_CONTEXT_FALLBACK", [{"citationId": 1, "filename": "small.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + result = inf_tools.build_rag_autoinject( + _convo(), {"thread_id": "t1", "context_length": 1200, "whole_doc": True} + ) + assert result is not None + assert _injected_text(result) == "TOPK_CONTEXT_FALLBACK" + + +def test_whole_doc_budget_reserves_image_parts(monkeypatch): + from core.rag import config + + monkeypatch.setattr(config, "WHOLE_DOC_MAX_TOKENS", 10_000) + scope = {"context_length": 7000, "response_headroom": 1000} + text_only = [{"role": "user", "content": [{"type": "text", "text": "summarize"}]}] + with_image = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}, + ], + } + ] + + assert ( + inf_tools._whole_doc_budget(scope, text_only) + - inf_tools._whole_doc_budget(scope, with_image) + == inf_tools._IMAGE_PART_TOKEN_ESTIMATE + ) + + +def test_build_rag_autoinject_server_kill_switch_blocks_whole_doc(rag_conn, monkeypatch): + # RAG_THREAD_WHOLE_DOC=0 stays authoritative; browser requests should not + # turn it back on by default. + from core.rag import config + + monkeypatch.setattr(config, "THREAD_WHOLE_DOC", False) + _add_doc(rag_conn, store.thread_scope("t1"), "d1", "doc.pdf", "h1", ["small body"]) + monkeypatch.setattr( + tool, + "search_for_autoinject", + lambda **kw: (_ for _ in ()).throw(AssertionError("retrieval should not run")), + ) + assert ( + inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "autoinject": False}) is None + ) + + +def test_whole_document_context_budgets_rendered_wrappers(rag_conn): + # Many tiny chunks add wrapper overhead beyond raw chunk token counts; budget + # the rendered prompt, not just stored text. + texts = ["x" for _ in range(120)] + _add_doc( + rag_conn, + store.thread_scope("t1"), + "d1", + "many-pages.pdf", + "h1", + texts, + tokens = [1 for _ in texts], + ) + assert tool.whole_document_context(scope_thread_id = "t1", max_tokens = 500) is None + + +def test_build_rag_autoinject_whole_doc_disabled_via_override(rag_conn, monkeypatch): + scope = store.thread_scope("t1") + _add_doc(rag_conn, scope, "d1", "doc.pdf", "h1", ["small body"]) + + sentinel = ("TOPK_TEXT", [{"citationId": 1, "filename": "doc.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + # whole_doc=False forces retrieval even though the doc fits. + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "whole_doc": False}) + assert result is not None + assert _injected_text(result) == "TOPK_TEXT" + + +def test_build_rag_autoinject_kb_scope_never_whole_doc(rag_conn, monkeypatch): + # A KB-only scope (no thread) goes through retrieval, never whole-doc. + kb_scope = store.kb_scope("K1") + _add_doc(rag_conn, kb_scope, "d1", "kb.pdf", "h1", ["kb body one", "kb body two"]) + + sentinel = ("KB_RETRIEVAL_TEXT", [{"citationId": 1, "filename": "kb.pdf", "text": "x"}]) + monkeypatch.setattr(tool, "search_for_autoinject", lambda **kw: sentinel) + + result = inf_tools.build_rag_autoinject(_convo(), {"kb_id": "K1"}) + assert result is not None + assert _injected_text(result) == "KB_RETRIEVAL_TEXT" + + +def test_whole_document_context_thread_scope_only(rag_conn): + # A project corpus chunk is never whole-doc injected, even with a thread attachment. + _add_doc(rag_conn, store.thread_scope("t1"), "td", "thread.txt", "h1", ["thread attachment"]) + _add_doc(rag_conn, store.project_scope("p1"), "pd", "project.txt", "h2", ["project corpus"]) + text, sources = tool.whole_document_context(scope_thread_id = "t1", max_tokens = 6000) + assert "thread attachment" in text + assert "project corpus" not in text + assert {s["filename"] for s in sources} == {"thread.txt"} + + +def test_build_rag_autoinject_appends_project_retrieval(rag_conn, monkeypatch): + # Project chat: thread attachment whole-doc'd AND project sources retrieved, merged. + _add_doc( + rag_conn, + store.thread_scope("t1"), + "td", + "thread.txt", + "h1", + ["thread chunk one", "thread chunk two"], + ) + proj = ( + "PROJ", + [ + { + "citationId": 1, + "chunkId": "pj:0", + "documentId": "pj", + "filename": "project.txt", + "page": None, + "text": "project passage zeta", + "score": 0.91, + } + ], + ) + captured = {} + + def fake_search(**kw): + captured.update(kw) + return proj + + monkeypatch.setattr(tool, "search_for_autoinject", fake_search) + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1", "project_id": "p1"}) + injected = _injected_text(result) + # Whole thread attachment AND the project passage are both injected. + assert "thread chunk one" in injected + assert "thread chunk two" in injected + assert "project passage zeta" in injected + # The companion retrieval was scoped to the project only (not thread or KB). + assert captured.get("scope_project_id") == "p1" + assert captured.get("scope_thread_id") is None + assert captured.get("scope_kb_id") is None + # Citation ids are sequential across the merged set: thread 1,2 then project 3. + assert ' whole-doc injection ──────── + + +def test_real_ingestion_feeds_whole_document(rag_conn, stub_embeddings, tmp_path): + """Drive the real ingestion worker on a multi-paragraph file, then confirm whole-doc + injection splices the entire document, not just retrieved chunks.""" + from core.rag import ingestion + + scope = store.thread_scope("t1") + body = ( + "# Quarterly Report\n\n" + + ("Revenue rose across every region this period. " * 40) + + "\n\nThe unique closing marker is xyzzy-sentinel for the final page. " * 40 + ) + src = tmp_path / "report.md" + src.write_text(body, encoding = "utf-8") + + document_id = store.create_document( + rag_conn, + scope = scope, + filename = "report.md", + sha256 = "sha-e2e", + thread_id = "t1", + status = "pending", + stored_path = str(src), + ) + job_id = ingestion._new_job(rag_conn, document_id, scope) + ingestion._run(job_id, document_id, scope, str(src), None) + + doc = store.get_document(rag_conn, document_id) + assert doc["status"] == "completed" + assert doc["num_chunks"] >= 2 # the doc chunked into multiple pieces + + result = inf_tools.build_rag_autoinject(_convo(), {"thread_id": "t1"}) + assert result is not None + injected = _injected_text(result) + # Opening and ending both present -> the whole file reached the model. + assert "Revenue rose" in injected + assert "xyzzy-sentinel" in injected + # Every stored chunk is represented as a numbered block. + assert injected.count(" void; setRagAutoInject: (value: RagAutoInject) => void; setRagAutoInjectMinScore: (score: number) => void; + setRagOcrScanned: (enabled: boolean) => void; + setRagCaptionFigures: (enabled: boolean) => void; setToolStatus: (status: string | null) => void; setGeneratingStatus: (status: string | null) => void; setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void; @@ -1077,6 +1091,8 @@ export const useChatRuntimeStore = create((set, get) => ({ DEFAULT_RAG_AUTOINJECT_MIN_SCORE, { min: 0, max: 1 }, ), + ragOcrScanned: loadBool(CHAT_RAG_OCR_KEY, DEFAULT_RAG_OCR), + ragCaptionFigures: loadBool(CHAT_RAG_CAPTION_KEY, DEFAULT_RAG_CAPTION), toolStatus: null, generatingStatus: null, activeDiffusionCanvas: null, @@ -1498,6 +1514,16 @@ export const useChatRuntimeStore = create((set, get) => ({ ); return { ragAutoInjectMinScore }; }), + setRagOcrScanned: (ragOcrScanned) => + set(() => { + saveBool(CHAT_RAG_OCR_KEY, ragOcrScanned); + return { ragOcrScanned }; + }), + setRagCaptionFigures: (ragCaptionFigures) => + set(() => { + saveBool(CHAT_RAG_CAPTION_KEY, ragCaptionFigures); + return { ragCaptionFigures }; + }), setToolStatus: (toolStatus) => set({ toolStatus }), setActiveDiffusionCanvas: (activeDiffusionCanvas) => set({ activeDiffusionCanvas }), diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index b8391e31e7..412f35d1a0 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -350,6 +350,9 @@ export interface OpenAIChatCompletionsRequest { mode: "hybrid" | "lexical" | "dense"; autoinject?: boolean; autoinject_min_score?: number; + + whole_doc?: boolean; + context_length?: number; }; auto_heal_tool_calls?: boolean; max_tool_calls_per_message?: number; diff --git a/studio/frontend/src/features/rag/api/rag-api.ts b/studio/frontend/src/features/rag/api/rag-api.ts index b1215d2d46..20800230b5 100644 --- a/studio/frontend/src/features/rag/api/rag-api.ts +++ b/studio/frontend/src/features/rag/api/rag-api.ts @@ -39,9 +39,17 @@ async function ragRequest( return json as T; } -async function ragUpload(path: string, file: File): Promise { +async function ragUpload( + path: string, + file: File, + ocr?: boolean, + caption?: boolean, +): Promise { const form = new FormData(); form.append("file", file); + // Per-upload overrides for the vision passes; omitted -> backend config default. + if (ocr !== undefined) form.append("ocr", String(ocr)); + if (caption !== undefined) form.append("caption", String(caption)); // No Content-Type: let the browser set the multipart boundary. const response = await authFetch(`${RAG_BASE}${path}`, { method: "POST", @@ -103,10 +111,14 @@ export async function listKnowledgeBaseDocuments( export function uploadKnowledgeBaseDocument( kbId: string, file: File, + ocr?: boolean, + caption?: boolean, ): Promise { return ragUpload( `/knowledge-bases/${encodeURIComponent(kbId)}/documents`, file, + ocr, + caption, ); } @@ -122,8 +134,15 @@ export async function listThreadDocuments( export function uploadThreadDocument( threadId: string, file: File, + ocr?: boolean, + caption?: boolean, ): Promise { - return ragUpload(`/threads/${encodeURIComponent(threadId)}/documents`, file); + return ragUpload( + `/threads/${encodeURIComponent(threadId)}/documents`, + file, + ocr, + caption, + ); } export async function listProjectDocuments( @@ -138,8 +157,15 @@ export async function listProjectDocuments( export function uploadProjectDocument( projectId: string, file: File, + ocr?: boolean, + caption?: boolean, ): Promise { - return ragUpload(`/projects/${encodeURIComponent(projectId)}/documents`, file); + return ragUpload( + `/projects/${encodeURIComponent(projectId)}/documents`, + file, + ocr, + caption, + ); } // Cached "does this project have indexed sources?" probe so the chat adapter can diff --git a/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx b/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx index 103ed5c8e8..be791b8e6d 100644 --- a/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx +++ b/studio/frontend/src/features/rag/components/retrieval-settings-section.tsx @@ -9,6 +9,7 @@ import { SelectValue, } from "@/components/ui/select"; import { Slider } from "@/components/ui/slider"; +import { Switch } from "@/components/ui/switch"; import { ToggleGroup, ToggleGroupItem } from "@/components/ui/toggle-group"; import { Tooltip, @@ -109,6 +110,12 @@ export function RetrievalSettingsSection() { const setRagAutoInjectMinScore = useChatRuntimeStore( (s) => s.setRagAutoInjectMinScore, ); + const ragOcrScanned = useChatRuntimeStore((s) => s.ragOcrScanned); + const setRagOcrScanned = useChatRuntimeStore((s) => s.setRagOcrScanned); + const ragCaptionFigures = useChatRuntimeStore((s) => s.ragCaptionFigures); + const setRagCaptionFigures = useChatRuntimeStore( + (s) => s.setRagCaptionFigures, + ); return (
@@ -202,6 +209,51 @@ export function RetrievalSettingsSection() { format={(v) => v.toFixed(2)} />
+ +
+
+ + OCR scanned pages + + Read text off scanned or image-only PDF pages with the loaded + model's vision, at upload time, so picture-only documents become + searchable. Needs a vision model; pages with a text layer are + unaffected. + + + + Transcribe image-only PDF pages when attaching. + +
+ +
+ +
+
+ + Describe figures & charts + + Caption PDF figures, charts, tables and diagrams at upload with the + loaded model's vision, so their content becomes searchable. Needs a + vision model; adds vision calls for detected figures. + + + + Read charts and diagrams when attaching. + +
+ +
); } diff --git a/studio/frontend/src/features/rag/components/use-rag-documents.ts b/studio/frontend/src/features/rag/components/use-rag-documents.ts index 488d9ed6d0..8e756b2782 100644 --- a/studio/frontend/src/features/rag/components/use-rag-documents.ts +++ b/studio/frontend/src/features/rag/components/use-rag-documents.ts @@ -2,6 +2,12 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { useCallback, useEffect, useRef, useState } from "react"; +import { useChatRuntimeStore } from "@/features/chat"; + +import { + CHAT_RAG_CAPTION_KEY, + CHAT_RAG_OCR_KEY, +} from "@/features/chat/stores/chat-runtime-store"; import { toast } from "@/lib/toast"; import { deleteDocument, @@ -45,13 +51,19 @@ export function useRagDocuments( }, [documents]); // documentId -> signature; forgotten on delete, cleared on scope change. const sigByDocId = useRef>(new Map()); - const sigAttached = useCallback( - (sig: string) => { - for (const s of sigByDocId.current.values()) if (s === sig) return true; - return false; - }, - [], - ); + // Skip a re-selected file only if a matching doc is healthy or still indexing. A doc + // that completed with 0 chunks is re-ingestable (e.g. a scan attached before a vision + // model loaded); the backend re-ingests on the same hash, so let it through. + const sigBlocksReupload = useCallback((sig: string) => { + const ids = new Set(); + for (const [id, s] of sigByDocId.current) if (s === sig) ids.add(id); + if (ids.size === 0) return false; + const docs = documentsRef.current.filter((d) => ids.has(d.id)); + if (docs.length === 0) return false; // sig tracked but doc gone -> allow re-upload + return docs.some( + (d) => d.status !== "completed" || (d.numChunks ?? 0) > 0, + ); + }, []); // True while upload() runs, so the scope-change effect can tell a real switch // from lazy thread materialization mid-upload (which must not reset). const uploadInFlightRef = useRef(false); @@ -82,7 +94,11 @@ export function useRagDocuments( const controller = new AbortController(); trackedJobs.current.set(jobId, controller); - const finish = (status: DocumentStatus, error?: string | null) => { + const finish = ( + status: DocumentStatus, + error?: string | null, + numChunks?: number | null, + ) => { if (status === "failed") { // Drop the chip rather than show "Failed"; warn via toast. sigByDocId.current.delete(documentId); @@ -91,7 +107,14 @@ export function useRagDocuments( description: error ?? "Indexing failed", }); } else { - patchDoc(documentId, { status, error: null, progress: 1 }); + // Record numChunks so re-selecting this file dedups (vs a 0-chunk doc, which + // stays re-ingestable); the SSE "complete" frame carries it. + patchDoc(documentId, { + status, + error: null, + progress: 1, + ...(numChunks != null ? { numChunks } : {}), + }); } trackedJobs.current.delete(jobId); }; @@ -105,7 +128,7 @@ export function useRagDocuments( progress: ev.progress ?? null, }); } else if (ev.type === "complete") { - finish("completed"); + finish("completed", null, ev.num_chunks); return; } else if (ev.type === "error") { finish("failed", ev.error ?? "Indexing failed"); @@ -121,6 +144,7 @@ export function useRagDocuments( ? "failed" : "completed", job.error, + job.numChunks, ); } catch { if (controller.signal.aborted) { @@ -132,7 +156,8 @@ export function useRagDocuments( for (let i = 0; i < 600; i++) { if (controller.signal.aborted) break; const job = await getJob(jobId); - if (job.status === "completed") return finish("completed"); + if (job.status === "completed") + return finish("completed", null, job.numChunks); if (job.status === "failed") { return finish("failed", job.error ?? "Indexing failed"); } @@ -231,12 +256,21 @@ export function useRagDocuments( tempId: string, ) => { try { + // Send vision-pass overrides only after the user has explicitly set them; + // otherwise backend env defaults own the ingest policy. + const state = useChatRuntimeStore.getState(); + const hasLocal = (key: string) => + typeof window !== "undefined" && window.localStorage.getItem(key) !== null; + const ocr = hasLocal(CHAT_RAG_OCR_KEY) ? state.ragOcrScanned : undefined; + const caption = hasLocal(CHAT_RAG_CAPTION_KEY) + ? state.ragCaptionFigures + : undefined; const result = activeScope.type === "kb" - ? await uploadKnowledgeBaseDocument(activeScope.kbId, file) + ? await uploadKnowledgeBaseDocument(activeScope.kbId, file, ocr, caption) : activeScope.type === "project" - ? await uploadProjectDocument(activeScope.projectId, file) - : await uploadThreadDocument(activeScope.threadId, file); + ? await uploadProjectDocument(activeScope.projectId, file, ocr, caption) + : await uploadThreadDocument(activeScope.threadId, file, ocr, caption); sigByDocId.current.set(result.documentId, fileSignature(file)); if (seenIds.has(result.documentId)) { setDocuments((rows) => rows.filter((row) => row.id !== tempId)); @@ -286,7 +320,7 @@ export function useRagDocuments( // one look like nothing happened. Dedup re-selections up front. const fresh: Array<{ tempId: string; file: File }> = []; for (const file of Array.from(files)) { - if (sigAttached(fileSignature(file))) { + if (sigBlocksReupload(fileSignature(file))) { toast.info(`${file.name} is already indexed - skipping`); continue; } @@ -332,7 +366,7 @@ export function useRagDocuments( uploadInFlightRef.current = false; } }, - [scope, uploadOne, sigAttached], + [scope, uploadOne, sigBlocksReupload], ); const remove = useCallback( diff --git a/studio/frontend/src/features/rag/types/rag.ts b/studio/frontend/src/features/rag/types/rag.ts index d01854f5a7..1277500ae6 100644 --- a/studio/frontend/src/features/rag/types/rag.ts +++ b/studio/frontend/src/features/rag/types/rag.ts @@ -39,6 +39,7 @@ export interface IndexJob { stage?: string | null; progress?: number | null; error?: string | null; + numChunks?: number | null; } /** One SSE frame from /jobs/{jobId}/events. */ @@ -47,6 +48,7 @@ export interface JobEvent { stage?: string | null; progress?: number | null; error?: string | null; + num_chunks?: number | null; } /** Coords 0..1, top-left origin. */ From 0a3e5a3172fc058680deccb6a94669dac2d487e9 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Tue, 30 Jun 2026 07:52:03 -0700 Subject: [PATCH 41/49] Studio: quick eject from the model selector (#6654) * Studio: quick eject from the model selector Add a one-click eject shortcut to the loaded-model pill so users do not have to open the picker to unload a model. - The loaded-status indicator shows a green checkmark at rest and swaps to a red eject icon on pill hover, with an "Eject model" tooltip. Clicking it ejects without opening the picker. - On Device tab now uses the placeholder "Search local models" instead of "Search Unsloth models". - The picker's "Eject model" button uses medium font weight. * Studio: drop unused group/eject marker class on the eject control * Studio: make the inline eject control valid HTML The eject shortcut was a focusable span (role/tabIndex) nested inside the trigger button. A button's content model forbids focusable descendants, so make it a plain decorative span (aria-hidden, no role/tabIndex) that keeps the mouse shortcut. Keyboard and screen-reader users eject via the picker's "Eject model" button. * Studio: disable the inline eject shortcut on touch devices On touch (no hover) the red eject icon and title tooltip never reveal, so tapping the loaded pill could unload the model with no visible affordance. Add [@media(hover:none)]:pointer-events-none so taps fall through to the trigger and open the picker; touch users eject from the picker instead. --------- Co-authored-by: shimmyshimmer Co-authored-by: Wasim Yousef Said --- .../assistant-ui/model-selector.tsx | 56 ++++++++++++++++--- .../assistant-ui/model-selector/pickers.tsx | 8 ++- 2 files changed, 55 insertions(+), 9 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 8938f4dda7..1fddf077ba 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -14,6 +14,7 @@ import { isCustomProviderType } from "@/features/chat/external-providers"; import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { cn } from "@/lib/utils"; import { + CheckmarkCircle02Icon, CloudIcon, DashboardSquare01Icon, Download01Icon, @@ -146,6 +147,7 @@ function ModelSelectorTrigger({ size = "default", className, dataTour, + onEject, }: { currentModel?: ModelOption; isLoaded: boolean; @@ -154,6 +156,7 @@ function ModelSelectorTrigger({ size?: "sm" | "default" | "lg"; className?: string; dataTour?: string; + onEject?: () => void; }) { return ( @@ -161,12 +164,15 @@ function ModelSelectorTrigger({ type="button" data-tour={dataTour} className={cn( - "unsloth-model-selector-trigger flex min-w-0 items-center gap-2 transition-colors", + "unsloth-model-selector-trigger group/trigger flex min-w-0 items-center gap-2 transition-colors", + // Suppress the pill's hover background while the eject hit area is + // hovered, so only the dot's own circle reacts. variant === "outline" && - "rounded-full border border-border/60 hover:bg-[#ececec] dark:hover:bg-[#2d2e32]", + "rounded-full border border-border/60 hover:bg-[#ececec] has-[[data-eject-hit]:hover]:!bg-transparent dark:hover:bg-[#2d2e32]", variant === "ghost" && - "rounded-full hover:bg-[#ececec] dark:hover:bg-[#2d2e32]", - variant === "muted" && "rounded-full bg-muted hover:bg-muted/80", + "rounded-full hover:bg-[#ececec] has-[[data-eject-hit]:hover]:!bg-transparent dark:hover:bg-[#2d2e32]", + variant === "muted" && + "rounded-full bg-muted hover:bg-muted/80 has-[[data-eject-hit]:hover]:!bg-muted", // More left padding than right; the chevron is pulled close to the // label (below) so the trigger reads balanced around the text. size === "sm" && "h-8 pl-3 pr-1.5 text-xs", @@ -175,9 +181,44 @@ function ModelSelectorTrigger({ className, )} > - {isLoaded && ( - - )} + {isLoaded && + (onEject ? ( + // Loaded status doubles as a mouse eject shortcut: green checkmark + // at rest, red eject icon on pill hover, click to eject. A plain + // span (no role/tabIndex) keeps it out of the trigger button's + // content model, which forbids focusable descendants. Keyboard and + // screen-reader users eject via the picker's "Eject model" button. + // aria-hidden marks it decorative; stopPropagation stops the + // popover from toggling. On touch (no hover) the eject icon and + // tooltip never reveal, so pointer-events-none disables the + // shortcut there and taps open the picker instead of ejecting. + event.stopPropagation()} + onClick={(event) => { + event.stopPropagation(); + onEject(); + }} + // Hit area larger than the icon, with a hover circle. Negative + // margin keeps the icon in the dot's original spot. + className="-m-1 flex size-5 shrink-0 cursor-pointer items-center justify-center rounded-full transition-colors hover:bg-black/10 dark:hover:bg-white/10 [@media(hover:none)]:pointer-events-none" + > + + + + ) : ( + + ))} {currentModel?.icon ? ( {currentModel.icon} @@ -644,6 +685,7 @@ export function ModelSelector({ size={size} className={className} dataTour={triggerDataTour} + onEject={onEject ? handleEject : undefined} /> setQuery(event.target.value)} - placeholder="Search Unsloth models" + placeholder={ + section === "downloaded" + ? "Search local models" + : "Search Unsloth models" + } data-model-picker-search-input={true} className="field-soft h-9 border-0 pl-8 pr-8" /> @@ -3447,7 +3451,7 @@ export function HubModelPicker({