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 01/31] 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 02/31] 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 03/31] 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 04/31] 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 05/31] 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 06/31] 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 07/31] 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 08/31] 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 09/31] 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 10/31] 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 11/31] =?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 12/31] 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 13/31] 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 14/31] 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 15/31] 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 16/31] 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 17/31] 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 18/31] 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 19/31] 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 20/31] 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 21/31] 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 22/31] (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 23/31] 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 24/31] 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 25/31] 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 26/31] 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 27/31] 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 28/31] 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 29/31] 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 30/31] 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 31/31] 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